[main] Muon fp8 param gather for decoupled layout - #5479
Conversation
ae3d338 to
1fc496b
Compare
|
/ok to test 1fc496b |
Wohox
changed the title
Wohox
changed the title
915b27a to
9a91868
Compare
…int fix Makes `use_distributed_optimizer` a per-buffer property so the LayerWise (Muon) optimizer can drop the persistent `dp_size * max(shard_load)` padding from its long-lived param and grad buffers. LayerWise-managed (Muon 2D-matrix) buffers get a compact no-padding DDP layout and locally disable DistributedOptimizer semantics (all-reduce gradients, whole-param ping-pong ownership, `allgather_params` param sync), while sibling buffers (embeddings, biases, layernorm) keep the standard byte-level DistributedOptimizer path. The decision is baked into each buffer's own `ddp_config` via `dataclasses.replace`, so bucket groups inherit it; a bucket group runs one collective type, so `partition_buckets` asserts that any buckets it merges agree. This layout is now the **default**: `use_layer_wise_param_layout` defaults to `False` on both `DistributedDataParallelConfig` and `OptimizerConfig`, and `--use-layer-wise-param-layout` opts back **into** the padded shard-aligned layout (e.g. for bit-for-bit comparison against older runs). Checkpoint fix, which belongs with this layout but was missed when it first landed: * Grad-buffer range maps are filtered to the params the optimizer instance actually owns. On this layout a DistOpt's buffers also carry buckets owned by the LayerWise child; unfiltered, the DistOpt builds ranges and state for params it does not own, duplicating what LayerWise already saves. * `sharded_param_state_dp_reshardable` save/load skip buckets this optimizer owns no param of. Membership is checked per param rather than via `params[0]`, so buckets mixing owned and LayerWise-managed params are handled. * A DP rank can own a shard that is entirely padding yet still overlaps `[0, numel_unpadded)`; such a shard now synthesizes a padding ShardedTensor from a captured fp32 template so global coverage holds, and writes it back into `state` (the branch rebinds `bucket_state`, so without the write-back the shard was dropped). * The original strict `empty bucket encountered` assert is kept for pure DistOpt runs; empty shards are only expected when LayerWise-managed params co-exist. * `fully_sharded_model_space` is asserted unsupported on this layout rather than failing obscurely later. The `*_param_layout` functional cases now pass `--use-layer-wise-param-layout` explicitly, since they exist to cover the padded layout and the default moved. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: Pingtian Li <pingtianl@nvidia.com>
Lets Muon persist and sync its matrix parameters as FP8 (mxfp8 on Blackwell, blockwise on
Hopper) instead of BF16, removing the standing BF16 parameter copy. Combined with the
compact layout this closes the Muon-vs-Adam memory gap while staying bit-for-bit equal to
the fp8-param-gather-OFF run.
FP8 tensors cannot be all-gathered directly (mxfp8's row/col block scales cannot be derived
from one another), so the whole-param transport rides BF16 and requantizes on landing:
stage each owned param from its fp32 master to bf16 (`_stage_param_to_bf16` -- a
high-precision source, not a lossy fp8 dequant), uneven all-gather-v, then requantize the
gathered bf16 into *every* rank's fp8 `param.data` so all ranks hold `Q(bf16(master))`;
`post_all_gather_processing` rebuilds the fp8 columnwise/transpose. The overlap
(`start_param_sync`/`finish_param_sync`) and non-overlap (`allgather_params`) paths share
one copy-back helper, and the staging decision is persisted per bucket so transport dtype
and copy-back cannot disagree. Params are dispatched by transport dtype: fp8 and bf16 ride
the bf16-staged helper, while native fp32 params (`keep_in_fp32`, e.g. the DeepSeek-V4 CSA
`ape`) are gathered in fp32 -- the bf16 path would silently downcast them.
Also in this commit:
* Rank-independent ping-pong ownership. `allgather_params` needs every DP rank to agree on
each param's single owner; `numel` alone is not a total order, so ownership is keyed by a
canonical `(chunk_idx, buffer_idx, global_start_index)` identity.
* iter-0 master parity: the fp32 master is seeded from the high-precision pre-quantization
init, and the master->model copy routes through bf16. Gathered fp8 params are tagged and
skipped there, since the all-gather's requantize already wrote `Q(bf16(master))`;
non-gathered fp8 params (MoE experts at `expt_dp == 1`) still get their copy.
* Non-owned params drop TE's high-precision init copy, which every DP rank would otherwise
retain for ~(dp-1)/dp of the LayerWise matrix params for the whole run.
* `force_sync` finalizes a pending LayerWise gather it would otherwise bypass, which would
leave stale `param.data` and pollute the reused `grad_data`.
* Single all-reduce buffer: keying fp8 Muon grads by `torch.uint8` split their gradients
into two all-reduce buffers where OFF has one; NCCL's fp32 accumulation order is
buffer-size sensitive, so the split diverged ~1 ULP from OFF and Newton-Schulz amplified
it into visible loss drift. fp8 Muon grads now key to their bf16 logical dtype and share
one buffer. A consequence is that a bucket can hold both fp8 and bf16 params, so the
staging decision scans the whole bucket rather than `params_list[0]`.
* Validation: fp8 param gather is accepted only on this layout, requires
`fp8_recipe in {mxfp8, blockwise}` (fp4 rejected outright), mxfp8 additionally requires
`--reuse-grad-buf-for-mxfp8-param-ag`, and the layout requires
`num_distributed_optimizer_instances == 1`.
Golden values for the `dist_dist_muon` ckpt-resume cases are refreshed at this tip
(8xH100 EP8, 2x GB200 EP8, 1x GB200 EP4), regenerated with CI's own harness on the
common_pile dataset. Each reported 'Exact: FAILED / APPROXIMATE: PASSED' beforehand -- a
~1 ULP reordering from the rank-independent ownership sort, not a regression; peak
relative delta on lm loss is ~1.6e-4.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Signed-off-by: Pingtian Li <pingtianl@nvidia.com>
9a91868 to
b4ba39f
Compare
sharded_param_state_dp_reshardable synthesizes optimizer-state padding for shards that lie entirely in a bucket's inter-param padding. That synthesis needs the state's key set and dtypes, which a rank owning no param in the bucket cannot observe locally. Derive the template per (gbuf_idx, dtype, bucket_idx) and reconcile it across the DP group with one all_gather_object. Per bucket is the axis that matters: dist_checkpointing enforces one dtype per KEY, and each bucket is its own key, so a template sampled from one bucket is not authoritative for another. The agreement assert is likewise per bucket, so buckets no rank sampled stay absent instead of aborting the save. The template cannot be derived from config: the save path runs the state through get_unscaled_state (which upcasts bf16/fp16/fp8 state to fp32 and returns int16 for store_param_remainders), and the key set comes from whatever optimizer.state holds. `step` is excluded -- it becomes a LocalNonpersistentObject and never a ShardedTensor. Device stays out of the gathered template and out of the compare. It is rank-local: under --optimizer-cpu-offload HybridDeviceOptimizer walks each rank's own shard list, so the CPU/GPU cutoff lands on a different param on every rank and a cross-rank compare would abort a legitimate save. dist_checkpointing validates dtype and shape but never device. Also reject fully_sharded_model_space up front for the decoupled compact layout. That format is independently non-functional for every DistributedOptimizer here -- it sets flattened_range on every non-factory param and ShardedTensor rejects that unconditionally -- so this assert does not fix the format; it makes the compact case fail early with an actionable message instead of dying per-param inside mapping.py. Narrow on purpose. Signed-off-by: Pingtian Li <pingtianl@nvidia.com>
_allgather_helper read params_list[0][0] to pick the device and dtype for the collective. The per-rank lists come from the ping-pong assignment and the dtype split in _dispatch, either of which can leave rank 0's slot empty while other ranks own params, so this raised IndexError. `any(owned ...)` in the caller does not prevent it. Take the first non-empty list instead, mirroring the fp8 twin. The same shape appears in the expert-parallel guards, which decide "there are no expert params anywhere" from expt_dp_params_list[0] alone and then disable the expert all-gather entirely. Both sharding paths happen to place the first expert param in slot 0 (ping-pong starts at expt_dp_idx 0; the layout path gives the first param of every expert bucket shard_id 0), so the two forms are equivalent for every reachable configuration today and this fixes no observable bug. It is a defensive change: it states the intended invariant, and unlike the helper it would have failed silently rather than loudly. Signed-off-by: Pingtian Li <pingtianl@nvidia.com>
…sampling
Three of this test's comparisons were passing for the wrong reason.
_is_quantized probed for a `dequantize` attribute, but torch.Tensor.dequantize
is defined on every tensor (it returns self.to(float32) for dense ones), so the
predicate was always True, _snapshot_params always returned {}, and the param
comparison was asserting {} == {}. Key off is_float8tensor and assert the
snapshot is non-empty.
_run_steps then snapshotted with include_quantized=True and fed that to the
ON-vs-OFF comparison, where OFF holds plain bf16 and ON holds Q(bf16(master)):
a dequantized fp32 view of the latter differs from the former by the
quantization step by construction. Drop it. ON legitimately skips its fp8
params while OFF keeps them as bf16, so ON's key set is a subset; pin the
dropped set at step 0 and require it to be unchanged and non-empty, rather than
accepting any shrinkage.
Under overlap_param_gather the param all-gather is deferred to the next forward
pre-hook, so right after step() the buffer still holds pre-update values and
there is no well-defined instant at which ON and OFF can be compared. Finalize
before sampling -- restaging the masters first, since the LayerWise buckets
alias the grad buffer that backward has just filled. Then re-arm: the forced
sync leaves param_gather_dispatched=True, which would turn every subsequent
forward pre-hook into a no-op and silently drop the deferred-gather coverage
this parametrization exists for. Assert the pre-hook is armed so that regression
cannot recur silently.
Finally, the force_sync test compared its two grad_data snapshots for equality,
but both paths end in grad_data.zero_(), so on the happy path that is zeros
against zeros; state the post-condition directly. Its param snapshots were also
asymmetric (one side included quantized storage, the other did not), which made
the key-set assert pass only because both were empty.
Signed-off-by: Pingtian Li <pingtianl@nvidia.com>
setup_model_and_optimizer never set use_layer_wise_param_layout, so the field was only ever whatever argparse defaulted to and no test could select the layout. Thread it through both helpers, and guard it against ddp_use_layer_wise -- the flag DDP actually receives -- rather than against the optimizer-name derived use_layer_wise. The old guard admitted optimizer='dist_muon' with dist_opt=False and use_layer_wise_param_layout=True, which builds DDP without a layout while the optimizer expects the shard-aligned one. test_layer_wise_optimizer_save_load then gains a LayerWise arm that actually reaches the layout: dist_opt=True and use_param_layout=True (the DDP-level routing switch), with use_layer_wise_param_layout=False selecting the compact side. It is the only fully_reshardable coverage of that layout above tp=pp=1 and the only one with grad_reduce_in_fp32 left at its default. The padded shard-aligned layout is deliberately out of scope. Both arms must request a sharding type now that both build a real DistributedOptimizer. The default fully_sharded_model_space is unusable: the compact layout is rejected up front, and an ordinary DistOpt dies deeper in replace() with "ShardedTensor.flattened_range is not supported". dp_reshardable synthesizes padding with torch.empty, so two saves of identical state differ in the uninitialized bytes and the terminal check_equal would flake. fully_reshardable has neither problem. The plain arm is kept as an in-file control -- test_optimizer.py already covers a plain bf16 DistOpt under fully_reshardable, and harder resharding cases besides -- and is labelled as such. Signed-off-by: Pingtian Li <pingtianl@nvidia.com>
746ebd8 to
0604c63
Compare
…am-gather-main-merged # Conflicts: # megatron/core/distributed/param_and_grad_buffer.py
…am-gather-main-merged # Conflicts: # megatron/training/arguments.py
|
/claude strict-review |
| help='If set, initialize with fake distributed process group and all distributed communication operations will be skipped. \ | ||
| This is quite useful for profiling memory usage of distributed training with just one GPU. \ | ||
| Setting WORLD_SIZE and RANK to the specific values for target distribtued scale.') | ||
| group.add_argument( |
There was a problem hiding this comment.
do we need to add this? isn't the config struct automatically populated into the args?
There was a problem hiding this comment.
I checked this more carefully: this flag is not auto-generated today. Neither OptimizerConfig nor DistributedDataParallelConfig is registered with ArgumentGroupFactory; get_megatron_optimizer_config() and get_megatron_ddp_config() only copy attributes that already exist on the argparse namespace.
Since this PR removes the old --no-use-layer-wise-param-layout flag and introduces the opt-in --use-layer-wise-param-layout form, this explicit parser entry is still required. I’ll keep it.
| per_chunk_layouts = [None] * n | ||
| if DP is DDP: | ||
| if use_layer_wise_distributed_optimizer and use_layer_wise_param_layout: | ||
| if use_layer_wise_distributed_optimizer: |
There was a problem hiding this comment.
why do we need to change this?
There was a problem hiding this comment.
This change makes ddp_config.use_layer_wise_param_layout the single source of truth for selecting the layout.
When the LayerWise optimizer is enabled, LayerWiseDistributedOptimizer.compute_full_param_layout() needs to run for both values: True builds the shard-aligned padded layout for LayerWise-managed parameters, while False builds the compact no-padding layout for those parameters and retains the standard DistOpt layout for sibling parameters such as embeddings, biases, and layer norms.
Therefore, skipping layout computation when the flag is False would also lose the required per-buffer routing and layout information.
| mxfp8 columnwise can't be derived from rowwise, so force columnwise before copy_ (TE rebuilds | ||
| both directions from the bf16); blockwise/Float8 columnwise is a lossless transpose. | ||
| """ | ||
| if is_mxfp8tensor(model_p): |
There was a problem hiding this comment.
this won't account for single weight, so we should either support it or ban it properly in the transformer config level
There was a problem hiding this comment.
Good catch. This copy-back path currently handles plain Float8/MXFP8 parameters, but it does not handle the Transformer Engine single-grouped-weight GroupedTensor representation correctly.
I’ll explicitly reject moe_single_grouped_weight with decoupled LayerWise FP8 parameter gather in config validation, unless we add a GroupedTensor-aware copy-back path and corresponding test coverage.
| main_param = getattr(p, "main_param", None) | ||
| if main_param is not None: | ||
| return main_param.detach().to(torch.bfloat16) | ||
| if is_float8tensor(p): |
There was a problem hiding this comment.
is_float8tensor is actually very vague, it could mean many fp8 recipes, so I suggest being a bit more explicit here ( if it's fp8 blockwise recipe, then fp8 blockwise, or mxfp8 tensor, etc.)
There was a problem hiding this comment.
is_float8tensor is intentionally a broad abstraction over Transformer Engine quantized tensor implementations, but I agree that it is too broad as the capability check in this path.
The LayerWise parameter-gather path currently supports only MXFP8 and blockwise FP8, so I’ll make that restriction explicit in the predicate/config validation rather than relying on the generic classification.
| if main_param is not None: | ||
| return main_param.detach().to(torch.bfloat16) | ||
| if is_float8tensor(p): | ||
| return dequantize_fp8_tensor(p).detach().to(torch.bfloat16) |
There was a problem hiding this comment.
this seems like a numerical lossy operation, when will _stage_param_to_bf16 be used?
There was a problem hiding this comment.
For this Muon path, every locally owned source parameter must have an FP32 main_param, so the intended operation is FP32 master → BF16. This is intentionally lossy relative to FP32, but it matches the FP8-param-gather-OFF baseline, which produces Q(bf16(master)).
The FP8-dequant fallback is not expected to execute in this path. I’ll narrow the helper or assert the main_param invariant so that it cannot silently take the lossier fallback.
There was a problem hiding this comment.
_stage_param_to_bf16 is used by the two alternative LayerWise parameter-sync implementations: LayerWiseDistributedOptimizer.allgather_params() for the non-overlap path and _ParamAndGradBucketGroup.start_param_sync() for the overlap path. Only one of these paths is used for a given sync configuration.
| arguments layer sets this flag and resets ``use_distributed_optimizer`` to False so | ||
| that the standard distributed-optimizer path is not triggered.""" | ||
|
|
||
| use_layer_wise_param_layout: bool = False |
There was a problem hiding this comment.
can you elaborate what does layout mean here?
There was a problem hiding this comment.
Here, layout means the mapping of each parameter into the contiguous DDP buffer: its start/end indices, bucket assignment, and the padding and shard boundaries used by that buffer.
For LayerWise-managed parameters, the padded layout places every whole parameter inside one DP shard, aligns parameter starts, and pads the shards to the same size. The compact layout instead uses the normal no-padding DDP ordering and keeps whole-parameter ownership in the LayerWise optimizer’s ping-pong assignment.
The layout always determines the grad-buffer placement. It also determines the param-buffer placement when that buffer uses DistributedOptimizer; compact non-DistOpt LayerWise buffers do not allocate a DDP param buffer.
| @@ -1110,12 +1131,46 @@ def _copy_model_grads_to_main_grads(self): | |||
| model_param.grad = model_param.main_grad | |||
|
|
|||
| def _copy_main_params_to_model_params(self): | |||
There was a problem hiding this comment.
this name is too similar with _copy_main_params_to_param_buffer and why do we need this in the first place?
There was a problem hiding this comment.
_copy_main_params_to_model_params() is the existing Float16OptimizerWithFloat16Params hook; the method name predates this change and pairs with _copy_model_params_to_main_params().
_copy_main_params_to_param_buffer() has a different role: it is the DistributedOptimizer/MXFP8 staging helper that writes FP32 master shards into the DDP parameter buffer before parameter all-gather.
A compact non-DistOpt LayerWise child does not own that byte-sharded parameter buffer, so after the optimizer step it must use _copy_main_params_to_model_params() instead. The new branch also ensures FP8 parameters are written as Q(bf16(master)) and skips parameters that will already be written by the LayerWise gather.
Signed-off-by: Pingtian Li <pingtianl@nvidia.com>
Signed-off-by: Pingtian Li <pingtianl@nvidia.com>
Signed-off-by: Pingtian Li <pingtianl@nvidia.com>
| return data | ||
|
|
||
| @classmethod | ||
| def _filter_gbuf_range_map( |
Signed-off-by: Pingtian Li <pingtianl@nvidia.com>
Signed-off-by: Pingtian Li <pingtianl@nvidia.com>
There was a problem hiding this comment.
Sorry, I don't follow the changes in this file.
Could you move these changes to a separate PR and code-comment the why better?
https://google.github.io/eng-practices/review/developer/small-cls.html
There was a problem hiding this comment.
There are multiple places that call the high_precision_init_value method, this is to reduce the duplicated code.
There was a problem hiding this comment.
I'm still confused. The change in this file made the code strictly harder to read -- what does this dedup buy us here?
6 participants
What does this PR do ?
Adds FP8 parameter gather to the Muon optimizer on the compact decoupled LayerWise DDP layout
(
--fp8-param-gather), so Muon persists and syncs its matrix parameters as FP8 instead of BF16 —bringing its memory footprint down to Adam's level while staying bit-for-bit equal to the
fp8-param-gather-OFF run.
Main-branch counterpart of the
devstack #5388 + #5470 (both merged).Compact-layout checkpoint support is intentionally out of scope for this PR.
Self-contained: it supersedes #5391, which carried the layout half on
mainand is nowclosed into this one. See Parity with the dev stack below.
Builds on the decoupled compact LayerWise layout (dev #5388 / main #5391), included here. Previously the non-DistOpt Muon path could
only hold and all-gather its parameters in BF16. With this change Muon's matrix params are persisted
as FP8 (mxfp8 on Blackwell, blockwise on Hopper), removing the standing BF16 parameter copy; combined
with the no-padding layout from #5388 this closes the Muon-vs-Adam memory gap.
What's new
1. An FP8-aware whole-parameter all-gather for the non-DistOpt Muon path. Muon's matrix buffers
own whole params via ping-pong assignment and sync through
allgather_params(not DistOpt'sbyte-shard reduce-scatter/all-gather). FP8 tensors can't be all-gathered directly — mxfp8's row/col
block scales can't be derived from one another — so the transport rides BF16 and requantizes on
landing: stage each owned param from its fp32 master → bf16 (
_stage_param_to_bf16, ahigh-precision source, not a lossy fp8 dequant) → uneven all-gather-v of the bf16 whole-params →
requantize the gathered bf16 into every rank's fp8
param.data(
copy_back_gathered_bf16_into_fp8_param) so all ranks holdQ(bf16(master))(== OFF / Adam) →post_all_gather_processingrebuilds the fp8 columnwise/transpose.Both copy-back sites — the overlap path (
start_param_sync/finish_param_sync) and thenon-overlap path (
allgather_params) — share one helper,_layerwise_copy_back_gathered_params(bucket, local_rank, fp8_staged=...). The staging decision ispersisted on the bucket as
bucket.layerwise_fp8_stagedatstart_param_synctime and read back bythe (possibly deferred) copy-back, so transport dtype and copy-back can never disagree. Plain bf16
(fp8pg OFF) collapses back to the original code.
Params are dispatched by transport dtype: fp8 and bf16 ride the bf16-staged helper, while native
fp32 params (weights marked
keep_in_fp32, e.g. the DeepSeek-V4 CSAape) are gathered in fp32 —routing them through the bf16-staged path would silently downcast them, and mixing fp32 with bf16 in
one flatten is invalid. For a pure-bf16 model the native group is empty and this collapses to the
original single-helper dispatch.
2.
use_distributed_optimizerbecomes a per-buffer property. On the decoupled layout,LayerWise-managed (Muon 2D-matrix) buffers locally disable DistOpt — compact no-padding layout,
all-reduce gradients, whole-param ping-pong ownership — while sibling buffers (embeddings, biases,
layernorm) keep the standard byte-level
DistributedOptimizerpath.The decision has a single source of truth: it is baked into each buffer's own
ddp_configviadataclasses.replace(...), so bucket groups inherit it throughbuffer.ddp_configrather thanthrough a separately threaded attribute. A bucket group runs a single collective type
(reduce-scatter for DistOpt buffers, all-reduce otherwise), so the merging paths in
partition_buckets(the fp8-merge branch,force_single_bucket_group) assert that the buckets theycombine agree, and the no-fp8 branch keeps one bucket per group so they never mix. When every bucket
agrees — every non-decoupled configuration — this collapses to exactly the previous grouping.
3. Rank-independent ping-pong ownership.
allgather_paramsrequires every DP rank to agree oneach param's single owner.
numelalone isn't a total order (a stable sort tie-breaks equal-numelparams by insertion order → different owners per rank), so ownership is now keyed by a canonical
identity
(chunk_idx, buffer_idx, global_start_index)from model construction, with an assert thatevery Muon param has one.
4. iter-0 master parity, and no double-write of gathered fp8 params. The fp32 master is seeded
from each param's high-precision pre-quantization init (
get_high_precision_init_val), not the lossyfp8 dequant, so fp8pg ON and OFF hold an identical master at iter 0 (matching DistOpt).
The master→model copy routes through bf16 (
Q(bf16(master))), gated by_layer_wise_non_distopt_child— a tagLayerWiseDistributedOptimizersets on theFloat16OptimizerWithFloat16Paramschildren it wraps. Gathered fp8 params are tagged_layer_wise_fp8_gatheredand skipped there, because the all-gather's requantize already writesQ(bf16(master))into theirparam.data; non-gathered fp8 params (e.g. MoE experts atexpt_dp == 1, which the all-gather skips) are untagged and still get their copy here.Routing those Float16 children away from the DDP param buffer in
step_with_ready_gradsis handledby the existing
hasattr(self, "_copy_main_params_to_param_buffer")probe — onlyDistributedOptimizerowns that byte-shard buffer — so this PR deliberately does not define araising stub of that method on
Float16OptimizerWithFloat16Params, which would make the probe alwayssucceed.
5. Memory: drop the non-owned high-precision init copies.
shard_params()removes non-ownedparams from the local optimizer groups, so the Float16 wrapping only clears TE's high-precision init
copy (a full-size CPU tensor per fp8 param) for locally owned params. Without an explicit sweep
every DP rank would retain ~
(dp-1)/dpof the LayerWise matrix params' bf16 CPU copies for the wholerun. The sweep is scoped to LayerWise-managed params — sibling DistOpt params must keep their init
value until their own optimizer's master creation consumes it.
6.
force_syncfinalizes a pending LayerWise gather. A pending LayerWise gather normally landsin
finish_param_sync(the forward pre-hook), whichforce_syncbypasses.start_param_syncnowfinalizes it explicitly — otherwise every rank keeps stale
param.dataand the gather payloadpollutes
grad_data(which the LayerWise path reuses as the all-gather receive buffer).7. Validation guards (
arguments.py): FP8 param gather is supported only on the decoupled layoutand is rejected on the padded LayerWise layout. On decouple it requires
fp8_recipe ∈ {mxfp8, blockwise};fp4_param_gatheris rejected unconditionally (the LayerWise gather routes buckets byis_float8tensor, so an NVFP4 param would silently take the raw flatten path). mxfp8 additionallyrequires
--reuse-grad-buf-for-mxfp8-param-ag, and the layout requiresnum_distributed_optimizer_instances == 1, since the non-DistOpt Muon buffers all-reduce within asingle instance and partial DistOpt would under-reduce Muon gradients across the DP domain.
8. Single all-reduce buffer (a correctness fix this combination exposed). With fp8pg ON,
LayerWise FP8 grads were keyed by
torch.uint8ingroup_params_for_buffers, splitting Muon'sgradients into two all-reduce buffers (uint8 + bf16) where OFF has one. Both reduce in fp32, but
NCCL's fp32 accumulation order is buffer-size sensitive, so the two-buffer order diverged ~1 ULP from
OFF on the DP-reduced dense FP8 2D weights — amplified by Muon's Newton–Schulz into a visible loss
divergence. Fix: a new
merge_layerwise_fp8_gradsflag (decouple-only) keys FP8 Muon grads by theirbf16 logical dtype so FP8 + bf16 grads share one fp32 all-reduce buffer, matching OFF. Two
all-reduces collapse to one (no extra communication); NVFP4's packed param buffer is untouched.
9. Mixed fp8/bf16 buckets (a direct consequence of 8). Once fp8 Muon grads key to their bf16
logical dtype, one buffer — and therefore one bucket — can hold both fp8 and bf16 params (e.g. an
fp8 matrix next to a non-quantized MoE router / DSA indexer / mHC weight). The bf16-staged transport
handles both dtypes, so the staging decision is
any(is_float8tensor(p) for p in bucket.params_list)rather than a peek at
params_list[0]: a bf16-first mixed bucket would otherwise be routed into theraw
_flatten_dense_tensors()path and crash on the MXFP8.view(-1).Correctness — bitwise ON == OFF
Deterministic,
clip_grad=0. Compared per-iterlm loss, forward output, per-param reducedmain_grad, fp32 master, and bf16 model param.(Before the single-buffer fix, loss diverged at it12 on H100 / it3 on GB200; per-param pre-reduce
grads were already identical, isolating the divergence to the DP all-reduce.) With
clip_grad=0,Adam (DistOpt byte-shard, not split) was already bitwise.
clip_grad>0adds a separatelayout-dependent global grad-norm fp32 sum term, so the
clip_grad=1.0benchmarking recipes are notbitwise — bit-exactness is established at
clip_grad=0.Memory & throughput
GB200 DP64 (mxfp8, deepseek_v4_pro_proxy 8L, TP1/PP1/EP64, SL4K):
FP8 param persistence −10.8 GB (identical on Muon and Adam); compact-layout padding removal −26.2 GB;
decouple+fp8pg vs padded Muon −37 GB (~26%), and ~11 GB below Adam-no-fp8pg. H100 dev proxy (DP8,
blockwise) reproduces the mechanism: fp8pg −1.1 GB (35,460 vs 36,594 MB); padding
0.0%vs0.6%.Refactor integrity: Muon padded / fp8pg OFF vs the pre-refactor DSv4+padded-Muon baseline is
142,858 vs 142,857 MB (1 MB) — the padded path is left equivalent.
Throughput (GB200 wall-clock, comparable within branch): Adam < Muon decouple < Muon padded. Muon
is ~5–6% slower than Adam (Newton–Schulz + uneven all-gather-v); fp8 param-gather itself adds ~1%;
the compact layout is ~4.8% faster than padded.
Parity with the dev stack
The FP8/layout source and test commits from the
devside are carried here:force_syncfinalize / non-owned init cleanup / fp4 guard, native-fp32 gather dispatch, added MoE +force_synccoverageThe FP8/layout feature symbols were checked symbol-by-symbol against the dev stack.
Main-only adaptations, all intentional:
ChainedOptimizer.prepare_model_params_for_param_sync()(main'spurpose-built API, added after
devforked) instead of iteratingchained_optimizers;is_grouped_mxfp8tensor()widening from [Main] Numerical fix for moe single grouped weight with fp8 fp4 primary weight and grad norm spikes #5487, and theLayerWise all-gather reads main's cached
self.dp_cp/self.expt_dpgroups from Enforce that the number of optimizer shards used in layout computation is the same used during the training iteration #6048;Golden values regenerated. The
dist_dist_muonckpt-resume goldens are refreshed on thisbranch, using CI's own harness (
tests/functional_tests/shell_test_utils/run_ci_test.sh) on thecommon_pile CI dataset and matching hardware, each running the full ckpt-resume flow (both phases
to iteration 100):
..._dist_dist_muondev_dgx_h100..._dist_dist_muondev_dgx_gb200..._dist_dist_muon_1nodedev_dgx_gb200Every run reported Exact comparison: FAILED / APPROXIMATE test: PASSED before the refresh — a
reordering, not a regression; peak
|rel Δ|on lm loss is ~1.6e-4. (dev's refreshed goldens arenot transferable:
mainanddevhave independently drifted baselines for these cases, ~1e-3across 94–97 of 100 iters, the same order as the change itself.) They were re-verified in compare
mode at the final tip: all three pass with zero exact-comparison failures.
The two
*_param_layoutcases now pass--use-layer-wise-param-layoutexplicitly, since thecompact layout became the default and they exist to cover the padded one; their goldens are
unchanged.
Issue tracking
For PRs from open-source community contributors:
Linked issue: Supersedes #5391 (the main-branch prerequisite, closed into this PR)
Contribution process
Pre-checks
Unit test:
tests/unit_tests/test_muon_decouple_fp8_param_gather.py— bitwise ON-vs-OFF checks on thedecoupled layout (
muon+use_distributed_optimizer→LayerWiseDistributedOptimizer), assertingper-step loss / forward output /
main_grad/ fp32 master / bf16 param are bitwise-identical underdeterministic_mode. Three cases, parametrized overfp8_recipe ∈ {blockwise, mxfp8}:test_on_vs_off_bitwise_identical(× overlap {OFF, ON}),test_moe_on_vs_off_bitwise_identical(× overlap × EP), and
test_force_sync_finalizes_pending_layerwise_gather. blockwise runs on Hopperonly and mxfp8 on Blackwell and newer, so each platform runs 7 and skips 7. Both halves verified
against the CI
Dockerfile.ci.devimage:test_muon_decouple_fp8_param_gather.py7 passed, 7 skipped(blockwise; mxfp8 skipped)test_muon_decouple_fp8_param_gather.py7 passed, 7 skipped(mxfp8; blockwise skipped)Code review
Feel free to message or comment @NVIDIA/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.
Dev-branch counterpart: #5470 (merged)
🤖 Generated with Claude Code