[Dev] Align decoupled Muon FP8 parameter-gather paths with main - #6015
[Dev] Align decoupled Muon FP8 parameter-gather paths with main#6015Wohox wants to merge 2 commits into
Conversation
94362fc to
56f71b5
Compare
593ce3b to
da968db
Compare
Four defects found while reviewing this PR. All four originate in the dev-side commits this PR ports (NVIDIA#6015 `56f71b517`/`593ce3b10`, NVIDIA#5470 `0866f3511`), so they apply to `dev` as well. 1. `sharded_param_state_fs_model_space` rejected every plain `ShardedTensor`. The guard asserted `isinstance(sharded_metadata, ShardedTensorFactory)` per param, using "is a factory" as a proxy for "is the decoupled LayerWise layout" -- unrelated properties. Only gated-MLP fc1 yields a factory, so this aborted `fully_sharded_model_space` for essentially every param of every ordinary DistributedOptimizer run (including resume from checkpoints whose metadata selects that format), with a misleading Muon message, and made the pre-existing two-branch handling below it dead code. Rejected once up front on the real condition instead, leaving the common path untouched. 2. `_allgather_helper` indexed rank 0's list blindly (`params_list[0][0].device`). Its fp8 twin already guards this, but the native-fp32 dispatch added by NVIDIA#5470 feeds it a per-rank split in which rank 0's list is empty whenever rank 0 owns no `keep_in_fp32` param -- an `IndexError` on every rank, on the first `allgather_params()`. Mirrored the fp8 twin's lookup and its no-owner early return. 3. The padding template was rebuilt per save from a hardcoded Adam key set and `main_params_dtype`. Lion (and Muon's `muon_scalar_optimizer`) have no `exp_avg_sq`, and `'param'`'s dtype is not `main_params_dtype`, so a rank owning nothing in any bucket emitted padding whose keys/dtypes disagreed with the real shards its peers wrote under the same key. That matters even though the padding is discarded on load: it is a real `ShardedTensor` tiling the same `(gbuf_world_numel_unpadded,)` global tensor, and `dist_checkpointing/validation.py` requires one dtype per key, so the save aborts. Record the key set and dtypes once in `__init__` (`_ckpt_pad_dtypes`) from `optimizer_state_keys` / `_get_state_key_dtype` -- the same helpers every other call site uses. They are a property of the optimizer and its config, hence identical on every DP rank, so no cross-DP reconciliation is needed. Where a rank does hold state we cross-check the two agree, so a wrong derivation surfaces on that rank instead of as a checkpointing dtype error on whichever rank happened to own nothing. 4. A 101-character line in `megatron/core` would fail the pylint `C0301` pre-commit hook (`max-line-length=100`, `files: ^megatron/core/.*`). Reflowed. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Pingtian Li <pingtianl@nvidia.com>
1433ad0 to
c89716a
Compare
Wohox
changed the title
|
/claude strict-review |
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
There was a problem hiding this comment.
Strict review passed — no significant issues found.
Traced the full data flow across all five changed files:
distrib_optimizer.py (main fix) — The store-back line gbuf_range_map_for_all_buckets[bucket_idx] = bucket_state correctly re-binds the synthesized padding shard into the returned state; without it the shard is dropped and DCP global-plan validation fails (chunks_volume < tensor_volume). Verified the synthesis coverage math by hand: the ShardedTensor covers exactly [world_shard_start, world_shard_start + pad_len) within [0, numel_unpadded), and the downstream padding-insertion loop adds no spurious trailing padding.
Collective-safety / backward compat — The all_gather_object is gated on use_layer_wise_distributed_optimizer, a config value identical on every DP rank, so all ranks take the same branch and an ordinary DistributedOptimizer keeps its communication-free save/load path. The relaxed empty bucket assert only changes behavior in a previously-crashing case (empty state + empty param_map); ordinary DistOpt still hits the strict param_map assertion. self.data_parallel_group is an existing instance attribute, not a new global-process-group read.
layer_wise_optimizer.py — _allgather_helper now derives device/dtype from the first non-empty per-rank list instead of indexing rank 0 blindly; this is symmetric across ranks (same partition on all), so the early return is collective-safe and mirrors the fp8 twin. The not any(...) expert-parallel guards correctly generalize the [0]-only emptiness check.
Mandatory unused-variable check — _optimizer_model_params, pad_templates/pad_device, and all new test params have real read/use paths.
Tests — Genuinely convert vacuous {} == {} asserts into meaningful ones and add store-back regression coverage.
LGTM.
85caba9 to
4516fe2
Compare
|
/ok to test 4516fe2 |
4516fe2 to
d05dd1f
Compare
|
/ok to test d05dd1f |
|
/ok to test 7db50e2 |
Wohox
changed the title
Batch MXFP8 copy-back through existing quantizers, share conversion helpers across DDP and MCore FSDP, and tighten validation plus offload-overlap handling. Make gather ownership checks robust to empty rank slots and add coverage for copy-back, overlap, and quantized parameter comparisons. Co-authored-by: Shiqing Fan <shiqingf@nvidia.com> Signed-off-by: Shiqing Fan <shiqingf@nvidia.com> Signed-off-by: Pingtian Li <pingtianl@nvidia.com>
abbe881 to
a3db0b5
Compare
Persist exact intra-parameter padding ranges in checkpoint content metadata so DCP validation can distinguish intentional sparse coverage from missing optimizer state. Preserve multi-bucket indices during load and cover trailing plus intra-parameter padding with value-level round-trip tests. Signed-off-by: Pingtian Li <pingtianl@nvidia.com>
a3db0b5 to
7610db8
Compare
1 participant
Summary
This PR is now the
dev-branch alignment counterpart of the decoupled Muon FP8parameter-gather work in #5479 on
main. It keeps the original #6015 LayerWisecheckpoint/param-gather fixes and ports the subsequent implementation cleanups and
performance work so the two branches use the same primitives and validation boundaries.
The refactor also incorporates the batched MXFP8 copy-back optimization from #6094.
megatron/training/training.pyremains unchanged; the shared behavior lives in the DDP,FP8, and optimizer layers that own it.
Changes
FP8 parameter staging and copy-back
instead of routing every parameter through
aten::copy_/__torch_dispatch__.allocating one BF16 tensor per parameter and flattening the temporary copies.
single-parameter compatibility wrapper.
initialization value in
pop_high_precision_init_val, shared by DDP, DistributedOptimizer, LayerWise Optimizer, Float16 Optimizer, and Megatron FSDP call sites.
does not opt into the LayerWise gather topology.
Validation and routing
blockwise FP8) instead of treating every generic Float8/GroupedTensor representation
as supported.
--moe-single-grouped-weightfor the compact LayerWise FP8 gather because thebatched copy-back does not support Transformer Engine GroupedTensor storage.
child-level pre-forward requirement does not imply that every bucket owned by that
child has already been dispatched.
Existing #6015 fixes retained
dp_reshardablebucket shards in the decoupled compactLayerWise configuration.
first real parameter.
Checkpoint redesign boundary
The existing checkpoint implementation is retained by this PR; the cleaner
padding-free redesign is intentionally not implemented here.
The follow-up design has two relevant rank-local cases (Distributed Optimizer ownership
is already resolved before these buffers are constructed, so an "unowned bucket" case
is not part of the design):
numel_unpaddedis trailing DP-alignment padding and should beomitted from save and load.
dense bucket checkpoint extent. It should be omitted by making the coverage contract
padding-aware, not by synthesizing checkpoint tensors. The change must preserve
explicit bucket identity and construct the corresponding load request without relying
on list position, so multi-bucket checkpoints round-trip by value.
This must be a scoped coverage rule for declared padding regions; globally disabling DCP
coverage validation would also accept missing optimizer state and is not safe.
Test plan
H100:
tests/unit_tests/test_muon_decouple_fp8_param_gather.py: 8 passed, 9 skippedper rank.
tests/unit_tests/test_fp8_utils.py: 2 passed, 3 skipped per rank.tests/unit_tests/distributed/test_fp8_param.py: 10 passed, 12 skipped per rank.tests/unit_tests/distributed/mfsdp_v1/test_mfsdp_fully_shard.py:197 passed, 125 skipped, 6 xfailed per rank. The six results are expected XFAILs,
not failures introduced by this PR.
tests/unit_tests/dist_checkpointing/test_layer_wise_optimizer.py:292 passed, 161 skipped per rank.
GB200:
tests/unit_tests/test_fp8_utils.py: 5 passed per rank, including the real MXFP8batched copy-back path.
numerical drift also reproduces on the pre-refactor
b9086f1a7baseline, and thepending-gather test calls
finish_grad_sync()after manually waiting the parametergather handle. Neither is caused by this refactor.
Formatting and static checks:
tools/autoformat.shscope passed Black, isort, pylint, and ruff.git diff --checkpassed.Related