diff --git a/megatron/core/optimizer/distrib_optimizer.py b/megatron/core/optimizer/distrib_optimizer.py index b388161a610..2ea3ae89766 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.""" @@ -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,11 +3010,43 @@ 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 - ): + # 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 + if not _bucket_is_managed_by_layer_wise_optimizer( + bucket, default_for_untagged=False + ) + ] + if not distopt_buckets: + # Entire group is LayerWise-owned; LayerWise will sync it. continue - model_chunk._start_bucket_group_param_sync(bucket_group, force_sync=False) + + 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, + 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/layer_wise_optimizer.py b/megatron/core/optimizer/layer_wise_optimizer.py index e8e173ffe06..a3e45ae50c4 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,88 @@ 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._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 + ``_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 + — 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. + + 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. + """ + # 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 +485,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:`_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)): opt = optimizers[i] if isinstance(opt, (Float16OptimizerWithFloat16Params, FP32Optimizer)): @@ -408,12 +500,52 @@ 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 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 + # 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 and muon's NS step diverges from the + # ``fp8_param_gather=False`` baseline. + _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:`_copy_main_params_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 @@ -718,23 +850,118 @@ def count_zeros(self): ) 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. + 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. """ 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] - ): + 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 + 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 _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 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) + 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: @@ -747,17 +974,34 @@ 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._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/core/optimizer/optimizer.py b/megatron/core/optimizer/optimizer.py index 44ae593718c..1569ad605c4 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.""" @@ -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,48 @@ 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/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/arguments.py b/megatron/training/arguments.py index dcafa9e8139..63ad87fe236 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.""" @@ -1075,6 +1075,49 @@ 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.' + ) + + # 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/megatron/training/training.py b/megatron/training/training.py index 3930cc46a21..fe70d75e44d 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. @@ -3684,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) 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..be5064e1aff --- /dev/null +++ b/tests/unit_tests/test_muon_mxfp8_fp8_param_gather.py @@ -0,0 +1,507 @@ +# 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 copy +import gc +import os +import sys + +import pytest +import torch +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 +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 +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.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 + +_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): + 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) + + +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. + """ + snapshot = {} + 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 + 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 + 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 + 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: + """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' + 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() + destroy_global_vars() + destroy_num_microbatches_calculator() + gc.collect() + + 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) + layer_spec = get_gpt_layer_with_transformer_engine_spec() + return GPTModel( + config=transformer_config, + transformer_layer_spec=layer_spec, + vocab_size=args.vocab_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, + pg_collection=config_kwargs.get("pg_collection"), + vp_stage=config_kwargs.get("vp_stage"), + ) + + 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'] + args = parse_args() + args.num_layers = 2 + args.vocab_size = 128 + 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 + 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.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' + 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 + # ``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 + # them explicitly. + args.fp8 = "e4m3" + args.fp8_recipe = fp8_recipe + args.fp8_param_gather = 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) + 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 _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) + + 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__}" + ) + return args, gpt_model, optimizer + + 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() + 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, + 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 + + 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, 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( + 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_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. + """ + 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_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, + overlap_param_gather=overlap_param_gather, + ) + _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): + _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}" + ) + + 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]: + _assert_tensor_equal( + grads_on[step][name], + grads_off[step][name], + 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}", + ) + + 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 common_master_names: + _assert_tensor_equal( + masters_on[step][name], + masters_off[step][name], + f"master mismatch at step {step} for {name}", + )