build: update development PyTorch image to 26.09

[balasaajay]

Update the development CI image from NVIDIA PyTorch 26.08 to 26.09 while retaining the merged Transformer Engine 2.20 release line. Validate the image update with functional FlashAttention pinned to generation 2.

Build and runtime changes

• Align the GitHub image pin and both GitLab development architecture pins at 26.09. Retain TE source 6ea2a74a9e98c99e6d7b164a33775cc457520027 (2.20.2+6ea2a74a) and cuDNN frontend 1.29.0.
• Keep DeepGEMM ea2b7805 and libdw-dev for the PyTorch/CUDA C++20 and FP8-header build compatibility fixes. Keep Mamba's C++20 build-flag backport without changing its runtime source pin.
• Revert this PR's FA4 Quack patch, Docker patch application, Hopper SplitKV/small-head/LSE workarounds and paged-FA4 replay test. Remove its direct CuTe/Quack dependency additions and restore the base TileLang 0.1.12 / TVM FFI 0.1.11 overrides. Preserve the base repository's existing FA4 support and other dependency constraints.
• Apply PR #7800: replace the build-specific CUDA graph RNG registration guard with a cached runtime capture probe and include its regression test. Retain clamped-SwiGLU backward compatibility and replay coverage, SSM decode-padding correction and the FSDP checkpoint None guard.

Test configuration

Set NVTE_FLASH_ATTN_V4=0 and pass --flash-attention-version 2 in MCore functional training/inference, inference-performance and both custom inference-server launchers. flash_attention_version is a model argument, not an environment variable; the explicit argument makes the pin effective. This selects the generation when FlashAttention is used and preserves each test's existing fused/local/FlashInfer/FlashMLA backend configuration.

The determinism profiling and MIMO checkpoint round-trip launchers also explicitly set NVTE_FLASH_ATTN_V4=0 and --flash-attention-version 2. These two custom paths bypass the shared training launcher; five added environment/argument lines complete the requested pin without changing their assertions, tolerances, test selection or golden values.

Set NVTE_FLASH_ATTN_V4=0 in both H100/GB200 unit recipes, covering latest and legacy tags. Generic shared GPT/hybrid inference configurations explicitly default flash_attention_version to 2; callers can still select another generation. The separate static-engine, hybrid-prefix, chunked-prefill, MTP and text-generation-controller fixture builders are also pinned to FA2. These generic builders otherwise selected unsupported FA4 paged attention and caused assertions followed by rank-divergence timeouts. The startup environment export alone does not control Megatron's direct attention dispatch, and unit initialization intentionally resets TE environment variables. Dedicated backend tests retain their selections. Temporarily mark only TestDynamicInferenceNVLS::test_batch_invariant_prefill_matches_full_forward as flaky_in_dev, tracked in #7958; this batch-invariant case requires FA3/FA4 and cannot be forced onto FA2. Its numerical cause remains unisolated. Test bodies, assertions, tolerances and all other test selections are preserved.

Restore the two MTP inference logprob references and the NanoV3 batch-128 performance reference to the PR base. Re-enable 12 GitHub inference cases (10 H100, 2 GB200) previously quarantined during the FA4-era numerical investigation for FA2 validation. That re-enable preserves GitLab MR/nightly selection. Keep the independently recorded baseline, NCCL and unit-state quarantines. Re-enable the four H100 and two one-node GB200 CP2 GitHub cases tracked in #7733, as requested after discussion with MCore; all six now pass five full FA2 repeats against their existing exact goldens. No assertion, tolerance or comparator is relaxed.

Functional golden references

Refresh eight development references from completed training in the current FA2 configuration: three Muon hardware/case combinations, H100 MFSDP v2 CP2, and four GB200 nightly CP2 variants. Loss changes remain below 0.01% at every step. The gradient-zero updates are explicitly approved after checking repeatability, including individual-step changes above the earlier 1% limit. All selected loss/zero trajectories are finite and identical across three observed attempts; these are retries of the first comparison, not five completed repeats. Checkpoint-resume tails for the Muon cases also match their uninterrupted trajectories.

