[Optimizer] HybridDeviceOptimizer resumes from a checkpoint without a second set of pinned copies, and exactly

[tigist-far]

Summary

Resuming a run that keeps its optimizer on the host (optimizer_cpu_offload) from a checkpoint with optimizer state failed three ways in HybridDeviceOptimizer's load_state_dict path, all on the same code path and all found on the hosted 2-node Nemotron 3 Super SFT run on 2026-09-09:

#	Symptom on the 120B run	Cause in hybrid_optimizer.py
1	KeyError: tensor(...) in the post-load hook, before the first step	_update_fp32_params_by_new_state indexed param_to_fp32_param for every state entry; under the distributed optimizer the parameters are fp32 shards, which have no fp32 clone
2	CUDA OOM during the load (178 GiB GPUs full)	the pre-load hook only swapped parameters that have an fp32 clone, so torch cast the loaded state of every offloaded parameter to the GPU (~90 GB per rank) before the post-load hook moved it back to the host
3	Worker pod OOM-killed at the next save: pinned host memory 930 → 1640 GB per node after the load	the post-load hook rebuilt the sub-optimizers (_init_sub_optimizers): a fresh pinned CPU copy of every parameter and fresh pinned gradient buffers, while the loaded state's master_param entries kept the previous copies alive

The fix keeps the sub-optimizers, their pinned copies and gradient buffers across a load and only re-syncs parameter groups and state; both hooks swap through param_to_inner_param / inner_param_to_orig_param (an fp32 clone, a pinned CPU copy, or the parameter itself) so the loaded state is cast on the device it belongs to; and the state sync copies a loaded master_param into the live inner parameter and re-points the entry at it.

A second, smaller commit fixes a fourth gap the exactness test exposed: TransformerEngine's FusedAdam keeps its Adam step in param_group["step"], which the hybrid optimizer neither saved nor restored, so after any load the GPU-resident half restarted its bias correction at step 1 (parameters diverged by ~7e-4 after one step while the offloaded half was bit-exact). The sub-optimizers' group steps are now mirrored into the HDO's own param groups, which the checkpoint already carries, and handed back after a load. dummy_step, the step path and the checkpoint layout are otherwise untouched; checkpoints without the new key load as before.

How a load runs now

sequenceDiagram
    participant DO as DistributedOptimizer.load_state_dict
    participant T as torch Optimizer.load_state_dict
    participant H as HybridDeviceOptimizer hooks
    DO->>H: pre-load hook
    Note over H: swap every group param for its inner param<br/>(fp32 clone / pinned CPU copy / itself)
    H->>T: cast state to the inner params' dtype and device
    Note over T: offloaded state stays on the CPU
    T->>H: post-load hook
    Note over H: swap back to the original params<br/>keep sub-optimizers, copies, gradient buffers<br/>re-sync param groups and state
    H->>H: copy a loaded master_param into the live inner param, re-point the entry

Loading

Before this change the post-load hook's last step was _init_sub_optimizers(), i.e. the right-hand column of the table above.

Areas changed

Area	Files	What changed
megatron/core/optimizer/cpu_offloading/hybrid_optimizer.py	1	pre/post load hooks swap through the inner-parameter maps; the post-load hook no longer rebuilds the sub-optimizers; _update_fp32_params_by_new_state syncs every inner parameter from master_param and tolerates parameters without a clone; the sub-optimizers' per-group step (FusedAdam keeps it there) is mirrored into the HDO's param groups after each step so it survives a checkpoint
tests/unit_tests/test_optimizer_cpu_offloading.py	1	test_load_state_dict_reuses_the_cpu_copies_and_resumes_exactly: after load_state_dict, inner parameters, sub-optimizers and the gradient-buffer map are the same objects; state tensors live on the inner parameters' devices; two further steps match the saved-from optimizer bit for bit — fp32 and bf16 parameters, offload fractions 0.5 and 1.0. On the previous code the fp32 cases raise the KeyError and the bf16 cases fail the identity checks

Measurements behind the change

From the hosted run (16 ranks, TP 8 × CP 2 × EP 8, optimizer fully offloaded, pods limited to 3020 GB):

	cold start	resumed with optimizer state (old code)
pinned host memory per pod after settling	930 GB	1640 GB
pinned memory mapped per rank	116 GB	190 GB
anonymous memory per pod	1.11 TB	1.11 TB

The extra ~500 GB per pod is 8 ranks × (31 GB of parameter copies + 31 GB of gradient buffers), i.e. one full set; the sampler saw pinned memory climb 1141 → 1640 GB across the load phase. With the fix the resumed run's floor returns to the cold start's.

Testing

Run on flamingo through megatron-bridge-internal's make test-shell-remote (image megatron-bridge:7f31fe20f, 2 × H100, this branch rsynced over the vendored 3rdparty/Megatron-LM, torch 2.13, TE 2.16):

Run	Result
test_load_state_dict_reuses_the_cpu_copies_and_resumes_exactly, this branch	4 passed (fp32 and bf16 × offload 0.5 and 1.0)
the same test against the previous hybrid_optimizer.py (mutation check)	4 failed: KeyError for the fp32 cases, the identity assertions for the bf16 cases
tests/unit_tests/test_optimizer_cpu_offloading.py, whole module	76 passed
Megatron-Bridge tests/unit_tests/training/test_checkpointing.py	173 passed

The fork's pre-commit hooks (black 26.3.0, pylint 3.2.6, isort 5.13.2) pass on both files.

On the hosted run the runtime equivalent of the first commit (nemotron hybrid_optimizer_resume, applied after the load) brought the resumed run's pinned host memory from 1640 GB back to 894 GB per pod.