Add chunked GPU optimizer-state offloading

[yanring]

• I, the PR author, have personally reviewed every line of this PR.

What does this PR do ?

Makes optimizer state stop being a GPU-memory problem: Adam moments and
master weights live in pinned CPU memory and only come back to GPU when they
are actually being used. In chunked mode, the optimizer step and the
first-step state init keep roughly one chunk of state resident instead of
all of it, and checkpoint save reads the CPU copies directly without
rematerializing anything on GPU.

The problem

Adam state is one of the largest GPU-memory consumers in training, and it is
only touched during the optimizer step. The rest of the time it sits there,
taking memory away from activations, which is exactly what long-context and
memory-heavy post-training runs are short of. The lifecycle also has hidden
spikes: state init materializes everything at once, and checkpoint save used
to bring every offloaded tensor back to GPU, so runs can OOM at the first
save even when the steady state fits.

Megatron has two mechanisms in this space today, and neither covers this:

• --optimizer-cpu-offload (HybridDeviceOptimizer) moves the Adam compute
  itself to CPU. It saves the memory but pays with a much slower step.
• --offload-optimizer-states ([Dev] [Reapply] Optimizer State and Master Weight Offloading #2987) had the right idea: keep Adam on GPU
  and move only the storage to CPU between steps. But its training-loop
  wiring was lost in a later main-to-dev sync merge, so today the flag builds
  an offloader that nothing ever calls (grep for offload_states() callers:
  there are none).

What this PR does

Revives #2987 and finishes the idea: optimizer state lives on CPU and GPU
residency is bounded at every point where the implementation can bound it.
Two places still temporarily materialize the full state on GPU: the
unchunked step (chunk-numel 0, the restored #2987 behavior, which reloads
everything before stepping) and an optimizer checkpoint load, which fills
the GPU tensors before the load path re-offloads them.

• Per-step offloading (restored wiring). After each optimizer step the
  moments and master weights copy D2H on a side stream and their GPU storage
  is freed. The H2D reload for the next step is kicked off during grad
  finalization, so it overlaps the grad reduce-scatter. Adam itself still
  runs on GPU at full speed; only the residency moves.
• Chunked optimizer step (new, --offload-optimizer-states-chunk-numel N).
  Full reload still needs the whole state on GPU during the step. In chunked
  mode the step walks groups of whole local parameter shards packed up to
  roughly N elements: reload one chunk, run FusedAdam on it, offload it
  back, move on. Chunks split only between parameters, so a single shard
  larger than N forms its own chunk and the effective residency bound is
  max(N, largest local shard). The lazy first-step state init follows the
  same bound, so even step 1 never holds the full state.
• No init spike in chunked mode. With chunked offloading and
  mcore-managed master weights, the fp32 masters are built directly in
  pinned CPU buffers instead of being materialized on GPU first; TE-managed
  optimizer state is initialized lazily chunk by chunk during the first
  step. Unchunked offload keeps [Dev] [Reapply] Optimizer State and Master Weight Offloading #2987's behavior (build on GPU, then
  offload).
• No save spike, and a real crash fix. Checkpoint save reads the
  offloaded CPU copies directly instead of reloading everything to GPU.
  While making that work we hit an actual illegal memory access: TE
  FusedAdam.state_dict() unscales every state tensor, which launches
  kernels on freed (zero-storage) tensors; the async IMA surfaces at the
  next device sync and can silently corrupt the checkpoint being written.
  Save now packs the common state dict with the base
  torch.optim.Optimizer.state_dict(), per-param values go through a read
  gate that returns the CPU copies, and any raw read of released storage
  asserts loudly instead of corrupting data.
• Load without surprises. Released states are reloaded before a mid-run
  optimizer load writes into them and re-offloaded afterwards;
  --no-load-optim / --finetune runs offload before the first checkpoint
  load.
• Works with FP8 primary weights. Chunked mode quantizes updated master
  shards back to FP8 per chunk under --fp8-param-gather, FP8 load targets
  can stage on CPU to avoid a full-model dequantization peak, and a new
  ddp_config.preserve_fp8_columnwise switch stops the eagerly recreated
  columnwise copies from sitting resident through checkpoint save: outside
  CUDA graphs TE creates them lazily in the next forward instead. The eager
  path stays on automatically under CUDA graphs, which need stable storage
  pointers. Note this changes the default for non-CUDA-graph FP8 runs from
  eager to lazy.
• Skipped steps stay clean. An inf/nan skip synchronizes the in-flight
  H2D copies through a new MegatronOptimizer.on_step_skipped() hook instead
  of leaving them pending into the next iteration.

Numerics are unchanged by design: state dtypes, the fp32 grad buffer, and
the Adam math are all untouched. Offloading only changes where the bytes
live between uses, which is also why the equivalence checks below can demand
exact matches rather than tolerances.

Constraints

• Distributed optimizer + TE FusedAdam only (same envelope as [Dev] [Reapply] Optimizer State and Master Weight Offloading #2987).
• Rejected combinations, all fail-loud at argument validation:
  --optimizer-cpu-offload (pick one strategy), --async-save (save reads
  the offloaded CPU copies; an async writer would race the next step's
  offloads overwriting those same buffers), --optimizer-cuda-graph (a
  captured step would replay on freed storage pointers), and
  --rl-offload-optimizer-during-inference (its raw optimizer reads would
  touch released storage). Chunked mode additionally rejects
  --fp4-param-gather and --reuse-grad-buf-for-mxfp8-param-ag.
• exp_avg / exp_avg_sq dtypes must be fp32 or bf16: checkpoint save reads
  the CPU bytes as-is, and TE stores fp16/fp8 states scaled.
• With TE-managed master weights (precision-aware optimizer without
  fp8/ds-fp8 param handling), main_params_dtype must be fp32 and
  --store-param-remainders must be off, for the same byte-exact-save
  reason.

Cost

The D2H/H2D traffic overlaps grad reduce-scatter and next-step compute in
full-reload mode. Chunked mode serializes per-chunk copies with the chunk
updates, so the optimizer step gets somewhat slower in exchange for the
residency bound; we run production with chunk sizes around 1e9 elements
where this cost is small relative to a training step.

Validation

This is the offload path we run in production training. Correctness there is
gated by an equivalence harness: a no-offload baseline and an offload run
with identical seed, data, and topology, FP8 param gather enabled, required
to produce exactly equal per-step loss and grad norm, while the offload
run's peak allocated memory must come in strictly below the baseline (an
offload that silently no-ops fails the harness). Checkpoint save/load/resume
is exercised the same way.

In this repo, tests/unit_tests/test_optimizer_state_offloading.py is
extended from the #2987 suite (all 6 original cases kept) with
chunked/unchunked equivalence, checkpoint save/load round-trip, memory
residency, FP8 blockwise, and step-skip coverage: 24 passed / 2 skipped
(both by-design opt-outs) on 1 node x 8 GPUs, torch 2.11 + TE 2.17; enabling
the env-gated FP8 integration test gives 25 passed / 1 skipped. black /
isort / ruff pass with the repo configs.

Issue tracking

Linked issue:

Pre-checks

• I have added relevant unit tests
• I have added relevant functional tests
• I have added proper typing to my code
• I have added relevant documentation
• I have run the autoformatter on my PR

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

[yanring]

cc @hxbai