Use the original generated full-precision artifact values. Update loss and gradient zeros where they failed; MFSDP v2 CP2 changes loss only. The two legacy GB200 Muon references require a complete canonical export to satisfy the full-precision validator, also migrating their ungated memory/timing fields. Those fields do not change these recipes’ pass/fail checks and are not a performance acceptance result. This golden refresh preserves unrelated blocks in the other six files and leaves test assertions, tolerances and recipe selection unchanged.

Signed means below are 100 × mean((old-new)/old); positive means lower new values. Muon loss/zeros/memory use 100 steps and timing 99; other cases use 50 steps and timing 49. No compared loss/zero samples are missing, non-finite or omitted for a zero denominator.

Case / hardware	Loss mean	Gradient-zero mean	Allocated-memory mean	Peak-memory mean	Timing mean
gpt3_mcore_te_mfsdp_v2_cp2 / H100	+0.00006624%	Unchanged	Unchanged	Unchanged	Unchanged
gpt3_mcore_te_tp2_pp2_cp2 / GB200	-0.00001952%	-0.33642704%	Unchanged	Unchanged	Unchanged
gpt3_mcore_te_tp2_pp2_cp2_calculate_per_token_loss / GB200	-0.00001952%	-0.33642704%	Unchanged	Unchanged	Unchanged
gpt3_mcore_te_tp2_pp2_cp2_etp4_calculate_per_token_loss_dp_last / GB200	-0.00001952%	-0.33642704%	Unchanged	Unchanged	Unchanged
gpt3_mcore_te_tp2_pp2_cp2_etp4_dp_last / GB200	-0.00001952%	-0.33642704%	Unchanged	Unchanged	Unchanged
gpt3_moe_mcore_te_ep8_resume_torch_dist_dist_muon / GB200	-0.00045020%	-0.37829370%	-0.05266490%	-0.01377754%	+24.32124745%
gpt3_moe_mcore_te_ep8_resume_torch_dist_muon / GB200	+0.00007861%	-0.57728597%	+0.06129428%	+0.07394847%	+13.54333840%
gpt3_moe_mcore_te_ep8_resume_torch_dist_muon / H100	-0.00013851%	-0.11381538%	Unchanged	Unchanged	Unchanged

Maximum individual-step gradient-zero changes are 8.9825% for GB200 Muon, 10.7590% for GB200 distributed Muon, 11% for H100 Muon, and 12.5397% for the four nightly CP2 variants. Their means above are not substitutes for per-step review. The legacy GB200 timing fields include noisy warmup samples and were migrated only as part of the complete original export.

Retain the separately authorized one-node GB200 DeepSeek loss refresh; its existing updated reference passed five full FA2 repeats. The six GitHub CP2 cases documented below use different hardware/case identities from the four GB200 nightly variants refreshed here.

CUDA graph RNG registration compatibility

Apply PR #7800 at head 2edf97fa42dc590621a81ce0a54f0f4fa0e19671. Its cached runtime probe detects whether a CUDA graph rejects RNG operations on an unregistered generator. This replaces the exact NGC 26.08 build exemption, since PyTorch version strings do not reliably distinguish lazy-registration behavior. Calls made during an active capture conservatively request registration and defer probing.

Both ported files match #7800's source exactly. The included regression passed through torch.distributed.run --standalone --nproc-per-node 1 on RTX 6000 Ada with both NGC 26.09 / PyTorch 2.14.0a0+b2c75dd062 (registration required) and NGC 26.08 / PyTorch 2.14.0a0+4fdf77b940 (registration unnecessary). Black, isort, Ruff, Pylint and whitespace validation pass. This is focused local GPU validation, not H100/GB200 acceptance.

CP2 validation and golden references

