Selective-FP32 runtime and native-FP32 optimizer checkpoint lifecycle

[FDecaYed]

Selective-FP32 runtime and native-FP32 optimizer checkpoint lifecycle

Summary

Preserve explicitly marked FP32 parameters while converting the rest of a
module, provide an FP8-disabled execution context, and make optimizer
checkpoint mapping correct for interleaved low-precision and native-FP32
parameters.

Scope and non-goals

• Add the selective-FP32 marking, conversion, and execution primitives.
• Preserve native-FP32 parameters in optimizer sharded state dictionaries.
• Do not add mHC, DSv4 attention, model construction, or recipe wiring.

Provenance

This local recut is reconstructed from the frozen main baseline
bb5647a9bdd0, not a replay of historical commits. It references the DSv4
tracking PR #5795 and consolidates the complete contracts from #5929 and
#5930. Original DSv4 integration credit: @hxbai in #5795.

Dependencies

None beyond the frozen baseline used for this local recut. This is a reusable
prerequisite outside the nine-PR DSv4-specific series.

Tests

• Selective marking, FP16/BF16 conversion, and the FP8-disabled context.
• Interleaved BF16/FP32/BF16 optimizer-state mapping and Adam state keys.

Publication

Publish in the first prerequisite round, rebased onto then-current main, so
the GitHub changeset contains only this slice. This is not a stacked-review PR.

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

[FDecaYed]

/ok to test a2e06de

[FDecaYed]

/ok to test 4283485

[FDecaYed]

/ok to test 985cb27

[janEbert]

The inner for-loop needs to guard each iteration on whether param.requires_grad, according to the logic that assigns to self.float16_groups and self.fp32_from_fp32_groups, right? Otherwise we may iterate float16_params too early and for unrelated parameters.
I think the design is a bit fragile with the zip potentially ending early and the next throwing a hard-to-understand error.

[FDecaYed]

requires_grad guard would still be unsafe here because the FP32 main copies are detached and therefore also have requires_grad=False.

I replaced the positional iterator with an explicit FP32-main-parameter-ID to model-parameter map. Frozen and native-FP32 parameters remain mapped to themselves, while the fp32_from_fp16 checkpoint IDs are now restricted to actual main copies. I also added a regression covering interleaved trainable BF16 / frozen BF16 / native FP32 / trainable BF16 parameters and verifying the resulting optimizer state IDs.

[janEbert]

Thanks! Also, adding strict=True is great.

[janEbert]

/ok to test 8ae07a9

[asolergi-nv]

/claude review

[claude]

Light review. The optimizer state-id remapping in sharded_state_dict looks correct (ids now enumerate self.optimizer.param_groups instead of assuming every inner param is an fp32 main copy, and float16_param_ids_per_group pairs fp32_from_fp16_params back to the right ids), and the param_to_fp32_param.get() guard in hybrid_optimizer.py matches the other call sites that already used .get()/in. Two things:

1. mark_keep_in_fp32 does not survive a device move for buffers.

megatron/core/transformer/module.py:437 documents the marker as applying to "a parameter or buffer", but the marker is a plain Python attribute and nn.Module._apply treats the two differently:

• Parameters: _apply goes through compute_should_use_set_data then torch._has_compatible_shallow_copy_type, which is true for a CPU-to-CUDA move (both dense), so it does param.data = param_applied and the nn.Parameter object — and therefore keep_in_fp32 — is preserved.
• Buffers: _apply does self._buffers[key] = fn(buf) unconditionally, so buf.cuda() returns a new tensor and keep_in_fp32 is silently dropped.

This matters because megatron/training/training.py moves the model to GPU before wrapping — model_module.cuda(...) at training.py:1802, then Float16Module(config, model_module) at training.py:1807. A buffer marked at construction time (i.e. on CPU, the normal case with use_cpu_initialization) loses its marker during .cuda() and gets cast to BF16/FP16 anyway, silently, with no error. The new tests only exercise the parameter path (Fp32MarkedToyNet.scale, DummyModule.fp32_param), so this is not covered.

Either narrow the docstring contract to parameters only, noting that nn.Module._apply replaces buffers with new tensor objects on device moves and drops the marker; or make the marker survive buffer replacement (re-tag inside the conversion, or key off buffer names rather than a tensor attribute) and add a buffer case to test_keep_in_fp32_params.

2. get_fp8_disabled_context nullcontext() branches are untested.

test_get_fp8_disabled_context_uses_disabled_te_context parametrizes only the two enabled paths (fp8="hybrid" giving fp8_autocast, fp8_param=True giving fp8_model_init). The nullcontext() early returns are the branches most likely to regress from a mis-edited boolean, and they are cheap to cover: one case with all four of fp8/fp4/fp8_param/fp4_param falsy asserting isinstance(result, nullcontext) for both is_init=False and is_init=True, plus an fp4-only case confirming FP4 reaches fp8_autocast (TE uses fp8_autocast for FP4 too, per the comment at fp4_utils.py:273).

[svcnvidia-nemo-ci]

🔄 Merge queue validation started!

You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/31433705373

[svcnvidia-nemo-ci]

🔄 Merge queue validation started!

You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/31438325213