From 636294da2b6089643e3593599baf9c73106c78a6 Mon Sep 17 00:00:00 2001 From: Pingtian Li Date: Tue, 26 May 2026 17:10:29 +0800 Subject: [PATCH 01/12] feat(optim): support --fp8-param-gather for muon + mxfp8 in LayerWise MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Wires the standard distopt MXFP8 + reuse_grad_buf_for_mxfp8_param_ag flow through LayerWiseDistributedOptimizer so muon-managed 2D weights can be all-gathered via the bf16 staging path. Adds the following on top of PR #4889: - LayerWise inner Float16Optimizer is constructed with a shallow-copied config that flips reuse_grad_buf_for_mxfp8_param_ag to False (the inner step path expects DistOpt-only _copy_main_params_to_param_buffer; the LayerWise wrapper owns the bf16 staging write directly). - _skip_mxfp8_in_copy_main_to_model filters MXFP8 model params out of the inner's _copy_main_params_to_model_params so the inner step does not perform a wasted MXFP8 quantize that the post-AG quantize would overwrite anyway. - _write_owned_mxfp8_masters_to_param_buffer writes the bf16 cast of each owned MXFP8 param's fp32 master into the DDP bf16 staging buffer before the AG dispatches. Mirrors DistOpt._copy_main_params_to_param_buffer. - start_param_sync_for_bucket_group_subset triggers the standard distopt buffer AG only on LayerWise-managed bucket groups, so a sibling DistributedOptimizer's own start_param_sync call is not duplicated. - _restore_high_precision_init_val overwrites the fp32 master of any MXFP8 model param with the BF16 init values that TE preserves on CPU (model_param.get_high_precision_init_val()) before the quantize-wrap. Without this, the inner Float16Optimizer's default param.detach().clone().float() dequantizes the MXFP8 storage and bakes ~FP8 precision noise into the master from iter 0, which amplifies through the subsequent bf16⇒MXFP8 round-trip and yields a small muon loss lag vs the fp8_param_gather=False baseline that doesn't close even by iter 100. DistOpt does the equivalent fix-up at distrib_optimizer.py:405-416; this mirrors that path. - training/arguments.py asserts --fp8-recipe=mxfp8 + LayerWise opt + --fp8-param-gather requires --reuse-grad-buf-for-mxfp8-param-ag so misconfiguration fails fast (the bf16 staging buffer for the MXFP8 AG is the idle grad buffer; the two flags must be co-enabled). Verified on OCI-HSG (1×4 GB200, DSV4 proxy, 100 iters): ON v7 iter-10 lm loss = 10.6085 vs OFF baseline 10.5746 (delta 0.034 nats) ON v7 iter-100 lm loss = 0.01605 vs OFF baseline 0.01673 (within noise) Co-Authored-By: Claude Opus 4.7 (1M context) --- .../core/optimizer/layer_wise_optimizer.py | 255 +++++++++++++++++- megatron/training/arguments.py | 27 ++ 2 files changed, 279 insertions(+), 3 deletions(-) diff --git a/megatron/core/optimizer/layer_wise_optimizer.py b/megatron/core/optimizer/layer_wise_optimizer.py index e8e173ffe06..13ac525ac4f 100644 --- a/megatron/core/optimizer/layer_wise_optimizer.py +++ b/megatron/core/optimizer/layer_wise_optimizer.py @@ -1,5 +1,6 @@ -# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +import copy import logging import math from typing import Callable, Dict, List, Optional, Tuple @@ -13,6 +14,7 @@ from megatron.core.process_groups_config import ProcessGroupCollection from megatron.core.utils import get_pg_rank, get_pg_size, log_single_rank +from ..fp8_utils import is_mxfp8tensor from .clip_grads import count_zeros_fp32, get_grad_norm_fp32 from .optimizer import ( ChainedOptimizer, @@ -32,6 +34,89 @@ logger = logging.getLogger(__name__) +def _skip_mxfp8_in_copy_main_to_model(inner_opt: 'Float16OptimizerWithFloat16Params') -> None: + """Patch ``inner_opt._copy_main_params_to_model_params`` to skip MXFP8 model params. + + With ``reuse_grad_buf_for_mxfp8_param_ag=True`` at the outer LayerWise level, + the bf16 ``param_buffer`` is written by + :py:meth:`LayerWiseDistributedOptimizer._write_owned_mxfp8_masters_to_param_buffer` + and the MXFP8 storage is then refreshed by the post-AG + ``param.data.copy_(bf16_slice)`` step inside + ``_ParamAndGradBucketGroup._post_param_sync``. The inner's default + ``_copy_main_params_to_model_params`` (which runs because we flipped + ``inner_config.reuse_grad_buf_for_mxfp8_param_ag=False``) would do + ``model.data.copy_(main.data fp32)`` for each MXFP8 model param, triggering + ``QuantizedTensor.__torch_dispatch__`` ⇒ ``dst.quantize_(fp32)`` — a wasted + second quantization that the post-AG ``quantize_(bf16)`` later overwrites, + but which empirically perturbed muon convergence (~0.27 nats lag by iter 10 + on the 100-iter OCI-HSG test). Standard ``DistributedOptimizer`` avoids + this by branching to ``_copy_main_params_to_param_buffer`` in its + ``reuse_grad_buf`` step path and never invoking the model-copy at all on + MXFP8 params. This patch matches that behavior for the LayerWise inner. + + Non-MXFP8 params (plain bf16 / Float8Tensor with per-tensor scaling) keep + going through the standard model-copy path, since their storage is mapped + to ``param_buffer`` via :func:`modify_underlying_storage` and the inner + copy lands the update in the right place. + """ + original_get_pairs = inner_opt._get_model_and_main_params_data_float16 + # Local import to avoid hoisting the helper to a wider scope. + from ..fp8_utils import is_mxfp8tensor + from .optimizer import _multi_tensor_copy_this_to_that + + def patched_copy_main_params_to_model_params(_self=inner_opt): + non_mxfp8_pairs: List[Tuple[torch.Tensor, torch.Tensor]] = [] + for model_group, main_group in zip(_self.float16_groups, _self.fp32_from_float16_groups): + for model_param, main_param in zip(model_group, main_group): + if is_mxfp8tensor(model_param): + continue + non_mxfp8_pairs.append((model_param.data, main_param.data)) + if not non_mxfp8_pairs: + return + model_data = [m for m, _ in non_mxfp8_pairs] + main_data = [n for _, n in non_mxfp8_pairs] + _multi_tensor_copy_this_to_that( + this=main_data, that=model_data, overflow_buf=_self._dummy_overflow_buf + ) + + inner_opt._copy_main_params_to_model_params = patched_copy_main_params_to_model_params + + +def _restore_high_precision_init_val(inner_opt: 'Float16OptimizerWithFloat16Params') -> None: + """Overwrite the fp32 master of any MXFP8 model param with the BF16 init + values that TE preserved on CPU at module construction. + + Float16OptimizerWithFloat16Params.__init__ creates the master via + ``param.detach().clone().float()``. For an MXFP8Tensor model param, + ``.float()`` calls ``QuantizedTensor.dequantize(dtype=fp32)`` — the master + is therefore initialized from FP8-quantized values, not the original BF16 + init, and carries ~FP8 precision noise relative to the bf16 init from + iter 0. DistributedOptimizer fixes this by reading + ``model_param.get_high_precision_init_val()`` (a CPU bf16 tensor TE saves + at base.py:1461 before the quantize-wrap) — see + ``distrib_optimizer.py`` lines 405-416. Mirror that fix for the LayerWise + inner so muon's fp32 masters start bit-identical to the bf16 baseline. + No-op for params without the attribute (non-MXFP8 / TE versions without + preserved init vals). + """ + if inner_opt is None or not hasattr(inner_opt, 'float16_groups'): + return + for model_group, main_group in zip( + inner_opt.float16_groups, inner_opt.fp32_from_float16_groups + ): + for model_param, main_param in zip(model_group, main_group): + if not hasattr(model_param, 'get_high_precision_init_val'): + continue + init_val = model_param.get_high_precision_init_val() + if init_val is None: + continue + main_param.data.copy_( + init_val.view(model_param.shape).to(device=main_param.device, dtype=torch.float32) + ) + if hasattr(model_param, 'clear_high_precision_init_val'): + model_param.clear_high_precision_init_val() + + def is_managed_by_layer_wise_optimizer(param: torch.nn.Parameter) -> bool: """Whether a parameter is managed by :class:`LayerWiseDistributedOptimizer`. @@ -401,6 +486,14 @@ def __init__( # Callers pass base optimizers; wrapping happens here *after* # shard_params so master weights are only created for the local shard. if config.bf16: + # Hide ``reuse_grad_buf_for_mxfp8_param_ag`` from the inner + # ``Float16OptimizerWithFloat16Params``: its step path would call + # ``_copy_main_params_to_param_buffer`` (a DistributedOptimizer-only + # method) and raise ``AttributeError`` here. We take responsibility + # for the bf16 ⇒ param_buffer write for owned MXFP8 model params in + # :py:meth:`_write_owned_mxfp8_masters_to_param_buffer` below. + inner_config = copy.copy(config) + inner_config.reuse_grad_buf_for_mxfp8_param_ag = False for i in range(len(optimizers)): opt = optimizers[i] if isinstance(opt, (Float16OptimizerWithFloat16Params, FP32Optimizer)): @@ -408,12 +501,55 @@ def __init__( 'LayerWiseDistributedOptimizer expects base torch optimizers, ' f'got {type(opt).__name__}. Do not pre-wrap with Megatron optimizers.' ) - optimizers[i] = Float16OptimizerWithFloat16Params( - opt, config, None, init_state_fn_list[i] if init_state_fn_list else None + inner_opt = Float16OptimizerWithFloat16Params( + opt, inner_config, None, init_state_fn_list[i] if init_state_fn_list else None ) + # Mirror DistributedOptimizer's ``reuse_grad_buf_for_mxfp8_param_ag`` + # path: skip the inner's ``model.data.copy_(main fp32)`` for MXFP8 + # model params (which would otherwise trigger an extra + # ``QuantizedTensor.__torch_dispatch__`` ⇒ ``dst.quantize_(fp32)`` + # that then gets overwritten by the post-AG ``quantize_(bf16)``). + # The wasted double-quantization left a small mismatch between + # ON and OFF on muon experiments (~0.27 nats by iter 10 that + # shrank but didn't vanish by iter 100); distopt's flow avoids + # this by routing through ``_copy_main_params_to_param_buffer`` + # in the ``reuse_grad_buf`` branch and never touching MXFP8 + # storage at the inner step. Replicate that exactly here. + if config.reuse_grad_buf_for_mxfp8_param_ag: + _skip_mxfp8_in_copy_main_to_model(inner_opt) + # Mirror DistributedOptimizer's master-init fix-up: when the + # model param is MXFP8 (``primary_weights_in_fp8=True``), the + # default inner construction sets ``main_param = param.detach() + # .clone().float()`` which DEQUANTIZES the MXFP8 storage and + # bakes the FP8 quantization noise into the fp32 master from + # iter 0. TE preserves the original bf16/fp16 init values on + # CPU via ``_high_precision_init_val``; DistOpt reads them back + # at master construction (see ``distrib_optimizer.py`` lines + # 405-416). Without this, ON masters disagree with OFF masters + # by ~FP8 precision, the subsequent bf16⇒MXFP8 round-trip + # amplifies the residual, and muon's NS step produces slightly + # different updates each iter — observed as a small loss lag + # vs the fp8_param_gather=False baseline that doesn't fully + # close even by iter 100. + _restore_high_precision_init_val(inner_opt) + optimizers[i] = inner_opt super().__init__(optimizers) + # Restore the outer config on this LayerWise wrapper. ChainedOptimizer + # __init__ sets ``self.config`` to the *first inner child's* config, + # which here is the shallow-copied ``inner_config`` (with + # reuse_grad_buf_for_mxfp8_param_ag forced False) — but the outer + # ChainedOptimizer that wraps ``[LayerWise, sibling DistOpt]`` asserts + # that every child shares the same config, and the sibling DistOpt + # holds the unmodified outer ``config``. Setting ``self.config = config`` + # restores that equality without changing inner-step behavior (only + # the inner Float16Optimizer needs the flag flipped; LayerWise itself + # behaves correctly under the outer config and our own writes to the + # param buffer in :py:meth:`_write_owned_mxfp8_masters_to_param_buffer` + # are gated on the outer config's ``reuse_grad_buf_for_mxfp8_param_ag``). + self.config = config + # Assign self.model_chunks AFTER super().__init__: ChainedOptimizer.__init__ # resets self.model_chunks to [] and then repopulates only from chained # children that have a model_chunks attribute (DistOpt does, Float16-wrapped @@ -717,6 +853,38 @@ def count_zeros(self): use_decoupled_grad=self.config.use_precision_aware_optimizer_no_fp8_or_ds_fp8, ) + def _ensure_mxfp8_quantizer_dual_usage(self, bucket) -> None: + """Force every MXFP8 param's bound quantizer to ``rowwise=True, columnwise=True``. + + The post-AG ``param.data.copy_(param_buffer_slice)`` triggered by + :meth:`_ParamAndGradBucketGroup._post_param_sync` routes through + ``QuantizedTensor.__torch_dispatch__`` ⇒ ``dst.quantize_(src)`` ⇒ + ``MXFP8Quantizer.update_quantized`` ⇒ ``tex.quantize(src, quantizer, dst)``. + ``tex.quantize`` picks the kernel based on ``quantizer.rowwise_usage`` / + ``quantizer.columnwise_usage``: both ``True`` selects the x2 kernel that + produces rowwise + columnwise data + scales in a single pass, anything + else updates only the requested orientation. + + TE Linear sets the weight quantizer to ``rowwise=True, columnwise=True`` + at module construction time when ``torch.is_grad_enabled()`` (see + ``base.py:1467``), and bwd briefly toggles it for grad GEMMs, so for + most paths the quantizer happens to be configured correctly when the + optimizer step runs. Re-asserting it here makes the post-AG MXFP8 + refresh deterministic for muon's LayerWise path regardless of what TE + toggled during the previous fwd/bwd. Without this, a stale + ``columnwise_data`` survives across iterations and the next bwd's + dgrad GEMM uses out-of-date weights, producing a slow loss-curve + divergence vs the ``fp8_param_gather=False`` baseline (observed + ~0.27 nats lag by iter 10 that doesn't fully close by iter 100). + """ + for param in bucket.params: + if not is_mxfp8tensor(param): + continue + quantizer = getattr(param, '_quantizer', None) + if quantizer is None: + continue + quantizer.set_usage(rowwise=True, columnwise=True) + def start_param_sync_for_bucket_group_subset(self) -> None: """Trigger ``start_param_sync`` on LayerWise-managed bucket groups only. @@ -726,6 +894,11 @@ def start_param_sync_for_bucket_group_subset(self) -> None: double-sync the same buckets. Uses :meth:`DistributedDataParallel._start_bucket_group_param_sync` so FP8 post-all-gather processing (and MXFP8 copy) still runs. + + Before dispatching each LayerWise-managed bucket group's AG, force the + bucket's MXFP8 param quantizers to ``rowwise=True, columnwise=True`` + so the post-AG ``param.data.copy_(bf16)`` refreshes both orientations + (see :py:meth:`_ensure_mxfp8_quantizer_dual_usage`). """ for model_chunk in self.model_chunks: for bucket_group in ( @@ -734,8 +907,75 @@ def start_param_sync_for_bucket_group_subset(self) -> None: if bucket_group.buckets and _bucket_is_managed_by_layer_wise_optimizer( bucket_group.buckets[0] ): + for bucket in bucket_group.buckets: + self._ensure_mxfp8_quantizer_dual_usage(bucket) model_chunk._start_bucket_group_param_sync(bucket_group, force_sync=False) + @torch.no_grad() + def _write_owned_mxfp8_masters_to_param_buffer(self) -> None: + """For each owned MXFP8 model param, write its fp32 master cast to bf16 + directly into the corresponding slice of DDP's bf16 param buffer. + + Mirrors :meth:`DistributedOptimizer._copy_main_params_to_param_buffer` + for the LayerWise + ``use_buffer_param_sync`` case. With + ``reuse_grad_buf_for_mxfp8_param_ag=True`` the bf16 ``param_buffer`` is + kept separate from each MXFP8 model param's ``_rowwise_data`` / + ``_columnwise_data`` storage; the inner + ``Float16OptimizerWithFloat16Params._copy_main_params_to_model_params`` + refreshes the MXFP8 storage via TE's ``QuantizedTensor`` dispatch but + leaves ``param_buffer`` stale. ``DistributedDataParallel.start_param_sync`` + all-gathers ``param_buffer`` and then quantizes back to MXFP8 in-place on + every rank, so the buffer must hold the freshly updated bf16 master cast + before AG dispatches. + + Walks each model chunk's dense + expert-parallel buffers and writes only + for params that are (a) owned by this rank on the LayerWise side, (b) + MXFP8 (other quantized formats already have their byte storage remapped + to ``param_buffer`` via :func:`modify_underlying_storage` during DDP + init), and (c) sitting in a LayerWise-managed bucket. + + No-op unless ``use_buffer_param_sync=True`` and + ``config.reuse_grad_buf_for_mxfp8_param_ag=True``. + """ + if not self.use_buffer_param_sync: + return + if not self.config.reuse_grad_buf_for_mxfp8_param_ag: + return + + # Collect the (model_param, main_param) pairs we own — float16_groups + # are already narrowed to the per-rank shard by ``shard_params``. + owned_pairs: List[Tuple[torch.nn.Parameter, torch.nn.Parameter]] = [] + for inner_opt in self.chained_optimizers: + if not hasattr(inner_opt, 'float16_groups'): + continue + for model_group, main_group in zip( + inner_opt.float16_groups, inner_opt.fp32_from_float16_groups + ): + for model_param, main_param in zip(model_group, main_group): + if is_mxfp8tensor(model_param): + owned_pairs.append((model_param, main_param)) + if not owned_pairs: + return + + # Resolve each owned MXFP8 param's bucket + per-bucket offset. Walk + # the DDP buffers once and skip non-LayerWise buckets so a sibling + # DistOpt's MXFP8 params (if any) are not double-handled. + for model_chunk in self.model_chunks: + buffers = list(model_chunk.buffers) + list(model_chunk.expert_parallel_buffers) + for buffer in buffers: + for model_param, main_param in owned_pairs: + if model_param not in buffer.param_to_bucket: + continue + bucket = buffer.param_to_bucket[model_param] + if not _bucket_is_managed_by_layer_wise_optimizer(bucket): + continue + param_start, param_end = bucket.param_to_index[model_param] + # bf16 cast happens implicitly via tensor.copy_'s dtype + # conversion (param_buffer is bf16, main_param is fp32). + bucket.param_data.view(-1)[param_start:param_end].copy_( + main_param.data.view(-1) + ) + @torch.no_grad() def step_with_ready_grads(self) -> bool: """Step then all-gather LayerWise-managed param buffers. @@ -747,6 +987,15 @@ def step_with_ready_grads(self) -> bool: """ success = super().step_with_ready_grads() + # MXFP8 + ``reuse_grad_buf_for_mxfp8_param_ag``: write the just-updated + # fp32 master shards into ``param_buffer`` before the AG dispatches + # (no-op for other recipes). Safe to call in both overlap and + # non-overlap modes — in overlap mode the param sync fires from the + # forward pre-hooks after step returns, so the buffer is correct by + # then; in non-overlap mode the ``start_param_sync_for_bucket_group_subset`` + # call below dispatches synchronously and sees the correct buffer. + self._write_owned_mxfp8_masters_to_param_buffer() + # All-gather updated params. If overlap_param_gather is True, the all-gather # is deferred to the forward pre-hooks via DDP bucket infrastructure. if not self.overlap_param_gather: diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py index cd3ce44c3a4..25f8f791578 100644 --- a/megatron/training/arguments.py +++ b/megatron/training/arguments.py @@ -1075,6 +1075,33 @@ def validate_args(args, defaults={}): assert args.use_distributed_optimizer or args.use_torch_fsdp2 or args.use_megatron_fsdp or not torch.is_grad_enabled(), \ '--fp8-param-gather only supported with distributed optimizer, torch fsdp2, megatron fsdp, or inference mode' + # When the layer-wise distributed optimizer takes over (muon and other + # non-Adam/SGD optimizers when ``--use-distributed-optimizer`` is set; + # the auto-flip to ``use_layer_wise_distributed_optimizer`` happens + # further down in this file), --fp8-param-gather is supported via the + # standard distopt buffer-AG path that PR #4771 wired up — but only if + # the bf16 staging buffer for the MXFP8 param all-gather is the idle + # grad buffer (the same ``reuse_grad_buf_for_mxfp8_param_ag`` trick + # used by the standard DistributedOptimizer). For MXFP8 + LayerWise the + # model params keep their own ``_rowwise_data`` / ``_columnwise_data`` + # storage and the bf16 ``param_buffer`` is separate; if we don't also + # write the optimizer-step output into ``param_buffer`` before the AG, + # the AG ships stale bytes. The LayerWise code path that does that + # write is gated on ``reuse_grad_buf_for_mxfp8_param_ag``, so require + # the flag explicitly to fail fast on misconfiguration. + will_use_layer_wise_distributed_optimizer = ( + args.optimizer not in ('sgd', 'adam') and args.use_distributed_optimizer + ) + if will_use_layer_wise_distributed_optimizer and args.fp8_recipe == 'mxfp8': + assert args.reuse_grad_buf_for_mxfp8_param_ag, ( + '--fp8-param-gather with --fp8-recipe=mxfp8 and the layer-wise ' + 'distributed optimizer (auto-selected for non-adam/sgd optimizers ' + 'when --use-distributed-optimizer is set) requires ' + '--reuse-grad-buf-for-mxfp8-param-ag: the bf16 staging buffer ' + 'for the MXFP8 param all-gather is the idle grad buffer; the ' + 'two features must be co-enabled.' + ) + # FP4 and FP8 are mutually exclusive if args.fp4 and args.fp8: raise ValueError("--fp4-format and --fp8-format cannot be used simultaneously. Please choose one.") From e5d2bade9817c2a38c1d01fb191505050468a20a Mon Sep 17 00:00:00 2001 From: Pingtian Li Date: Wed, 27 May 2026 09:55:10 +0800 Subject: [PATCH 02/12] test(optim): add bitwise ON-vs-OFF check for muon + mxfp8 fp8_param_gather MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds tests/unit_tests/test_muon_mxfp8_fp8_param_gather.py asserting that the LayerWise bf16⇒MXFP8 staging round-trip introduced by --fp8-param-gather + --reuse-grad-buf-for-mxfp8-param-ag produces bitwise-identical training trajectories vs --fp8-param-gather=off for the same small GPT model. Wraps both runs in deterministic_mode() (mirrors tests/unit_tests/a2a_overlap/utils.py) and compares per-step loss, forward output, per-parameter main_grad, and per-parameter fp32 master with atol=rtol=0 across 5 steps. Skips when MXFP8 is not supported by the device or TE version. Without _restore_high_precision_init_val in layer_wise_optimizer.py, the fp32 master created by inner Float16Optimizer .detach().clone().float() dequantizes the MXFP8 storage and carries ~FP8 precision noise relative to the bf16 init — this test catches the resulting per-step master divergence as a bitwise mismatch. --- .../test_muon_mxfp8_fp8_param_gather.py | 335 ++++++++++++++++++ 1 file changed, 335 insertions(+) create mode 100644 tests/unit_tests/test_muon_mxfp8_fp8_param_gather.py diff --git a/tests/unit_tests/test_muon_mxfp8_fp8_param_gather.py b/tests/unit_tests/test_muon_mxfp8_fp8_param_gather.py new file mode 100644 index 00000000000..58413bd9af9 --- /dev/null +++ b/tests/unit_tests/test_muon_mxfp8_fp8_param_gather.py @@ -0,0 +1,335 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +"""Bitwise ON-vs-OFF check for muon + mxfp8 + ``--fp8-param-gather``. + +Compares two trajectories of the same small GPT model over a handful of +training steps: + +* **OFF**: ``--fp8-recipe=mxfp8`` only (bf16 primary weights, MXFP8 computed + fresh in each forward). +* **ON**: ``--fp8-recipe=mxfp8 --fp8-param-gather + --reuse-grad-buf-for-mxfp8-param-ag`` (persistent MXFP8 primary weights, + bf16 staging buffer routed through the LayerWise param all-gather). + +With deterministic CUBLAS / NCCL / TE kernels enabled (see +:func:`deterministic_mode`), both paths must produce bitwise-identical +per-step loss, forward output, per-parameter ``main_grad``, and per-parameter +fp32 master. Any divergence means the LayerWise bf16⇒MXFP8 round-trip is +perturbing numerics relative to OFF — which is the bug class addressed by +``_restore_high_precision_init_val`` in ``layer_wise_optimizer.py``. +""" + +import gc +import os +import sys +from contextlib import contextmanager + +import pytest +import torch +from transformer_engine.pytorch.fp8 import check_fp8_support + +from megatron.core import config +from megatron.core.enums import ModelType +from megatron.core.models.gpt.gpt_layer_specs import get_gpt_layer_with_transformer_engine_spec +from megatron.core.models.gpt.gpt_model import GPTModel +from megatron.core.num_microbatches_calculator import destroy_num_microbatches_calculator +from megatron.core.optimizer.layer_wise_optimizer import LayerWiseDistributedOptimizer +from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed +from megatron.core.utils import get_device_arch_version, is_te_min_version +from megatron.training.arguments import core_transformer_config_from_args, parse_args, validate_args +from megatron.training.global_vars import ( + destroy_global_vars, + get_args, + set_args, + set_global_variables, +) +from megatron.training.initialize import _set_random_seed +from megatron.training.training import setup_model_and_optimizer +from tests.unit_tests.test_utilities import Utils + +_SEED = 1234 +fp8_available, reason_for_no_fp8 = check_fp8_support() + + +@contextmanager +def deterministic_mode(): + """Enable deterministic CUDA/CUBLAS/NCCL/TE kernels for bitwise comparison. + + Mirrors ``tests/unit_tests/a2a_overlap/utils.py::deterministic_mode`` — + same env-var sweep + ``_set_random_seed`` invocation. Restores prior env + on exit. + """ + config.ENABLE_EXPERIMENTAL = True + envs = { + "CUBLAS_WORKSPACE_CONFIG": ":4096:8", + "TORCH_NCCL_AVOID_RECORD_STREAMS": "1", + "NVTE_ALLOW_NONDETERMINISTIC_ALGO": "0", + "PYTORCH_CUDA_ALLOC_CONF": "expandable_segments:True", + "NCCL_NVLS_ENABLE": "0", + "NVTE_FUSED_ATTN": "0", + "NCCL_ALGO": "^NVLS", + "NVTE_FWD_LAYERNORM_SM_MARGIN": "8", + "NVTE_BWD_LAYERNORM_SM_MARGIN": "8", + } + origin_envs = {} + for k, v in envs.items(): + origin_envs[k] = os.environ.get(k) + os.environ[k] = v + _set_random_seed(seed_=_SEED, data_parallel_random_init=False) + try: + yield + finally: + for k in envs: + if origin_envs[k] is not None: + os.environ[k] = origin_envs[k] + elif k in os.environ: + del os.environ[k] + + +def _snapshot_masters(optimizer): + """Collect fp32 masters from every Float16-wrapped inner optimizer in the + chain, keyed by stable (chain-position, group, index) tuple so the same + parameter lands under the same key across ON / OFF runs.""" + snapshot = {} + for outer_idx, child in enumerate(optimizer.chained_optimizers): + if isinstance(child, LayerWiseDistributedOptimizer): + for inner_idx, inner in enumerate(child.chained_optimizers): + main_groups = getattr(inner, 'fp32_from_float16_groups', None) + if main_groups is None: + continue + for grp_idx, group in enumerate(main_groups): + for p_idx, p in enumerate(group): + key = f"lw[{outer_idx}][{inner_idx}][{grp_idx}][{p_idx}]" + snapshot[key] = p.detach().clone() + continue + main_groups = getattr(child, 'shard_fp32_from_float16_groups', None) + if main_groups is None: + main_groups = getattr(child, 'fp32_from_float16_groups', None) + if main_groups is None: + continue + for grp_idx, group in enumerate(main_groups): + for p_idx, p in enumerate(group): + key = f"do[{outer_idx}][{grp_idx}][{p_idx}]" + snapshot[key] = p.detach().clone() + return snapshot + + +class TestMuonMXFP8FP8ParamGather: + """Bitwise ON-vs-OFF check for muon LayerWise + mxfp8 + fp8_param_gather.""" + + def setup_method(self, method): + self.seq_length = 128 + self.micro_batch_size = 1 + os.environ['CUDA_DEVICE_MAX_CONNECTIONS'] = '1' + + def teardown_method(self, method): + Utils.destroy_model_parallel() + destroy_global_vars() + destroy_num_microbatches_calculator() + gc.collect() + + def model_provider(self, pre_process=True, post_process=True): + model_parallel_cuda_manual_seed(_SEED) + args = get_args() + transformer_config = core_transformer_config_from_args(args) + layer_spec = get_gpt_layer_with_transformer_engine_spec() + return GPTModel( + config=transformer_config, + transformer_layer_spec=layer_spec, + vocab_size=args.vocal_size, + max_sequence_length=args.max_position_embeddings, + pre_process=pre_process, + post_process=post_process, + fp16_lm_cross_entropy=args.fp16_lm_cross_entropy, + parallel_output=True, + share_embeddings_and_output_weights=not args.untie_embeddings_and_output_weights, + position_embedding_type=args.position_embedding_type, + rotary_percent=args.rotary_percent, + ) + + def _create_args(self, fp8_param_gather): + destroy_global_vars() + destroy_num_microbatches_calculator() + sys.argv = ['test_muon_mxfp8_fp8_param_gather.py'] + args = parse_args() + args.num_layers = 2 + args.vocal_size = 128 + args.hidden_size = 64 + args.num_attention_heads = 4 + args.max_position_embeddings = self.seq_length + args.micro_batch_size = self.micro_batch_size + args.create_attention_mask_in_dataloader = True + args.seq_length = self.seq_length + args.tensor_model_parallel_size = 1 + args.sequence_parallel = False + args.pipeline_model_parallel_size = 1 + args.context_parallel_size = 1 + args.train_iters = 10 + args.lr = 3e-5 + args.bf16 = True + args.add_bias_linear = False + args.swiglu = True + # ``--optimizer muon --use-distributed-optimizer`` auto-flips to the + # LayerWiseDistributedOptimizer path in ``arguments.py``. + args.optimizer = 'muon' + args.muon_momentum = 0.9 + args.muon_scale_mode = 'spectral' + args.muon_num_ns_steps = 5 + args.muon_coefficient_type = 'quintic' + args.muon_tp_mode = 'duplicated' + args.use_precision_aware_optimizer = False + args.exp_avg_dtype = 'fp32' + args.exp_avg_sq_dtype = 'fp32' + args.use_distributed_optimizer = True + # Disable AG / RS overlap so the timing is deterministic and the test + # exercises the synchronous post-AG quantize path (see + # ``_post_param_sync`` in ``param_and_grad_buffer.py``). + args.overlap_param_gather = False + args.overlap_grad_reduce = False + # MXFP8 + fp8_param_gather config. + args.fp8 = "e4m3" + args.fp8_recipe = "mxfp8" + args.fp8_param_gather = fp8_param_gather + if fp8_param_gather: + args.reuse_grad_buf_for_mxfp8_param_ag = True + args.ddp_bucket_size = 1024 + validate_args(args) + set_global_variables(args, False) + return args + + def _build_batch(self): + data = list(range(self.seq_length)) + input_ids = torch.tensor(data, dtype=torch.int64).repeat((self.micro_batch_size, 1)).cuda() + labels = 1 + torch.tensor(data, dtype=torch.int64).repeat((self.micro_batch_size, 1)).cuda() + position_ids = ( + torch.tensor(data, dtype=torch.int64).repeat((self.micro_batch_size, 1)).cuda() + ) + attention_mask = torch.ones( + (self.micro_batch_size, 1, self.seq_length, self.seq_length), dtype=bool + ).cuda() + loss_mask = torch.ones(self.seq_length).repeat((self.micro_batch_size, 1)).cuda() + return input_ids, labels, position_ids, attention_mask, loss_mask + + def _run_steps(self, fp8_param_gather, num_steps): + """Build model + LayerWise optimizer, run ``num_steps`` deterministic + training steps, and return per-step snapshots of loss, forward output, + per-parameter ``main_grad`` (before step), and per-parameter fp32 + master (after step).""" + args = self._create_args(fp8_param_gather=fp8_param_gather) + set_args(args) + torch.manual_seed(_SEED) + Utils.initialize_model_parallel(tensor_model_parallel_size=1, expert_model_parallel_size=1) + + input_ids, labels, position_ids, attention_mask, loss_mask = self._build_batch() + gpt_model, optimizer, _ = setup_model_and_optimizer( + self.model_provider, ModelType.encoder_or_decoder + ) + assert len(gpt_model) == 1 + # Muon + use_distributed_optimizer must auto-promote to LayerWise. + assert isinstance(optimizer.chained_optimizers[0], LayerWiseDistributedOptimizer), ( + "muon + use_distributed_optimizer should route to " + "LayerWiseDistributedOptimizer; got " + f"{type(optimizer.chained_optimizers[0]).__name__}" + ) + + losses, outputs, grads_per_step, masters_per_step = [], [], [], [] + + for _ in range(num_steps): + gpt_model[0].zero_grad_buffer() + optimizer.zero_grad() + gpt_model[0].set_is_first_microbatch() + output = gpt_model[0].forward( + input_ids=input_ids, + position_ids=position_ids, + attention_mask=attention_mask, + labels=labels, + loss_mask=loss_mask, + ) + loss = output.mean() + loss.backward() + if args.overlap_grad_reduce: + gpt_model[0].finish_grad_sync() + + # Snapshot main_grad before the optimizer step zeros / overwrites it. + grad_snapshot = { + name: p.main_grad.detach().clone() + for name, p in gpt_model[0].named_parameters() + if p.main_grad is not None + } + + update_successful, _, _ = optimizer.step() + assert update_successful + + masters_per_step.append(_snapshot_masters(optimizer)) + grads_per_step.append(grad_snapshot) + losses.append(loss.detach().clone()) + outputs.append(output.detach().clone()) + + return losses, outputs, grads_per_step, masters_per_step + + @pytest.mark.skipif( + get_device_arch_version() < 10, reason="MXFP8 requires Blackwell architecture or newer" + ) + @pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8) + @pytest.mark.skipif( + not is_te_min_version("2.3.0.dev0"), reason="TE 2.3.0.dev0 is required for MXFP8" + ) + def test_on_vs_off_bitwise_identical(self): + """fp8_param_gather=ON must produce bitwise-identical loss, forward + output, per-parameter gradient, and per-parameter fp32 master vs + fp8_param_gather=OFF for muon + mxfp8 over multiple training steps.""" + num_steps = 5 + + with deterministic_mode(): + losses_off, outputs_off, grads_off, masters_off = self._run_steps( + fp8_param_gather=False, num_steps=num_steps + ) + losses_on, outputs_on, grads_on, masters_on = self._run_steps( + fp8_param_gather=True, num_steps=num_steps + ) + + assert len(losses_on) == len(losses_off) == num_steps + + for step in range(num_steps): + torch.testing.assert_close( + losses_on[step], + losses_off[step], + atol=0, + rtol=0, + msg=lambda m, s=step: f"loss mismatch at step {s}: {m}", + ) + torch.testing.assert_close( + outputs_on[step], + outputs_off[step], + atol=0, + rtol=0, + msg=lambda m, s=step: f"output mismatch at step {s}: {m}", + ) + + assert set(grads_on[step].keys()) == set(grads_off[step].keys()), ( + f"grad parameter set mismatch at step {step}: " + f"on={sorted(grads_on[step].keys())} " + f"off={sorted(grads_off[step].keys())}" + ) + for name in grads_on[step]: + torch.testing.assert_close( + grads_on[step][name], + grads_off[step][name], + atol=0, + rtol=0, + msg=lambda m, s=step, n=name: f"grad mismatch at step {s} for {n}: {m}", + ) + + assert set(masters_on[step].keys()) == set(masters_off[step].keys()), ( + f"master parameter set mismatch at step {step}: " + f"on={sorted(masters_on[step].keys())} " + f"off={sorted(masters_off[step].keys())}" + ) + for name in masters_on[step]: + torch.testing.assert_close( + masters_on[step][name], + masters_off[step][name], + atol=0, + rtol=0, + msg=lambda m, s=step, n=name: f"master mismatch at step {s} for {n}: {m}", + ) From d1ae16a0ac669466e0e01d3d265f908c1ef93fa0 Mon Sep 17 00:00:00 2001 From: pingtianl Date: Wed, 27 May 2026 01:29:04 -0700 Subject: [PATCH 03/12] Fix deferred DistOpt param sync for LayerWise MXFP8 Ensure sibling DistributedOptimizer buckets inside mixed bucket groups are synchronized when LayerWise Muon uses MXFP8 param gather with grad-buffer reuse. --- megatron/core/optimizer/distrib_optimizer.py | 25 +- megatron/core/optimizer/optimizer.py | 44 ++- .../test_muon_mxfp8_fp8_param_gather.py | 313 ++++++++++++------ 3 files changed, 277 insertions(+), 105 deletions(-) diff --git a/megatron/core/optimizer/distrib_optimizer.py b/megatron/core/optimizer/distrib_optimizer.py index b388161a610..d487477837e 100644 --- a/megatron/core/optimizer/distrib_optimizer.py +++ b/megatron/core/optimizer/distrib_optimizer.py @@ -3003,11 +3003,28 @@ def start_param_sync_for_bucket_group_subset(self) -> None: ): if not bucket_group.buckets: continue - if _bucket_is_managed_by_layer_wise_optimizer( - bucket_group.buckets[0], default_for_untagged=False - ): + distopt_buckets = [ + bucket + for bucket in bucket_group.buckets + if not _bucket_is_managed_by_layer_wise_optimizer( + bucket, default_for_untagged=False + ) + ] + if not distopt_buckets: continue - model_chunk._start_bucket_group_param_sync(bucket_group, force_sync=False) + + if len(distopt_buckets) == len(bucket_group.buckets): + model_chunk._start_bucket_group_param_sync(bucket_group, force_sync=False) + else: + distopt_bucket_group = type(a ra m)( + distopt_buckets, + bucket_group.ddp_config, + bucket_group.intra_distributed_optimizer_instance_group, + bucket_group.intra_distributed_optimizer_instance_size, + ) + model_chunk._start_bucket_group_param_sync( + distopt_bucket_group, force_sync=False + ) @torch.no_grad() def step_with_ready_grads(self) -> bool: diff --git a/megatron/core/optimizer/optimizer.py b/megatron/core/optimizer/optimizer.py index ddc3dd8620e..a1d40b973b9 100644 --- a/megatron/core/optimizer/optimizer.py +++ b/megatron/core/optimizer/optimizer.py @@ -1311,7 +1311,13 @@ def _enable_deferred_mxfp8_param_sync(self) -> List[Tuple[Any, Any]]: deferred_bucket_groups = [] deferred_bucket_group_ids = set() - for optimizer in self.chained_optimizers: + def iter_distributed_optimizers(optimizer): + if isinstance(optimizer, DistributedOptimizer): + yield optimizer + for child in getattr(optimizer, 'chained_optimizers', []): + yield from iter_distributed_optimizers(child) + + for optimizer in iter_distributed_optimizers(self): if not isinstance(optimizer, DistributedOptimizer): continue @@ -1322,23 +1328,47 @@ def _enable_deferred_mxfp8_param_sync(self) -> List[Tuple[Any, Any]]: ): if not bucket_group.buckets: continue - if _bucket_is_managed_by_layer_wise_optimizer( - bucket_group.buckets[0], default_for_untagged=False - ): - continue bucket_group_id = id(bucket_group) if bucket_group_id in deferred_bucket_group_ids: continue + distopt_buckets = [ + bucket + for bucket in bucket_group.buckets + if not _bucket_is_managed_by_layer_wise_optimizer( + bucket, default_for_untagged=False + ) + ] + if not distopt_buckets: + continue + deferred_bucket_group_ids.add(bucket_group_id) - deferred_bucket_groups.append((model_chunk, bucket_group)) + if len(distopt_buckets) == len(bucket_group.buckets): + deferred_bucket_groups.append((model_chunk, bucket_group)) + else: + deferred_bucket_groups.append( + ( + model_chunk, + type(bucket_group)( + distopt_buckets, + bucket_group.ddp_config, + bucket_group.intra_distributed_optimizer_instance_group, + bucket_group.intra_distributed_optimizer_instance_size, + ), + ) + ) return deferred_bucket_groups def _disable_deferred_mxfp8_param_sync(self) -> None: """Disable deferred DistOpt param sync.""" - for optimizer in self.chained_optimizers: + def iter_optimizers(optimizer): + yield optimizer + for child in getattr(optimizer, 'chained_optimizers', []): + yield from iter_optimizers(child) + + for optimizer in iter_optimizers(self): if hasattr(optimizer, '_defer_param_sync'): optimizer._defer_param_sync = False diff --git a/tests/unit_tests/test_muon_mxfp8_fp8_param_gather.py b/tests/unit_tests/test_muon_mxfp8_fp8_param_gather.py index 58413bd9af9..0d59b195ece 100644 --- a/tests/unit_tests/test_muon_mxfp8_fp8_param_gather.py +++ b/tests/unit_tests/test_muon_mxfp8_fp8_param_gather.py @@ -22,20 +22,20 @@ import gc import os import sys -from contextlib import contextmanager +import copy import pytest import torch from transformer_engine.pytorch.fp8 import check_fp8_support -from megatron.core import config from megatron.core.enums import ModelType from megatron.core.models.gpt.gpt_layer_specs import get_gpt_layer_with_transformer_engine_spec from megatron.core.models.gpt.gpt_model import GPTModel from megatron.core.num_microbatches_calculator import destroy_num_microbatches_calculator from megatron.core.optimizer.layer_wise_optimizer import LayerWiseDistributedOptimizer from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed -from megatron.core.utils import get_device_arch_version, is_te_min_version +from megatron.core.utils import is_te_min_version +from megatron.training.utils import get_device_arch_version from megatron.training.arguments import core_transformer_config_from_args, parse_args, validate_args from megatron.training.global_vars import ( destroy_global_vars, @@ -43,75 +43,157 @@ set_args, set_global_variables, ) -from megatron.training.initialize import _set_random_seed from megatron.training.training import setup_model_and_optimizer +from tests.unit_tests.a2a_overlap.utils import deterministic_mode from tests.unit_tests.test_utilities import Utils _SEED = 1234 fp8_available, reason_for_no_fp8 = check_fp8_support() -@contextmanager -def deterministic_mode(): - """Enable deterministic CUDA/CUBLAS/NCCL/TE kernels for bitwise comparison. +def _clone_optimizer_state_value(value): + if torch.is_tensor(value): + return value.detach().clone() + if isinstance(value, dict): + return {k: _clone_optimizer_state_value(v) for k, v in value.items()} + if isinstance(value, list): + return [_clone_optimizer_state_value(v) for v in value] + if isinstance(value, tuple): + return tuple(_clone_optimizer_state_value(v) for v in value) + return copy.deepcopy(value) - Mirrors ``tests/unit_tests/a2a_overlap/utils.py::deterministic_mode`` — - same env-var sweep + ``_set_random_seed`` invocation. Restores prior env - on exit. + +def _iter_leaf_optimizers(optimizer): + for child in getattr(optimizer, 'chained_optimizers', []): + yield from _iter_leaf_optimizers(child) + if not hasattr(optimizer, 'chained_optimizers'): + yield optimizer + + +def _snapshot_masters(model): + """Collect fp32 masters keyed by model parameter name. + + Optimizer-internal group layouts can legitimately differ between + fp8_param_gather ON / OFF, especially for MXFP8 tensors that may be split + into internal quantized members. The model parameter name is the stable + identity we want to compare. """ - config.ENABLE_EXPERIMENTAL = True - envs = { - "CUBLAS_WORKSPACE_CONFIG": ":4096:8", - "TORCH_NCCL_AVOID_RECORD_STREAMS": "1", - "NVTE_ALLOW_NONDETERMINISTIC_ALGO": "0", - "PYTORCH_CUDA_ALLOC_CONF": "expandable_segments:True", - "NCCL_NVLS_ENABLE": "0", - "NVTE_FUSED_ATTN": "0", - "NCCL_ALGO": "^NVLS", - "NVTE_FWD_LAYERNORM_SM_MARGIN": "8", - "NVTE_BWD_LAYERNORM_SM_MARGIN": "8", - } - origin_envs = {} - for k, v in envs.items(): - origin_envs[k] = os.environ.get(k) - os.environ[k] = v - _set_random_seed(seed_=_SEED, data_parallel_random_init=False) - try: - yield - finally: - for k in envs: - if origin_envs[k] is not None: - os.environ[k] = origin_envs[k] - elif k in os.environ: - del os.environ[k] - - -def _snapshot_masters(optimizer): - """Collect fp32 masters from every Float16-wrapped inner optimizer in the - chain, keyed by stable (chain-position, group, index) tuple so the same - parameter lands under the same key across ON / OFF runs.""" snapshot = {} - for outer_idx, child in enumerate(optimizer.chained_optimizers): - if isinstance(child, LayerWiseDistributedOptimizer): - for inner_idx, inner in enumerate(child.chained_optimizers): - main_groups = getattr(inner, 'fp32_from_float16_groups', None) - if main_groups is None: + for name, param in model.named_parameters(): + main_param = getattr(param, 'main_param', None) + if main_param is None: + continue + snapshot[name] = main_param.detach().clone() + return snapshot + + +def _snapshot_model_params(model): + return { + name: param.detach().clone() + for name, param in model.named_parameters() + if not _is_quantized_param(param) + } + + +def _is_quantized_param(param): + return hasattr(param, 'dequantize') or hasattr(param.data, 'dequantize') + + +def _assert_tensor_equal(actual, expected, message): + if torch.equal(actual, expected): + return + diff = (actual.float() - expected.float()).abs() + max_index = int(diff.argmax().item()) if diff.numel() > 0 else 0 + actual_flat = actual.detach().reshape(-1) + expected_flat = expected.detach().reshape(-1) + raise AssertionError( + f"{message}: max_diff={diff.max().item()}, " + f"max_index={max_index}, " + f"actual_at_max={actual_flat[max_index].item() if actual_flat.numel() else None}, " + f"expected_at_max={expected_flat[max_index].item() if expected_flat.numel() else None}, " + f"actual_dtype={actual.dtype}, expected_dtype={expected.dtype}, " + f"actual={actual}, expected={expected}" + ) + + +def _snapshot_optimizer_states(model, optimizer): + param_to_name = {param: name for name, param in model.named_parameters()} + states = {} + for leaf_optimizer in _iter_leaf_optimizers(optimizer): + torch_optimizer = getattr(leaf_optimizer, 'optimizer', None) + if torch_optimizer is None: + continue + for name, param in model.named_parameters(): + main_param = getattr(param, 'main_param', None) + if main_param is None or main_param not in torch_optimizer.state: + continue + states[name] = _clone_optimizer_state_value(torch_optimizer.state[main_param]) + + for model_group, main_group in zip( + getattr(leaf_optimizer, 'model_float16_groups', []), + getattr(leaf_optimizer, 'shard_fp32_from_float16_groups', []), + ): + for model_param, main_param in zip(model_group, main_group): + name = param_to_name.get(model_param) + if name is None or main_param is None or main_param not in torch_optimizer.state: continue - for grp_idx, group in enumerate(main_groups): - for p_idx, p in enumerate(group): - key = f"lw[{outer_idx}][{inner_idx}][{grp_idx}][{p_idx}]" - snapshot[key] = p.detach().clone() + states[name] = _clone_optimizer_state_value(torch_optimizer.state[main_param]) + return states + + +def _snapshot_initial_state(model, optimizer): + return { + 'model_params': { + name: param.detach().clone() for name, param in model.named_parameters() + }, + 'masters': _snapshot_masters(model), + 'optimizer_states': _snapshot_optimizer_states(model, optimizer), + } + + +def _restore_optimizer_states(model, optimizer, optimizer_states): + param_to_name = {param: name for name, param in model.named_parameters()} + for leaf_optimizer in _iter_leaf_optimizers(optimizer): + torch_optimizer = getattr(leaf_optimizer, 'optimizer', None) + if torch_optimizer is None: continue - main_groups = getattr(child, 'shard_fp32_from_float16_groups', None) - if main_groups is None: - main_groups = getattr(child, 'fp32_from_float16_groups', None) - if main_groups is None: + for name, param in model.named_parameters(): + main_param = getattr(param, 'main_param', None) + if main_param is None or name not in optimizer_states: + continue + torch_optimizer.state[main_param] = _clone_optimizer_state_value(optimizer_states[name]) + + for model_group, main_group in zip( + getattr(leaf_optimizer, 'model_float16_groups', []), + getattr(leaf_optimizer, 'shard_fp32_from_float16_groups', []), + ): + for model_param, main_param in zip(model_group, main_group): + name = param_to_name.get(model_param) + if name is None or main_param is None or name not in optimizer_states: + continue + torch_optimizer.state[main_param] = _clone_optimizer_state_value( + optimizer_states[name] + ) + + +@torch.no_grad() +def _restore_initial_state(model, optimizer, initial_state): + source_model_params = initial_state['model_params'] + source_masters = initial_state['masters'] + + for name, param in model.named_parameters(): + if name not in source_model_params: continue - for grp_idx, group in enumerate(main_groups): - for p_idx, p in enumerate(group): - key = f"do[{outer_idx}][{grp_idx}][{p_idx}]" - snapshot[key] = p.detach().clone() - return snapshot + param.data.copy_(source_model_params[name].to(device=param.device)) + + optimizer.reload_model_params() + + for name, param in model.named_parameters(): + main_param = getattr(param, 'main_param', None) + if main_param is not None and name in source_masters: + main_param.data.copy_(source_masters[name].to(device=main_param.device)) + + _restore_optimizer_states(model, optimizer, initial_state['optimizer_states']) class TestMuonMXFP8FP8ParamGather: @@ -121,6 +203,11 @@ def setup_method(self, method): self.seq_length = 128 self.micro_batch_size = 1 os.environ['CUDA_DEVICE_MAX_CONNECTIONS'] = '1' + Utils.initialize_model_parallel( + tensor_model_parallel_size=1, + pipeline_model_parallel_size=1, + expert_model_parallel_size=1, + ) def teardown_method(self, method): Utils.destroy_model_parallel() @@ -128,7 +215,7 @@ def teardown_method(self, method): destroy_num_microbatches_calculator() gc.collect() - def model_provider(self, pre_process=True, post_process=True): + def model_provider(self, pre_process=True, post_process=True, **config_kwargs): model_parallel_cuda_manual_seed(_SEED) args = get_args() transformer_config = core_transformer_config_from_args(args) @@ -145,6 +232,8 @@ def model_provider(self, pre_process=True, post_process=True): share_embeddings_and_output_weights=not args.untie_embeddings_and_output_weights, position_embedding_type=args.position_embedding_type, rotary_percent=args.rotary_percent, + pg_collection=config_kwargs.get("pg_collection"), + vp_stage=config_kwargs.get("vp_stage"), ) def _create_args(self, fp8_param_gather): @@ -154,7 +243,8 @@ def _create_args(self, fp8_param_gather): args = parse_args() args.num_layers = 2 args.vocal_size = 128 - args.hidden_size = 64 + args.hidden_size = 128 + args.ffn_hidden_size = 256 args.num_attention_heads = 4 args.max_position_embeddings = self.seq_length args.micro_batch_size = self.micro_batch_size @@ -166,9 +256,13 @@ def _create_args(self, fp8_param_gather): args.context_parallel_size = 1 args.train_iters = 10 args.lr = 3e-5 + args.clip_grad = 0.0 args.bf16 = True args.add_bias_linear = False args.swiglu = True + args.hidden_dropout = 0.0 + args.attention_dropout = 0.0 + args.attention_backend = "unfused" # ``--optimizer muon --use-distributed-optimizer`` auto-flips to the # LayerWiseDistributedOptimizer path in ``arguments.py``. args.optimizer = 'muon' @@ -210,17 +304,11 @@ def _build_batch(self): loss_mask = torch.ones(self.seq_length).repeat((self.micro_batch_size, 1)).cuda() return input_ids, labels, position_ids, attention_mask, loss_mask - def _run_steps(self, fp8_param_gather, num_steps): - """Build model + LayerWise optimizer, run ``num_steps`` deterministic - training steps, and return per-step snapshots of loss, forward output, - per-parameter ``main_grad`` (before step), and per-parameter fp32 - master (after step).""" + def _build_model_and_optimizer(self, fp8_param_gather): args = self._create_args(fp8_param_gather=fp8_param_gather) set_args(args) torch.manual_seed(_SEED) - Utils.initialize_model_parallel(tensor_model_parallel_size=1, expert_model_parallel_size=1) - input_ids, labels, position_ids, attention_mask, loss_mask = self._build_batch() gpt_model, optimizer, _ = setup_model_and_optimizer( self.model_provider, ModelType.encoder_or_decoder ) @@ -231,8 +319,15 @@ def _run_steps(self, fp8_param_gather, num_steps): "LayerWiseDistributedOptimizer; got " f"{type(optimizer.chained_optimizers[0]).__name__}" ) + return args, gpt_model, optimizer - losses, outputs, grads_per_step, masters_per_step = [], [], [], [] + def _run_steps(self, args, gpt_model, optimizer, num_steps): + """Run ``num_steps`` deterministic training steps and return per-step + snapshots of loss, forward output, per-parameter ``main_grad`` (before + step), and per-parameter fp32 master (after step).""" + input_ids, labels, position_ids, attention_mask, loss_mask = self._build_batch() + + losses, outputs, grads_per_step, masters_per_step, params_per_step = [], [], [], [], [] for _ in range(num_steps): gpt_model[0].zero_grad_buffer() @@ -260,12 +355,13 @@ def _run_steps(self, fp8_param_gather, num_steps): update_successful, _, _ = optimizer.step() assert update_successful - masters_per_step.append(_snapshot_masters(optimizer)) + params_per_step.append(_snapshot_model_params(gpt_model[0])) + masters_per_step.append(_snapshot_masters(gpt_model[0])) grads_per_step.append(grad_snapshot) losses.append(loss.detach().clone()) outputs.append(output.detach().clone()) - return losses, outputs, grads_per_step, masters_per_step + return losses, outputs, grads_per_step, masters_per_step, params_per_step @pytest.mark.skipif( get_device_arch_version() < 10, reason="MXFP8 requires Blackwell architecture or newer" @@ -281,29 +377,49 @@ def test_on_vs_off_bitwise_identical(self): num_steps = 5 with deterministic_mode(): - losses_off, outputs_off, grads_off, masters_off = self._run_steps( - fp8_param_gather=False, num_steps=num_steps + off_args, off_model, off_optimizer = self._build_model_and_optimizer( + fp8_param_gather=False ) - losses_on, outputs_on, grads_on, masters_on = self._run_steps( - fp8_param_gather=True, num_steps=num_steps + initial_state = _snapshot_initial_state(off_model[0], off_optimizer) + on_args, on_model, on_optimizer = self._build_model_and_optimizer( + fp8_param_gather=True ) + _restore_initial_state(on_model[0], on_optimizer, initial_state) + + losses_off, outputs_off, grads_off, masters_off, params_off = [], [], [], [], [] + losses_on, outputs_on, grads_on, masters_on, params_on = [], [], [], [], [] + + for _ in range(num_steps): + off_step = self._run_steps(off_args, off_model, off_optimizer, 1) + on_step = self._run_steps(on_args, on_model, on_optimizer, 1) + + for dst, src in zip( + (losses_off, outputs_off, grads_off, masters_off, params_off), + off_step, + ): + dst.extend(src) + for dst, src in zip( + (losses_on, outputs_on, grads_on, masters_on, params_on), + on_step, + ): + dst.extend(src) + + del off_model, on_model, off_optimizer, on_optimizer + gc.collect() + torch.cuda.empty_cache() assert len(losses_on) == len(losses_off) == num_steps for step in range(num_steps): - torch.testing.assert_close( + _assert_tensor_equal( losses_on[step], losses_off[step], - atol=0, - rtol=0, - msg=lambda m, s=step: f"loss mismatch at step {s}: {m}", + f"loss mismatch at step {step}", ) - torch.testing.assert_close( + _assert_tensor_equal( outputs_on[step], outputs_off[step], - atol=0, - rtol=0, - msg=lambda m, s=step: f"output mismatch at step {s}: {m}", + f"output mismatch at step {step}", ) assert set(grads_on[step].keys()) == set(grads_off[step].keys()), ( @@ -312,24 +428,33 @@ def test_on_vs_off_bitwise_identical(self): f"off={sorted(grads_off[step].keys())}" ) for name in grads_on[step]: - torch.testing.assert_close( + _assert_tensor_equal( grads_on[step][name], grads_off[step][name], - atol=0, - rtol=0, - msg=lambda m, s=step, n=name: f"grad mismatch at step {s} for {n}: {m}", + f"grad mismatch at step {step} for {name}", + ) + + assert set(params_on[step].keys()) == set(params_off[step].keys()), ( + f"model parameter set mismatch at step {step}: " + f"on={sorted(params_on[step].keys())} " + f"off={sorted(params_off[step].keys())}" + ) + for name in params_on[step]: + _assert_tensor_equal( + params_on[step][name], + params_off[step][name], + f"model parameter mismatch after step {step} for {name}", ) - assert set(masters_on[step].keys()) == set(masters_off[step].keys()), ( - f"master parameter set mismatch at step {step}: " + common_master_names = set(masters_on[step].keys()) & set(masters_off[step].keys()) + assert common_master_names, ( + f"no common local fp32 masters to compare at step {step}: " f"on={sorted(masters_on[step].keys())} " f"off={sorted(masters_off[step].keys())}" ) - for name in masters_on[step]: - torch.testing.assert_close( + for name in common_master_names: + _assert_tensor_equal( masters_on[step][name], masters_off[step][name], - atol=0, - rtol=0, - msg=lambda m, s=step, n=name: f"master mismatch at step {s} for {n}: {m}", + f"master mismatch at step {step} for {name}", ) From 20b0e43fcc62edfb236024bb3c2362476ae3b13e Mon Sep 17 00:00:00 2001 From: Pingtian Li Date: Wed, 27 May 2026 16:52:09 +0800 Subject: [PATCH 04/12] fix(optim): block blockwise FP8 + scope-down comments + skip blockwise in test Three follow-ups on the muon LayerWise FP8 param-gather path: 1. Add hard assertion in `arguments.py` that blocks `--fp8-param-gather` with any recipe other than mxfp8 when the optimizer auto-promotes to `LayerWiseDistributedOptimizer`. The bf16-staging + post-AG quantize round-trip is wired up only for MXFP8's `_rowwise_data` / `_columnwise_data` storage; using blockwise FP8 today would silently gather stale storage. Fail fast at arg-parse time. 2. Strip experiment-specific data from new comments in `layer_wise_optimizer.py` (per-iter loss numbers, cluster names, absolute nat lags). Keep the *what* and *why* of each helper so the intent stays clear without hard-coding ephemeral run details. 3. Make the unit test parametrize over `fp8_recipe in [mxfp8, blockwise]` and explicitly `pytest.skip` blockwise. Today the LayerWise path only supports mxfp8; the parametrization documents that surface while the skip keeps blockwise from accidentally running and producing a misleading bitwise mismatch. Also fix the typo `type(a ra m)(` -> `type(bucket_group)(` in `distrib_optimizer.py:3019` (was a paste artifact in the prior commit). --- megatron/core/optimizer/distrib_optimizer.py | 4 +- .../core/optimizer/layer_wise_optimizer.py | 32 +++++----- megatron/core/optimizer/optimizer.py | 3 +- megatron/training/arguments.py | 18 +++++- .../test_muon_mxfp8_fp8_param_gather.py | 58 ++++++++++--------- 5 files changed, 66 insertions(+), 49 deletions(-) diff --git a/megatron/core/optimizer/distrib_optimizer.py b/megatron/core/optimizer/distrib_optimizer.py index d487477837e..bf6ac015e4b 100644 --- a/megatron/core/optimizer/distrib_optimizer.py +++ b/megatron/core/optimizer/distrib_optimizer.py @@ -1,4 +1,4 @@ -# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved. +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. """Megatron distributed optimizer.""" @@ -3016,7 +3016,7 @@ def start_param_sync_for_bucket_group_subset(self) -> None: if len(distopt_buckets) == len(bucket_group.buckets): model_chunk._start_bucket_group_param_sync(bucket_group, force_sync=False) else: - distopt_bucket_group = type(a ra m)( + distopt_bucket_group = type(bucket_group)( distopt_buckets, bucket_group.ddp_config, bucket_group.intra_distributed_optimizer_instance_group, diff --git a/megatron/core/optimizer/layer_wise_optimizer.py b/megatron/core/optimizer/layer_wise_optimizer.py index 13ac525ac4f..f29f46040d2 100644 --- a/megatron/core/optimizer/layer_wise_optimizer.py +++ b/megatron/core/optimizer/layer_wise_optimizer.py @@ -47,10 +47,10 @@ def _skip_mxfp8_in_copy_main_to_model(inner_opt: 'Float16OptimizerWithFloat16Par ``inner_config.reuse_grad_buf_for_mxfp8_param_ag=False``) would do ``model.data.copy_(main.data fp32)`` for each MXFP8 model param, triggering ``QuantizedTensor.__torch_dispatch__`` ⇒ ``dst.quantize_(fp32)`` — a wasted - second quantization that the post-AG ``quantize_(bf16)`` later overwrites, - but which empirically perturbed muon convergence (~0.27 nats lag by iter 10 - on the 100-iter OCI-HSG test). Standard ``DistributedOptimizer`` avoids - this by branching to ``_copy_main_params_to_param_buffer`` in its + second quantization that the post-AG ``quantize_(bf16)`` later overwrites + — and which perturbs muon convergence relative to the + ``fp8_param_gather=False`` baseline. Standard ``DistributedOptimizer`` + avoids this by branching to ``_copy_main_params_to_param_buffer`` in its ``reuse_grad_buf`` step path and never invoking the model-copy at all on MXFP8 params. This patch matches that behavior for the LayerWise inner. @@ -509,12 +509,12 @@ def __init__( # model params (which would otherwise trigger an extra # ``QuantizedTensor.__torch_dispatch__`` ⇒ ``dst.quantize_(fp32)`` # that then gets overwritten by the post-AG ``quantize_(bf16)``). - # The wasted double-quantization left a small mismatch between - # ON and OFF on muon experiments (~0.27 nats by iter 10 that - # shrank but didn't vanish by iter 100); distopt's flow avoids - # this by routing through ``_copy_main_params_to_param_buffer`` - # in the ``reuse_grad_buf`` branch and never touching MXFP8 - # storage at the inner step. Replicate that exactly here. + # The wasted double-quantization perturbs muon convergence + # vs the ``fp8_param_gather=False`` baseline. distopt's flow + # avoids this by routing through + # ``_copy_main_params_to_param_buffer`` in the + # ``reuse_grad_buf`` branch and never touching MXFP8 storage + # at the inner step. Replicate that exactly here. if config.reuse_grad_buf_for_mxfp8_param_ag: _skip_mxfp8_in_copy_main_to_model(inner_opt) # Mirror DistributedOptimizer's master-init fix-up: when the @@ -526,11 +526,8 @@ def __init__( # CPU via ``_high_precision_init_val``; DistOpt reads them back # at master construction (see ``distrib_optimizer.py`` lines # 405-416). Without this, ON masters disagree with OFF masters - # by ~FP8 precision, the subsequent bf16⇒MXFP8 round-trip - # amplifies the residual, and muon's NS step produces slightly - # different updates each iter — observed as a small loss lag - # vs the fp8_param_gather=False baseline that doesn't fully - # close even by iter 100. + # by ~FP8 precision and muon's NS step diverges from the + # ``fp8_param_gather=False`` baseline. _restore_high_precision_init_val(inner_opt) optimizers[i] = inner_opt @@ -873,9 +870,8 @@ def _ensure_mxfp8_quantizer_dual_usage(self, bucket) -> None: refresh deterministic for muon's LayerWise path regardless of what TE toggled during the previous fwd/bwd. Without this, a stale ``columnwise_data`` survives across iterations and the next bwd's - dgrad GEMM uses out-of-date weights, producing a slow loss-curve - divergence vs the ``fp8_param_gather=False`` baseline (observed - ~0.27 nats lag by iter 10 that doesn't fully close by iter 100). + dgrad GEMM uses out-of-date weights, producing a loss-curve + divergence vs the ``fp8_param_gather=False`` baseline. """ for param in bucket.params: if not is_mxfp8tensor(param): diff --git a/megatron/core/optimizer/optimizer.py b/megatron/core/optimizer/optimizer.py index a1d40b973b9..8c123463363 100644 --- a/megatron/core/optimizer/optimizer.py +++ b/megatron/core/optimizer/optimizer.py @@ -1,4 +1,4 @@ -# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved. +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. """Megatron optimizer.""" @@ -1363,6 +1363,7 @@ def iter_distributed_optimizers(optimizer): def _disable_deferred_mxfp8_param_sync(self) -> None: """Disable deferred DistOpt param sync.""" + def iter_optimizers(optimizer): yield optimizer for child in getattr(optimizer, 'chained_optimizers', []): diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py index 25f8f791578..fe8a46b4df9 100644 --- a/megatron/training/arguments.py +++ b/megatron/training/arguments.py @@ -1,4 +1,4 @@ -# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. """Megatron arguments.""" @@ -1102,6 +1102,22 @@ def validate_args(args, defaults={}): 'two features must be co-enabled.' ) + # Only MXFP8 is supported by the LayerWise FP8 param-gather path + # today; the bf16-staging + post-AG quantize round-trip is wired up + # exclusively for MXFP8's ``_rowwise_data`` / ``_columnwise_data`` + # storage. Block other FP8 recipes (e.g. blockwise) explicitly so + # the configuration fails fast instead of silently using stale + # storage at all-gather time. + if will_use_layer_wise_distributed_optimizer and args.fp8_recipe != 'mxfp8': + raise AssertionError( + f'--fp8-param-gather with --fp8-recipe={args.fp8_recipe} is not ' + 'supported by the layer-wise distributed optimizer ' + '(auto-selected for non-adam/sgd optimizers when ' + '--use-distributed-optimizer is set); only --fp8-recipe=mxfp8 ' + 'is supported. Disable --fp8-param-gather or switch to ' + '--fp8-recipe=mxfp8.' + ) + # FP4 and FP8 are mutually exclusive if args.fp4 and args.fp8: raise ValueError("--fp4-format and --fp8-format cannot be used simultaneously. Please choose one.") diff --git a/tests/unit_tests/test_muon_mxfp8_fp8_param_gather.py b/tests/unit_tests/test_muon_mxfp8_fp8_param_gather.py index 0d59b195ece..4223efd26f8 100644 --- a/tests/unit_tests/test_muon_mxfp8_fp8_param_gather.py +++ b/tests/unit_tests/test_muon_mxfp8_fp8_param_gather.py @@ -19,10 +19,10 @@ ``_restore_high_precision_init_val`` in ``layer_wise_optimizer.py``. """ +import copy import gc import os import sys -import copy import pytest import torch @@ -35,7 +35,6 @@ from megatron.core.optimizer.layer_wise_optimizer import LayerWiseDistributedOptimizer from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed from megatron.core.utils import is_te_min_version -from megatron.training.utils import get_device_arch_version from megatron.training.arguments import core_transformer_config_from_args, parse_args, validate_args from megatron.training.global_vars import ( destroy_global_vars, @@ -44,6 +43,7 @@ set_global_variables, ) from megatron.training.training import setup_model_and_optimizer +from megatron.training.utils import get_device_arch_version from tests.unit_tests.a2a_overlap.utils import deterministic_mode from tests.unit_tests.test_utilities import Utils @@ -143,9 +143,7 @@ def _snapshot_optimizer_states(model, optimizer): def _snapshot_initial_state(model, optimizer): return { - 'model_params': { - name: param.detach().clone() for name, param in model.named_parameters() - }, + 'model_params': {name: param.detach().clone() for name, param in model.named_parameters()}, 'masters': _snapshot_masters(model), 'optimizer_states': _snapshot_optimizer_states(model, optimizer), } @@ -236,7 +234,7 @@ def model_provider(self, pre_process=True, post_process=True, **config_kwargs): vp_stage=config_kwargs.get("vp_stage"), ) - def _create_args(self, fp8_param_gather): + def _create_args(self, fp8_param_gather, fp8_recipe="mxfp8"): destroy_global_vars() destroy_num_microbatches_calculator() sys.argv = ['test_muon_mxfp8_fp8_param_gather.py'] @@ -280,11 +278,14 @@ def _create_args(self, fp8_param_gather): # ``_post_param_sync`` in ``param_and_grad_buffer.py``). args.overlap_param_gather = False args.overlap_grad_reduce = False - # MXFP8 + fp8_param_gather config. + # FP8 + fp8_param_gather config. Only ``mxfp8`` is wired through the + # LayerWise bf16-staging + post-AG quantize round-trip; other recipes + # are blocked by ``arguments.py`` and the parametrized test skips + # them explicitly. args.fp8 = "e4m3" - args.fp8_recipe = "mxfp8" + args.fp8_recipe = fp8_recipe args.fp8_param_gather = fp8_param_gather - if fp8_param_gather: + if fp8_param_gather and fp8_recipe == "mxfp8": args.reuse_grad_buf_for_mxfp8_param_ag = True args.ddp_bucket_size = 1024 validate_args(args) @@ -304,8 +305,8 @@ def _build_batch(self): loss_mask = torch.ones(self.seq_length).repeat((self.micro_batch_size, 1)).cuda() return input_ids, labels, position_ids, attention_mask, loss_mask - def _build_model_and_optimizer(self, fp8_param_gather): - args = self._create_args(fp8_param_gather=fp8_param_gather) + def _build_model_and_optimizer(self, fp8_param_gather, fp8_recipe="mxfp8"): + args = self._create_args(fp8_param_gather=fp8_param_gather, fp8_recipe=fp8_recipe) set_args(args) torch.manual_seed(_SEED) @@ -363,6 +364,7 @@ def _run_steps(self, args, gpt_model, optimizer, num_steps): return losses, outputs, grads_per_step, masters_per_step, params_per_step + @pytest.mark.parametrize("fp8_recipe", ["mxfp8", "blockwise"]) @pytest.mark.skipif( get_device_arch_version() < 10, reason="MXFP8 requires Blackwell architecture or newer" ) @@ -370,19 +372,29 @@ def _run_steps(self, args, gpt_model, optimizer, num_steps): @pytest.mark.skipif( not is_te_min_version("2.3.0.dev0"), reason="TE 2.3.0.dev0 is required for MXFP8" ) - def test_on_vs_off_bitwise_identical(self): + def test_on_vs_off_bitwise_identical(self, fp8_recipe): """fp8_param_gather=ON must produce bitwise-identical loss, forward output, per-parameter gradient, and per-parameter fp32 master vs - fp8_param_gather=OFF for muon + mxfp8 over multiple training steps.""" + fp8_param_gather=OFF for muon + mxfp8 over multiple training steps. + + Only mxfp8 is wired up for the LayerWise FP8 param-gather path; + other recipes (e.g. blockwise) are blocked by ``arguments.py`` and + skipped here so the test surface reflects what is supported. + """ + if fp8_recipe != "mxfp8": + pytest.skip( + f"--fp8-recipe={fp8_recipe} is not supported with the LayerWise " + "FP8 param-gather path; only mxfp8 is wired up." + ) num_steps = 5 with deterministic_mode(): off_args, off_model, off_optimizer = self._build_model_and_optimizer( - fp8_param_gather=False + fp8_param_gather=False, fp8_recipe=fp8_recipe ) initial_state = _snapshot_initial_state(off_model[0], off_optimizer) on_args, on_model, on_optimizer = self._build_model_and_optimizer( - fp8_param_gather=True + fp8_param_gather=True, fp8_recipe=fp8_recipe ) _restore_initial_state(on_model[0], on_optimizer, initial_state) @@ -394,13 +406,11 @@ def test_on_vs_off_bitwise_identical(self): on_step = self._run_steps(on_args, on_model, on_optimizer, 1) for dst, src in zip( - (losses_off, outputs_off, grads_off, masters_off, params_off), - off_step, + (losses_off, outputs_off, grads_off, masters_off, params_off), off_step ): dst.extend(src) for dst, src in zip( - (losses_on, outputs_on, grads_on, masters_on, params_on), - on_step, + (losses_on, outputs_on, grads_on, masters_on, params_on), on_step ): dst.extend(src) @@ -411,15 +421,9 @@ def test_on_vs_off_bitwise_identical(self): assert len(losses_on) == len(losses_off) == num_steps for step in range(num_steps): + _assert_tensor_equal(losses_on[step], losses_off[step], f"loss mismatch at step {step}") _assert_tensor_equal( - losses_on[step], - losses_off[step], - f"loss mismatch at step {step}", - ) - _assert_tensor_equal( - outputs_on[step], - outputs_off[step], - f"output mismatch at step {step}", + outputs_on[step], outputs_off[step], f"output mismatch at step {step}" ) assert set(grads_on[step].keys()) == set(grads_off[step].keys()), ( From 919d21a4012d1d12f8b669798c4651c564680a54 Mon Sep 17 00:00:00 2001 From: Pingtian Li Date: Wed, 27 May 2026 16:54:49 +0800 Subject: [PATCH 05/12] test(optim): skip muon mxfp8 fp8_param_gather test when MXFP8 unsupported Add an explicit `check_mxfp8_support()` skipif on top of the existing arch / TE-version / fp8-available guards so the test exits cleanly on environments where MXFP8 is not actually available (e.g. pre-Blackwell GPUs, sm_120+ which TE currently rejects, or builds without MXFP8 kernels), with TE's own diagnostic string as the skip reason. --- tests/unit_tests/test_muon_mxfp8_fp8_param_gather.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/tests/unit_tests/test_muon_mxfp8_fp8_param_gather.py b/tests/unit_tests/test_muon_mxfp8_fp8_param_gather.py index 4223efd26f8..ec59f43ef72 100644 --- a/tests/unit_tests/test_muon_mxfp8_fp8_param_gather.py +++ b/tests/unit_tests/test_muon_mxfp8_fp8_param_gather.py @@ -26,7 +26,7 @@ import pytest import torch -from transformer_engine.pytorch.fp8 import check_fp8_support +from transformer_engine.pytorch.fp8 import check_fp8_support, check_mxfp8_support from megatron.core.enums import ModelType from megatron.core.models.gpt.gpt_layer_specs import get_gpt_layer_with_transformer_engine_spec @@ -49,6 +49,7 @@ _SEED = 1234 fp8_available, reason_for_no_fp8 = check_fp8_support() +mxfp8_available, reason_for_no_mxfp8 = check_mxfp8_support() def _clone_optimizer_state_value(value): @@ -365,6 +366,7 @@ def _run_steps(self, args, gpt_model, optimizer, num_steps): return losses, outputs, grads_per_step, masters_per_step, params_per_step @pytest.mark.parametrize("fp8_recipe", ["mxfp8", "blockwise"]) + @pytest.mark.skipif(not mxfp8_available, reason=reason_for_no_mxfp8) @pytest.mark.skipif( get_device_arch_version() < 10, reason="MXFP8 requires Blackwell architecture or newer" ) From a6610329b27c749dfc4e88e3b022bcf575c8c373 Mon Sep 17 00:00:00 2001 From: Pingtian Li <158665726+Wohox@users.noreply.github.com> Date: Thu, 28 May 2026 13:34:33 +0800 Subject: [PATCH 06/12] Update tests/unit_tests/test_muon_mxfp8_fp8_param_gather.py Co-authored-by: claude[bot] <209825114+claude[bot]@users.noreply.github.com> --- tests/unit_tests/test_muon_mxfp8_fp8_param_gather.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/unit_tests/test_muon_mxfp8_fp8_param_gather.py b/tests/unit_tests/test_muon_mxfp8_fp8_param_gather.py index ec59f43ef72..3ce8a38941c 100644 --- a/tests/unit_tests/test_muon_mxfp8_fp8_param_gather.py +++ b/tests/unit_tests/test_muon_mxfp8_fp8_param_gather.py @@ -222,7 +222,7 @@ def model_provider(self, pre_process=True, post_process=True, **config_kwargs): return GPTModel( config=transformer_config, transformer_layer_spec=layer_spec, - vocab_size=args.vocal_size, + vocab_size=args.vocab_size, max_sequence_length=args.max_position_embeddings, pre_process=pre_process, post_process=post_process, From 1c142814171a905514d07323fa25e250c393c2fa Mon Sep 17 00:00:00 2001 From: Pingtian Li <158665726+Wohox@users.noreply.github.com> Date: Thu, 28 May 2026 13:34:49 +0800 Subject: [PATCH 07/12] Update tests/unit_tests/test_muon_mxfp8_fp8_param_gather.py Co-authored-by: claude[bot] <209825114+claude[bot]@users.noreply.github.com> --- tests/unit_tests/test_muon_mxfp8_fp8_param_gather.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/unit_tests/test_muon_mxfp8_fp8_param_gather.py b/tests/unit_tests/test_muon_mxfp8_fp8_param_gather.py index 3ce8a38941c..fd75c931c68 100644 --- a/tests/unit_tests/test_muon_mxfp8_fp8_param_gather.py +++ b/tests/unit_tests/test_muon_mxfp8_fp8_param_gather.py @@ -241,7 +241,7 @@ def _create_args(self, fp8_param_gather, fp8_recipe="mxfp8"): sys.argv = ['test_muon_mxfp8_fp8_param_gather.py'] args = parse_args() args.num_layers = 2 - args.vocal_size = 128 + args.vocab_size = 128 args.hidden_size = 128 args.ffn_hidden_size = 256 args.num_attention_heads = 4 From 0d7fa16b76ea15cd20f278fa2db27db4cd83cc36 Mon Sep 17 00:00:00 2001 From: Pingtian Li Date: Thu, 28 May 2026 13:52:44 +0800 Subject: [PATCH 08/12] fix(optim): per-bucket LayerWise filter + drop unused local Address review comments on PR #4987: 1. ``LayerWise.start_param_sync_for_bucket_group_subset`` now filters buckets per-bucket instead of checking only ``buckets[0]``, and builds a sub-group containing just the LayerWise-managed buckets when the group is mixed-ownership. ``partition_buckets`` Case 3 (FP8 present, ``reduce_scatter_with_fp32_accumulation=False``) merges non-FP8 DistOpt-managed bf16 buckets (biases, layernorms) into the last FP8 bucket group; without per-bucket filtering, LayerWise dispatched AG for the entire group while the sibling DistOpt also synced those same bf16 buckets via its own per-bucket filter, double-gathering them. Mirrors the per-bucket pattern that ``d1ae16a0a`` applied to the DistOpt deferred-sync path. 2. Drop the unused ``original_get_pairs = inner_opt ._get_model_and_main_params_data_float16`` line in ``_skip_mxfp8_in_copy_main_to_model``. Leftover from development; never referenced. --- .../core/optimizer/layer_wise_optimizer.py | 55 +++++++++++++------ 1 file changed, 38 insertions(+), 17 deletions(-) diff --git a/megatron/core/optimizer/layer_wise_optimizer.py b/megatron/core/optimizer/layer_wise_optimizer.py index f29f46040d2..dc4c8dd196a 100644 --- a/megatron/core/optimizer/layer_wise_optimizer.py +++ b/megatron/core/optimizer/layer_wise_optimizer.py @@ -59,7 +59,6 @@ def _skip_mxfp8_in_copy_main_to_model(inner_opt: 'Float16OptimizerWithFloat16Par to ``param_buffer`` via :func:`modify_underlying_storage` and the inner copy lands the update in the right place. """ - original_get_pairs = inner_opt._get_model_and_main_params_data_float16 # Local import to avoid hoisting the helper to a wider scope. from ..fp8_utils import is_mxfp8tensor from .optimizer import _multi_tensor_copy_this_to_that @@ -882,30 +881,52 @@ def _ensure_mxfp8_quantizer_dual_usage(self, bucket) -> None: quantizer.set_usage(rowwise=True, columnwise=True) def start_param_sync_for_bucket_group_subset(self) -> None: - """Trigger ``start_param_sync`` on LayerWise-managed bucket groups only. + """Trigger ``start_param_sync`` only on LayerWise-managed buckets. Walks each model chunk's dense + expert-parallel bucket groups and - skips any group not managed by LayerWise, so a sibling - :class:`DistributedOptimizer`'s own ``start_param_sync`` call does not - double-sync the same buckets. Uses - :meth:`DistributedDataParallel._start_bucket_group_param_sync` so FP8 - post-all-gather processing (and MXFP8 copy) still runs. - - Before dispatching each LayerWise-managed bucket group's AG, force the - bucket's MXFP8 param quantizers to ``rowwise=True, columnwise=True`` - so the post-AG ``param.data.copy_(bf16)`` refreshes both orientations - (see :py:meth:`_ensure_mxfp8_quantizer_dual_usage`). + filters **per-bucket** so a sibling :class:`DistributedOptimizer`'s own + ``start_param_sync`` call does not double-sync any bucket. Mixed + bucket groups can occur in ``partition_buckets`` Case 3 (FP8 present, + ``reduce_scatter_with_fp32_accumulation=False``), which appends + non-FP8 DistOpt-managed bf16 buckets (biases, layernorms) into the + last FP8 bucket group; checking only ``buckets[0]`` and dispatching + AG on the whole group would AG those DistOpt-managed bf16 buckets a + second time. Mirrors the per-bucket filter applied to + :meth:`DistributedOptimizer.start_param_sync_for_bucket_group_subset`. + Uses :meth:`DistributedDataParallel._start_bucket_group_param_sync` + so FP8 post-all-gather processing (and MXFP8 copy) still runs. + + Before dispatching each LayerWise-managed bucket group's AG, force + the bucket's MXFP8 param quantizers to ``rowwise=True, + columnwise=True`` so the post-AG ``param.data.copy_(bf16)`` refreshes + both orientations (see + :py:meth:`_ensure_mxfp8_quantizer_dual_usage`). """ for model_chunk in self.model_chunks: for bucket_group in ( model_chunk.bucket_groups + model_chunk.expert_parallel_bucket_groups ): - if bucket_group.buckets and _bucket_is_managed_by_layer_wise_optimizer( - bucket_group.buckets[0] - ): - for bucket in bucket_group.buckets: - self._ensure_mxfp8_quantizer_dual_usage(bucket) + if not bucket_group.buckets: + continue + lw_buckets = [ + bucket + for bucket in bucket_group.buckets + if _bucket_is_managed_by_layer_wise_optimizer(bucket) + ] + if not lw_buckets: + continue + for bucket in lw_buckets: + self._ensure_mxfp8_quantizer_dual_usage(bucket) + if len(lw_buckets) == len(bucket_group.buckets): model_chunk._start_bucket_group_param_sync(bucket_group, force_sync=False) + else: + lw_bucket_group = type(bucket_group)( + lw_buckets, + bucket_group.ddp_config, + bucket_group.intra_distributed_optimizer_instance_group, + bucket_group.intra_distributed_optimizer_instance_size, + ) + model_chunk._start_bucket_group_param_sync(lw_bucket_group, force_sync=False) @torch.no_grad() def _write_owned_mxfp8_masters_to_param_buffer(self) -> None: From 70c5487dc9b12d607f161254f9726abd21654740 Mon Sep 17 00:00:00 2001 From: Pingtian Li Date: Fri, 29 May 2026 10:21:54 +0800 Subject: [PATCH 09/12] doc(optim): explain per-bucket filter + mixed-group sub-group construction MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Address review comment on PR #4987: the per-bucket filter and the ``type(bucket_group)(distopt_buckets, …)`` sub-group construction added in ``DistributedOptimizer.start_param_sync_for_bucket_group_subset`` had no inline justification. Without that context it isn't obvious why we walk one bucket at a time instead of dispatching the whole group, or when the synthesized sub-group path is taken. Extends the docstring to call out the ``partition_buckets`` Case 3 scenario where the last FP8 bucket group holds both DistOpt-managed bf16 buckets (biases / layernorms) and LayerWise-managed FP8 buckets, and adds inline comments on each of the three branches (entire group LayerWise-owned, pure DistOpt group, mixed group) describing what is dispatched and why the synthesized sub-group inherits the parent's ddp_config / DP group / DP world size so the AG collective lands on the same comm. Pure documentation; no behaviour change. --- megatron/core/optimizer/distrib_optimizer.py | 40 +++++++++++++++----- 1 file changed, 31 insertions(+), 9 deletions(-) diff --git a/megatron/core/optimizer/distrib_optimizer.py b/megatron/core/optimizer/distrib_optimizer.py index bf6ac015e4b..2ea3ae89766 100644 --- a/megatron/core/optimizer/distrib_optimizer.py +++ b/megatron/core/optimizer/distrib_optimizer.py @@ -2982,15 +2982,22 @@ def copy_group_params(model_groups, shard_main_groups): copy_group_params(self.model_fp32_groups, self.shard_fp32_groups) def start_param_sync_for_bucket_group_subset(self) -> None: - """Trigger ``start_param_sync`` on DistOpt-managed bucket groups only. - - Walks each model chunk's DDP bucket groups and skips those tagged - ``is_managed_by_layer_wise_optimizer=True`` (so a sibling - :class:`LayerWiseDistributedOptimizer` does not double-sync the same - buckets). When no LayerWise tagging is present every bucket group is - included — matching the previous ``model_chunk.start_param_sync()`` - behaviour. Uses :meth:`DistributedDataParallel._start_bucket_group_param_sync` - so FP8 post-all-gather processing (and MXFP8 copy) still runs. + """Trigger ``start_param_sync`` only on DistOpt-managed buckets. + + Filters DDP bucket groups **per-bucket** so a sibling + :class:`LayerWiseDistributedOptimizer` (e.g. the Muon LayerWise + optimizer) does not double-sync any bucket. When no LayerWise tagging + is present every bucket group is included — matching the previous + ``model_chunk.start_param_sync()`` behaviour. Mixed-ownership bucket + groups (some buckets DistOpt-managed, some LayerWise-managed) can + appear in ``param_and_grad_buffer.partition_buckets`` Case 3 (FP8 + present, ``reduce_scatter_with_fp32_accumulation=False``), which + appends non-FP8 DistOpt-managed bf16 buckets (biases, layernorms) + into the last FP8 bucket group. For those groups we synthesize a + bucket group that contains only the DistOpt buckets so the + LayerWise side's all-gather and ours stay disjoint. Uses + :meth:`DistributedDataParallel._start_bucket_group_param_sync` so + FP8 post-all-gather processing (and MXFP8 copy) still runs. """ # Deferred import: layer_wise_optimizer's compute_full_param_layout # lazily imports DistributedOptimizer, so importing the helper at @@ -3003,6 +3010,9 @@ def start_param_sync_for_bucket_group_subset(self) -> None: ): if not bucket_group.buckets: continue + # Restrict to buckets *not* claimed by LayerWise — those will + # be all-gathered by ``LayerWiseDistributedOptimizer. + # start_param_sync_for_bucket_group_subset`` instead. distopt_buckets = [ bucket for bucket in bucket_group.buckets @@ -3011,11 +3021,23 @@ def start_param_sync_for_bucket_group_subset(self) -> None: ) ] if not distopt_buckets: + # Entire group is LayerWise-owned; LayerWise will sync it. continue if len(distopt_buckets) == len(bucket_group.buckets): + # Pure DistOpt group — dispatch the original bucket + # group as-is so all per-bucket-group state on the DDP + # side (handles, gather lists, etc.) is reused. model_chunk._start_bucket_group_param_sync(bucket_group, force_sync=False) else: + # Mixed group (Case 3 above): build a thin transient + # bucket group containing only the DistOpt buckets and + # dispatch that. The synthesized group reuses the + # parent's ddp_config / DP process group / DP world + # size so the AG collective lands on the same comm. + # The LayerWise side will independently sync the + # LayerWise-tagged buckets from the same parent group + # via its mirror of this routine. distopt_bucket_group = type(bucket_group)( distopt_buckets, bucket_group.ddp_config, From 477dc5a40cc3a51f6228dd6e0e6e4019981f0124 Mon Sep 17 00:00:00 2001 From: Pingtian Li Date: Tue, 2 Jun 2026 16:19:38 +0800 Subject: [PATCH 10/12] fix(layer_wise): support fp8_param_gather for muon+mxfp8 in overlap_param_gather mode MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit PR 4987 enabled fp8_param_gather for the LayerWise (muon) optimizer with --fp8-recipe mxfp8 + --reuse-grad-buf-for-mxfp8-param-ag, but only validated overlap_param_gather=False. With overlap_param_gather=True every muon-managed 2D weight froze at init (loss flat ~12.5, grad norm pinned ~205) while Adam (standard DistributedOptimizer) was unaffected. Root cause: reuse_grad_buf_for_mxfp8_param_ag aliases the bf16 param-AG staging buffer onto the grad buffer. In overlap mode the LayerWise param all-gather + post-AG copy into the MXFP8 param.data is deferred to the next forward's pre-hook, but zero_grad_buffer() runs at the start of that iteration and zeroes the shared buffer first. The deferred AG then ships a zeroed staging buffer and the post-AG copy overwrites every muon param.data with zeros. The standard DistributedOptimizer avoids this because train_step re-stages its masters via _copy_main_params_to_param_buffer() AFTER zero_grad_buffer() — but that loop was gated to isinstance(DistributedOptimizer) and skipped the LayerWise optimizer, making the bug muon-only. Fix: give LayerWiseDistributedOptimizer a duck-typed _copy_main_params_to_param_buffer() that re-writes the owned MXFP8 masters into the staging buffer post-zero, and widen the train_step post-zero loop to isinstance(opt, (DistributedOptimizer, LayerWiseDistributedOptimizer)) so both optimizers are handled by a single polymorphic call. Verified on DSv4 Flash Proxy / GB200x16 / TP1 PP2 EP8 / mxfp8 / muon: overlap + fp8_param_gather now descends matching the fp8pg-off and non-overlap baselines, with normal grad norms. The deterministic bitwise ON-vs-OFF unit test (test_muon_mxfp8_fp8_param_gather.py) passes. --- .../core/optimizer/layer_wise_optimizer.py | 94 +++++++------------ megatron/training/training.py | 9 +- 2 files changed, 44 insertions(+), 59 deletions(-) diff --git a/megatron/core/optimizer/layer_wise_optimizer.py b/megatron/core/optimizer/layer_wise_optimizer.py index dc4c8dd196a..a3e45ae50c4 100644 --- a/megatron/core/optimizer/layer_wise_optimizer.py +++ b/megatron/core/optimizer/layer_wise_optimizer.py @@ -39,7 +39,7 @@ def _skip_mxfp8_in_copy_main_to_model(inner_opt: 'Float16OptimizerWithFloat16Par With ``reuse_grad_buf_for_mxfp8_param_ag=True`` at the outer LayerWise level, the bf16 ``param_buffer`` is written by - :py:meth:`LayerWiseDistributedOptimizer._write_owned_mxfp8_masters_to_param_buffer` + :py:meth:`LayerWiseDistributedOptimizer._copy_main_params_to_param_buffer` and the MXFP8 storage is then refreshed by the post-AG ``param.data.copy_(bf16_slice)`` step inside ``_ParamAndGradBucketGroup._post_param_sync``. The inner's default @@ -490,7 +490,7 @@ def __init__( # ``_copy_main_params_to_param_buffer`` (a DistributedOptimizer-only # method) and raise ``AttributeError`` here. We take responsibility # for the bf16 ⇒ param_buffer write for owned MXFP8 model params in - # :py:meth:`_write_owned_mxfp8_masters_to_param_buffer` below. + # :py:meth:`_copy_main_params_to_param_buffer` below. inner_config = copy.copy(config) inner_config.reuse_grad_buf_for_mxfp8_param_ag = False for i in range(len(optimizers)): @@ -542,7 +542,7 @@ def __init__( # restores that equality without changing inner-step behavior (only # the inner Float16Optimizer needs the flag flipped; LayerWise itself # behaves correctly under the outer config and our own writes to the - # param buffer in :py:meth:`_write_owned_mxfp8_masters_to_param_buffer` + # param buffer in :py:meth:`_copy_main_params_to_param_buffer` # are gated on the outer config's ``reuse_grad_buf_for_mxfp8_param_ag``). self.config = config @@ -849,37 +849,6 @@ def count_zeros(self): use_decoupled_grad=self.config.use_precision_aware_optimizer_no_fp8_or_ds_fp8, ) - def _ensure_mxfp8_quantizer_dual_usage(self, bucket) -> None: - """Force every MXFP8 param's bound quantizer to ``rowwise=True, columnwise=True``. - - The post-AG ``param.data.copy_(param_buffer_slice)`` triggered by - :meth:`_ParamAndGradBucketGroup._post_param_sync` routes through - ``QuantizedTensor.__torch_dispatch__`` ⇒ ``dst.quantize_(src)`` ⇒ - ``MXFP8Quantizer.update_quantized`` ⇒ ``tex.quantize(src, quantizer, dst)``. - ``tex.quantize`` picks the kernel based on ``quantizer.rowwise_usage`` / - ``quantizer.columnwise_usage``: both ``True`` selects the x2 kernel that - produces rowwise + columnwise data + scales in a single pass, anything - else updates only the requested orientation. - - TE Linear sets the weight quantizer to ``rowwise=True, columnwise=True`` - at module construction time when ``torch.is_grad_enabled()`` (see - ``base.py:1467``), and bwd briefly toggles it for grad GEMMs, so for - most paths the quantizer happens to be configured correctly when the - optimizer step runs. Re-asserting it here makes the post-AG MXFP8 - refresh deterministic for muon's LayerWise path regardless of what TE - toggled during the previous fwd/bwd. Without this, a stale - ``columnwise_data`` survives across iterations and the next bwd's - dgrad GEMM uses out-of-date weights, producing a loss-curve - divergence vs the ``fp8_param_gather=False`` baseline. - """ - for param in bucket.params: - if not is_mxfp8tensor(param): - continue - quantizer = getattr(param, '_quantizer', None) - if quantizer is None: - continue - quantizer.set_usage(rowwise=True, columnwise=True) - def start_param_sync_for_bucket_group_subset(self) -> None: """Trigger ``start_param_sync`` only on LayerWise-managed buckets. @@ -895,12 +864,6 @@ def start_param_sync_for_bucket_group_subset(self) -> None: :meth:`DistributedOptimizer.start_param_sync_for_bucket_group_subset`. Uses :meth:`DistributedDataParallel._start_bucket_group_param_sync` so FP8 post-all-gather processing (and MXFP8 copy) still runs. - - Before dispatching each LayerWise-managed bucket group's AG, force - the bucket's MXFP8 param quantizers to ``rowwise=True, - columnwise=True`` so the post-AG ``param.data.copy_(bf16)`` refreshes - both orientations (see - :py:meth:`_ensure_mxfp8_quantizer_dual_usage`). """ for model_chunk in self.model_chunks: for bucket_group in ( @@ -915,8 +878,6 @@ def start_param_sync_for_bucket_group_subset(self) -> None: ] if not lw_buckets: continue - for bucket in lw_buckets: - self._ensure_mxfp8_quantizer_dual_usage(bucket) if len(lw_buckets) == len(bucket_group.buckets): model_chunk._start_bucket_group_param_sync(bucket_group, force_sync=False) else: @@ -929,21 +890,30 @@ def start_param_sync_for_bucket_group_subset(self) -> None: model_chunk._start_bucket_group_param_sync(lw_bucket_group, force_sync=False) @torch.no_grad() - def _write_owned_mxfp8_masters_to_param_buffer(self) -> None: - """For each owned MXFP8 model param, write its fp32 master cast to bf16 - directly into the corresponding slice of DDP's bf16 param buffer. - - Mirrors :meth:`DistributedOptimizer._copy_main_params_to_param_buffer` - for the LayerWise + ``use_buffer_param_sync`` case. With - ``reuse_grad_buf_for_mxfp8_param_ag=True`` the bf16 ``param_buffer`` is - kept separate from each MXFP8 model param's ``_rowwise_data`` / + def _copy_main_params_to_param_buffer(self) -> None: + """Write each owned MXFP8 fp32 master, cast to bf16, into DDP's bf16 + param buffer. + + Duck-typed twin of + :meth:`DistributedOptimizer._copy_main_params_to_param_buffer` so the + ``train_step`` post-``zero_grad_buffer`` loop can treat both optimizers + uniformly (``isinstance(opt, (DistributedOptimizer, + LayerWiseDistributedOptimizer))``) and call this method polymorphically. + + With ``reuse_grad_buf_for_mxfp8_param_ag=True`` the bf16 ``param_buffer`` + is kept separate from each MXFP8 model param's ``_rowwise_data`` / ``_columnwise_data`` storage; the inner ``Float16OptimizerWithFloat16Params._copy_main_params_to_model_params`` refreshes the MXFP8 storage via TE's ``QuantizedTensor`` dispatch but leaves ``param_buffer`` stale. ``DistributedDataParallel.start_param_sync`` all-gathers ``param_buffer`` and then quantizes back to MXFP8 in-place on every rank, so the buffer must hold the freshly updated bf16 master cast - before AG dispatches. + before the AG dispatches. Because the staging buffer is aliased onto the + grad buffer, ``zero_grad_buffer()`` zeroes it at the start of every + iteration, so ``train_step`` re-stages here (after the zero, before the + deferred forward-pre-hook AG); the in-step call from + :meth:`step_with_ready_grads` keeps the non-overlap and first-iteration + paths correct. Walks each model chunk's dense + expert-parallel buffers and writes only for params that are (a) owned by this rank on the LayerWise side, (b) @@ -1011,19 +981,27 @@ def step_with_ready_grads(self) -> bool: # forward pre-hooks after step returns, so the buffer is correct by # then; in non-overlap mode the ``start_param_sync_for_bucket_group_subset`` # call below dispatches synchronously and sees the correct buffer. - self._write_owned_mxfp8_masters_to_param_buffer() + self._copy_main_params_to_param_buffer() # All-gather updated params. If overlap_param_gather is True, the all-gather # is deferred to the forward pre-hooks via DDP bucket infrastructure. + # + # NOTE for the MXFP8 + ``reuse_grad_buf_for_mxfp8_param_ag`` + overlap case: + # the ``_copy_main_params_to_param_buffer`` above writes into the + # bf16 staging buffer that is aliased onto the grad buffer, which the next + # iteration's ``zero_grad_buffer()`` zeroes BEFORE the deferred forward + # pre-hook AG runs. The masters are therefore re-staged post-zero from + # ``train_step`` via :meth:`_copy_main_params_to_param_buffer` (the same + # call ``train_step`` makes for the standard ``DistributedOptimizer``); + # the write above is harmless but only authoritative for the first + # iteration (before the forward pre-hook is enabled). if not self.overlap_param_gather: if self.use_buffer_param_sync: # Model params are views into the DDP param buffer - # (ddp_config.use_distributed_optimizer=True). The optimizer step - # already copied updated fp32 main params → bf16 model params (= - # buffer views), so the buffer is up-to-date. Trigger the standard - # buffer all-gather, but only for LayerWise-managed bucket groups - # so a sibling DistributedOptimizer's own ``start_param_sync`` call - # is not duplicated for the same buckets. + # (ddp_config.use_distributed_optimizer=True). Trigger the buffer + # all-gather, but only for LayerWise-managed bucket groups so a + # sibling DistributedOptimizer's own ``start_param_sync`` is not + # duplicated for the same buckets. self.start_param_sync_for_bucket_group_subset() else: self.allgather_params() diff --git a/megatron/training/training.py b/megatron/training/training.py index 3930cc46a21..6321dc22e62 100644 --- a/megatron/training/training.py +++ b/megatron/training/training.py @@ -2214,7 +2214,14 @@ def train_step(forward_step_func, data_iterator, model, optimizer, opt_param_sch full_cg_captured = FullCudaGraphWrapper.cuda_graph.get("training") is not None if forward_pre_hook_enabled or full_cg_captured: for optim_instance in optimizer.chained_optimizers: - if isinstance(optim_instance, DistributedOptimizer): + # Both the standard DistributedOptimizer and the LayerWise + # (muon) optimizer keep their MXFP8 masters in a staging + # buffer aliased onto the grad buffer that zero_grad_buffer() + # just zeroed; re-stage them so the deferred forward pre-hook + # all-gather ships fresh weights instead of the zeroed buffer. + if isinstance( + optim_instance, (DistributedOptimizer, LayerWiseDistributedOptimizer) + ): optim_instance._copy_main_params_to_param_buffer() # Forward pass. From cfc7f3101429569c1fc296f5be3e8a240d155695 Mon Sep 17 00:00:00 2001 From: Pingtian Li Date: Wed, 3 Jun 2026 09:32:29 +0800 Subject: [PATCH 11/12] test(optim): cover overlap_param_gather path in muon mxfp8 fp8_param_gather test Parametrize the bitwise ON-vs-OFF test over overlap_param_gather in {False, True}. The True case is the regression guard for the frozen-loss bug on the deferred forward-pre-hook all-gather path: the bf16 staging buffer is aliased onto the grad buffer that zero_grad_buffer() zeroes each iteration, so without re-staging the masters post-zero the all-gather ships zeros and the muon-managed weights never update. _run_steps now mirrors train_step's post-zero_grad_buffer re-stage (_copy_main_params_to_param_buffer on each DistributedOptimizer / LayerWiseDistributedOptimizer) when reuse_grad_buf_for_mxfp8_param_ag and overlap_param_gather are set; DDP auto-registers the forward pre-hook when overlap_param_gather is True. --overlap-param-gather requires --overlap-grad-reduce, so the two are co-enabled. Co-Authored-By: Claude Opus 4.8 (1M context) --- .../test_muon_mxfp8_fp8_param_gather.py | 63 +++++++++++++++---- 1 file changed, 52 insertions(+), 11 deletions(-) diff --git a/tests/unit_tests/test_muon_mxfp8_fp8_param_gather.py b/tests/unit_tests/test_muon_mxfp8_fp8_param_gather.py index fd75c931c68..be5064e1aff 100644 --- a/tests/unit_tests/test_muon_mxfp8_fp8_param_gather.py +++ b/tests/unit_tests/test_muon_mxfp8_fp8_param_gather.py @@ -32,6 +32,7 @@ from megatron.core.models.gpt.gpt_layer_specs import get_gpt_layer_with_transformer_engine_spec from megatron.core.models.gpt.gpt_model import GPTModel from megatron.core.num_microbatches_calculator import destroy_num_microbatches_calculator +from megatron.core.optimizer.distrib_optimizer import DistributedOptimizer from megatron.core.optimizer.layer_wise_optimizer import LayerWiseDistributedOptimizer from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed from megatron.core.utils import is_te_min_version @@ -235,7 +236,7 @@ def model_provider(self, pre_process=True, post_process=True, **config_kwargs): vp_stage=config_kwargs.get("vp_stage"), ) - def _create_args(self, fp8_param_gather, fp8_recipe="mxfp8"): + def _create_args(self, fp8_param_gather, fp8_recipe="mxfp8", overlap_param_gather=False): destroy_global_vars() destroy_num_microbatches_calculator() sys.argv = ['test_muon_mxfp8_fp8_param_gather.py'] @@ -274,11 +275,16 @@ def _create_args(self, fp8_param_gather, fp8_recipe="mxfp8"): args.exp_avg_dtype = 'fp32' args.exp_avg_sq_dtype = 'fp32' args.use_distributed_optimizer = True - # Disable AG / RS overlap so the timing is deterministic and the test - # exercises the synchronous post-AG quantize path (see - # ``_post_param_sync`` in ``param_and_grad_buffer.py``). - args.overlap_param_gather = False - args.overlap_grad_reduce = False + # ``overlap_param_gather=False`` exercises the synchronous post-AG + # quantize path (``start_param_sync_for_bucket_group_subset`` called from + # ``step_with_ready_grads``); ``True`` exercises the deferred forward + # pre-hook all-gather path (the more important / production path), where + # the bf16 staging buffer is re-staged post-``zero_grad_buffer`` in + # ``_run_steps`` exactly as ``train_step`` does. ``--overlap-param-gather`` + # requires ``--overlap-grad-reduce`` (see ``arguments.py``), so the two + # are co-enabled. + args.overlap_param_gather = overlap_param_gather + args.overlap_grad_reduce = overlap_param_gather # FP8 + fp8_param_gather config. Only ``mxfp8`` is wired through the # LayerWise bf16-staging + post-AG quantize round-trip; other recipes # are blocked by ``arguments.py`` and the parametrized test skips @@ -306,8 +312,14 @@ def _build_batch(self): loss_mask = torch.ones(self.seq_length).repeat((self.micro_batch_size, 1)).cuda() return input_ids, labels, position_ids, attention_mask, loss_mask - def _build_model_and_optimizer(self, fp8_param_gather, fp8_recipe="mxfp8"): - args = self._create_args(fp8_param_gather=fp8_param_gather, fp8_recipe=fp8_recipe) + def _build_model_and_optimizer( + self, fp8_param_gather, fp8_recipe="mxfp8", overlap_param_gather=False + ): + args = self._create_args( + fp8_param_gather=fp8_param_gather, + fp8_recipe=fp8_recipe, + overlap_param_gather=overlap_param_gather, + ) set_args(args) torch.manual_seed(_SEED) @@ -334,6 +346,22 @@ def _run_steps(self, args, gpt_model, optimizer, num_steps): for _ in range(num_steps): gpt_model[0].zero_grad_buffer() optimizer.zero_grad() + + # Mirror ``train_step``: with ``reuse_grad_buf_for_mxfp8_param_ag`` the + # bf16 staging buffer is aliased onto the grad buffer that + # ``zero_grad_buffer()`` just zeroed, so re-stage the masters before + # the deferred forward-pre-hook all-gather (registered by DDP when + # ``overlap_param_gather=True``) ships them — otherwise the AG would + # gather a zeroed buffer and the muon-managed weights would freeze. + # No-op when overlap is off (the in-step ``step_with_ready_grads`` + # write + synchronous AG cover that path). + if args.reuse_grad_buf_for_mxfp8_param_ag and args.overlap_param_gather: + for optim_instance in optimizer.chained_optimizers: + if isinstance( + optim_instance, (DistributedOptimizer, LayerWiseDistributedOptimizer) + ): + optim_instance._copy_main_params_to_param_buffer() + gpt_model[0].set_is_first_microbatch() output = gpt_model[0].forward( input_ids=input_ids, @@ -365,6 +393,7 @@ def _run_steps(self, args, gpt_model, optimizer, num_steps): return losses, outputs, grads_per_step, masters_per_step, params_per_step + @pytest.mark.parametrize("overlap_param_gather", [False, True]) @pytest.mark.parametrize("fp8_recipe", ["mxfp8", "blockwise"]) @pytest.mark.skipif(not mxfp8_available, reason=reason_for_no_mxfp8) @pytest.mark.skipif( @@ -374,11 +403,19 @@ def _run_steps(self, args, gpt_model, optimizer, num_steps): @pytest.mark.skipif( not is_te_min_version("2.3.0.dev0"), reason="TE 2.3.0.dev0 is required for MXFP8" ) - def test_on_vs_off_bitwise_identical(self, fp8_recipe): + def test_on_vs_off_bitwise_identical(self, fp8_recipe, overlap_param_gather): """fp8_param_gather=ON must produce bitwise-identical loss, forward output, per-parameter gradient, and per-parameter fp32 master vs fp8_param_gather=OFF for muon + mxfp8 over multiple training steps. + Run for both ``overlap_param_gather`` settings: ``False`` exercises the + synchronous post-AG quantize path, ``True`` the deferred forward + pre-hook all-gather path. The overlap=True case is the regression guard + for the frozen-loss bug — the bf16 staging buffer is aliased onto the + grad buffer that ``zero_grad_buffer()`` zeroes each iteration, so without + the post-zero re-stage the deferred AG would ship zeros and the + muon-managed weights would never update. + Only mxfp8 is wired up for the LayerWise FP8 param-gather path; other recipes (e.g. blockwise) are blocked by ``arguments.py`` and skipped here so the test surface reflects what is supported. @@ -392,11 +429,15 @@ def test_on_vs_off_bitwise_identical(self, fp8_recipe): with deterministic_mode(): off_args, off_model, off_optimizer = self._build_model_and_optimizer( - fp8_param_gather=False, fp8_recipe=fp8_recipe + fp8_param_gather=False, + fp8_recipe=fp8_recipe, + overlap_param_gather=overlap_param_gather, ) initial_state = _snapshot_initial_state(off_model[0], off_optimizer) on_args, on_model, on_optimizer = self._build_model_and_optimizer( - fp8_param_gather=True, fp8_recipe=fp8_recipe + fp8_param_gather=True, + fp8_recipe=fp8_recipe, + overlap_param_gather=overlap_param_gather, ) _restore_initial_state(on_model[0], on_optimizer, initial_state) From 79dba142014444f94c04d6d46802830bada0e0f3 Mon Sep 17 00:00:00 2001 From: Pingtian Li Date: Wed, 3 Jun 2026 10:09:04 +0800 Subject: [PATCH 12/12] fix(optim): re-stage LayerWise MXFP8 masters on the eval + paged-stash paths The training-loop re-stage (post-zero_grad_buffer _copy_main_params_to_param_buffer) was widened to LayerWiseDistributedOptimizer, but two structurally identical sites were missed: - training.py eval block: with reuse_grad_buf_for_mxfp8_param_ag + overlap_param_gather, the pre-eval zero_grad_buffer() zeroes the staging buffer aliased onto the grad buffer, then only DistributedOptimizer params were re-staged. disable_forward_pre_hook( param_sync=True) would then all-gather zeroed buffers for muon-managed MXFP8 weights during eval. - paged_stash.py _try_copy_main_params: reachable when moe_expert_rank_capacity_factor is set together with reuse_grad_buf_for_mxfp8_param_ag + overlap_param_gather (copy_main_params is gated on exactly those flags); LayerWise was skipped, leaving its staging buffer zeroed. Both now re-stage LayerWiseDistributedOptimizer too. _copy_main_params_to_param_buffer self-guards on use_buffer_param_sync + reuse_grad_buf_for_mxfp8_param_ag, so the call is a no-op for configs that don't need it. Co-Authored-By: Claude Opus 4.8 (1M context) --- megatron/core/transformer/moe/paged_stash.py | 9 ++++++++- megatron/training/training.py | 8 +++++++- 2 files changed, 15 insertions(+), 2 deletions(-) diff --git a/megatron/core/transformer/moe/paged_stash.py b/megatron/core/transformer/moe/paged_stash.py index 135b4804ad1..29623fbc4b2 100644 --- a/megatron/core/transformer/moe/paged_stash.py +++ b/megatron/core/transformer/moe/paged_stash.py @@ -11,6 +11,7 @@ from megatron.core._rank_utils import log_single_rank from megatron.core.full_cuda_graph import FullCudaGraphWrapper from megatron.core.optimizer.distrib_optimizer import DistributedOptimizer +from megatron.core.optimizer.layer_wise_optimizer import LayerWiseDistributedOptimizer from megatron.core.utils import get_attr_wrapped_model logger = logging.getLogger(__name__) @@ -1362,7 +1363,13 @@ def prepare_for_rerun(self, is_training=True): if self.copy_main_params: def _try_copy_main_params(opt): - if isinstance(opt, DistributedOptimizer) and hasattr( + if isinstance(opt, LayerWiseDistributedOptimizer): + # The LayerWise (muon) optimizer also stages its MXFP8 masters in + # the buffer aliased onto the grad buffer zeroed above; its + # _copy_main_params_to_param_buffer self-guards on + # use_buffer_param_sync + reuse_grad_buf_for_mxfp8_param_ag. + opt._copy_main_params_to_param_buffer() + elif isinstance(opt, DistributedOptimizer) and hasattr( opt, 'shard_fp32_from_float16_groups' ): opt._copy_main_params_to_param_buffer() diff --git a/megatron/training/training.py b/megatron/training/training.py index 6321dc22e62..fe70d75e44d 100644 --- a/megatron/training/training.py +++ b/megatron/training/training.py @@ -3691,10 +3691,16 @@ def trace_handler(p): if args.reuse_grad_buf_for_mxfp8_param_ag and args.overlap_param_gather: # disable_forward_pre_hook(param_sync=True) below force-syncs params for eval. # Copy the main params to param buffer before the forced AllGather. + # Both the standard DistributedOptimizer and the LayerWise (muon) + # optimizer stage their MXFP8 masters in the buffer aliased onto the + # grad buffer that zero_grad_buffer() just zeroed; re-stage both so the + # forced eval all-gather ships fresh weights instead of the zeroed buffer. for model_chunk in model: model_chunk.zero_grad_buffer() for optim_instance in optimizer.chained_optimizers: - if isinstance(optim_instance, DistributedOptimizer): + if isinstance( + optim_instance, (DistributedOptimizer, LayerWiseDistributedOptimizer) + ): optim_instance._copy_main_params_to_param_buffer() if should_disable_forward_pre_hook(args): disable_forward_pre_hook(model)