Signed head a7a6534eec8c8b0968810b7d307c198e749a333d contains the #7800 port and re-enables exactly six CP2 cases from #7733: four H100 variants and two one-node GB200 aliases. Recipe-parser validation confirms exactly these six additions to GitHub L1, with GitLab MR/nightly selection unchanged.

All six cases passed five complete 50-step repeats each using current source a7a6534ee and FA2: 30 official comparator passes, including exact and approximate loss/gradient-zero checks. Independently decoded all 30 original TensorBoard event files and verified that their full-precision loss and zero-count arrays match the generated JSON and existing references exactly. Runtime source, effective FA2 selection and disabled FA4 were verified from the original logs.

No golden-file edits are needed for these six GitHub CP2 cases. All 12 gated metric blocks match the existing references, including every individual step. The finite/full-precision checker passes for all six files. No assertion, tolerance or comparator was changed, and timing/memory blocks were preserved.

Signed differences below are 100 × mean((old-new)/old) over 50 shared steps. Positive means the new value is lower; all means and maximum individual-step differences here are exactly zero.

Case / environment	Repeats × steps	lm loss	num-zeros	Original JET pipeline
gpt3_mcore_te_tp2_pp2_cp2_1node / gb200	5 × 50	0.000000%	0.000000%	72301798
gpt3_mcore_te_tp2_pp2_cp2_calculate_per_token_loss_1node / gb200	5 × 50	0.000000%	0.000000%	72301807
gpt3_mcore_te_tp2_pp2_cp2 / h100	5 × 50	0.000000%	0.000000%	72301809
gpt3_mcore_te_tp2_pp2_cp2_calculate_per_token_loss / h100	5 × 50	0.000000%	0.000000%	72301783
gpt3_mcore_te_tp2_pp2_cp2_etp4_calculate_per_token_loss_dp_last / h100	5 × 50	0.000000%	0.000000%	72301797
gpt3_mcore_te_tp2_pp2_cp2_etp4_dp_last / h100	5 × 50	0.000000%	0.000000%	72301826

The targeted workloads used verified immutable dependency images while building/testing the explicitly selected current source. H100 used 72272726-amd64@sha256:33073b99eaade9f718bfe0410fae9cba3ef4c1e7f9d2d12a56c3b7347a7668a1, built at d842b7a44; its effective AMD64 dependencies match this revision (subsequent image changes add an ARM64-only MoK stage). GB200 used 72281847-arm64@sha256:86c021b904caf5c031f78059ed3f001a69ac718efe19bfadc3a954fe8adcbef5, built at 090200d9e, with byte-identical dependency/build inputs. Wrapper build traces preserve this image provenance. Earlier setup and scheduler failures occurred before training and are excluded from completed-repeat counts.

Current full CI status

Signed head 89fde7cdec0878005da42ba4c2aec08eff9eac25 adds only two H100 GitLab recipe-line changes to the previously validated 880404c7e revision. The distributed Adam and distributed Muon EP8 cases are already quarantined on GitHub; extend their quarantine to GitLab mr, and also mr-slim for distributed Muon. Both pass ten regular/resume pairs in each of the two scheduled main baselines, but repeatedly fail on the PR with unchanged model configurations and H100 references.

H100 case	Largest absolute per-step loss change	Largest per-step gradient-zero change	Existing approximate checks	Tracker
Distributed Adam	0.00686620%	1,858.51635% (step 99: 53,638 → 1,050,509)	Loss passes; zeros fail at 8/100 steps	#7201
Distributed Muon	0.01060030%	9.70565% (step 27: 1,257 → 1,379)	Both pass	#7734

These two H100 references are outside the approved eight-reference exception. Preserve their goldens, test code and tolerances. Recipe-parser validation confirms only these GitLab exclusions; GitHub selection, GB200 cases, the passing regular-Muon case and baseline-failing cases remain unchanged. The commit is signed and signed off, and whitespace validation passes.

