perf(recipe): tune Nemotron Omni H100 fine-tuning

[cuichenx]

What does this PR do ?

Adds validated one-node H100 fine-tuning variants for Nemotron 3 Nano Omni, including lower-precision optimizer state, HybridEP for 4K CORD-v2 SFT/PEFT, and packed CP2 all-to-all execution for 8K SFT. It also adds an explicit, default-off option that makes Transformer Engine precision-aware optimizer checkpoints saveable and reloadable under tight GPU-memory margins without changing their portable values or dtypes.

The promoted 4K full-SFT recipe now uses MBS4, fused attention, Transformer Engine cross entropy, fused weighted squared-ReLU, router fusion, and selective MoE/layernorm recompute. These are execution changes only: the real image-text data, objective, GBS64, trained parameters, and natural routing semantics are preserved.

Changelog

• Add self-contained eight-H100 CORD-v2 full-SFT, packed 8K SFT, and LoRA PEFT recipes while preserving real image-text data, batch semantics, objectives, trained parameters, and natural MoE routing.
• Store BF16 main gradients and Adam moments with scaled FP16 optimizer main parameters and no parameter remainders, replacing the prior FP32 optimizer state in the promoted H100 variants.
• Fail closed during optimizer setup if Transformer Engine's effective main/moment dtypes or remainder state differ from the requested Bridge precision contract.
• Use HybridEP with safe zero-weight uneven-dispatch padding for 4K SFT/PEFT; retain all-to-all for packed CP2 after a controlled forward-progress screen.
• Add the default-off stage_precision_aware_optimizer_state_on_cpu checkpoint option. Enable it only in the two full-SFT H100 recipes whose complete optimizer checkpoints otherwise exhaust device memory; reject non-torch_dist use and leave PEFT and portable recipes disabled.
• Add focused recipe and checkpoint tests and update the model verification card with exact-commit, immutable-revision 100-step evidence.

Validation

All promoted runs used the public nvcr.io/nvidia/nemo:26.06 container, immutable model and CORD-v2 revisions, eight H100 GPUs, and real image-text samples.

The optimized 4K full-SFT run used exact clean Bridge commit 9b69b47a321de14b0453c230dd130e18ae0f00f7 and pinned Megatron-Core commit 24bad8e677d22625d86ef2a54c9506b6e4992c93. Runtime validation confirmed both dense and expert Transformer Engine FusedAdam instances as scaled FP16 main parameters, BF16 main gradients/moments, and no parameter remainders on every rank.

Workload	Topology	Steps	Loss	Final-10 ms	Final-10 TFLOP/s/GPU	Peak sampled memory
4K full SFT	TP2/CP1/EP8, HybridEP	100	1.281308 → 0.03330838	3,603.67	197.10	76,349 MiB
4K LoRA PEFT	TP2/CP1/EP8, HybridEP	100	1.282941 → 0.02634745	11,621.42	61.25	44,873 MiB
8K packed full SFT	TP4/CP2/EP8, all-to-all	100	1.251364 → 0.01020405	15,577.90	93.73	75,197 MiB

The optimized 4K run improves final-ten model throughput by 77.3% over the prior exact 111.15 TFLOP/s/GPU evidence. It recorded zero skipped/NaN iterations and saved complete model, optimizer, scheduler, RNG, and train-state checkpoints at steps 50 and 100. A separate direct continuation loaded the step-50 distributed optimizer and train state, completed steps 51–100 with zero skipped/NaN iterations, and reproduced final-ten performance at 3,604.14 ms / 197.07 TFLOP/s/GPU. Across the 50 aligned resumed steps, mean absolute LM-loss drift was 0.00050938 and maximum absolute drift was 0.00155107 under natural HybridEP routing.

The final opt-in implementation was rechecked at exact clean Bridge commit e88515d35c4efb4c4bd1a4d9a0ae3c4f865f0c9d: the enabled 4K full-SFT recipe saved a complete distributed step-51 optimizer/RNG checkpoint (about 445.8 GB), and a fresh process loaded that checkpoint and completed step 52 with finite loss/gradient and zero skipped/NaN iterations. This follow-up changes only whether CPU staging is selected; it does not change training arithmetic or the 100-step performance configuration above.

Controlled full-SFT screens rejected alternatives that were slower, memory-heavier, unsupported for this architecture, or incompatible with the reduced-precision optimizer contract. In particular, fused residual RMSNorm was noise-level faster but consumed about 4 GiB more sampled peak memory; fused QKV+RoPE does not support the real multimodal forward shape; grouped shared-expert GEMM requires SwiGLU rather than this model's squared-ReLU; and capturable optimizer graphing requires FP32 master weights.

Additional checks:

• 230 focused and adjacent recipe/optimizer/checkpoint unit tests passed for the implementation; an additional 38 checkpoint-config tests passed.
• The final flattened/tuned recipe passes all 19 targeted Nemotron Omni recipe tests.
• uv run pre-commit run --all-files passed.
• Model verification card validator passed.
• Public-diff privacy scan and git diff --check passed.
• All commits have valid GPG signatures and DCO sign-offs.

GitHub Actions CI

This is a draft PR. Final CI can be triggered against the published head commit.

Before your PR is "Ready for review"

Pre checks:

• Read and followed the contributor guidelines.
• Added focused recipe and checkpoint tests.
• Updated the model verification card.
• Optional-component impact: no new optional dependency or import surface.

Additional Information

No linked issue.

[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.

[cuichenx]

/ok to test d6fb26b

[cuichenx]

/ok to test cd0a475

[cuichenx]

/ok to test 93ccec7

[kamran-nvidia]

/ok to test d7da085

[kamran-nvidia]

/ok to test 1cf8d6d