dist_ckpt: add --stream-ckpt-dequant to fix OOM on large FP8/MXFP8 loads with --fp8-param-gather - #4451
Conversation
|
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. |
|
/ok to test add6edc |
|
/ok to test 088b812 |
|
/claude review |
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
| for _ in range(mcore_sh_ten.prepend_axis_num): | ||
| assert ten.size(0) == 1 |
There was a problem hiding this comment.
Bug: since ten is no longer mutated inside the loop, ten.size(0) is checked repeatedly instead of validating each prepended axis. If prepend_axis_num > 1 and a non-leading prepended axis isn't actually 1, the assertion silently passes and view produces a wrong shape.
The original code worked because ten = ten[0] peeled off axis 0 each iteration, so the next ten.size(0) checked the next axis. With the view-based approach the loop variable needs to advance the axis index:
| for _ in range(mcore_sh_ten.prepend_axis_num): | |
| assert ten.size(0) == 1 | |
| for i in range(mcore_sh_ten.prepend_axis_num): | |
| assert ten.size(i) == 1 |
|
/ok to test a13ba55 |
|
@asolergi-nv High-level question: why do we want to gate this feature behind the |
|
This has been validated in megatron-inference by @santhnm2 . We are able to run ultra on 4 GB200s in 8-bit precision. |
|
/ok to test 1964c9b |
Resolve conflict in megatron/training/arguments.py: keep the branch's tensorwise stream_ckpt_dequant safety check and drop the FlashInfer inference-backend import validation, which main removed in NVIDIA#5791 ("Prevent FlashInfer sampling from running with CUDA graphs"). The FlashInfer argument itself is unchanged; only the relocated runtime validation block is dropped. Signed-off-by: Antoni-Joan Solergibert <asolergibert@nvidia.com>
|
/ok to test 89e4b96 |
|
🔄 Merge queue validation started! You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/30273509109 |
|
🔄 Merge queue validation started! You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/30286225199 |
|
🔄 Merge queue validation started! You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/30295594142 |
|
🔄 Merge queue validation started! You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/30299352591 |
|
🔄 Merge queue validation started! You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/30306309809 |
|
🔄 Merge queue validation started! You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/30311011972 |
|
🔄 Merge queue validation started! You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/30335393420 |
|
🔄 Merge queue validation started! You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/30337947510 |
9 participants
What does this PR do ?
Problem
Loading a distributed checkpoint on a large model trained with
--fp8-param-gather(or--fp4-param-gather) can OOM on GPU during the load itself — not during training. Root cause:force_all_tensors_to_non_fp8inmegatron/core/dist_checkpointing/serialization.pywalks the entire sharded state dict before the checkpoint read starts and, for everyQuantizedTensordestination (Float8 / MXFP8 / blockwise FP8 / NVFP4), allocates a fresh BF16 scratch tensor of the same shape. For an N-element quantized weight set that's a transient 2·N bytes of GPU memory sitting idle until the read finishes — and for large models, 2·N exceeds what's left on device.A secondary OOM was surfaced by the first fix:
_unwrap_pyt_sharded_tensorusedten[0]to strip prepended singleton axes, which onMXFP8Tensorfalls through toQuantizedTensor.__torch_dispatch__and dequantizes the entire tensor to BF16 just to select the leading index — re-creating the same OOM per shard after a successful load.Solution
Move the dequantize from an upfront bulk pass into PyTorch DCP's per-tensor
LoadPlannerhooks.MCoreLoadPlanner.resolve_tensornow detectsQuantizedTensordestinations, allocates a single BF16 scratch per tensor, returns it as the load destination, and commit_tensor copy-quantizes it back into the original quantized tensor. Only one scratch is live at any moment — peak transient GPU overhead drops from2·sum(numel)to~2·max(numel). Delayed-scaling amax_history pollution is prevented via a per-tensor snapshot/restore around the quantize-on-copy.The
_unwrap_pyt_sharded_tensorOOM is fixed by using view() (which all TE quantized tensor classes implement natively without dequantizing) instead of ten[0]/squeeze.Default behaviour is unchanged. The streaming path is gated behind a new opt-in flag
--stream-ckpt-dequant(defaultFalse); without it, the load path is byte-identical to today. A warning is emitted when--fp8-param-gatheror--fp4-param-gatheris used without--stream-ckpt-dequantand a checkpoint load is configured, pointing users at the escape hatch.Tests
New file
tests/unit_tests/dist_checkpointing/test_stream_ckpt_dequant.py:test_fp8_save_load_content_equivalence— FP8 round-trip with streaming on and off; asserts loaded values are bit-equivalent within FP8 quantization tolerance.test_fp8_amax_history_not_polluted— seeds a sentinel into the delayed-scaling quantizer's amax, runs a streaming load, asserts the amax is exactly restored (no amax_history pollution).test_mxfp8_save_load_content_equivalence— MXFP8 round-trip both paths; incidentally exercises the_unwrap_pyt_sharded_tensorview fix (without it this test OOMs onten[0]dequantize fallback).test_plain_tensor_untouched_by_streaming_path— plain BF16 tensors must round-trip losslessly under both paths; guards against the streaming branch breaking non-quantized loads.test_default_is_off— asserts the new stream_ckpt_dequant default is False on both TorchDistLoadShardedStrategy and MCoreLoadPlanner, preserving legacy behaviour.test_planner_state_cleanup_after_load— asserts the planner's per-tensor tracking dict is empty after a streaming load completes, confirming scratch tensors are GC-eligible and no state leaks.test_streaming_flag_forwards_through_fpsl_wrapper— asserts FullyParallelLoadStrategyWrapper.stream_ckpt_dequant forwards correctly from the wrapped base strategy.Issue tracking
For PRs from open-source community contributors:
Linked issue:
Contribution process
Pre-checks
Code review
Feel free to message or comment the @mcore-oncall to help accelerate your merge into main. The less complex your PR is, the faster it will be approved and merged!
All PRs start as draft. If you open a non-draft PR, it will be automatically converted to draft.
Step 1: Mark PR as "Ready for Review"
.github/CODEOWNERS.Final Review might get declined if these requirements are not fulfilled.
Step 2: Final Review
For PRs that change
megatron/core, once all expert reviewers have approved, theFinal Reviewlabel is applied automatically and final reviewers are assigned.For PRs outside
megatron/core, this step is skipped.Step 3: Approved
Once all required reviewers have approved, the
Approvedlabel is applied automatically.Merge
Any member of mcore-engineers will be able to merge your PR.
For MRs into `dev` branch
The proposed review process for `dev` branch is under active discussion.MRs are mergable after one approval by either
eharper@nvidia.comorzijiey@nvidia.com.