Current-head validation, as of 2026-10-08 17:11 UTC:

• Signed head 5b38ce0d2c92bfe59be235f37400889937e437af adds the five custom-launcher configuration lines described above. Black, isort, Ruff, Pylint, shell syntax and whitespace checks pass. Both modified custom paths now pass their actual contracts: the deterministic/non-deterministic profiling ratio is 1.070, below 1.35, and MIMO passes its single checkpoint round-trip comparison. Their FA2/NVTE4-off configuration is verified through the executed source and available arguments; this is not a claim about a particular kernel implementation. MIMO captures successful subprocess output, so its imported TE version is not independently dumped.
• GitHub CICD 37787649936 tests merge 3396df0792b25951089cb02da9a24c9370779df3, with parents main 91655aa78007e74e55c963f99861b8a97368c024 and head 5b38ce0d2. Both native builds and all 42 unit buckets passed and were checked against their original results. GDN2 required one automatic launcher recovery; the MoE bucket passed its 830 selected cases on all eight ranks without a retry. All 110 functional cases (80 H100, 30 GB200) are accepted: 100 completed their five-repeat contracts, with 715 canonical comparator passes, and ten custom/smoke cases passed their own contracts. The full workflow and unit-test coverage report completed successfully. Earlier-head results are not credited to this run.
• GitLab MR 72440710 and nightly 72440728 use fork source 5b38ce0d2. All native images built successfully, and the actual generated matrices are verified: 158 MR cases and 167 nightly cases, excluding GB300. Standard numerical cases use five repeats; custom contracts retain their existing behavior. The default time limit is 2,700 seconds, with the existing 5,400-second override for one nightly Nemotron recipe. LTS pins are unchanged.
• All eight refreshed references passed on this head: 55 canonical comparisons and 3,500 training steps, including all five repetitions and applicable resume checks. Independent full-precision artifact audits verified finite loss and applicable gradient-zero arrays, exact agreement with the adopted references, and no worker restarts in these eight workloads. Original runtime logs bind these results to MCore 5b38ce0d2, TE 2.20.2+6ea2a74a, and the FA2 configuration. This numerical result does not establish performance or memory equivalence.
• All ten CP2 case/platform combinations tracked by #7733 are now validated: four H100 GitHub cases, two one-node GB200 GitHub aliases and four two-node GB200 nightly cases. Their original artifacts contain 50 passing comparisons, 2,500 steps and 5,000 finite loss/gradient-zero samples, exact across each case's five repeats and current reference. The six GitHub cases retain their existing references; the four nightly references use the approved updates above. The CI-failure tracker is closed with its historical drift/memory evidence and causal uncertainty preserved.
• Nightly is complete: 160 accepted results and 7 failed recipes. The 160 results comprise 158 numerical cases and two intentional smokes, with 1,125 comparator passes. Six failures match the failing metric categories in the latest completed scheduled nightly baseline. Two-node DeepSeek MFSDPv1 has an additional PR exact-loss mismatch while its baseline fails timing; current approximate loss and timing checks pass. It remains separately tracked in #7735. All seven baseline-failing recipes remain enabled and unchanged as requested. Successful automatic recoveries are retained in the audit, including six final workloads that needed a TCPStore port restart; they do not establish a fleet repair.
• MR is complete: 151 accepted results, one additional green wrapper with a hidden baseline failure, and six failed recipes. Five failed recipes also fail in the latest completed scheduled baseline. The sixth, H100 GPTOSS, has a token trajectory identical to one observed in an older scheduled baseline; the latest scheduled baseline has no numerical result, so it does not establish current numerical equivalence. The green moe_perf wrapper contains the known baseline RandomSTE.generator error and is excluded from acceptance. The GRPO throughput wrapper ultimately succeeded after automatic recovery; its earlier attempts failed the symmetric timing gate with median iteration times 15.33% and 19.98% lower than the unchanged reference. Those failures remain in the record, and no reference or tolerance was changed.

