Released TransformerEngine in Dockerfile

[Oblynx]

The Dockerfile (=> also dev env) is currently using a specific commit of TransformerEnginer, and adding a patch on top manually. The patch was written in summer of 2025 and there's been multiple releases of TransformerEngine since then. Although I can't fully track the patch from the github issues/PRs, there's a good chance that the latest TE release (2.11) has included it.

Can we use the TE v2.11?

[Oblynx]

This TransformerEngine PR seems to solve the issue we patch in bionemo-framework: NVIDIA/TransformerEngine#1963

Merged July 22.

[pstjohn]

Hi @Oblynx, sorry for the slow reply! We've largely moved development over to the bionemo-recipes subfolder, where our Dockerfiles are much more stripped down, and we live closer to the bleeding pytorch / TransformerEngine versions. We may remove some of the top-level framework container code in the next few months.

What are you currently using that container for? Would your use case be covered by some of the work we're doing in recipes?

[Oblynx]

Hi @pstjohn thanks for your reply! I'm experimenting with Evo2 embeddings. I didn't see any mention of Evo2 in your recipes project. Do you think your Dockerfiles from recipes would work with Evo2, have you been developing like that lately?

[pstjohn]

we're actively working on https://github.com/NVIDIA/bionemo-framework/tree/main/bionemo-recipes/recipes/evo2_megatron, but let me defer to @jstjohn who is leading that effort

[jstjohn]

Yes, you should be able to start from that linked evo2 recipe, and bump the pytorch image to one that has the TE version you want. We're pretty up to date already.

[Oblynx]

Oh wow, your experiments on the README look like a great improvement! My goal is to run inference with the 40b1M model using as few GPUs as possible, and in fact with the nemo2 implementation so far I had no luck scaling the CP. I was constantly hitting OOM, despite using a couple of H200.

I'll try this version when I get the chance. Do you need any help with the packaging/tooling?