Summary

--fp8-param-gather + --reuse-grad-buf-for-mxfp8-param-ag (mxfp8) was already wired for the standard DistributedOptimizer, but the muon LayerWiseDistributedOptimizer was missing the pieces needed to take that same path. This PR fills those gaps so --fp8-param-gather works for muon-managed 2D weights when --fp8-recipe=mxfp8, on top of #4889's muon-aware bucket alignment, with bitwise parity against the --fp8-param-gather=off baseline.

Design principle: reuse the DistOpt machinery instead of growing a parallel LayerWise path. The bf16-staging write is exposed as a duck-typed _copy_main_params_to_param_buffer() so train_step drives both optimizers through one code path, and the per-bucket param-sync filtering mirrors DistOpt's exactly.

Background: with this feature the model params are kept as persistent MXFP8 tensors; the bf16 param_buffer is aliased onto the idle grad buffer; the all-gather ships bf16 and then quantizes back to MXFP8 in place on every rank. Every gap below is a consequence of the muon/LayerWise path not yet handling some part of that lifecycle that DistOpt already did.

The gaps fall into four areas.

Area 1 — Master-weight numerics for persistent FP8 params

These two gaps made the muon loss drift from the fp8_param_gather=off baseline. Both exist only because the params are now persistent MXFP8 tensors.

1a — fp32 master was initialized from dequantized MXFP8.
The inner Float16OptimizerWithFloat16Params.__init__ builds the master via param.detach().clone().float(). For an MXFP8 param, .float() dequantizes, so the master is born from FP8-rounded values and carries ~FP8 noise from iter 0. DistOpt avoids this by reading model_param.get_high_precision_init_val() (the bf16 init TE preserves on CPU). Without the fix, the iter-10 lm-loss gap is ≈0.27 nats.
→ filled by _restore_high_precision_init_val (mirrors distrib_optimizer.py:405-416).

1b — the inner step re-quantized MXFP8 params redundantly.
_copy_main_params_to_model_params would run model.data.copy_(main_fp32) on each MXFP8 model param, triggering an extra QuantizedTensor quantize_(fp32) that the post-AG quantize_(bf16) then overwrites — wasted work that also perturbs muon convergence. DistOpt's reuse_grad_buf step path never touches MXFP8 storage at the inner step.
→ filled by _skip_mxfp8_in_copy_main_to_model (+ a shallow-copied inner_config that hides reuse_grad_buf_for_mxfp8_param_ag from the inner optimizer; self.config is restored after super().__init__ so the outer ChainedOptimizer's shared-config assertion still holds).

Area 2 — Staging the bf16 param buffer before the all-gather

With reuse_grad_buf_for_mxfp8_param_ag the bf16 param_buffer is separate from the MXFP8 storage, and the all-gather ships param_buffer. Two gaps left it holding the wrong bytes.

2a — the updated master was never written into the staging buffer.
The optimizer step updates the fp32 master, but nothing copied it (cast to bf16) into param_buffer before the AG, so the AG shipped stale bytes.
→ filled by _copy_main_params_to_param_buffer — a duck-typed twin of DistributedOptimizer._copy_main_params_to_param_buffer, so train_step can call it polymorphically over (DistributedOptimizer, LayerWiseDistributedOptimizer). Self-guards on use_buffer_param_sync + reuse_grad_buf_for_mxfp8_param_ag, so it is a no-op for configs that don't need it.

2b — the staging buffer is zeroed before the deferred all-gather (overlap path).
Because the staging buffer is aliased onto the grad buffer, the next iteration's zero_grad_buffer() zeroes it after the in-step write but before the deferred forward-pre-hook all-gather. Under --overlap-param-gather the AG therefore shipped a zeroed buffer and the muon-managed weights never updated (loss frozen). The fix is not to disable overlap, but to re-stage post-zero on every site that force-syncs through the param buffer:

• train_step (training loop) — the same hook DistOpt already uses;
• the eval block (disable_forward_pre_hook(param_sync=True) force-syncs params for eval);
• paged_stash _try_copy_main_params (reachable when moe_expert_rank_capacity_factor is set with reuse_grad_buf + overlap).
  All three now re-stage LayerWiseDistributedOptimizer too, not just DistributedOptimizer.

Area 3 — Disjoint param-sync when LayerWise and DistOpt coexist

muon manages 2D weights via LayerWise; everything else goes through a sibling DistributedOptimizer (#4771). They share the same DDP buffers, so the gap is making each optimizer all-gather only its own buckets.

3a — per-bucket ownership / mixed bucket groups.
A bucket group can be mixed-ownership: in partition_buckets Case 3 (FP8 present, reduce_scatter_with_fp32_accumulation=False) non-FP8 DistOpt-managed bf16 buckets (biases, layernorms) get appended into the last FP8 bucket group. Checking only buckets[0] and dispatching AG on the whole group double-syncs the foreign buckets. Both optimizers now filter per-bucket and, for mixed groups, synthesize a thin transient bucket group containing only the buckets they own. Applied symmetrically in three places: DistributedOptimizer.start_param_sync_for_bucket_group_subset, LayerWiseDistributedOptimizer.start_param_sync_for_bucket_group_subset, and the deferred-sync path in optimizer.py (which also recurses the chained-optimizer tree so nested DistOpt instances are reached).

Area 4 — Fail-fast configuration guards

The remaining gap was silent misconfiguration: combinations that would quietly all-gather stale storage.

4a — mxfp8 + LayerWise + --fp8-param-gather now requires --reuse-grad-buf-for-mxfp8-param-ag (the two must be co-enabled).
4b — non-mxfp8 + LayerWise + --fp8-param-gather is explicitly rejected (only MXFP8's _rowwise/_columnwise round-trip is wired up today).

Deliberately NOT in this PR

• Quantizer dual-usage re-assert. An earlier revision re-asserted quantizer.set_usage(rowwise=True, columnwise=True) before each AG, guessing TE's bwd toggled columnwise_usage. Running the bitwise test in deterministic_mode() without that re-assert still reproduces atol=rtol=0 ON-vs-OFF parity, proving it unnecessary; the helper was removed to keep the LayerWise surface minimal.
• Blockwise FP8 + LayerWise + --fp8-param-gather — rejected at arg-parse (4b) and skipped in the test; needs a separate _rowwise/_columnwise layout + post-AG path.

Test plan

• tests/unit_tests/test_muon_mxfp8_fp8_param_gather.py::test_on_vs_off_bitwise_identical — builds the same small GPT twice (fp8_param_gather=False/True), restores the ON run's init state from the OFF snapshot so both start bit-identical, runs 5 muon steps in deterministic_mode(), and asserts atol=rtol=0 on loss, forward output, per-parameter main_grad, fp32 master, and (non-MXFP8) model params. Parametrized over overlap_param_gather ∈ {False, True} — the True case is the regression guard for Area 2b. Skips automatically when MXFP8 is unsupported (pre-Blackwell, sm_120+, TE without MXFP8). Verified on GB200: [mxfp8-False] and [mxfp8-True] both pass; blockwise variants skip by design.
• End-to-end convergence smoke (DSV4 proxy, GB200, muon, e4m3 mxfp8): ON tracks the OFF baseline (iter-10 within ≈0.03 nats; ≈0.27 nats without Area 1a). Overlap ON: before Area 2b the loss was frozen; after, the overlap run tracks the non-overlap run.

Changes by file

🤖 Generated with Claude Code