The checkout-only DNS failure passed after the one scoped wrapper retry already recorded, completing five comparisons and 250 training steps. Other service-managed and worker recoveries are retained separately. The intermittent setup/compilation trackers #7742, #7745, and port-collision tracker #7920 remain open. Required human CODEOWNERS approvals are still pending.

Completed validation before the scope-only change: GitHub run 37740317698 completed successfully on head 880404c7e and tested merge a834439d06f682810337ac87d9c44664d5281f90 (main parent 159b443c). All 42 unit buckets and 110 functional jobs (80 H100, 30 GB200) passed. The functional originals contain 715 canonical comparator passes: 430 regular, 215 resume and 70 inference. One hundred jobs completed their five-repeat contracts; ten custom/smoke jobs were checked against their own contracts and are not represented as five-repeat numerical tests.

The four H100 inference buckets passed 3,976 distinct cases on all eight ranks, including all 50 cases addressed by the FA2 fixture pins. The Transformer bucket passed 1,272 distinct cases across its intended pipeline stages. The three GB200 unit buckets passed 220 distinct cases on all four ranks. Existing skips/deselections are outside acceptance. GDN2 parallel needed one automatic launcher recovery matching #7742. Five functional jobs recovered from worker TCPStore port collisions tracked in #7920; all their training and comparisons completed, with no functional launcher/manual retries. The intentional fault-injection case is recorded separately.

All seven priority functional cases (six re-enabled GitHub CP2 variants plus H100 MFSDPv2 CP2) also passed independent full-precision raw-value audits across all five repeats. Both canonical H100 and GB200 MTP cases passed all five inference comparisons with FA2 and the current references; #7730 is closed for that resolved configuration. This is not a claim that the historical FA4 numerical cause is understood. Runtime functional logs on both hardware lanes report TE 2.20.2+6ea2a74a.

The preceding exact-head GitLab MR 72369725 and nightly 72369841 accepted all eight refreshed references, completing 3,500 training steps and 55 canonical comparator passes over five repetitions per case. Retained finite, full-precision loss and applicable gradient-zero arrays match the adopted references across repetitions and checkpoint resumes. Earlier checkout-only failures are tracked separately. The two repeatedly failing H100 MR cases are being stopped after their evidence was retained; neither is counted as passing.

Baseline-failing GitLab recipes remain unchanged as requested. That includes the two-node DeepSeek MFSDPv1 nightly recipe: its baseline fails timing only, while the PR additionally has up to 2.86894% per-step loss drift (the approximate loss check passes). This extra loss mismatch exceeds the strict sub-1%-at-every-step update limit and remains tracked separately in #7735; it is not classified as a shared baseline loss failure. Original MoE performance logs also retain the baseline RandomSTE.generator error behind a successful wrapper, so wrapper success is not treated as performance acceptance.

[copy-pr-bot]

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

[balasaajay]

/ok to test dac0cff

[balasaajay]

/ok to test 6ebe61d

[balasaajay]

/ok to test e9f01a0

[balasaajay]

/ok to test 2f16cb1

[balasaajay]

/ok to test b73d6cb

[balasaajay]

/ok to test 25b5f63

[Phlip79]

Same here.

[balasaajay]

This patch fixes native-extension build failures because our pinned Mamba explicitly uses C++17 while the new PyTorch ATen headers require C++20; it changes only four compiler flags. No published Mamba release currently contains the fix, and state-spaces/mamba#1000 remains unmerged, so we retain the patch until a suitable upstream revision is adopted and validated.

[Phlip79]

Can you please add this comment to the actual file?

[balasaajay]

/ok to test 64977df

[balasaajay]

/ok to test d842b7a

[balasaajay]

/ok to test 090200d

[cspades]

LGTM for the inference test failure resolution specifically

flash_attention_version=2,