perf: speed up TransformerEngine source builds

[joyang-nv]

What does this PR do ?

Reduces cold TransformerEngine source-build time by limiting CUDA compilation to the H100/B200 architectures currently targeted by the image, increasing TE-native parallelism, enabling nvcc worker threads, and selecting only the PyTorch framework. All build knobs remain overridable.

It also makes the GB200/arm64 Apptainer source installation resilient to transient GitHub rate limiting while downloading build dependencies.

Why these TransformerEngine settings?

NVTE_CUDA_ARCHS=90;100

The pinned TransformerEngine release defaults to building CUDA 13 code for 75;80;89;90;100;120. The NeMo-RL image currently sets TORCH_CUDA_ARCH_LIST="9.0 10.0" for H100 and B200, so compiling the four additional architecture families adds substantial source-build time without benefiting those targets. TransformerEngine internally expands architecture 100 into the required generic and Blackwell-specific variants for architecture-specific sources.

This is the main performance optimization, but it has a compatibility tradeoff: this source-built TransformerEngine will not support A100, L40, or other omitted architectures. If the release image must remain broadly portable, the architecture restriction should instead be supplied only to the H100/GB200 CI builds as an overridable build argument.

NVTE_FRAMEWORK=pytorch

TransformerEngine can auto-detect installed frameworks and build bindings for PyTorch, JAX, or Paddle. NeMo-RL and Megatron use the PyTorch bindings, so setting this explicitly makes the build deterministic and prevents an environment-dependent extra framework build. This is primarily a correctness/reproducibility guard; it is expected to contribute much less to the speedup than restricting CUDA architectures.

Apptainer dependency-download reliability

The GB200/arm64 image build also exposed a transient failure while Apptainer downloaded an unauthenticated GitHub patch:

16.98 + curl -f -L -sS -o ./e2fsprogs-246.patch https://github.com/tytso/e2fsprogs/pull/246.patch
17.09 curl: (22) The requested URL returned error: 429

docker/install_apptainer.sh now retries Apptainer's dependency downloader up to five times with exponential backoff (5s, 10s, 20s, and 40s between attempts). Apptainer cleans partial dependency files before each attempt; if all attempts fail, the script preserves the final nonzero status and fails the image build.

Issues

None.

Usage

No user-facing API changes. Docker builds may override NVTE_BUILD_MAX_JOBS, NVTE_BUILD_THREADS_PER_JOB, NVTE_CUDA_ARCHS, or NVTE_FRAMEWORK as build arguments.

Performance comparison

The before and after measurements use the same CI build path and the same Building transformer-engine / Built transformer-engine timestamps.

Measurement	Before optimization	After optimization	Change
TransformerEngine source build	4,393.1s (73m13.1s)	1,260.6s (21m00.6s)	3,132.5s saved (71.3% reduction, 3.49x faster)

Before optimization

The baseline CI log reports:

#130 219.7    Building transformer-engine @ git+https://github.com/NVIDIA/TransformerEngine.git@42b840051647eef89761a16dfdff87e82bb253ab
#130 427.9       Built fast-hadamard-transform @ git+https://github.com/Dao-AILab/fast-hadamard-transform.git@e7706faf8d1c3b9f241e36860640ad1dac644ede
#130 630.6       Built mamba-ssm @ git+https://github.com/state-spaces/mamba.git@a14b1dff0454a3bc27d9eb31355dc01e4b2490ec
#130 831.5       Built causal-conv1d @ git+https://github.com/Dao-AILab/causal-conv1d@4f6ae4e26ae5fe8af9372f8d312ab25cc4595223
#130 4612.8      Built transformer-engine @ git+https://github.com/NVIDIA/TransformerEngine.git@42b840051647eef89761a16dfdff87e82bb253ab

uv builds these source packages concurrently; the three other native packages are sibling builds, not TransformerEngine phases. They all finish by 831.5s, while TransformerEngine continues until 4612.8s. TransformerEngine is therefore the critical path and remains active for 3,781.3s (63m01.3s) after the final listed sibling build completes.

After optimization

The optimized CI log reports:

#130 219.9    Building transformer-engine @ git+https://github.com/NVIDIA/TransformerEngine.git@42b840051647eef89761a16dfdff87e82bb253ab
#130 435.0       Built fast-hadamard-transform @ git+https://github.com/Dao-AILab/fast-hadamard-transform.git@e7706faf8d1c3b9f241e36860640ad1dac644ede
#130 634.6       Built mamba-ssm @ git+https://github.com/state-spaces/mamba.git@a14b1dff0454a3bc27d9eb31355dc01e4b2490ec
#130 846.1       Built causal-conv1d @ git+https://github.com/Dao-AILab/causal-conv1d@4f6ae4e26ae5fe8af9372f8d312ab25cc4595223
#130 1480.5      Built transformer-engine @ git+https://github.com/NVIDIA/TransformerEngine.git@42b840051647eef89761a16dfdff87e82bb253ab

The sibling package timings remain similar, while TransformerEngine's source-build interval falls from 4,393.1s to 1,260.6s. After the final listed sibling build completes, the remaining TransformerEngine critical-path time falls from 63m01.3s to 10m34.4s.

Validation

Additional local measurements:

• Standalone cold TransformerEngine wheel build: 11m44s.
• End-to-end cold-TE uv sync --link-mode symlink --frozen --extra mcore --no-install-project: 18m00s, including dependency downloads, recursive TE checkout/submodules, and concurrent native dependency builds.
• BuildKit --call=outline validation passed and exposed the intended argument defaults.
• bash -n docker/install_apptainer.sh passed.
• git diff --check passed.

Before your PR is "Ready for review"

Pre checks:

• Read and followed the contributor guidelines.
• No new tests are necessary for this Docker build-configuration change; cold source-build validation is documented above.
• Unit and functional test suites were not run locally; CI:L0 will run the applicable image, unit, lint, and documentation checks.
• No user-facing documentation changes are necessary; the settings are documented inline and remain overridable.

Additional Information

The existing global uv cache continues to reuse completed TE wheels. A separate persistent ccache flow was intentionally avoided because it would not improve the first cold build and would require additional registry-cache lifecycle complexity for ephemeral CI builders.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[joyang-nv]

/ok to test d193bcb

[joyang-nv]

/ok to test a801ec5

[joyang-nv]

/ok to test 26abf6d

[joyang-nv]

Root cause of the uv lock --check failure:

• b0edc6e7 (chore: Bump Gym (#3326)) advanced the Gym submodule from d67ad661 to 3e0f307d.
• The new Gym revision added dependencies, but the root uv.lock was not regenerated. The late full-workspace lock check in this PR correctly detected that mismatch.
• 26abf6df regenerates the lockfile with CI-pinned uv 0.11.18; uv lock --check now passes for all 516 packages.

This failure was therefore unrelated to the TransformerEngine optimization or Apptainer retry changes.
@yfw for CC.