From 3a73182cf557bbedf8ff257d12820411eed19a65 Mon Sep 17 00:00:00 2001 From: Pingtian Li Date: Wed, 29 Jul 2026 21:38:38 +0800 Subject: [PATCH 01/12] Support decoupled compact LayerWise DDP layout for Muon, with checkpoint fix Makes `use_distributed_optimizer` a per-buffer property so the LayerWise (Muon) optimizer can drop the persistent `dp_size * max(shard_load)` padding from its long-lived param and grad buffers. LayerWise-managed (Muon 2D-matrix) buffers get a compact no-padding DDP layout and locally disable DistributedOptimizer semantics (all-reduce gradients, whole-param ping-pong ownership, `allgather_params` param sync), while sibling buffers (embeddings, biases, layernorm) keep the standard byte-level DistributedOptimizer path. The decision is baked into each buffer's own `ddp_config` via `dataclasses.replace`, so bucket groups inherit it; a bucket group runs one collective type, so `partition_buckets` asserts that any buckets it merges agree. This layout is now the **default**: `use_layer_wise_param_layout` defaults to `False` on both `DistributedDataParallelConfig` and `OptimizerConfig`, and `--use-layer-wise-param-layout` opts back **into** the padded shard-aligned layout (e.g. for bit-for-bit comparison against older runs). Checkpoint fix, which belongs with this layout but was missed when it first landed: * Grad-buffer range maps are filtered to the params the optimizer instance actually owns. On this layout a DistOpt's buffers also carry buckets owned by the LayerWise child; unfiltered, the DistOpt builds ranges and state for params it does not own, duplicating what LayerWise already saves. * `sharded_param_state_dp_reshardable` save/load skip buckets this optimizer owns no param of. Membership is checked per param rather than via `params[0]`, so buckets mixing owned and LayerWise-managed params are handled. * A DP rank can own a shard that is entirely padding yet still overlaps `[0, numel_unpadded)`; such a shard now synthesizes a padding ShardedTensor from a captured fp32 template so global coverage holds, and writes it back into `state` (the branch rebinds `bucket_state`, so without the write-back the shard was dropped). * The original strict `empty bucket encountered` assert is kept for pure DistOpt runs; empty shards are only expected when LayerWise-managed params co-exist. * `fully_sharded_model_space` is asserted unsupported on this layout rather than failing obscurely later. The `*_param_layout` functional cases now pass `--use-layer-wise-param-layout` explicitly, since they exist to cover the padded layout and the default moved. Co-Authored-By: Claude Opus 5 (1M context) Signed-off-by: Pingtian Li --- .../distributed_data_parallel_config.py | 6 + .../core/distributed/param_and_grad_buffer.py | 123 ++++++++-- megatron/core/optimizer/distrib_optimizer.py | 126 +++++++++- .../core/optimizer/layer_wise_optimizer.py | 24 +- megatron/core/optimizer/optimizer_config.py | 7 + megatron/training/arguments.py | 36 ++- megatron/training/training.py | 38 +-- .../model_config.yaml | 1 - .../model_config.yaml | 1 - .../model_config.yaml | 1 + .../model_config.yaml | 1 + .../test_layer_wise_optimizer.py | 225 +++++++++++++++++- tests/unit_tests/dist_checkpointing/utils.py | 85 ++++++- 13 files changed, 607 insertions(+), 67 deletions(-) diff --git a/megatron/core/distributed/distributed_data_parallel_config.py b/megatron/core/distributed/distributed_data_parallel_config.py index 96925ce120c..7b6845e4f56 100644 --- a/megatron/core/distributed/distributed_data_parallel_config.py +++ b/megatron/core/distributed/distributed_data_parallel_config.py @@ -36,6 +36,12 @@ class DistributedDataParallelConfig: enabled. Defaults to 1, which means DistOpt is across entire DP domain. """ + use_layer_wise_param_layout: bool = False + """Layer-wise (Muon) optimizer only. When True, LayerWise-managed buffers use + the shard-aligned padded LayerWise param layout. When False (default), the compact + decoupled layout is selected instead. + """ + check_for_nan_in_grad: bool = False """ If true, check for NaNs and Infs in gradients _before_ communication collective. diff --git a/megatron/core/distributed/param_and_grad_buffer.py b/megatron/core/distributed/param_and_grad_buffer.py index 50fa566d1b6..35eeff100bf 100644 --- a/megatron/core/distributed/param_and_grad_buffer.py +++ b/megatron/core/distributed/param_and_grad_buffer.py @@ -1,5 +1,6 @@ # Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved. +import dataclasses import fnmatch import functools import logging @@ -1062,6 +1063,17 @@ def __init__( self.gradient_scaling_factor = gradient_scaling_factor self.nccl_ub = nccl_ub + self._is_layer_wise_buffer = bool( + self.params and getattr(self.params[0], "is_managed_by_layer_wise_optimizer", False) + ) + # Bake the per-buffer DistOpt decision into this buffer's ddp_config (single source of + # truth; bucket groups inherit it): a LayerWise (Muon) buffer on the compact decoupled + # layout disables DistributedOptimizer, while sibling buffers keep the model-level setting. + if self._is_layer_wise_buffer and not getattr( + self.ddp_config, "use_layer_wise_param_layout", True + ): + self.ddp_config = dataclasses.replace(self.ddp_config, use_distributed_optimizer=False) + # Data structures to store underlying buckets and relevant indexing data. self.buckets = [] self.param_to_bucket = {} # Param -> bucket mapping. @@ -1112,6 +1124,27 @@ def __init__( # nvfp4_packed_numel_unpadded is already set by _compute_nvfp4_packed_layout. assert self.numel_unpadded <= self.numel + + # Diagnostic: log persistent buffer size vs. unpadded payload so the cost of any + # optimizer-driven padding (e.g. the LayerWise shard-aligned ``dp_size * max(shard_load)`` + # layout) is visible per buffer. Emit at INFO only when it is interesting — a + # LayerWise-managed buffer or one that actually carries padding — and DEBUG otherwise, so + # ordinary (zero-padding) buffers do not spam non-experimental runs. + _padding = self.numel - self.numel_unpadded + _pad_frac = _padding / max(self.numel_unpadded, 1) + log_on_each_pipeline_stage( + logger, + logging.INFO if (self._is_layer_wise_buffer or _padding > 0) else logging.DEBUG, + f"ParamAndGradBuffer layout: param_dtype={self.param_dtype} " + f"grad_dtype={self.grad_dtype} dp_world_size={self.data_parallel_world_size} " + f"layerwise={self._is_layer_wise_buffer} " + f"distopt={self.ddp_config.use_distributed_optimizer} " + f"numel={self.numel} numel_unpadded={self.numel_unpadded} " + f"padding={_padding} ({_pad_frac:.1%})", + tp_group=self.tp_group, + dp_cp_group=self.dp_cp_group, + ) + if self.has_nvfp4_params: assert self.nvfp4_packed_numel_unpadded <= self.nvfp4_packed_numel if self.ddp_config.use_distributed_optimizer: @@ -1703,22 +1736,72 @@ def partition_buckets( assert fp8_buffer is None fp8_buffer = buffer - # Case 1: Put all buckets into a single bucket group if force_single_bucket_group is True. + # A bucket group performs a single collective type (reduce-scatter for DistOpt buffers, + # all-reduce otherwise), so buckets merged into one group must agree on the effective + # per-buffer ``use_distributed_optimizer``. The decoupled LayerWise layout + # (``use_layer_wise_param_layout=False``) gives LayerWise (Muon) buffers + # ``use_distributed_optimizer=False`` while sibling buffers keep True; the no-fp8 Case 2 below + # keeps every bucket in its own group so they never mix, but the merging Cases 1/3 must assert + # consistency. + _ddp_config = buffers[0].ddp_config + _decouple = not getattr(_ddp_config, "use_layer_wise_param_layout", True) + + def _bucket_distopt(bucket): + """This bucket's effective ``use_distributed_optimizer``.""" + is_lw = bool( + bucket.params_list + and getattr(bucket.params_list[0], "is_managed_by_layer_wise_optimizer", False) + ) + if _decouple and is_lw: + return False + return _ddp_config.use_distributed_optimizer + + def _merged_use_distributed_optimizer(merge_buckets): + values = {_bucket_distopt(bucket) for bucket in merge_buckets} + assert len(values) == 1, ( + "Cannot merge buckets with differing effective use_distributed_optimizer into one " + "bucket group. This happens when the decoupled LayerWise layout " + "(use_layer_wise_param_layout=False) mixes LayerWise (all-reduce) and non-LayerWise " + "(reduce-scatter) buffers under a merging bucketing strategy (e.g. the fp8 merge " + "path). Disable bucket merging for the decoupled LayerWise path." + ) + return values.pop() + + # Case 1: Put all buckets into a single bucket group if force_single_bucket_group is True + # (e.g. disable_bucketing / non-first VPP chunks). A bucket group performs a single + # collective type, so when the decoupled LayerWise layout (use_layer_wise_param_layout=False) + # mixes LayerWise (all-reduce, non-DistOpt) and non-LayerWise (reduce-scatter, DistOpt) + # buffers in one chunk, we cannot + # merge them into a single group. Split by the effective per-bucket use_distributed_optimizer + # instead, preserving order. When all buckets agree (the non-decoupled case) this collapses + # to exactly one group, identical to the previous behavior. if force_single_bucket_group: - buckets = [] - ddp_config = buffers[0].ddp_config data_parallel_group = buffers[0].data_parallel_group data_parallel_world_size = buffers[0].data_parallel_world_size + ordered_distopt_values = [] + buckets_by_distopt = {} + # buffer.ddp_config already carries the per-buffer use_distributed_optimizer. + ddp_config_by_distopt = {} for buffer in buffers: - assert ddp_config == buffer.ddp_config assert data_parallel_group == buffer.data_parallel_group assert data_parallel_world_size == buffer.data_parallel_world_size - buckets.extend(buffer.buckets) - - bucket_group = _ParamAndGradBucketGroup( - buckets, ddp_config, data_parallel_group, data_parallel_world_size - ) - return [bucket_group] + distopt = buffer.ddp_config.use_distributed_optimizer + ddp_config_by_distopt.setdefault(distopt, buffer.ddp_config) + for bucket in buffer.buckets: + if distopt not in buckets_by_distopt: + buckets_by_distopt[distopt] = [] + ordered_distopt_values.append(distopt) + buckets_by_distopt[distopt].append(bucket) + + return [ + _ParamAndGradBucketGroup( + buckets_by_distopt[distopt], + ddp_config_by_distopt[distopt], + data_parallel_group, + data_parallel_world_size, + ) + for distopt in ordered_distopt_values + ] if fp8_buffer is None: # Case 2: When there is no fp8 buffer in the input buffers, let each bucket group have @@ -1737,11 +1820,12 @@ def partition_buckets( return bucket_groups else: # Case 3: When using fp8 params, merge all non-fp8 buckets into the last fp8 bucket group. - non_fp8_buckets = [] + # Track each non-fp8 bucket with its buffer's (authoritative) ddp_config. + non_fp8_buckets = [] # list of (bucket, ddp_config) for buffer in buffers: if buffer.param_dtype != torch.uint8: for bucket in buffer.buckets: - non_fp8_buckets.append(bucket) + non_fp8_buckets.append((bucket, buffer.ddp_config)) bucket_groups = [] for bucket in fp8_buffer.buckets: @@ -1755,17 +1839,17 @@ def partition_buckets( bucket_groups.append( _ParamAndGradBucketGroup( [bucket], - buffer.ddp_config, + fp8_buffer.ddp_config, buffer.data_parallel_group, buffer.data_parallel_world_size, ) ) if non_fp8_buckets: - for non_fp8_bucket in non_fp8_buckets: + for non_fp8_bucket, non_fp8_ddp_config in non_fp8_buckets: bucket_groups.append( _ParamAndGradBucketGroup( [non_fp8_bucket], - buffer.ddp_config, + non_fp8_ddp_config, buffer.data_parallel_group, buffer.data_parallel_world_size, ) @@ -1773,14 +1857,19 @@ def partition_buckets( continue # Skip the default bucket group creation below else: - group_buckets = [bucket] + non_fp8_buckets + group_buckets = [bucket] + [b for b, _ in non_fp8_buckets] else: # The first N-1 bucket groups. group_buckets = [bucket] + # Merged buckets must share the fp8 group's effective use_distributed_optimizer. + assert ( + _merged_use_distributed_optimizer(group_buckets) + == fp8_buffer.ddp_config.use_distributed_optimizer + ) bucket_groups.append( _ParamAndGradBucketGroup( group_buckets, - buffer.ddp_config, + fp8_buffer.ddp_config, buffer.data_parallel_group, buffer.data_parallel_world_size, ) diff --git a/megatron/core/optimizer/distrib_optimizer.py b/megatron/core/optimizer/distrib_optimizer.py index 782d3263570..50b2cf00f93 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.""" @@ -8,7 +8,7 @@ from collections import ChainMap from dataclasses import replace from logging import getLogger -from typing import Any, Callable, Dict, List, Optional, Tuple +from typing import Any, Callable, Dict, List, Optional, Set, Tuple import torch import torch.nn.functional @@ -255,6 +255,26 @@ def _build_model_gbuf_range(cls, param_and_grad_buffer: _ParamAndGradBuffer, buc return data + @classmethod + def _filter_gbuf_range_map( + cls, gbuf_range_map: Dict, optimizer_params: Set[torch.nn.Parameter] + ) -> Dict: + """Filter grad-buffer range maps to the params owned by this optimizer instance.""" + return { + dtype: [ + { + **range_map, + "param_map": { + param: param_range + for param, param_range in range_map["param_map"].items() + if param in optimizer_params + }, + } + for range_map in range_maps + ] + for dtype, range_maps in gbuf_range_map.items() + } + @classmethod def _build_gbuf_range_map(cls, param_and_grad_buffer: _ParamAndGradBuffer): """Builds a map between parameters and their ranges in the grad buffer. @@ -724,6 +744,13 @@ def __init__( for model_idx, buffers in self.per_model_buffers.items(): self.per_model_bucket_groups[model_idx] = partition_buckets(buffers) + optimizer_params = { + param for param_group in self.optimizer.param_groups for param in param_group['params'] + } + # Model params this optimizer owns, captured before the fp32-master swap below. Used by + # sharded_param_state_dp_reshardable / load to skip buckets it owns no param of, and to + # distinguish decoupled LayerWise setups (empty shards expected) from pure DistOpt. + self._optimizer_model_params = optimizer_params self.gbuf_ranges = [] self.per_bucket_numel = [] self.per_bucket_numel_unpadded = [] @@ -743,7 +770,9 @@ def __init__( ] } ) - self.gbuf_ranges.append(self._build_gbuf_range_map(buffer)) + self.gbuf_ranges.append( + self._filter_gbuf_range_map(self._build_gbuf_range_map(buffer), optimizer_params) + ) self.model_param_gbuf_map = self._build_model_param_gbuf_map(self.gbuf_ranges) # Add main_param field to each parameter. We will use this fp32 copy to compute @@ -1863,6 +1892,36 @@ def sharded_param_state_dp_reshardable( data_parallel_world_size = self.data_parallel_group.size() state = self.get_parameter_state_dp_reshardable() + + # fp32 optimizer-state {key: (dtype, device)} captured before the loop below mutates + # ``state``, so an empty shard can synthesize valid padding ShardedTensors. + pad_template = None + for _g in range(len(self.gbuf_ranges)): + for _bs_all in state[_g].values(): + for _bs in _bs_all: + if _bs: + pad_template = { + k: (v.dtype, v.device) + for k, v in _bs[0].items() + if isinstance(v, torch.Tensor) + } + break + if pad_template is not None: + break + if pad_template is not None: + break + + # A rank that owns nothing in any bucket has no local sample. The optimizer-state keys + # and their dtypes are config-determined and identical on every DP rank, so derive the + # template from config instead of gathering it across DP -- no communication needed. + if pad_template is None: + _pad_device = torch.cuda.current_device() + pad_template = { + 'param': (self.config.main_params_dtype, _pad_device), + 'exp_avg': (self.config.exp_avg_dtype, _pad_device), + 'exp_avg_sq': (self.config.exp_avg_sq_dtype, _pad_device), + } + # per_bucket_numel metadata is saved separately for each TPxPP domain. for per_bucket_key in ('per_bucket_numel', 'per_bucket_numel_unpadded'): key = ( @@ -1894,9 +1953,53 @@ def sharded_param_state_dp_reshardable( f'.gbuf_idx_{gbuf_idx}.dtype_{dtype}.bucket_idx_{bucket_idx}' ) - # The global ckpt tensors must be fully covered. - # We add extra empty padding if necessary - assert bucket_state, 'empty bucket encountered' + # Skip buckets this optimizer owns no param of. In the decoupled compact + # LayerWise (Muon) layout our buffers also hold buckets whose params the + # LayerWise child owns; their state is saved there and every shard here is + # empty. Checking membership (vs the params[0] tag) also handles buckets that + # mix owned and LayerWise-managed params. + bucket = self.buffers[gbuf_idx].buckets[bucket_idx] + if not any(p in self._optimizer_model_params for p in bucket.params_list): + continue + + # bucket_state is built 1:1 from this rank's param_map, so an empty state is + # legitimate only when the rank owns no param range in this bucket (a small + # bucket + 64-element shard alignment can put a whole shard in padding). If it + # does own a range yet the state is empty, the state was lost -- keep the + # original strict check. + if not bucket_state: + assert not self.gbuf_ranges[gbuf_idx][dtype][bucket_idx]['param_map'], ( + f'empty dp_reshardable state for {sharded_bucket_key} but this rank ' + 'owns param ranges in the bucket (optimizer state lost)' + ) + world_shard_start = data_parallel_rank * gbuf_local_numel + if world_shard_start >= gbuf_world_numel_unpadded: + # Shard is entirely past the unpadded end (trailing padding); not saved. + continue + pad_len = min( + gbuf_local_numel, gbuf_world_numel_unpadded - world_shard_start + ) + # Synthesize the padding for this shard's overlap with [0, numel_unpadded), + # using pad_template for the correct fp32 keys/dtypes. The template is + # always available (config-derived above); a None here would mean the + # coverage we must emit is being dropped -- fail loudly rather than skip. + assert ( + pad_template is not None + ), f'no padding template for {sharded_bucket_key}; coverage would be dropped' + bucket_state = [ + { + **{ + k: torch.empty(pad_len, dtype=_dt, device=_dev) + for k, (_dt, _dev) in pad_template.items() + }, + 'gbuf_local_start': 0, + 'gbuf_local_end': pad_len, + 'padding': True, + } + ] + # Store the synthesized shard back into ``state``: this branch rebinds + # ``bucket_state`` to a new list, so without this the shard is dropped. + gbuf_range_map_for_all_buckets[bucket_idx] = bucket_state # Insert padding between parameter tensors to ensure full coverage as needed. all_pad_tensors = {} @@ -2015,6 +2118,11 @@ def _get_param_state_sharded_tensors(model_param, item_slice): f" Hint: {KEEP_VARS_HINT}" ) from e + assert isinstance(sharded_metadata, ShardedTensorFactory), ( + "fully_sharded_model_space is not supported for the decoupled compact " + "LayerWise (Muon) optimizer layout" + ) + # Set DP corresponding replica_id coordinate to 0. assert ( len(sharded_metadata.replica_id) == 3 @@ -2075,6 +2183,12 @@ def load_parameter_state_from_dp_reshardable(self, state_dict): assert len(gbuf_range_maps) == 1, "single dtype supported, for now." for dtype, gbuf_range_map_for_all_buckets in gbuf_range_maps.items(): for bucket_idx, gbuf_range_map in enumerate(gbuf_range_map_for_all_buckets): + # Skip buckets this optimizer owns no param of (see the save counterpart). + if not any( + p in self._optimizer_model_params + for p in self.buffers[gbuf_idx].buckets[bucket_idx].params_list + ): + continue bucket_state = state_dict[gbuf_idx][dtype][bucket_idx] bucket_state = [ bucket_state_elem diff --git a/megatron/core/optimizer/layer_wise_optimizer.py b/megatron/core/optimizer/layer_wise_optimizer.py index a132e1be3a1..b9f1b77e539 100644 --- a/megatron/core/optimizer/layer_wise_optimizer.py +++ b/megatron/core/optimizer/layer_wise_optimizer.py @@ -396,8 +396,15 @@ def compute_full_param_layout( :class:`FullParamLayout` with a :class:`PerBufferParamLayout` per buffer group. """ # Avoid a circular import: DistributedOptimizer imports LayerWise indirectly. + from ..distributed.param_and_grad_buffer import _compute_default_per_buffer_param_layout from .distrib_optimizer import DistributedOptimizer + # Decoupled layout (use_layer_wise_param_layout=False): LayerWise (Muon) buffers use a + # compact no-padding DDP layout (and locally disable DistributedOptimizer semantics in + # DDP), so they must NOT receive the shard-aligned ``dp_size * max(shard_load)`` padded + # layout here. Non-LayerWise buffers keep DistOpt's byte-level layout regardless. + decouple_ddp_layout = not ddp_config.use_layer_wise_param_layout + buffer_groups = group_params_for_buffers(params, ddp_config.grad_reduce_in_fp32) layouts = {} for buffer_key, (group_params, param_indices) in buffer_groups.items(): @@ -413,6 +420,15 @@ def compute_full_param_layout( # Dispatch per buffer: LayerWise (Muon) params get the shard-aligned # layout; non-LayerWise params (e.g. Adam-managed embeddings, biases) # get DistOpt's byte-level layout. + if buffer_key.is_managed_by_layer_wise_optimizer and decouple_ddp_layout: + # Compact no-padding layout (DDP treats this buffer as non-DistOpt). Attach + # param_indices so DDP's layout/grouping consistency check passes. + per_buffer_layout = _compute_default_per_buffer_param_layout( + group_params, bucket_size + ) + per_buffer_layout.param_indices = param_indices + layouts[buffer_key] = per_buffer_layout + continue if buffer_key.is_managed_by_layer_wise_optimizer: compute_per_buffer_layout = ( LayerWiseDistributedOptimizer._compute_per_buffer_param_layout @@ -444,6 +460,7 @@ def __init__( """ self.pg_collection = pg_collection + self.decouple_ddp_layout = not config.use_layer_wise_param_layout # The data-parallel groups this optimizer shards parameters over. Cached here so the # sharding, all-gather and broadcast paths read one attribute instead of reaching back @@ -467,7 +484,7 @@ def __init__( ) full_param_layouts = None - if model_chunks is not None: + if model_chunks is not None and not self.decouple_ddp_layout: full_param_layouts = [ chunk.full_param_layout for chunk in model_chunks @@ -475,12 +492,13 @@ def __init__( ] or None self.shard_params(optimizers, full_param_layouts) - # When a full_param_layout is available, ddp_config.use_distributed_optimizer + # When a full_param_layout is available (no decoupling), use_distributed_optimizer # is True and model params are views into the DDP param buffer. After the # optimizer step copies updated fp32 main params → bf16 model params, the # buffer is already up-to-date in-place. We can use DDP's buffer-based # all-gather (start_param_sync) instead of the flatten/unflatten allgather_params - # path. + # In the decouple path, Muon buffers are non-DistOpt and own + # whole params via ping-pong, so we use the legacy allgather_params path instead. self.use_buffer_param_sync = full_param_layouts is not None # Set up overlap param gather using DDP bucket infrastructure. diff --git a/megatron/core/optimizer/optimizer_config.py b/megatron/core/optimizer/optimizer_config.py index 24f9a032c47..b741815e62a 100644 --- a/megatron/core/optimizer/optimizer_config.py +++ b/megatron/core/optimizer/optimizer_config.py @@ -327,6 +327,13 @@ class OptimizerConfig: arguments layer sets this flag and resets ``use_distributed_optimizer`` to False so that the standard distributed-optimizer path is not triggered.""" + use_layer_wise_param_layout: bool = False + """Layer-wise (Muon) optimizer only. When True, LayerWise-managed buffers use + the shard-aligned padded LayerWise param layout. When False (default), the compact decoupled + layout is selected: LayerWise-managed (Muon) buffers use a compact no-padding DDP layout with + all-reduce gradients and legacy whole-param ping-pong ownership + ``allgather_params`` + param sync.""" + overlap_param_gather: bool = False """If true, overlap param all-gather with forward compute. This argument is intended to have the same value as the "overlap_param_gather" argument diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py index fd30984b38a..e16abf49d3a 100644 --- a/megatron/training/arguments.py +++ b/megatron/training/arguments.py @@ -1686,6 +1686,21 @@ def validate_args(args, defaults={}): assert args.ckpt_format in ["torch", "torch_dist"], "Emerging optimizer supports torch and torch_dist checkpoint format." + if args.use_layer_wise_distributed_optimizer: + assert not args.fp8_param_gather and not getattr(args, 'fp4_param_gather', False), ( + "Layer-wise (Muon) distributed optimizer does not support FP8/FP4 parameter gather " + "(fp8_param_gather / fp4_param_gather). Use fp8_param_gather=False (e.g. blockwise/" + "MXFP8 compute with parameters persisted in bf16)." + ) + if not args.use_layer_wise_param_layout: + assert args.num_distributed_optimizer_instances == 1, ( + "the decoupled compact LayerWise DDP layout (the default; pass " + "--use-layer-wise-param-layout for the padded layout) requires " + "num_distributed_optimizer_instances == 1: the non-DistOpt LayerWise (Muon) buffers " + "only all-reduce within a single optimizer instance, so partial DistOpt (>1 " + "instance) would under-reduce Muon gradients across the full data-parallel domain." + ) + # Make sure all functionality that requires Gloo process groups is disabled. if not args.use_gloo_process_groups: if args.use_distributed_optimizer: @@ -2996,17 +3011,16 @@ def _add_distributed_args(parser): dest='align_param_gather') group.add_argument('--use-distributed-optimizer', action='store_true', help='Use distributed optimizer.') - group.add_argument('--no-use-layer-wise-param-layout', - action='store_false', - dest='use_layer_wise_param_layout', - help='Opt out of the precomputed LayerWise param layout. When set, ' - 'falls back to the legacy LayerWise ping-pong path: all params ' - '(including non-Muon embeddings, biases, layernorm) live in a single ' - 'LayerWise buffer and the optimizer uses the allgather_params() codepath. ' - 'The default (precomputed layout) routes non-Muon params through a ' - 'separate DistributedOptimizer with byte-level sharding, which is faster ' - 'and uses less padding but produces different bf16 reduction ordering ' - 'and so will not match legacy-path loss curves bit-for-bit.') + group.add_argument('--use-layer-wise-param-layout', + action='store_true', + help='Opt INTO the padded shard-aligned LayerWise param layout. The default ' + 'is the compact decoupled layout, where LayerWise (Muon 2D) buffers use a ' + 'no-padding DDP layout and locally disable DistributedOptimizer (all-reduce ' + 'grads + whole-param ping-pong + allgather_params), while sibling buffers keep ' + 'the byte-level DistributedOptimizer; this avoids the persistent ' + 'dp_size * max(shard_load) padding. Pass this flag to restore the padded ' + 'layout (e.g. for bit-for-bit comparison; it uses a different bf16 reduction ' + 'ordering).') group.add_argument('--use-nccl-ub', action='store_true', dest='nccl_ub', help='Use the userbuffer registration for DP/FSDP communication buffers.' 'This option will reduce GPU SM usage for the DP/FSDP communication,' diff --git a/megatron/training/training.py b/megatron/training/training.py index 57b8efe9390..98d97a162db 100644 --- a/megatron/training/training.py +++ b/megatron/training/training.py @@ -1573,7 +1573,6 @@ def wrap_model_chunks_with_ddp( ddp_config, *, use_layer_wise_distributed_optimizer=False, - use_layer_wise_param_layout=True, DP=DDP, pg_collection=None, bucket_sizes=None, @@ -1584,12 +1583,13 @@ def wrap_model_chunks_with_ddp( Centralises the DDP-wrapping wiring shared between :func:`get_model` and unit tests. - For ``use_layer_wise_distributed_optimizer=True`` and ``use_layer_wise_param_layout=True``: - forces ``ddp_config.use_distributed_optimizer=True`` (mutated in place; needed - for reduce-scatter), and computes per-chunk shard-aligned layouts via - :meth:`LayerWiseDistributedOptimizer.compute_full_param_layout`. With - ``use_layer_wise_param_layout=False``, no layout is supplied and LayerWise falls back - to its legacy ``allgather_params`` sync path. + For ``use_layer_wise_distributed_optimizer=True``: forces + ``ddp_config.use_distributed_optimizer=True`` (mutated in place; needed for reduce-scatter) + and computes per-chunk layouts via + :meth:`LayerWiseDistributedOptimizer.compute_full_param_layout`. That method picks the + padded shard-aligned LayerWise layout or the compact decoupled layout per + ``ddp_config.use_layer_wise_param_layout`` (True → padded, False → compact, with LayerWise + buffers treated as non-DistOpt and synced via legacy ``allgather_params``). For non-layerwise with ``ddp_config.use_distributed_optimizer=True``: computes per-chunk byte-level layouts via @@ -1605,11 +1605,9 @@ def wrap_model_chunks_with_ddp( model_chunks: List of model chunks to wrap (un-DDP-wrapped). config: :class:`TransformerConfig`. ddp_config: :class:`DistributedDataParallelConfig`. Mutated in place when - ``use_layer_wise_distributed_optimizer=True`` and ``use_layer_wise_param_layout=True``. + ``use_layer_wise_distributed_optimizer=True``. Its ``use_layer_wise_param_layout`` + field selects the padded (default) vs compact decoupled LayerWise layout. use_layer_wise_distributed_optimizer: Whether the layerwise wiring runs. - use_layer_wise_param_layout: When ``use_layer_wise_distributed_optimizer=True``, - controls whether to compute and supply a shard-aligned param layout - to DDP. ``False`` keeps LayerWise on its legacy sync path. DP: The DDP class to construct (``DistributedDataParallel`` or an FSDP variant). pg_collection: Optional :class:`ProcessGroupCollection`. When provided, @@ -1632,13 +1630,18 @@ def wrap_model_chunks_with_ddp( # Compute per-chunk layouts (DDP only). per_chunk_layouts = [None] * n if DP is DDP: - if use_layer_wise_distributed_optimizer and use_layer_wise_param_layout: + if use_layer_wise_distributed_optimizer: + # LayerWise (Muon) optimizer. Force use_distributed_optimizer=True so sibling + # non-LayerWise buffers (embeddings, biases, layernorm) shard with the byte-level + # DistributedOptimizer layout, and tag params so DDP buffer grouping routes + # LayerWise-managed matrices (Muon's Newton-Schulz domain) to a separate buffer. + # The padded-vs-compact LayerWise layout decision is made inside + # compute_full_param_layout / _ParamAndGradBuffer from use_layer_wise_param_layout: + # by default (compact) LayerWise buffers get the no-padding layout and the per-buffer + # override flips use_distributed_optimizer off for them; with + # --use-layer-wise-param-layout they stay on the padded DistOpt layout. ddp_config.use_distributed_optimizer = True compute_layout = LayerWiseDistributedOptimizer.compute_full_param_layout - # Tag params so DDP buffer grouping routes LayerWise-managed matrices - # (Muon's Newton-Schulz domain) to a shard-aligned buffer and routes - # everything else (embeddings, biases, layernorm) to a separate - # DistOpt-style buffer. tag_params_for_buffer_routing(model_chunks) elif not use_layer_wise_distributed_optimizer and ddp_config.use_distributed_optimizer: compute_layout = DistributedOptimizer.compute_full_param_layout @@ -1884,9 +1887,6 @@ def build_model(): use_layer_wise_distributed_optimizer=getattr( args, 'use_layer_wise_distributed_optimizer', False ), - use_layer_wise_param_layout=getattr( - args, 'use_layer_wise_param_layout', True - ), DP=DP, pg_collection=pg_collection if args.use_megatron_fsdp else None, bucket_sizes=per_chunk_bucket_sizes, diff --git a/tests/functional_tests/test_cases/moe/gpt3_moe_mcore_te_ep8_resume_torch_dist_dist_muon/model_config.yaml b/tests/functional_tests/test_cases/moe/gpt3_moe_mcore_te_ep8_resume_torch_dist_dist_muon/model_config.yaml index fffad86a016..028074bd34f 100644 --- a/tests/functional_tests/test_cases/moe/gpt3_moe_mcore_te_ep8_resume_torch_dist_dist_muon/model_config.yaml +++ b/tests/functional_tests/test_cases/moe/gpt3_moe_mcore_te_ep8_resume_torch_dist_dist_muon/model_config.yaml @@ -66,5 +66,4 @@ MODEL_ARGS: --async-save: true --use-persistent-ckpt-worker: true --use-distributed-optimizer: true - --no-use-layer-wise-param-layout: true TEST_TYPE: ckpt-resume diff --git a/tests/functional_tests/test_cases/moe/gpt3_moe_mcore_te_ep8_resume_torch_dist_dist_muon_1node/model_config.yaml b/tests/functional_tests/test_cases/moe/gpt3_moe_mcore_te_ep8_resume_torch_dist_dist_muon_1node/model_config.yaml index 874e7dccf8d..40ac94eed9b 100644 --- a/tests/functional_tests/test_cases/moe/gpt3_moe_mcore_te_ep8_resume_torch_dist_dist_muon_1node/model_config.yaml +++ b/tests/functional_tests/test_cases/moe/gpt3_moe_mcore_te_ep8_resume_torch_dist_dist_muon_1node/model_config.yaml @@ -66,5 +66,4 @@ MODEL_ARGS: --async-save: true --use-persistent-ckpt-worker: true --use-distributed-optimizer: true - --no-use-layer-wise-param-layout: true TEST_TYPE: ckpt-resume diff --git a/tests/functional_tests/test_cases/moe/gpt3_moe_mcore_te_ep8_resume_torch_dist_dist_muon_param_layout/model_config.yaml b/tests/functional_tests/test_cases/moe/gpt3_moe_mcore_te_ep8_resume_torch_dist_dist_muon_param_layout/model_config.yaml index 028074bd34f..40f068cc90c 100644 --- a/tests/functional_tests/test_cases/moe/gpt3_moe_mcore_te_ep8_resume_torch_dist_dist_muon_param_layout/model_config.yaml +++ b/tests/functional_tests/test_cases/moe/gpt3_moe_mcore_te_ep8_resume_torch_dist_dist_muon_param_layout/model_config.yaml @@ -66,4 +66,5 @@ MODEL_ARGS: --async-save: true --use-persistent-ckpt-worker: true --use-distributed-optimizer: true + --use-layer-wise-param-layout: true TEST_TYPE: ckpt-resume diff --git a/tests/functional_tests/test_cases/moe/gpt3_moe_mcore_te_ep8_resume_torch_dist_dist_muon_param_layout_1node/model_config.yaml b/tests/functional_tests/test_cases/moe/gpt3_moe_mcore_te_ep8_resume_torch_dist_dist_muon_param_layout_1node/model_config.yaml index 40ac94eed9b..34a85516794 100644 --- a/tests/functional_tests/test_cases/moe/gpt3_moe_mcore_te_ep8_resume_torch_dist_dist_muon_param_layout_1node/model_config.yaml +++ b/tests/functional_tests/test_cases/moe/gpt3_moe_mcore_te_ep8_resume_torch_dist_dist_muon_param_layout_1node/model_config.yaml @@ -66,4 +66,5 @@ MODEL_ARGS: --async-save: true --use-persistent-ckpt-worker: true --use-distributed-optimizer: true + --use-layer-wise-param-layout: true TEST_TYPE: ckpt-resume diff --git a/tests/unit_tests/dist_checkpointing/test_layer_wise_optimizer.py b/tests/unit_tests/dist_checkpointing/test_layer_wise_optimizer.py index 42ef0a401ee..8be5353297d 100644 --- a/tests/unit_tests/dist_checkpointing/test_layer_wise_optimizer.py +++ b/tests/unit_tests/dist_checkpointing/test_layer_wise_optimizer.py @@ -1,4 +1,4 @@ -# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. from copy import deepcopy from functools import partial @@ -23,6 +23,7 @@ from megatron.core.utils import get_pg_size from megatron.training.arguments import parse_args from megatron.training.checkpointing import load_checkpoint, save_checkpoint +from megatron.training.utils import get_device_arch_version from tests.unit_tests.dist_checkpointing import ( TempNamedDir, init_basic_mock_args, @@ -253,6 +254,228 @@ def test_layer_wise_optimizer_save_load(self, tmp_path_dist_ckpt, tp, pp, bf16): check_equal(plain_sd_A, plain_sd_B) + # grad_reduce_in_fp32=True gives the mixed-dtype (bf16 param, fp32 grad) DistOpt sibling + # buffer of the decoupled Muon layout -- the case that fails 'Failed to validate global plan' + # in sharded_param_state_dp_reshardable on the real fp8 SFT save. + @pytest.mark.parametrize('grad_reduce_in_fp32', [False, True]) + @pytest.mark.parametrize('bf16', [True]) + def test_dp_reshardable_decouple_ckpt(self, tmp_path_dist_ckpt, bf16, grad_reduce_in_fp32): + """Save/load of the decoupled compact LayerWise (Muon) optimizer in ``dp_reshardable`` + format, for both the uniform (bf16, bf16) and the mixed-dtype (bf16, fp32) sibling + DistOpt buffers. Exercises ``sharded_param_state_dp_reshardable`` on the decouple path, + including the empty-bucket and global-plan-coverage handling. + """ + Utils.initialize_model_parallel(1, 1) # tp=pp=1 -> dp = world_size + metadata = {'distrib_optim_sharding_type': 'dp_reshardable'} + + def _build(seed): + return setup_model_and_optimizer( + seed=seed, + tp=1, + pp=1, + bf16=bf16, + dist_opt=True, + initialize_fn=initialize_gpt_model, + optimizer='dist_muon', + use_param_layout=True, + grad_reduce_in_fp32=grad_reduce_in_fp32, + ) + + with TempNamedDir( + tmp_path_dist_ckpt / 'test_layer_wise_dp_reshardable', sync=True + ) as ckpt_dir: + model_A, optimizer_A = _build(2) + model_sd = model_A[0].sharded_state_dict() + optim_sd = optimizer_A.sharded_state_dict(model_sd, metadata=metadata) + save(optim_sd, ckpt_dir) + + model_B, optimizer_B = _build(3) + model_sd_B = model_B[0].sharded_state_dict() + load_sd = optimizer_B.sharded_state_dict(model_sd_B, is_loading=True, metadata=metadata) + state_dict = load(load_sd, ckpt_dir) + optimizer_B.load_state_dict(state_dict) + Utils.destroy_model_parallel() + + @pytest.mark.parametrize('ep', [2, 4]) + def test_dp_reshardable_decouple_moe_ckpt(self, tmp_path_dist_ckpt, ep): + """dp_reshardable save/load round-trip of the decoupled compact LayerWise (Muon) + optimizer on an MoE model with expert parallelism and mixed-dtype (bf16 param / + fp32 grad) sibling DistOpt buffers. Single bucket per buffer (no ``ddp_bucket_size``), + so ownership is dense and the round-trip loads cleanly. + """ + Utils.initialize_model_parallel( + tensor_model_parallel_size=1, + pipeline_model_parallel_size=1, + expert_model_parallel_size=ep, + ) + metadata = {'distrib_optim_sharding_type': 'dp_reshardable'} + + def _build(seed): + return setup_moe_model_and_optimizer( + seed=seed, + tp=1, + pp=1, + ep=ep, + bf16=True, + dist_opt=True, + optimizer='dist_muon', + use_param_layout=True, + grad_reduce_in_fp32=True, + ) + + with TempNamedDir( + tmp_path_dist_ckpt / 'test_layer_wise_dp_reshardable_moe', sync=True + ) as ckpt_dir: + model_A, optimizer_A = _build(2) + model_sd = model_A[0].sharded_state_dict() + optim_sd = optimizer_A.sharded_state_dict(model_sd, metadata=metadata) + save(optim_sd, ckpt_dir) + + model_B, optimizer_B = _build(3) + model_sd_B = model_B[0].sharded_state_dict() + load_sd = optimizer_B.sharded_state_dict(model_sd_B, is_loading=True, metadata=metadata) + state_dict = load(load_sd, ckpt_dir) + optimizer_B.load_state_dict(state_dict) + Utils.destroy_model_parallel() + + def test_dp_reshardable_moe_synth_save(self, tmp_path_dist_ckpt): + """Regression for the empty-bucket-synth coverage gap in + ``DistributedOptimizer.sharded_param_state_dp_reshardable``. + + A small ``ddp_bucket_size`` splits the sibling DistOpt buffer of the decoupled + compact LayerWise (Muon) layout into many small buckets. Params are 64-element + aligned inside a bucket, so at DP=4 (local shard = 32) some DP rank owns a shard + that lies entirely inside inter-param padding while still overlapping + ``[0, gbuf_world_numel_unpadded)`` -- the empty-bucket-synth path. Set + ``DEBUG_DP_RESHARDABLE=1`` to see the per-bucket ``synth=True`` ``[DPRESH ...]`` + lines (this config fires it ~34x/rank-bucket). + + With the fix the synthesized padding ShardedTensor is stored back into the returned + ``state`` so the per-bucket global tensor is fully covered and torch DCP ``save`` + validates the global plan. Without the store-back line those shards are dropped and + ``save`` raises ``ValueError: Failed to validate global plan`` (chunks_volume < + tensor_volume). + + Note: this asserts the *save* half only. Extending it to a load round-trip trips a + separate, pre-existing ``dp_reshardable`` multi-bucket load defect (a + ``len(bucket_state) == len(param_map)`` mismatch in + ``load_parameter_state_from_dp_reshardable``) that also fires for a non-synth + multi-bucket config, so it is out of scope for this store-back fix. + """ + Utils.initialize_model_parallel( + tensor_model_parallel_size=1, + pipeline_model_parallel_size=1, + expert_model_parallel_size=2, + ) + metadata = {'distrib_optim_sharding_type': 'dp_reshardable'} + + with TempNamedDir(tmp_path_dist_ckpt / 'moe_synth_save', sync=True) as ckpt_dir: + model, optimizer = setup_moe_model_and_optimizer( + seed=2, + tp=1, + pp=1, + ep=2, + bf16=True, + dist_opt=True, + optimizer='dist_muon', + use_param_layout=True, + grad_reduce_in_fp32=True, + ddp_bucket_size=32, + ) + model_sd = model[0].sharded_state_dict() + optim_sd = optimizer.sharded_state_dict(model_sd, metadata=metadata) + # Fails with "Failed to validate global plan" if the synth store-back line + # in sharded_param_state_dp_reshardable is removed. + save(optim_sd, ckpt_dir) + Utils.destroy_model_parallel() + + # NOTE: 'fully_sharded_model_space' is intentionally NOT covered here. It is incompatible with + # the decoupled compact LayerWise (Muon) DistOpt layout regardless of precision: the sibling + # embedding / output_layer params produce a flattened_range ShardedTensor that + # dist_checkpointing rejects ('ShardedTensor.flattened_range is not supported'). This was + # confirmed to fail identically for bf16 and fp8 (so it is NOT fp8-specific), and it raises + # non-uniformly across DP ranks, so it cannot be asserted cleanly in a distributed test. It is + # a pre-existing limitation, independent of the dp_reshardable empty-bucket fix. + @pytest.mark.parametrize('fp8', [False, True]) + @pytest.mark.parametrize('sharding_type', ['dp_reshardable', 'fully_reshardable']) + def test_decouple_ckpt_roundtrip_values(self, tmp_path_dist_ckpt, sharding_type, fp8): + """Value-level save/load round-trip of the decoupled compact LayerWise (Muon) optimizer, + for the ``dp_reshardable`` and ``fully_reshardable`` sharding formats, for both bf16 and + quantized FP8 (MXFP8) model params. + + Correctness (not just "does not crash"): save optimizer A in ``sharding_type``, load it + into a differently seeded optimizer B, then assert B's optimizer state equals A's *by + value*. If the round-trip is faithful, B's fp32 master / exp_avg / exp_avg_sq must match + A's exactly. + + The comparison is done through a padding-free *canonical view*: both A's and (post-load) + B's in-memory state are re-serialized with ``fully_reshardable`` (model-centric, no + padding, deterministic) and compared bitwise via ``load_plain_tensors`` + ``check_equal``. + This is required because ``dp_reshardable``'s own bucket-space checkpoint serializes + inter-param / empty-shard padding as uninitialized ``torch.empty`` (values discarded on + load), so two saves of identical state differ in the padding bytes and cannot be compared + directly. ``fully_reshardable`` has no such padding, so it is a faithful canonical view of + the real optimizer state for both formats under test. + + The ``fp8`` axis covers the fp8 -> fp32-master path: FP8 changes model-param storage to a + Float8Tensor while optimizer state stays fp32, so the dequantize path + (``_is_distopt_quantized_param`` in distrib_optimizer.py) must round-trip. Both bf16 and + fp8 are exercised. + """ + # setup_moe_model_and_optimizer's fp8 path uses fp8_recipe='mxfp8', which needs + # Blackwell or newer; on older archs TE raises inside dequantize. Mirrors the arch + # guard in tests/unit_tests/test_muon_decouple_fp8_param_gather.py. + if fp8 and get_device_arch_version() < 10: + pytest.skip("mxfp8 requires Blackwell architecture or newer") + + from megatron.core.dist_checkpointing import load_plain_tensors + + Utils.initialize_model_parallel(1, 1) # tp=pp=1 -> dp = world_size + metadata = {'distrib_optim_sharding_type': sharding_type} + # Padding-free canonical view used to compare optimizer state by value. + canonical = {'distrib_optim_sharding_type': 'fully_reshardable'} + + def _build(seed): + kwargs = dict( + seed=seed, + tp=1, + pp=1, + bf16=True, + dist_opt=True, + initialize_fn=initialize_gpt_model, + optimizer='dist_muon', + use_param_layout=True, + grad_reduce_in_fp32=True, + ) + if fp8: + kwargs['fp8'] = True + return setup_model_and_optimizer(**kwargs) + + tag = f'{"fp8" if fp8 else "bf16"}_{sharding_type}' + with ( + TempNamedDir(tmp_path_dist_ckpt / f'{tag}_rt', sync=True) as rt_dir, + TempNamedDir(tmp_path_dist_ckpt / f'{tag}_A', sync=True) as canon_dir_A, + TempNamedDir(tmp_path_dist_ckpt / f'{tag}_B', sync=True) as canon_dir_B, + ): + # Save A in the format under test, load it into a differently seeded B. + model_A, optimizer_A = _build(2) + model_sd_A = model_A[0].sharded_state_dict() + save(optimizer_A.sharded_state_dict(model_sd_A, metadata=metadata), rt_dir) + + model_B, optimizer_B = _build(3) + model_sd_B = model_B[0].sharded_state_dict() + load_sd = optimizer_B.sharded_state_dict(model_sd_B, is_loading=True, metadata=metadata) + optimizer_B.load_state_dict(load(load_sd, rt_dir)) + + # Compare A vs post-load B by value, through the padding-free canonical view. + save(optimizer_A.sharded_state_dict(model_sd_A, metadata=canonical), canon_dir_A) + save(optimizer_B.sharded_state_dict(model_sd_B, metadata=canonical), canon_dir_B) + Utils.destroy_model_parallel() + + Utils.initialize_model_parallel(1, 1) + check_equal(load_plain_tensors(canon_dir_A), load_plain_tensors(canon_dir_B)) + Utils.destroy_model_parallel() + @pytest.mark.parametrize('tp', [1, 2, 4]) @pytest.mark.parametrize('pp', [1, 2, 4]) def test_layer_wise_optimizer_grad_norm(self, tp, pp): diff --git a/tests/unit_tests/dist_checkpointing/utils.py b/tests/unit_tests/dist_checkpointing/utils.py index 9a8b502eba0..e7716794b67 100644 --- a/tests/unit_tests/dist_checkpointing/utils.py +++ b/tests/unit_tests/dist_checkpointing/utils.py @@ -1,4 +1,4 @@ -# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. from functools import partial from typing import Any, Callable, Tuple, Union @@ -9,6 +9,7 @@ from megatron.core.dist_checkpointing.strategies.cached_metadata_filesystem_reader import ( CachedMetadataFileSystemReader, ) +from megatron.core.fp8_utils import is_float8tensor from megatron.core.models.gpt import GPTModel from megatron.core.models.gpt.gpt_layer_specs import ( get_gpt_layer_local_spec, @@ -28,7 +29,7 @@ def initialize_gpt_model( - pre_process=True, post_process=True, seed=0, use_glu=True, **config_kwargs + pre_process=True, post_process=True, seed=0, use_glu=True, fp8=False, **config_kwargs ): # These kwargs are passed through training.get_model for model construction, # but are not part of TransformerConfig; strip them before building config. @@ -45,11 +46,24 @@ def initialize_gpt_model( use_cpu_initialization=True, bf16=True, ) + if fp8: + # FP8 params only materialize with TransformerEngine layers, GPU init, and a + # hidden size that satisfies the MXFP8 32-element block-scaling alignment. The + # tiny (hidden=16, local-spec) default gives zero fp8 params. + default_config_kwargs.update( + hidden_size=128, + num_attention_heads=8, + use_cpu_initialization=False, + fp8='e4m3', + fp8_recipe='mxfp8', + fp8_param=True, + ) default_config_kwargs.update(**config_kwargs) transformer_config = TransformerConfig(**default_config_kwargs, gated_linear_unit=use_glu) + spec = get_gpt_layer_with_transformer_engine_spec() if fp8 else get_gpt_layer_local_spec() model = GPTModel( config=transformer_config, - transformer_layer_spec=get_gpt_layer_local_spec(), + transformer_layer_spec=spec, vocab_size=128, max_sequence_length=4, pre_process=pre_process, @@ -58,6 +72,10 @@ def initialize_gpt_model( with torch.no_grad(): for p in model.parameters(): + # Float8Tensor params own quantized storage; skip the plain random_ init + # (embeddings / layernorms remain plain tensors and are still randomized). + if is_float8tensor(p): + continue p.random_() return model @@ -191,6 +209,9 @@ def setup_model_and_optimizer( ep=1, etp=1, use_megatron_fsdp=False, + ddp_bucket_size=None, + grad_reduce_in_fp32=False, + fp8=False, ): optimizer_type = optimizer use_layer_wise = False @@ -225,10 +246,24 @@ def setup_model_and_optimizer( mock_args.megatron_fsdp_main_grads_dtype = None mock_args.megatron_fsdp_grad_comm_dtype = None mock_args.gradient_accumulation_fusion = False + mock_args.ddp_bucket_size = ddp_bucket_size + # grad_reduce_in_fp32 -> ddp_config grad_dtype=fp32 while params stay bf16, i.e. the + # mixed-dtype (bf16, fp32) gradient buffer / optimizer-state bucket. + mock_args.accumulate_allreduce_grads_in_fp32 = grad_reduce_in_fp32 mock_args.use_distributed_optimizer = ddp_use_dist_opt mock_args.use_layer_wise_distributed_optimizer = ddp_use_layer_wise if ddp_use_layer_wise: mock_args.optimizer = optimizer + # Only forward ``fp8`` to initialize_fn when enabled, so existing callers + # passing an initialize_fn without an ``fp8`` parameter keep working. + extra_init_kwargs = {} + if fp8: + # Make DDP build an fp8 param-gather buffer (MXFP8 reuses the grad buffer + # for the param all-gather). ``fp8_param`` on the TransformerConfig (set in + # ``initialize_gpt_model``) is what actually quantizes the model weights. + mock_args.fp8_param_gather = True + mock_args.reuse_grad_buf_for_mxfp8_param_ag = True + extra_init_kwargs['fp8'] = True model = get_model( partial( initialize_fn, @@ -240,9 +275,17 @@ def setup_model_and_optimizer( expert_model_parallel_size=ep, expert_tensor_parallel_size=etp, bf16=bf16, + **extra_init_kwargs, ) ) + if fp8: + # Guard against a silent config regression: if the fp8/TE path stops producing + # quantized weights the checkpoint round-trip would no longer exercise the + # fp8 -> fp32-master mapping we mean to test. + num_fp8_params = sum(1 for m in model for p in m.parameters() if is_float8tensor(p)) + assert num_fp8_params > 0, "fp8=True but no Float8Tensor params were created" + config = OptimizerConfig( bf16=bf16, params_dtype=torch.bfloat16 if bf16 else torch.float, @@ -341,6 +384,8 @@ def setup_moe_model_and_optimizer( use_glu=False, optimizer='adam', use_param_layout=False, + ddp_bucket_size=None, + grad_reduce_in_fp32=False, ): optimizer_type = optimizer use_layer_wise = False @@ -357,6 +402,17 @@ def setup_moe_model_and_optimizer( mock_args = parse_args(ignore_unknown_args=True) with mock.patch('megatron.training.training.get_args', new=lambda: mock_args): init_basic_mock_args(mock_args, tp, pp, bf16=bf16) + mock_args.ddp_bucket_size = ddp_bucket_size + # ``resolve_ddp_bucket_size`` (megatron/training/training.py) discards an explicit + # bucket_size unless overlap_grad_reduce is on -- otherwise the whole buffer is a + # single bucket. Turning overlap on (no backward is run here) lets ddp_bucket_size + # split the sibling DistOpt buffer into many small buckets so some DP rank owns a + # shard made only of inter-param alignment padding -> the empty-bucket-synth path. + if ddp_bucket_size is not None: + mock_args.overlap_grad_reduce = True + # grad_reduce_in_fp32 -> ddp_config grad_dtype=fp32 while params stay bf16, i.e. the + # mixed-dtype (bf16, fp32) gradient buffer / optimizer-state bucket. + mock_args.accumulate_allreduce_grads_in_fp32 = grad_reduce_in_fp32 mock_args.use_distributed_optimizer = ddp_use_dist_opt mock_args.use_layer_wise_distributed_optimizer = ddp_use_layer_wise if ddp_use_layer_wise: @@ -392,12 +448,25 @@ def setup_moe_model_and_optimizer( torch.manual_seed(seed + 1) model_parallel_cuda_manual_seed(seed + 1) + def _init_states(opt): + # In the decoupled compact LayerWise + DistOpt layout the top-level + # ChainedOptimizer wraps another ChainedOptimizer (LayerWise, which has no + # ``init_state_fn``) alongside a sibling DistOpt; recurse so the Muon + # Float16 sub-optimizers still get seeded, and skip optimizers without + # ``init_state_fn`` (DistOpt seeds its state elsewhere). + if isinstance(opt, ChainedOptimizer): + for child in opt.chained_optimizers: + _init_states(child) + return + if not hasattr(opt, 'init_state_fn'): + return + if not hasattr(opt, 'optimizer'): + opt.init_state_fn(opt) + else: + opt.init_state_fn(opt.optimizer) + if optimizer_type in ('muon', 'dist_muon'): - for opt in optimizer.chained_optimizers: - if not hasattr(opt, 'optimizer'): - opt.init_state_fn(opt) - else: - opt.init_state_fn(opt.optimizer) + _init_states(optimizer) else: for opt in optimizer.chained_optimizers: for group in opt.param_groups: From b4ba39fbbcd16b3f440b9154a060ae5961262387 Mon Sep 17 00:00:00 2001 From: Pingtian Li Date: Wed, 29 Jul 2026 23:28:53 +0800 Subject: [PATCH 02/12] Support FP8 parameter gather on the decoupled compact LayerWise layout Lets Muon persist and sync its matrix parameters as FP8 (mxfp8 on Blackwell, blockwise on Hopper) instead of BF16, removing the standing BF16 parameter copy. Combined with the compact layout this closes the Muon-vs-Adam memory gap while staying bit-for-bit equal to the fp8-param-gather-OFF run. FP8 tensors cannot be all-gathered directly (mxfp8's row/col block scales cannot be derived from one another), so the whole-param transport rides BF16 and requantizes on landing: stage each owned param from its fp32 master to bf16 (`_stage_param_to_bf16` -- a high-precision source, not a lossy fp8 dequant), uneven all-gather-v, then requantize the gathered bf16 into *every* rank's fp8 `param.data` so all ranks hold `Q(bf16(master))`; `post_all_gather_processing` rebuilds the fp8 columnwise/transpose. The overlap (`start_param_sync`/`finish_param_sync`) and non-overlap (`allgather_params`) paths share one copy-back helper, and the staging decision is persisted per bucket so transport dtype and copy-back cannot disagree. Params are dispatched by transport dtype: fp8 and bf16 ride the bf16-staged helper, while native fp32 params (`keep_in_fp32`, e.g. the DeepSeek-V4 CSA `ape`) are gathered in fp32 -- the bf16 path would silently downcast them. Also in this commit: * Rank-independent ping-pong ownership. `allgather_params` needs every DP rank to agree on each param's single owner; `numel` alone is not a total order, so ownership is keyed by a canonical `(chunk_idx, buffer_idx, global_start_index)` identity. * iter-0 master parity: the fp32 master is seeded from the high-precision pre-quantization init, and the master->model copy routes through bf16. Gathered fp8 params are tagged and skipped there, since the all-gather's requantize already wrote `Q(bf16(master))`; non-gathered fp8 params (MoE experts at `expt_dp == 1`) still get their copy. * Non-owned params drop TE's high-precision init copy, which every DP rank would otherwise retain for ~(dp-1)/dp of the LayerWise matrix params for the whole run. * `force_sync` finalizes a pending LayerWise gather it would otherwise bypass, which would leave stale `param.data` and pollute the reused `grad_data`. * Single all-reduce buffer: keying fp8 Muon grads by `torch.uint8` split their gradients into two all-reduce buffers where OFF has one; NCCL's fp32 accumulation order is buffer-size sensitive, so the split diverged ~1 ULP from OFF and Newton-Schulz amplified it into visible loss drift. fp8 Muon grads now key to their bf16 logical dtype and share one buffer. A consequence is that a bucket can hold both fp8 and bf16 params, so the staging decision scans the whole bucket rather than `params_list[0]`. * Validation: fp8 param gather is accepted only on this layout, requires `fp8_recipe in {mxfp8, blockwise}` (fp4 rejected outright), mxfp8 additionally requires `--reuse-grad-buf-for-mxfp8-param-ag`, and the layout requires `num_distributed_optimizer_instances == 1`. Golden values for the `dist_dist_muon` ckpt-resume cases are refreshed at this tip (8xH100 EP8, 2x GB200 EP8, 1x GB200 EP4), regenerated with CI's own harness on the common_pile dataset. Each reported 'Exact: FAILED / APPROXIMATE: PASSED' beforehand -- a ~1 ULP reordering from the rank-independent ownership sort, not a regression; peak relative delta on lm loss is ~1.6e-4. Co-Authored-By: Claude Opus 5 (1M context) Signed-off-by: Pingtian Li --- .../distributed/distributed_data_parallel.py | 14 +- .../core/distributed/param_and_grad_buffer.py | 209 ++-- megatron/core/fp8_utils.py | 27 +- megatron/core/optimizer/distrib_optimizer.py | 6 +- .../core/optimizer/layer_wise_optimizer.py | 224 ++++- megatron/core/optimizer/optimizer.py | 65 +- megatron/training/arguments.py | 51 +- megatron/training/training.py | 6 +- .../golden_values_dev_dgx_gb200.json | 934 ++++++++--------- .../golden_values_dev_dgx_h100.json | 940 +++++++++--------- .../golden_values_dev_dgx_gb200.json | 750 +++++++------- .../test_muon_decouple_fp8_param_gather.py | 453 +++++++++ 12 files changed, 2258 insertions(+), 1421 deletions(-) create mode 100644 tests/unit_tests/test_muon_decouple_fp8_param_gather.py diff --git a/megatron/core/distributed/distributed_data_parallel.py b/megatron/core/distributed/distributed_data_parallel.py index 68f020fb638..90d130dedd8 100644 --- a/megatron/core/distributed/distributed_data_parallel.py +++ b/megatron/core/distributed/distributed_data_parallel.py @@ -1,4 +1,4 @@ -# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved. +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. import logging from contextlib import contextmanager @@ -123,8 +123,16 @@ def __init__( param_to_name[param] = name all_params.append(param) - # Group parameters by (param_dtype, grad_dtype, is_expert_parallel). - buffer_groups = group_params_for_buffers(all_params, self.ddp_config.grad_reduce_in_fp32) + # Group parameters by (param_dtype, grad_dtype, is_expert_parallel). fp8 params key to + # uint8 (own buffer); partition_buckets later merges the small non-fp8 bucket groups into + # the fp8 group to aggregate their communication. + buffer_groups = group_params_for_buffers( + all_params, + self.ddp_config.grad_reduce_in_fp32, + merge_layerwise_fp8_grads=not getattr( + self.ddp_config, 'use_layer_wise_param_layout', True + ), + ) # Auto-compute layouts when using distributed optimizer but no layout was provided. # This maintains backward compatibility for callers that create DDP directly diff --git a/megatron/core/distributed/param_and_grad_buffer.py b/megatron/core/distributed/param_and_grad_buffer.py index 35eeff100bf..2faeb02b7e3 100644 --- a/megatron/core/distributed/param_and_grad_buffer.py +++ b/megatron/core/distributed/param_and_grad_buffer.py @@ -1,4 +1,4 @@ -# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved. +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. import dataclasses import fnmatch @@ -29,6 +29,8 @@ modify_nvfp4_rowwise_storage, ) from ..fp8_utils import ( + _stage_param_to_bf16, + copy_back_gathered_bf16_into_fp8_param, copy_tensor_to_quantized_param, is_float8tensor, is_grouped_mxfp8tensor, @@ -177,6 +179,36 @@ def wait(self): self.handles = None +@torch.no_grad() +def _layerwise_copy_back_gathered_params(bucket, local_rank: int, fp8_staged: bool = False) -> None: + """Copy each rank's gathered params from the bucket gather slots into model params (non-DistOpt + LayerWise overlap path). ``fp8_staged`` MUST match ``start_param_sync``'s staging decision. + + * bf16 (``fp8_staged=False``): unflatten against the params, ``copy_`` into non-owned + ``model_p.data`` (owned already hold the staged value). + * fp8 (``fp8_staged=True``): the all-gather rode bf16; requantize ALL ranks (owned included) + via ``copy_back_gathered_bf16_into_fp8_param`` so every owner holds ``Q(bf16(master))``. + + no_grad: in-place copy_ on a leaf param trips autograd's in-place guard. + """ + for idx, params in enumerate(bucket.layerwise_params_list): + if len(params) == 0: + continue + if fp8_staged: + templates = [torch.empty(p.shape, device="meta", dtype=torch.bfloat16) for p in params] + updated_params = _unflatten_dense_tensors(bucket.layerwise_gather_list[idx], templates) + for updated_p, model_p in zip(updated_params, params): + copy_back_gathered_bf16_into_fp8_param(model_p, updated_p) + continue + # bf16 transport: owned params already hold the staged bf16 value in their data, so only + # non-owned ranks need the copy. + if idx == local_rank: + continue + updated_params = _unflatten_dense_tensors(bucket.layerwise_gather_list[idx], params) + for updated_p, model_p in zip(updated_params, params): + model_p.data.copy_(updated_p) + + class _ParamAndGradBucketGroup: """ Put multiple buckets into a group so that their communications can be aggregated together. @@ -299,14 +331,41 @@ def reset(self): self.is_last_microbatch = True self.grad_reduce_finished = False + def _finalize_layerwise_param_sync(self): + """Copy gathered LayerWise (non-DistOpt) params back and release the reused grad buffer. + + Every path that completes a LayerWise param all-gather must run this before + ``_post_param_sync``: the gathered (possibly bf16-staged fp8) whole params sit in + ``bucket.layerwise_gather_list`` (views into ``grad_data``) until they are unflattened + into ``param.data``, and ``grad_data`` must be re-zeroed afterwards so the next + backward's accumulation into ``main_grad`` does not start from the gather payload. + """ + if self.ddp_config.use_distributed_optimizer: + return + for bucket in self.buckets: + if bucket.layerwise_gather_list is None: + continue + # Unflatten and copy gathered params for each rank (FP8-aware: see helper). + _layerwise_copy_back_gathered_params( + bucket, + self.intra_distributed_optimizer_instance_rank, + fp8_staged=getattr(bucket, 'layerwise_fp8_staged', False), + ) + bucket.layerwise_gather_list = None + # Zero out grad_data since it was reused as the all-gather + # receive buffer. Without this, accumulation into main_grad + # (a view into grad_data) would start from the result of the + # latest parameter all-gather instead of zero. + bucket.grad_data.zero_() + def _post_param_sync(self): """Run post-processing after param all-gather completes.""" if self.ddp_config.reuse_grad_buf_for_mxfp8_param_ag: for bucket in self.buckets: if bucket.param_data is None: - # LayerWise variable-size gather path: params are already updated via - # unflatten + copy_ in finish_param_sync, and there is no param_data - # buffer to copy back from. + # LayerWise variable-size gather path (incl. the non-DistOpt decoupled + # bucket): params are already updated via unflatten + copy_ / requantize in + # finish_param_sync, and there is no param_data buffer to copy back from. continue has_non_quantized_weight = False for param in bucket.params: @@ -395,6 +454,10 @@ def start_param_sync(self, force_sync: bool = False): if self.param_gather_handle is not None: self.param_gather_handle.wait() self.param_gather_handle = None + # A pending LayerWise gather normally lands in finish_param_sync (forward + # pre-hook), which force_sync bypasses -- finalize it here or every rank + # keeps stale ``param.data`` and the gather payload pollutes ``grad_data``. + self._finalize_layerwise_param_sync() self._post_param_sync() return else: @@ -422,10 +485,33 @@ def start_param_sync(self, force_sync: bool = False): local_rank = self.intra_distributed_optimizer_instance_rank group = self.intra_distributed_optimizer_instance_group layerwise_work_handles = [] + # Decouple fp8 param-gather: model params are Float8/MXFP8 but the all-gather rides bf16 + # — stage owned to bf16, gather, requantize on copy-back. Plain bf16 collapses to the + # original path. + decouple = not getattr(self.ddp_config, 'use_layer_wise_param_layout', True) for bucket in self.buckets: - # Use param dtype (e.g., bf16), NOT grad dtype (which may be - # fp32 when grad_reduce_in_fp32 is enabled). - param_dtype = bucket.params_list[0].dtype + # A decoupled LayerWise (Muon) bucket can MIX fp8 and bf16 params: + # merge_layerwise_fp8_grads keys fp8 Muon grads by their bf16 logical dtype so they + # share ONE buffer (hence bucket) with their bf16 siblings (e.g. an MoE router / + # DSA indexer / mHC weight that is not fp8-quantized). The bf16-staged path handles + # both dtypes, so a bucket holding ANY fp8 param must take it -- scanning only + # params_list[0] mis-routes a bf16-first mixed bucket into the raw + # _flatten_dense_tensors() path, which crashes on the MXFP8 .view(-1). + bucket_is_fp8 = bool( + decouple + and bucket.params_list + and any(is_float8tensor(p) for p in bucket.params_list) + ) + # TODO(perf, blockwise-only): blockwise could gather the owner's fp8 rowwise data + # (~2x less comm) + its small scale_inv and rebuild columnwise via transpose + # (Adam/DistOpt-style), instead of bf16. mxfp8 must stay on bf16: its row/col block + # scales cannot be derived from one another. + # + # Persist the staging decision so the copy-back (sync here, overlap in + # finish_param_sync) uses the same signal, keeping transport and copy-back in sync. + bucket.layerwise_fp8_staged = bucket_is_fp8 + # Transport dtype: bf16 for decouple fp8 param-gather; else the param's own dtype. + param_dtype = torch.bfloat16 if bucket_is_fp8 else bucket.params_list[0].dtype if max(bucket.layerwise_param_flat_sizes) == 0: bucket.layerwise_gather_list = None @@ -452,21 +538,30 @@ def start_param_sync(self, force_sync: bool = False): # Detach from autograd since start_param_sync may be called # during the forward pass where autograd is active. if local_size > 0: - # MXFP8 params can't be flattened (view(-1) unsupported); gather the - # fp32 master (param.main_param -> bf16), which the receive-side copy_ - # re-quantizes. Non-mxfp8 params flatten as-is. - src_params = [] - for p in bucket.layerwise_params_list[local_rank]: - if is_mxfp8tensor(p): - main_param = getattr(p, "main_param", None) - assert main_param is not None, ( - "LayerWise mxfp8 param sync needs param.main_param (fp32 " - "master) to stage the all-gather source; got None." - ) - src_params.append(main_param.to(param_dtype)) - else: - src_params.append(p) - flat_local_params = _flatten_dense_tensors(src_params).detach() + if bucket_is_fp8: + # Decoupled layout: stage fp32 master->bf16 (high-precision source), not + # lossy dequant(fp8). Copy-back requantizes every rank, owner included. + staged = [ + _stage_param_to_bf16(p) + for p in bucket.layerwise_params_list[local_rank] + ] + flat_local_params = _flatten_dense_tensors(staged) + else: + # Padded LayerWise layout: MXFP8 params can't be flattened (view(-1) + # unsupported); gather the fp32 master (param.main_param -> bf16), which + # the receive-side copy_ re-quantizes. Non-mxfp8 params flatten as-is. + src_params = [] + for p in bucket.layerwise_params_list[local_rank]: + if is_mxfp8tensor(p): + main_param = getattr(p, "main_param", None) + assert main_param is not None, ( + "LayerWise mxfp8 param sync needs param.main_param (fp32 " + "master) to stage the all-gather source; got None." + ) + src_params.append(main_param.to(param_dtype)) + else: + src_params.append(p) + flat_local_params = _flatten_dense_tensors(src_params).detach() local_slot_view.copy_(flat_local_params) bucket.layerwise_gather_list = gather_list @@ -480,23 +575,7 @@ def start_param_sync(self, force_sync: bool = False): self.param_gather_handle = _LayerwiseAllGatherHandle(layerwise_work_handles) else: # Synchronous: unflatten and copy gathered params immediately. - for bucket in self.buckets: - if bucket.layerwise_gather_list is None: - continue - for idx, params in enumerate(bucket.layerwise_params_list): - if len(params) == 0 or idx == local_rank: - continue - updated_params = _unflatten_dense_tensors( - bucket.layerwise_gather_list[idx], params - ) - for updated_p, model_p in zip(updated_params, params): - model_p.data.copy_(updated_p) - bucket.layerwise_gather_list = None - # Zero out grad_data since it was reused as the all-gather - # receive buffer. Without this, accumulation into main_grad - # (a view into grad_data) would start from the result of the - # latest parameter all-gather instead of zero. - bucket.grad_data.zero_() + self._finalize_layerwise_param_sync() self.param_gather_handle = None else: # Standard distributed optimizer path: use _coalescing_manager. @@ -567,29 +646,7 @@ def finish_param_sync(self, skip_next_bucket_dispatch: bool = False): else: self.next_param_gather_bucket_group.start_param_sync() - if not self.ddp_config.use_distributed_optimizer: - for bucket in self.buckets: - if bucket.layerwise_gather_list is None: - continue - # Unflatten and copy gathered params for each rank. - for idx, params in enumerate(bucket.layerwise_params_list): - # Skip local params and empty tensors. - if ( - len(params) == 0 - or idx == self.intra_distributed_optimizer_instance_rank - ): - continue - updated_params = _unflatten_dense_tensors( - bucket.layerwise_gather_list[idx], params - ) - for updated_p, model_p in zip(updated_params, params): - model_p.data.copy_(updated_p) - bucket.layerwise_gather_list = None - # Zero out grad_data since it was reused as the all-gather - # receive buffer. Without this, accumulation into main_grad - # (a view into grad_data) would start from the result of the - # latest parameter all-gather instead of zero. - bucket.grad_data.zero_() + self._finalize_layerwise_param_sync() self._post_param_sync() def start_grad_sync(self, force_all_reduce: Optional[bool] = False): @@ -837,6 +894,10 @@ def free_overlap_buffers(self): if self.param_gather_handle is not None: self.param_gather_handle.wait() self.param_gather_handle = None + # Finalize a pending LayerWise gather before dropping its receive views: + # discarding layerwise_gather_list here would leave stale ``param.data`` + # (checkpointed weights) and gather payload in ``grad_data``. + self._finalize_layerwise_param_sync() for bucket in self.buckets: bucket.layerwise_gather_list = None @@ -881,7 +942,9 @@ def register_grad_ready( def group_params_for_buffers( - params: List[torch.nn.Parameter], grad_reduce_in_fp32: bool + params: List[torch.nn.Parameter], + grad_reduce_in_fp32: bool, + merge_layerwise_fp8_grads: bool = False, ) -> Dict['BufferKey', Tuple[List[torch.nn.Parameter], List[int]]]: """Group parameters by buffer identity for buffer allocation. @@ -899,6 +962,8 @@ def group_params_for_buffers( Args: params: List of parameters to group. grad_reduce_in_fp32: Whether gradients are reduced in FP32. + merge_layerwise_fp8_grads: Decouple layout only — merge LayerWise (Muon) fp8 grads with + their bf16 siblings into one fp32 all_reduce buffer (see below). Returns: Dict mapping BufferKey to (params_list, param_indices). @@ -921,6 +986,15 @@ def group_params_for_buffers( param, 'is_managed_by_layer_wise_optimizer', False ) + # Decouple layout only: key fp8 Muon grads by their bf16 logical dtype so fp8 + bf16 grads + # share ONE fp32 all_reduce buffer; a split uint8/bf16 reduction diverges ~1 ULP from OFF. + if ( + merge_layerwise_fp8_grads + and is_float8tensor(param) + and is_managed_by_layer_wise_optimizer + ): + param_dtype = param.dtype + key = BufferKey( param_dtype, grad_dtype, is_expert_parallel, is_managed_by_layer_wise_optimizer ) @@ -1715,6 +1789,10 @@ def partition_buckets( has completed. This is because we need to wait for the non-fp8 params from the beginning layers to obtain their gradients. - Combining the non-fp8 bucket with the last fp8 bucket can help avoid this issue. + - A bucket group runs one collective type, so only buckets agreeing on the effective + per-buffer ``use_distributed_optimizer`` are merged; non-fp8 buckets with a different + value (the decouple-LayerWise path) go to their own group(s). When all buckets agree, + this collapses to the original behavior. Args: buffers (list): list of input buffers. @@ -1740,8 +1818,8 @@ def partition_buckets( # all-reduce otherwise), so buckets merged into one group must agree on the effective # per-buffer ``use_distributed_optimizer``. The decoupled LayerWise layout # (``use_layer_wise_param_layout=False``) gives LayerWise (Muon) buffers - # ``use_distributed_optimizer=False`` while sibling buffers keep True; the no-fp8 Case 2 below - # keeps every bucket in its own group so they never mix, but the merging Cases 1/3 must assert + # ``use_distributed_optimizer=False`` while sibling buffers keep True; the no-fp8 branch below + # keeps every bucket in its own group so they never mix, but the merging branches must assert # consistency. _ddp_config = buffers[0].ddp_config _decouple = not getattr(_ddp_config, "use_layer_wise_param_layout", True) @@ -1874,4 +1952,5 @@ def _merged_use_distributed_optimizer(merge_buckets): buffer.data_parallel_world_size, ) ) + return bucket_groups diff --git a/megatron/core/fp8_utils.py b/megatron/core/fp8_utils.py index 5411b676d83..d8a2a343fe5 100644 --- a/megatron/core/fp8_utils.py +++ b/megatron/core/fp8_utils.py @@ -1,4 +1,4 @@ -# Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved. +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. """Utility functions related to FP8 that are used throughout Megatron core""" @@ -268,6 +268,31 @@ def dequantize_fp8_tensor(fp8_tensor: torch.Tensor) -> torch.Tensor: return fp8_tensor.from_float8() +def copy_back_gathered_bf16_into_fp8_param(model_p: torch.Tensor, src_bf16: torch.Tensor) -> None: + """Requantize a gathered bf16 whole-param into an fp8 (Float8/MXFP8) model param in place. + + mxfp8 columnwise can't be derived from rowwise, so force columnwise before copy_ (TE rebuilds + both directions from the bf16); blockwise/Float8 columnwise is a lossless transpose. + """ + if is_mxfp8tensor(model_p): + quantizer = model_p.data._get_quantizer() + quantizer.set_usage(rowwise=True, columnwise=True) + model_p.data.copy_(src_bf16) + + +def _stage_param_to_bf16(p: torch.Tensor) -> torch.Tensor: + """Stage a param to a detached bf16 whole-param for fp8 param-gather transport. + + Prefer the fp32 master (high-precision source); else dequantize an fp8 param; else copy bf16. + """ + main_param = getattr(p, "main_param", None) + if main_param is not None: + return main_param.detach().to(torch.bfloat16) + if is_float8tensor(p): + return dequantize_fp8_tensor(p).detach().to(torch.bfloat16) + return p.detach().to(torch.bfloat16) + + def _resolve_callable_from_python_import_path(dotted_path: str): """Resolve a Python import path like 'pkg.mod.func' to a callable. diff --git a/megatron/core/optimizer/distrib_optimizer.py b/megatron/core/optimizer/distrib_optimizer.py index 50b2cf00f93..e5417e8318d 100644 --- a/megatron/core/optimizer/distrib_optimizer.py +++ b/megatron/core/optimizer/distrib_optimizer.py @@ -633,7 +633,11 @@ def compute_full_param_layout( Returns: FullParamLayout with a PerBufferParamLayout per buffer group. """ - buffer_groups = group_params_for_buffers(params, ddp_config.grad_reduce_in_fp32) + buffer_groups = group_params_for_buffers( + params, + ddp_config.grad_reduce_in_fp32, + merge_layerwise_fp8_grads=not getattr(ddp_config, 'use_layer_wise_param_layout', True), + ) layouts = {} for buffer_key, (group_params, param_indices) in buffer_groups.items(): if buffer_key.is_expert_parallel: diff --git a/megatron/core/optimizer/layer_wise_optimizer.py b/megatron/core/optimizer/layer_wise_optimizer.py index b9f1b77e539..624fde5fe7d 100644 --- a/megatron/core/optimizer/layer_wise_optimizer.py +++ b/megatron/core/optimizer/layer_wise_optimizer.py @@ -1,4 +1,4 @@ -# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. import logging import math @@ -14,6 +14,12 @@ 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 ( + _stage_param_to_bf16, + copy_back_gathered_bf16_into_fp8_param, + is_float8tensor, + post_all_gather_processing, +) from .clip_grads import count_zeros_fp32, get_grad_norm_fp32 from .optimizer import ( ChainedOptimizer, @@ -71,6 +77,16 @@ def _bucket_is_managed_by_layer_wise_optimizer(bucket, default_for_untagged: boo return param.is_managed_by_layer_wise_optimizer +def _param_sort_key(numel: int, identity: tuple) -> tuple: + """Rank-independent total-order key for ping-pong ownership: ``(numel, *canonical-identity)``. + + ``numel`` alone is not a total order (stable sort tie-breaks by rank-local insertion order), + so equal-numel params would get different owners across ranks; the canonical identity + ``(chunk_idx, buffer_idx, global_start_index)`` makes it total and identical on every rank. + """ + return (numel,) + tuple(identity) + + def tag_params_for_buffer_routing(model_chunks) -> None: """Tag every requires-grad param with ``is_managed_by_layer_wise_optimizer``. @@ -405,7 +421,13 @@ def compute_full_param_layout( # layout here. Non-LayerWise buffers keep DistOpt's byte-level layout regardless. decouple_ddp_layout = not ddp_config.use_layer_wise_param_layout - buffer_groups = group_params_for_buffers(params, ddp_config.grad_reduce_in_fp32) + # fp8 Muon grads key to uint8 (own buffer); partition_buckets later merges the non-fp8 + # bucket groups into the fp8 group to aggregate communication. + buffer_groups = group_params_for_buffers( + params, + ddp_config.grad_reduce_in_fp32, + merge_layerwise_fp8_grads=not getattr(ddp_config, 'use_layer_wise_param_layout', True), + ) layouts = {} for buffer_key, (group_params, param_indices) in buffer_groups.items(): if buffer_key.is_expert_parallel: @@ -421,8 +443,8 @@ def compute_full_param_layout( # layout; non-LayerWise params (e.g. Adam-managed embeddings, biases) # get DistOpt's byte-level layout. if buffer_key.is_managed_by_layer_wise_optimizer and decouple_ddp_layout: - # Compact no-padding layout (DDP treats this buffer as non-DistOpt). Attach - # param_indices so DDP's layout/grouping consistency check passes. + # Decouple path (incl. FP8 param-gather): compact no-padding layout (DDP treats this + # buffer as non-DistOpt). Attach param_indices so DDP's consistency check passes. per_buffer_layout = _compute_default_per_buffer_param_layout( group_params, bucket_size ) @@ -490,7 +512,28 @@ def __init__( for chunk in model_chunks if hasattr(chunk, 'full_param_layout') and chunk.full_param_layout is not None ] or None - self.shard_params(optimizers, full_param_layouts) + # Decouple path keeps whole-matrix ping-pong ownership (Newton-Schulz runs on whole + # matrices on one rank; param sync via ``allgather_params``). ``model_chunks`` lets the + # ping-pong fallback break equal-numel ties by a rank-independent identity (see below). + self.shard_params(optimizers, full_param_layouts, model_chunks) + + # Engage FP8 param sync automatically when the decouple-managed params are actually + # quantized (fp8_param_gather on + TE Float8/MXFP8 weights). Off -> plain bf16 path. + # Also tag the gathered fp8 params: the fp8 all-gather (``_allgather_helper_fp8``) + # requantizes bf16 -> each rank's fp8 ``param.data``, so the child optimizer's pre-gather + # fp8 copy-back into ``param.data`` is redundant for them and is skipped. Params in these + # per-rank lists are all-gathered (dp_cp / expt_dp size > 1 here); non-gathered fp8 params + # (e.g. expt_dp == 1 experts, which are absent from these lists) still need the copy-back. + self.use_fp8_param_sync = False + if self.decouple_ddp_layout: + for params_list in (self.dp_cp_params_list, self.expt_dp_params_list): + if not params_list: + continue + for per_rank in params_list: + for p in per_rank: + if is_float8tensor(p): + self.use_fp8_param_sync = True + p._layer_wise_fp8_gathered = True # When a full_param_layout is available (no decoupling), use_distributed_optimizer # is True and model params are views into the DDP param buffer. After the @@ -531,6 +574,27 @@ def __init__( optimizers[i] = Float16OptimizerWithFloat16Params( opt, config, None, init_state_fn_list[i] if init_state_fn_list else None ) + # Non-DistOpt LayerWise child has no byte-shard param buffer. Tag it so + # ``_copy_main_params_to_model_params`` routes fp8 weights through the bf16 + # round-trip (``Q(bf16(master))``), matching the fp8-param-gather-OFF baseline; + # the gather staging itself reuses the grad buffer. + optimizers[i]._layer_wise_non_distopt_child = True + + # shard_params() removed non-owned params from the local optimizer groups, so the + # Float16 wrapping above only clears the TE high-precision init copy (a full-size CPU + # tensor per fp8 param) for locally owned params. Without this sweep every DP rank + # retains ~(dp-1)/dp of the LayerWise matrix params' bf16 CPU copies for the whole + # run. The per-rank ownership lists cover all gathered params (owned entries were + # already cleared during master creation; TE's clear is a no-op then). Scoped to + # LayerWise-managed params only: sibling DistOpt params must keep their init val + # until their own optimizer's master creation consumes it. + for params_list in (self.dp_cp_params_list, self.expt_dp_params_list): + if not params_list: + continue + for per_rank_params in params_list: + for p in per_rank_params: + if hasattr(p, 'clear_high_precision_init_val'): + p.clear_high_precision_init_val() super().__init__(optimizers) @@ -548,7 +612,7 @@ def __init__( # This way each rank do some duplicated work but allgather_v is no longer needed # All current distopt optimization can also be potentially applied - def shard_params(self, optimizers, full_param_layouts=None): + def shard_params(self, optimizers, full_param_layouts=None, model_chunks=None): """Shard params across ranks according to the computed param layout. Each param's shard assignment is derived from the :class:`FullParamLayout` @@ -577,7 +641,7 @@ def shard_params(self, optimizers, full_param_layouts=None): if full_param_layouts is not None: self._shard_params_from_layout(optimizers, full_param_layouts, dp_cp_size, expt_dp_size) else: - self._shard_params_ping_pong(optimizers, dp_cp_size, expt_dp_size) + self._shard_params_ping_pong(optimizers, dp_cp_size, expt_dp_size, model_chunks) def _shard_params_from_layout(self, optimizers, full_param_layouts, dp_cp_size, expt_dp_size): """Derive shard assignments from the param layout.""" @@ -647,16 +711,44 @@ def _shard_params_from_layout(self, optimizers, full_param_layouts, dp_cp_size, if expt_dp_size == 1 or len(self.expt_dp_params_list[0]) == 0: self.expt_dp_params_list = None - def _shard_params_ping_pong(self, optimizers, dp_cp_size, expt_dp_size): - """Legacy ping-pong-by-numel shard assignment (no layout available). + def _build_param_sort_keys(self, model_chunks): + """Build ``{param: (chunk_idx, buffer_idx, global_start_index)}`` — a rank-independent key + for every requires-grad param. + + Both the chunk/buffer enumeration order and the ``param_index_map`` offsets come purely + from model construction (identical across DP ranks), so the key is the same on every rank. + Used to break equal-numel ties in ``_shard_params_ping_pong``. Returns ``None`` if no layout + info is available, so the caller falls back to legacy numel-only ordering. + """ + if model_chunks is None: + return None + identity: Dict[torch.nn.Parameter, tuple] = {} + for chunk_idx, chunk in enumerate(model_chunks): + buffers = list(getattr(chunk, 'buffers', [])) + list( + getattr(chunk, 'expert_parallel_buffers', []) + ) + for buffer_idx, buffer in enumerate(buffers): + param_index_map = getattr(buffer, 'param_index_map', None) + if param_index_map is None: + continue + for param, (global_start, _global_end, _bucket_id) in param_index_map.items(): + identity[param] = (chunk_idx, buffer_idx, global_start) + return identity or None + + def _shard_params_ping_pong(self, optimizers, dp_cp_size, expt_dp_size, model_chunks=None): + """Legacy ping-pong shard assignment (no layout available). Legacy: this method is a fallback for when no ``full_param_layout`` is provided. Once all call sites supply a layout, this can be removed in favor of :meth:`_shard_params_from_layout`. - List of parameters are sorted by numel and assigned to ranks in ping-pong style. - Example of 4 ranks and 10 parameters p0-p9 after sorting, then dp_cp_params_list - will be [[p0, p7, p8], [p1, p6, p9], [p2, p5], [p3, p4]]. + Parameters are sorted by a rank-independent TOTAL order and assigned ping-pong style. E.g. + 4 ranks, 10 params p0-p9 -> [[p0, p7, p8], [p1, p6, p9], [p2, p5], [p3, p4]]. + + CRITICAL: the sort key MUST be identical across DP ranks. ``numel`` alone is not (stable + sort tie-breaks equal-numel params by insertion order), which would give different owners + per rank -> params double-owned or zero-owned on the first step. So we tie-break by the + canonical identity ``(chunk_idx, buffer_idx, global_start_index)``. """ dp_cp_idx, expt_dp_idx = 0, 0 # Create ping-pong style loop so memory is more balanced. @@ -669,12 +761,25 @@ def _shard_params_ping_pong(self, optimizers, dp_cp_size, expt_dp_size): for optimizer in optimizers: param_groups += optimizer.param_groups - # Sort param in all groups by param numel and assign to each rank evenly. + # Sort param in all groups by a rank-independent TOTAL order, then assign to each rank. + identity = self._build_param_sort_keys(model_chunks) param_list = [] for group_index, group in enumerate(param_groups): for p in group["params"]: param_list.append((p, group_index)) - param_list.sort(key=lambda x: x[0].numel()) + if identity is not None: + # Total order: (numel, canonical-global-identity). Identical on every DP rank. + missing = [p for (p, _) in param_list if p not in identity] + assert not missing, ( + "ping-pong ownership requires a canonical identity for every Muon param, " + f"but {len(missing)} param(s) were not found in any model-chunk buffer's " + "param_index_map. Cannot guarantee identical ownership across ranks (the " + "allgather_params gather assumes every rank agrees on each param's single owner)." + ) + param_list.sort(key=lambda x: _param_sort_key(x[0].numel(), identity[x[0]])) + else: + # No layout info: keep the legacy numel-only ordering. + param_list.sort(key=lambda x: x[0].numel()) param_groups_this_rank = [[] for g in param_groups] # Assign params to rank in ping-pong style loop. @@ -758,8 +863,70 @@ def allgather_params(self) -> None: call sites supply a ``full_param_layout``, this can be removed — the standard distributed optimizer buffer all-gather (via ``start_param_sync``) replaces this flatten/unflatten path. + + Two transport variants share the same uneven (all-gather-v) shape: + + * **bf16** (``use_fp8_param_sync=False``): all-gather owned bf16 ``param.data``, copy_ into + non-owned params. + * **fp8** (``use_fp8_param_sync=True``): stage owned fp32 master->bf16, all-gather bf16, + requantize into EVERY rank's ``param.data`` (owned included) so all hold + ``Q(bf16(master))`` (== OFF/Adam). Then ``post_all_gather_processing`` rebuilds fp8 + columnwise/transpose (blockwise/Float8; mxfp8 noop since copy-back already forced it). """ + # FP8-aware variant: stage bf16, uneven all-gather bf16, requantize per rank. + def _allgather_helper_fp8(params_list, group): + # TODO(perf, blockwise-only): blockwise could gather the owner's fp8 rowwise data + # (~2x less comm) instead of bf16; mxfp8 must stay on bf16. See the matching TODO in + # ``_ParamAndGradBucketGroup.start_param_sync`` for the full rationale. + rank = get_pg_rank(group) + dp_size = get_pg_size(group) + # Device from any non-empty owned list (rank 0 may own zero params in the layout). + device = next((params[0].device for params in params_list if len(params) > 0), None) + if device is None: + # No rank owns any param in this buffer -> nothing to gather. + return + + # Stage fp32 master->bf16 (high-precision source), not lossy dequant(fp8). + owned = params_list[rank] + src = ( + _flatten_dense_tensors([_stage_param_to_bf16(p) for p in owned]) + if len(owned) > 0 + else torch.empty(0, device=device, dtype=torch.bfloat16) + ) + flat_sizes = [sum(p.numel() for p in params) for params in params_list] + if max(flat_sizes) == 0: + return + + gather_list = [] + for i in range(dp_size): + if i == rank: + gather_list.append(src) + else: + gather_list.append( + torch.empty(flat_sizes[i], device=device, dtype=torch.bfloat16) + ) + + torch.distributed.all_gather(gather_list, src, group=group) + + # Requantize the gathered bf16 into EVERY rank's params (owned included) so all ranks + # hold Q(bf16(master)), matching OFF/Adam. Unflatten by param shape (logical numel). + for idx, params in enumerate(params_list): + if len(params) == 0: + continue + templates = [ + torch.empty(p.shape, device="meta", dtype=torch.bfloat16) for p in params + ] + updated_params = _unflatten_dense_tensors(gather_list[idx], templates) + for updated_bf16, model_p in zip(updated_params, params): + copy_back_gathered_bf16_into_fp8_param(model_p, updated_bf16) + + # Rebuild fp8 columnwise/transpose after the gather (mirrors the overlap / DistOpt + # paths; blockwise/Float8 build it, mxfp8 is a noop). Else it'd be deferred to forward. + fp8_params = [p for params in params_list for p in params if is_float8tensor(p)] + if fp8_params: + post_all_gather_processing(fp8_params) + # helper function to flatten local params, all-gather, # unflatten and copy to model params def _allgather_helper(params_list, group): @@ -799,10 +966,35 @@ def _allgather_helper(params_list, group): if self.pg_collection is None: return + + def _dispatch(params_list, group): + # Split each rank's owned params by transport dtype. fp8 and bf16 params ride the + # bf16 transport (the fp8 helper stages master->bf16 / requantizes, a no-op in + # precision for bf16); native fp32 params (e.g. weights marked keep_in_fp32 such as + # the DeepSeek-V4 CSA ``ape``) must be gathered in fp32 -- routing them through the + # bf16-staged path would silently downcast them, and mixing fp32 with bf16 in one + # flatten is invalid. For a pure-bf16 model (no fp32 Muon params) the native group is + # empty and this collapses to the original single-helper dispatch. + staged = [ + [p for p in owned if is_float8tensor(p) or p.dtype != torch.float32] + for owned in params_list + ] + native = [ + [p for p in owned if not is_float8tensor(p) and p.dtype == torch.float32] + for owned in params_list + ] + if any(owned for owned in staged): + staged_helper = ( + _allgather_helper_fp8 if self.use_fp8_param_sync else _allgather_helper + ) + staged_helper(staged, group) + if any(owned for owned in native): + _allgather_helper(native, group) + if self.dp_cp_params_list: - _allgather_helper(self.dp_cp_params_list, self.dp_cp) + _dispatch(self.dp_cp_params_list, self.dp_cp) if self.expt_dp_params_list: - _allgather_helper(self.expt_dp_params_list, self.expt_dp) + _dispatch(self.expt_dp_params_list, self.expt_dp) @torch.no_grad() def broadcast_params(self): diff --git a/megatron/core/optimizer/optimizer.py b/megatron/core/optimizer/optimizer.py index dac16f4a2ee..0e28d62f974 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.""" @@ -46,6 +46,7 @@ optim_state_to_sharding_state, ) from ..dist_checkpointing.utils import add_prefix_for_sharding +from ..fp8_utils import copy_back_gathered_bf16_into_fp8_param, is_float8tensor from ..optimizer_param_scheduler import ParamGroupOverride as _ParamGroupOverride from ..transformer.module import param_is_not_shared from ..utils import log_single_rank @@ -678,6 +679,9 @@ def __init__( super().__init__(optimizer, config, init_state_fn) self.grad_scaler = grad_scaler + # Tagged True by LayerWiseDistributedOptimizer on its non-DistOpt children. + self._layer_wise_non_distopt_child = False + # None grad scaler is only supported for bf16. if self.grad_scaler is None: assert not self.config.fp16, 'fp16 expects a grad scaler.' @@ -795,8 +799,9 @@ def step_with_ready_grads(self) -> bool: if not self.is_stub_optimizer: # The reuse_grad_buf (fp8-param-gather) path stages master params into the DDP # param buffer, which only DistributedOptimizer owns. Optimizers without it - # (e.g. LayerWiseDistributedOptimizer's Float16 base opts) must instead copy - # master -> model params so the forward sees the update. + # (e.g. LayerWiseDistributedOptimizer's Float16 base opts, which are tagged + # ``_layer_wise_non_distopt_child``) must instead copy master -> model params so + # the forward sees the update. if self.config.reuse_grad_buf_for_mxfp8_param_ag and hasattr( self, "_copy_main_params_to_param_buffer" ): @@ -805,6 +810,10 @@ def step_with_ready_grads(self) -> bool: if not self.config.overlap_param_gather: self._copy_main_params_to_param_buffer() else: + # Non-DistOpt LayerWise children have no byte-shard param buffer to stage into + # (``_copy_main_params_to_param_buffer`` would raise), so even under reuse_grad_buf + # they copy fp32 master straight into model ``param.data`` (fp8 re-quantized in + # place, or bf16); the grad-buffer reuse applies to the later param-gather staging. self._copy_main_params_to_model_params() if timers is not None: @@ -1008,8 +1017,20 @@ def __init__( # float16 params: if param.type() in ['torch.cuda.HalfTensor', 'torch.cuda.BFloat16Tensor']: float16_params_this_group.append(param) - # Create a copy - main_param = param.detach().clone().float() + # Seed the fp32 master from the high-precision pre-quantization init + # for fp8 params (not the lossy fp8 dequant), matching DistOpt so + # fp8_param_gather ON/OFF hold an identical master at iter 0. + if hasattr(param, 'get_high_precision_init_val'): + main_param = ( + param.get_high_precision_init_val() + .detach() + .clone() + .to(param.device) + .float() + ) + param.clear_high_precision_init_val() + else: + main_param = param.detach().clone().float() # Copy tensor model parallel attributes. tensor_parallel.copy_tensor_model_parallel_attributes(main_param, param) tensor_parallel.copy_gtp_attributes(main_param, param) @@ -1110,12 +1131,46 @@ def _copy_model_grads_to_main_grads(self): model_param.grad = model_param.main_grad def _copy_main_params_to_model_params(self): + # Non-DistOpt LayerWise fp8: route master->model through bf16 (Q(bf16(master))) to match the + # fp8-param-gather-OFF baseline (a direct fp32->fp8 copy would write Q(fp32 master)). This + # also covers MoE expert weights at expt_dp==1, which are not gathered. + if self._layer_wise_non_distopt_child: + other_model_data, other_main_data = [], [] + 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_float8tensor(model_param): + # Gathered fp8 params get ``Q(bf16(master))`` written into ``param.data`` + # by the fp8 all-gather's requantize (``_allgather_helper_fp8``), which + # would overwrite this copy -- so skip it for them. Non-gathered fp8 params + # (e.g. MoE experts at expt_dp == 1, which the all-gather skips) are not + # tagged and still get their ``Q(bf16(master))`` written here. + if not getattr(model_param, '_layer_wise_fp8_gathered', False): + copy_back_gathered_bf16_into_fp8_param( + model_param, main_param.detach().to(torch.bfloat16) + ) + else: + other_model_data.append(model_param.data) + other_main_data.append(main_param.data) + if other_model_data: + _multi_tensor_copy_this_to_that( + this=other_main_data, + that=other_model_data, + overflow_buf=self._dummy_overflow_buf, + ) + return # Only needed for the float16 params. model_data, main_data = self._get_model_and_main_params_data_float16() _multi_tensor_copy_this_to_that( this=main_data, that=model_data, overflow_buf=self._dummy_overflow_buf ) + # NOTE: deliberately no ``_copy_main_params_to_param_buffer`` here. Only + # ``DistributedOptimizer`` owns the byte-shard param buffer, and + # ``MixedPrecisionOptimizer.step_with_ready_grads`` probes for the method with ``hasattr`` + # to route non-DistOpt optimizers (incl. LayerWise's Float16 children) to + # ``_copy_main_params_to_model_params`` instead. Defining a raising stub here would make + # that probe always succeed. + def _copy_model_params_to_main_params(self, state_dict=None): assert state_dict is None, "Initialize main params from state dict is not supported" # Only needed for the float16 params. diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py index e16abf49d3a..bc633e8c07b 100644 --- a/megatron/training/arguments.py +++ b/megatron/training/arguments.py @@ -1687,12 +1687,26 @@ def validate_args(args, defaults={}): if args.use_layer_wise_distributed_optimizer: - assert not args.fp8_param_gather and not getattr(args, 'fp4_param_gather', False), ( - "Layer-wise (Muon) distributed optimizer does not support FP8/FP4 parameter gather " - "(fp8_param_gather / fp4_param_gather). Use fp8_param_gather=False (e.g. blockwise/" - "MXFP8 compute with parameters persisted in bf16)." - ) if not args.use_layer_wise_param_layout: + # Decoupled compact LayerWise: fp8 parameter gather is supported via the FP8-aware + # whole-param all-gather. Only mxfp8/blockwise (fp4 out of scope); mxfp8 needs + # reuse_grad_buf. fp4 is rejected unconditionally -- the LayerWise gather routes + # buckets by is_float8tensor, so an NVFP4 param would silently take the raw + # flatten path. + assert not getattr(args, 'fp4_param_gather', False), ( + "Decoupled compact LayerWise DDP layout supports fp8 parameter gather only " + "(mxfp8 or blockwise); fp4_param_gather is out of scope." + ) + if args.fp8_param_gather: + assert args.fp8_recipe in ('mxfp8', 'blockwise'), ( + "fp8 parameter gather on the decoupled compact LayerWise DDP layout requires " + f"fp8_recipe in {{'mxfp8', 'blockwise'}}; got {args.fp8_recipe!r}." + ) + if args.fp8_recipe == 'mxfp8': + assert args.reuse_grad_buf_for_mxfp8_param_ag, ( + "mxfp8 + --fp8-param-gather on the decoupled compact LayerWise DDP layout " + "requires --reuse-grad-buf-for-mxfp8-param-ag (or use fp8_recipe='blockwise')." + ) assert args.num_distributed_optimizer_instances == 1, ( "the decoupled compact LayerWise DDP layout (the default; pass " "--use-layer-wise-param-layout for the padded layout) requires " @@ -1700,6 +1714,13 @@ def validate_args(args, defaults={}): "only all-reduce within a single optimizer instance, so partial DistOpt (>1 " "instance) would under-reduce Muon gradients across the full data-parallel domain." ) + else: + # Padded LayerWise param layout: fp8/fp4 parameter gather is not supported here. + assert not args.fp8_param_gather and not getattr(args, 'fp4_param_gather', False), ( + "Layer-wise (Muon) distributed optimizer with the padded param layout does not " + "support FP8/FP4 parameter gather. Drop --use-layer-wise-param-layout to get the " + "default decoupled compact layout, or set fp8_param_gather=False." + ) # Make sure all functionality that requires Gloo process groups is disabled. if not args.use_gloo_process_groups: @@ -3011,16 +3032,6 @@ def _add_distributed_args(parser): dest='align_param_gather') group.add_argument('--use-distributed-optimizer', action='store_true', help='Use distributed optimizer.') - group.add_argument('--use-layer-wise-param-layout', - action='store_true', - help='Opt INTO the padded shard-aligned LayerWise param layout. The default ' - 'is the compact decoupled layout, where LayerWise (Muon 2D) buffers use a ' - 'no-padding DDP layout and locally disable DistributedOptimizer (all-reduce ' - 'grads + whole-param ping-pong + allgather_params), while sibling buffers keep ' - 'the byte-level DistributedOptimizer; this avoids the persistent ' - 'dp_size * max(shard_load) padding. Pass this flag to restore the padded ' - 'layout (e.g. for bit-for-bit comparison; it uses a different bf16 reduction ' - 'ordering).') group.add_argument('--use-nccl-ub', action='store_true', dest='nccl_ub', help='Use the userbuffer registration for DP/FSDP communication buffers.' 'This option will reduce GPU SM usage for the DP/FSDP communication,' @@ -3075,6 +3086,16 @@ def _add_distributed_args(parser): help='If set, initialize with fake distributed process group and all distributed communication operations will be skipped. \ This is quite useful for profiling memory usage of distributed training with just one GPU. \ Setting WORLD_SIZE and RANK to the specific values for target distribtued scale.') + group.add_argument( + '--use-layer-wise-param-layout', + action='store_true', + help='Opt INTO the padded shard-aligned LayerWise param layout. The default is the compact ' + 'decoupled layout, where LayerWise (Muon 2D) buffers use a no-padding DDP layout and locally ' + 'disable DistributedOptimizer (all-reduce grads + whole-param ping-pong + allgather_params), ' + 'while sibling buffers keep the byte-level DistributedOptimizer; this avoids the persistent ' + 'dp_size * max(shard_load) padding. Pass this flag to restore the padded layout (e.g. for ' + 'bit-for-bit comparison; it uses a different bf16 reduction ordering).', + ) return parser diff --git a/megatron/training/training.py b/megatron/training/training.py index 98d97a162db..c87885cf233 100644 --- a/megatron/training/training.py +++ b/megatron/training/training.py @@ -1637,9 +1637,9 @@ def wrap_model_chunks_with_ddp( # LayerWise-managed matrices (Muon's Newton-Schulz domain) to a separate buffer. # The padded-vs-compact LayerWise layout decision is made inside # compute_full_param_layout / _ParamAndGradBuffer from use_layer_wise_param_layout: - # by default (compact) LayerWise buffers get the no-padding layout and the per-buffer - # override flips use_distributed_optimizer off for them; with - # --use-layer-wise-param-layout they stay on the padded DistOpt layout. + # with the default padded layout LayerWise buffers stay DistOpt; with + # the default (no --use-layer-wise-param-layout) they get the compact no-padding layout and + # per-buffer override flips use_distributed_optimizer off for them. ddp_config.use_distributed_optimizer = True compute_layout = LayerWiseDistributedOptimizer.compute_full_param_layout tag_params_for_buffer_routing(model_chunks) diff --git a/tests/functional_tests/test_cases/moe/gpt3_moe_mcore_te_ep8_resume_torch_dist_dist_muon/golden_values_dev_dgx_gb200.json b/tests/functional_tests/test_cases/moe/gpt3_moe_mcore_te_ep8_resume_torch_dist_dist_muon/golden_values_dev_dgx_gb200.json index 8ee325dec71..9389af9e845 100644 --- a/tests/functional_tests/test_cases/moe/gpt3_moe_mcore_te_ep8_resume_torch_dist_dist_muon/golden_values_dev_dgx_gb200.json +++ b/tests/functional_tests/test_cases/moe/gpt3_moe_mcore_te_ep8_resume_torch_dist_dist_muon/golden_values_dev_dgx_gb200.json @@ -6,104 +6,104 @@ "values": { "1": 10.94223, "2": 10.91037, - "3": 10.92694, - "4": 10.91781, - "5": 10.92105, - "6": 10.92688, + "3": 10.92695, + "4": 10.91755, + "5": 10.92096, + "6": 10.92734, "7": 10.92235, - "8": 10.92031, + "8": 10.9198, "9": 10.93219, - "10": 10.92611, - "11": 10.91797, - "12": 10.93443, - "13": 10.91919, - "14": 10.90942, - "15": 10.91807, - "16": 10.92132, - "17": 10.93511, - "18": 10.91657, - "19": 10.92834, - "20": 10.9022, - "21": 10.92588, - "22": 10.92401, - "23": 10.92945, - "24": 10.9046, - "25": 10.91421, - "26": 10.91452, - "27": 10.92959, - "28": 10.92001, - "29": 10.92055, - "30": 10.91213, - "31": 10.90591, - "32": 10.91832, - "33": 10.91666, - "34": 10.90319, - "35": 10.92638, - "36": 10.90801, - "37": 10.90602, - "38": 10.91644, - "39": 10.90971, - "40": 10.90394, - "41": 10.91095, - "42": 10.9147, - "43": 10.90761, - "44": 10.9166, - "45": 10.91017, - "46": 10.90603, - "47": 10.90082, - "48": 10.89309, - "49": 10.90441, - "50": 10.89335, - "51": 10.89814, - "52": 10.89698, - "53": 10.90087, - "54": 10.89136, - "55": 10.90606, - "56": 10.88392, - "57": 10.88847, - "58": 10.87993, - "59": 10.87881, - "60": 10.86754, - "61": 10.87005, - "62": 10.8627, - "63": 10.87072, - "64": 10.8622, - "65": 10.86482, - "66": 10.86786, - "67": 10.85568, - "68": 10.84834, - "69": 10.85356, - "70": 10.85766, - "71": 10.85007, - "72": 10.84966, - "73": 10.83916, - "74": 10.83516, - "75": 10.84065, - "76": 10.82979, - "77": 10.82946, - "78": 10.81214, - "79": 10.82951, - "80": 10.83397, - "81": 10.81606, - "82": 10.814, - "83": 10.8053, - "84": 10.78055, - "85": 10.78201, - "86": 10.79832, - "87": 10.79865, - "88": 10.80163, - "89": 10.78541, - "90": 10.78282, - "91": 10.78264, - "92": 10.78229, - "93": 10.75819, - "94": 10.76732, - "95": 10.76526, - "96": 10.73862, - "97": 10.73009, - "98": 10.75335, - "99": 10.76001, - "100": 10.73865 + "10": 10.9268, + "11": 10.91768, + "12": 10.93475, + "13": 10.91892, + "14": 10.90993, + "15": 10.91834, + "16": 10.92077, + "17": 10.93466, + "18": 10.91632, + "19": 10.92825, + "20": 10.90227, + "21": 10.92573, + "22": 10.92408, + "23": 10.92944, + "24": 10.90453, + "25": 10.91385, + "26": 10.91451, + "27": 10.92971, + "28": 10.91958, + "29": 10.92019, + "30": 10.91171, + "31": 10.90632, + "32": 10.91789, + "33": 10.91662, + "34": 10.90253, + "35": 10.92677, + "36": 10.9074, + "37": 10.9057, + "38": 10.91635, + "39": 10.9101, + "40": 10.9038, + "41": 10.91171, + "42": 10.91498, + "43": 10.90784, + "44": 10.91689, + "45": 10.91033, + "46": 10.9063, + "47": 10.9005, + "48": 10.89371, + "49": 10.90465, + "50": 10.89305, + "51": 10.89875, + "52": 10.89729, + "53": 10.90151, + "54": 10.89159, + "55": 10.90604, + "56": 10.88379, + "57": 10.88868, + "58": 10.88007, + "59": 10.87846, + "60": 10.86744, + "61": 10.86956, + "62": 10.86268, + "63": 10.87044, + "64": 10.86199, + "65": 10.86449, + "66": 10.8678, + "67": 10.85569, + "68": 10.84792, + "69": 10.85294, + "70": 10.8576, + "71": 10.84968, + "72": 10.84963, + "73": 10.83941, + "74": 10.83469, + "75": 10.84086, + "76": 10.82966, + "77": 10.82949, + "78": 10.81271, + "79": 10.82939, + "80": 10.83423, + "81": 10.81568, + "82": 10.81436, + "83": 10.80441, + "84": 10.78114, + "85": 10.78199, + "86": 10.79933, + "87": 10.79915, + "88": 10.80194, + "89": 10.78455, + "90": 10.78197, + "91": 10.78281, + "92": 10.78238, + "93": 10.75825, + "94": 10.768, + "95": 10.76521, + "96": 10.73841, + "97": 10.73064, + "98": 10.75338, + "99": 10.75991, + "100": 10.73847 } }, "num-zeros": { @@ -111,106 +111,106 @@ "end_step": 100, "step_interval": 1, "values": { - "1": 1375.0, - "2": 1150.0, - "3": 1244.0, - "4": 1230.0, - "5": 1180.0, - "6": 1258.0, - "7": 1543.0, - "8": 1272.0, - "9": 1271.0, - "10": 1218.0, - "11": 1244.0, - "12": 1236.0, - "13": 1482.0, - "14": 1391.0, - "15": 1219.0, - "16": 1316.0, - "17": 1338.0, - "18": 1226.0, - "19": 1215.0, - "20": 1217.0, - "21": 1455.0, - "22": 1319.0, - "23": 1299.0, - "24": 1568.0, - "25": 1277.0, - "26": 1358.0, - "27": 1283.0, - "28": 1237.0, - "29": 1292.0, - "30": 1296.0, - "31": 1381.0, - "32": 1446.0, - "33": 1366.0, - "34": 1431.0, - "35": 1354.0, - "36": 1339.0, - "37": 1314.0, - "38": 1323.0, - "39": 1248.0, - "40": 1309.0, - "41": 1226.0, - "42": 1071.0, - "43": 1309.0, - "44": 1338.0, - "45": 1422.0, - "46": 1327.0, - "47": 1384.0, - "48": 1428.0, - "49": 1396.0, - "50": 1203.0, - "51": 1246.0, - "52": 1284.0, - "53": 1320.0, - "54": 1199.0, - "55": 1320.0, - "56": 1196.0, - "57": 1223.0, - "58": 1267.0, - "59": 1323.0, - "60": 1308.0, - "61": 1206.0, - "62": 1398.0, - "63": 1300.0, - "64": 1319.0, - "65": 1458.0, - "66": 1153.0, - "67": 1380.0, - "68": 1189.0, - "69": 1365.0, - "70": 1255.0, - "71": 1106.0, - "72": 1332.0, - "73": 1292.0, - "74": 1261.0, - "75": 1308.0, - "76": 1345.0, - "77": 1285.0, - "78": 1234.0, - "79": 1257.0, - "80": 1223.0, - "81": 1407.0, - "82": 1040.0, - "83": 1258.0, - "84": 1231.0, - "85": 1183.0, - "86": 1280.0, - "87": 1159.0, - "88": 1288.0, - "89": 1191.0, - "90": 1184.0, - "91": 1153.0, - "92": 1118.0, - "93": 1335.0, - "94": 1318.0, - "95": 1138.0, - "96": 1283.0, - "97": 1298.0, - "98": 1258.0, - "99": 1306.0, - "100": 1404.0 + "1": 1356.0, + "2": 1146.0, + "3": 1276.0, + "4": 1214.0, + "5": 1218.0, + "6": 1222.0, + "7": 1561.0, + "8": 1321.0, + "9": 1328.0, + "10": 1214.0, + "11": 1229.0, + "12": 1292.0, + "13": 1401.0, + "14": 1366.0, + "15": 1136.0, + "16": 1296.0, + "17": 1316.0, + "18": 1274.0, + "19": 1192.0, + "20": 1216.0, + "21": 1415.0, + "22": 1404.0, + "23": 1285.0, + "24": 1569.0, + "25": 1257.0, + "26": 1366.0, + "27": 1280.0, + "28": 1198.0, + "29": 1307.0, + "30": 1349.0, + "31": 1372.0, + "32": 1383.0, + "33": 1359.0, + "34": 1402.0, + "35": 1301.0, + "36": 1242.0, + "37": 1414.0, + "38": 1229.0, + "39": 1247.0, + "40": 1344.0, + "41": 1286.0, + "42": 1050.0, + "43": 1301.0, + "44": 1358.0, + "45": 1359.0, + "46": 1216.0, + "47": 1334.0, + "48": 1390.0, + "49": 1428.0, + "50": 1202.0, + "51": 1247.0, + "52": 1275.0, + "53": 1352.0, + "54": 1230.0, + "55": 1197.0, + "56": 1137.0, + "57": 1268.0, + "58": 1349.0, + "59": 1228.0, + "60": 1361.0, + "61": 1173.0, + "62": 1368.0, + "63": 1238.0, + "64": 1359.0, + "65": 1421.0, + "66": 1159.0, + "67": 1417.0, + "68": 1143.0, + "69": 1331.0, + "70": 1241.0, + "71": 1150.0, + "72": 1426.0, + "73": 1314.0, + "74": 1243.0, + "75": 1343.0, + "76": 1407.0, + "77": 1281.0, + "78": 1288.0, + "79": 1271.0, + "80": 1189.0, + "81": 1353.0, + "82": 1133.0, + "83": 1240.0, + "84": 1236.0, + "85": 1243.0, + "86": 1282.0, + "87": 1182.0, + "88": 1345.0, + "89": 1172.0, + "90": 1233.0, + "91": 1220.0, + "92": 1161.0, + "93": 1369.0, + "94": 1380.0, + "95": 1177.0, + "96": 1244.0, + "97": 1398.0, + "98": 1211.0, + "99": 1238.0, + "100": 1359.0 } }, "mem-allocated-bytes": { @@ -223,101 +223,101 @@ "3": 993740800.0, "4": 993725952.0, "5": 993734144.0, - "6": 993718272.0, - "7": 993741312.0, - "8": 993709056.0, - "9": 993735168.0, - "10": 993729536.0, - "11": 993708544.0, + "6": 993719296.0, + "7": 993740800.0, + "8": 993710592.0, + "9": 993733632.0, + "10": 993730048.0, + "11": 993709568.0, "12": 993701888.0, - "13": 993785856.0, - "14": 993799168.0, - "15": 993735168.0, - "16": 993737728.0, - "17": 993748480.0, - "18": 993705984.0, - "19": 993765888.0, - "20": 993732096.0, - "21": 993782784.0, - "22": 993733120.0, - "23": 993708544.0, - "24": 993744896.0, - "25": 993698304.0, - "26": 993810432.0, - "27": 993738752.0, - "28": 993734144.0, - "29": 993771008.0, - "30": 993692672.0, - "31": 993728000.0, - "32": 993732608.0, - "33": 993720320.0, - "34": 993783296.0, - "35": 993718272.0, + "13": 993787392.0, + "14": 993797632.0, + "15": 993737216.0, + "16": 993738240.0, + "17": 993748992.0, + "18": 993710592.0, + "19": 993765376.0, + "20": 993731584.0, + "21": 993783296.0, + "22": 993733632.0, + "23": 993709056.0, + "24": 993745408.0, + "25": 993696256.0, + "26": 993809920.0, + "27": 993739264.0, + "28": 993735680.0, + "29": 993768960.0, + "30": 993694208.0, + "31": 993729024.0, + "32": 993733632.0, + "33": 993719808.0, + "34": 993782272.0, + "35": 993720832.0, "36": 993655808.0, - "37": 993730560.0, + "37": 993731072.0, "38": 993759744.0, - "39": 993702912.0, - "40": 993736704.0, - "41": 993739264.0, + "39": 993702400.0, + "40": 993738240.0, + "41": 993742336.0, "42": 993735168.0, - "43": 993736192.0, - "44": 993706496.0, - "45": 993702912.0, + "43": 993735680.0, + "44": 993705472.0, + "45": 993702400.0, "46": 993764352.0, "47": 993757696.0, - "48": 993702912.0, - "49": 993731072.0, - "50": 993682944.0, - "51": 993666560.0, - "52": 993688064.0, - "53": 993730048.0, - "54": 993729536.0, - "55": 993711616.0, - "56": 993721344.0, - "57": 993754112.0, - "58": 993792512.0, - "59": 993699840.0, - "60": 993713152.0, - "61": 993680896.0, - "62": 993712128.0, - "63": 993723904.0, - "64": 993727488.0, - "65": 993713152.0, - "66": 993652736.0, - "67": 993728000.0, + "48": 993703936.0, + "49": 993730048.0, + "50": 993681408.0, + "51": 993665024.0, + "52": 993688576.0, + "53": 993732608.0, + "54": 993730048.0, + "55": 993710592.0, + "56": 993720320.0, + "57": 993756672.0, + "58": 993790976.0, + "59": 993699328.0, + "60": 993711616.0, + "61": 993681920.0, + "62": 993711616.0, + "63": 993724416.0, + "64": 993729024.0, + "65": 993713664.0, + "66": 993652224.0, + "67": 993729536.0, "68": 993741824.0, "69": 993719808.0, "70": 993671168.0, "71": 993694720.0, - "72": 993763840.0, - "73": 993697792.0, - "74": 993695232.0, - "75": 993733632.0, + "72": 993765376.0, + "73": 993698304.0, + "74": 993693696.0, + "75": 993732608.0, "76": 993709568.0, "77": 993723904.0, - "78": 993714688.0, - "79": 993704448.0, - "80": 993724416.0, - "81": 993707520.0, - "82": 993683456.0, - "83": 993739264.0, - "84": 993730560.0, - "85": 993741824.0, - "86": 993717248.0, - "87": 993612800.0, - "88": 993669120.0, - "89": 993677824.0, - "90": 993700352.0, - "91": 993731072.0, - "92": 993679872.0, - "93": 993647616.0, - "94": 993690112.0, - "95": 993677824.0, - "96": 993651712.0, - "97": 993678336.0, - "98": 993742848.0, - "99": 993701888.0, - "100": 993733632.0 + "78": 993712128.0, + "79": 993704960.0, + "80": 993725440.0, + "81": 993709056.0, + "82": 993682944.0, + "83": 993739776.0, + "84": 993731584.0, + "85": 993740800.0, + "86": 993716224.0, + "87": 993611776.0, + "88": 993669632.0, + "89": 993675776.0, + "90": 993700864.0, + "91": 993730048.0, + "92": 993680384.0, + "93": 993649664.0, + "94": 993690624.0, + "95": 993676288.0, + "96": 993652224.0, + "97": 993679360.0, + "98": 993742336.0, + "99": 993704960.0, + "100": 993732608.0 } }, "mem-max-allocated-bytes": { @@ -337,94 +337,94 @@ "10": 3331357184.0, "11": 3331357184.0, "12": 3331357184.0, - "13": 3357274624.0, - "14": 3362478592.0, - "15": 3362478592.0, - "16": 3362478592.0, - "17": 3362478592.0, - "18": 3362478592.0, - "19": 3362478592.0, - "20": 3362478592.0, - "21": 3362478592.0, - "22": 3362478592.0, - "23": 3362478592.0, - "24": 3362478592.0, - "25": 3362478592.0, - "26": 3375542784.0, - "27": 3375542784.0, - "28": 3375542784.0, - "29": 3375542784.0, - "30": 3375542784.0, - "31": 3375542784.0, - "32": 3375542784.0, - "33": 3375542784.0, - "34": 3375542784.0, - "35": 3375542784.0, - "36": 3375542784.0, - "37": 3375542784.0, - "38": 3375542784.0, - "39": 3375542784.0, - "40": 3375542784.0, - "41": 3375542784.0, - "42": 3375542784.0, - "43": 3375542784.0, - "44": 3375542784.0, - "45": 3375542784.0, - "46": 3375542784.0, - "47": 3375542784.0, - "48": 3375542784.0, - "49": 3375542784.0, - "50": 3375542784.0, - "51": 3375542784.0, - "52": 3375542784.0, - "53": 3375542784.0, - "54": 3375542784.0, - "55": 3375542784.0, - "56": 3375542784.0, - "57": 3375542784.0, - "58": 3375542784.0, - "59": 3375542784.0, - "60": 3375542784.0, - "61": 3375542784.0, - "62": 3375542784.0, - "63": 3375542784.0, - "64": 3375542784.0, - "65": 3375542784.0, - "66": 3375542784.0, - "67": 3375542784.0, - "68": 3375542784.0, - "69": 3375542784.0, - "70": 3375542784.0, - "71": 3375542784.0, - "72": 3375542784.0, - "73": 3375542784.0, - "74": 3375542784.0, - "75": 3375542784.0, - "76": 3375542784.0, - "77": 3375542784.0, - "78": 3375542784.0, - "79": 3375542784.0, - "80": 3375542784.0, - "81": 3375542784.0, - "82": 3375542784.0, - "83": 3375542784.0, - "84": 3375542784.0, - "85": 3375542784.0, - "86": 3375542784.0, - "87": 3375542784.0, - "88": 3375542784.0, - "89": 3375542784.0, - "90": 3375542784.0, - "91": 3375542784.0, - "92": 3375542784.0, - "93": 3375542784.0, - "94": 3375542784.0, - "95": 3375542784.0, - "96": 3375542784.0, - "97": 3375542784.0, - "98": 3375542784.0, - "99": 3375542784.0, - "100": 3375542784.0 + "13": 3358355456.0, + "14": 3362544640.0, + "15": 3362544640.0, + "16": 3362544640.0, + "17": 3362544640.0, + "18": 3362544640.0, + "19": 3362544640.0, + "20": 3362544640.0, + "21": 3362544640.0, + "22": 3362544640.0, + "23": 3362544640.0, + "24": 3362544640.0, + "25": 3362544640.0, + "26": 3375151104.0, + "27": 3375151104.0, + "28": 3375151104.0, + "29": 3375151104.0, + "30": 3375151104.0, + "31": 3375151104.0, + "32": 3375151104.0, + "33": 3375151104.0, + "34": 3375151104.0, + "35": 3375151104.0, + "36": 3375151104.0, + "37": 3375151104.0, + "38": 3375151104.0, + "39": 3375151104.0, + "40": 3375151104.0, + "41": 3375151104.0, + "42": 3375151104.0, + "43": 3375151104.0, + "44": 3375151104.0, + "45": 3375151104.0, + "46": 3375151104.0, + "47": 3375151104.0, + "48": 3375151104.0, + "49": 3375151104.0, + "50": 3375151104.0, + "51": 3375151104.0, + "52": 3375151104.0, + "53": 3375151104.0, + "54": 3375151104.0, + "55": 3375151104.0, + "56": 3375151104.0, + "57": 3375151104.0, + "58": 3375151104.0, + "59": 3375151104.0, + "60": 3375151104.0, + "61": 3375151104.0, + "62": 3375151104.0, + "63": 3375151104.0, + "64": 3375151104.0, + "65": 3375151104.0, + "66": 3375151104.0, + "67": 3375151104.0, + "68": 3375151104.0, + "69": 3375151104.0, + "70": 3375151104.0, + "71": 3375151104.0, + "72": 3375151104.0, + "73": 3375151104.0, + "74": 3375151104.0, + "75": 3375151104.0, + "76": 3375151104.0, + "77": 3375151104.0, + "78": 3375151104.0, + "79": 3375151104.0, + "80": 3375151104.0, + "81": 3375151104.0, + "82": 3375151104.0, + "83": 3375151104.0, + "84": 3375151104.0, + "85": 3375151104.0, + "86": 3375151104.0, + "87": 3375151104.0, + "88": 3375151104.0, + "89": 3375151104.0, + "90": 3375151104.0, + "91": 3375151104.0, + "92": 3375151104.0, + "93": 3375151104.0, + "94": 3375151104.0, + "95": 3375151104.0, + "96": 3375151104.0, + "97": 3375151104.0, + "98": 3375151104.0, + "99": 3375151104.0, + "100": 3375151104.0 } }, "iteration-time": { @@ -433,105 +433,105 @@ "step_interval": 1, "values": { "1": "nan", - "2": 8.49271, - "3": 0.25688, - "4": 0.22124, - "5": 0.21824, - "6": 0.21083, - "7": 0.20712, - "8": 0.19956, - "9": 0.19815, - "10": 0.1926, - "11": 0.19351, - "12": 0.19256, - "13": 0.19281, - "14": 0.19127, - "15": 0.19154, - "16": 0.18548, - "17": 0.19195, - "18": 0.1983, - "19": 0.25528, - "20": 0.22052, - "21": 0.28285, - "22": 0.2177, - "23": 0.21469, - "24": 0.21162, - "25": 0.23116, - "26": 0.18774, - "27": 0.18611, - "28": 0.19088, - "29": 0.18301, - "30": 0.18331, - "31": 0.18247, - "32": 0.18431, - "33": 0.19054, - "34": 0.18773, - "35": 0.18089, - "36": 0.18111, - "37": 0.18092, - "38": 0.18475, - "39": 0.18645, - "40": 0.18212, - "41": 0.18174, - "42": 0.18552, - "43": 0.18528, - "44": 0.17967, - "45": 0.18292, - "46": 0.18062, - "47": 0.1801, - "48": 0.18752, - "49": 0.18553, - "50": 0.18368, - "51": 0.45252, - "52": 0.22591, - "53": 0.1812, - "54": 0.18363, - "55": 0.18724, - "56": 0.18336, - "57": 0.18169, - "58": 0.18373, - "59": 0.18345, - "60": 0.18069, - "61": 0.18631, - "62": 0.18358, - "63": 0.18134, - "64": 0.18375, - "65": 0.18154, - "66": 0.18594, - "67": 0.18109, - "68": 0.18792, - "69": 0.1783, - "70": 0.18347, - "71": 0.18189, - "72": 0.18531, - "73": 0.18202, - "74": 0.18041, - "75": 0.1821, - "76": 0.17952, - "77": 0.18016, - "78": 0.18287, - "79": 0.18241, - "80": 0.17827, - "81": 0.18204, - "82": 0.1888, - "83": 0.18477, - "84": 0.18326, - "85": 0.1834, - "86": 0.18213, - "87": 0.18585, - "88": 0.18155, - "89": 0.17993, - "90": 0.17981, - "91": 0.18275, - "92": 0.18312, - "93": 0.18, - "94": 0.18267, - "95": 0.18972, - "96": 0.18008, - "97": 0.18222, - "98": 0.1815, - "99": 0.17947, - "100": 0.18358 + "2": 5.90294, + "3": 0.26152, + "4": 0.21613, + "5": 0.20882, + "6": 0.20882, + "7": 0.20075, + "8": 0.19842, + "9": 0.20133, + "10": 0.1901, + "11": 0.1904, + "12": 0.18409, + "13": 0.1932, + "14": 0.18651, + "15": 0.183, + "16": 0.18296, + "17": 0.18867, + "18": 0.18622, + "19": 0.18922, + "20": 0.18043, + "21": 0.17865, + "22": 0.18905, + "23": 0.18161, + "24": 0.19104, + "25": 0.17529, + "26": 0.1875, + "27": 0.1744, + "28": 0.17705, + "29": 0.17844, + "30": 0.18877, + "31": 0.18617, + "32": 0.17614, + "33": 0.17845, + "34": 0.18037, + "35": 0.17481, + "36": 0.18838, + "37": 0.1829, + "38": 0.17985, + "39": 0.17982, + "40": 0.17707, + "41": 0.1781, + "42": 0.17821, + "43": 0.1774, + "44": 0.17516, + "45": 0.17739, + "46": 0.17181, + "47": 0.17523, + "48": 0.17835, + "49": 0.17671, + "50": 0.17428, + "51": 0.30857, + "52": 0.24281, + "53": 0.1875, + "54": 0.17344, + "55": 0.17332, + "56": 0.17908, + "57": 0.17446, + "58": 0.17629, + "59": 0.17871, + "60": 0.1818, + "61": 0.17728, + "62": 0.17708, + "63": 0.17819, + "64": 0.17691, + "65": 0.18031, + "66": 0.1753, + "67": 0.17847, + "68": 0.17835, + "69": 0.17499, + "70": 0.17444, + "71": 0.17275, + "72": 0.17357, + "73": 0.17775, + "74": 0.17772, + "75": 0.17724, + "76": 0.18592, + "77": 0.17383, + "78": 0.17675, + "79": 0.17717, + "80": 0.17934, + "81": 0.19001, + "82": 0.17417, + "83": 0.17929, + "84": 0.17939, + "85": 0.18325, + "86": 0.17452, + "87": 0.18009, + "88": 0.17535, + "89": 0.17776, + "90": 0.17555, + "91": 0.18803, + "92": 0.17793, + "93": 0.17961, + "94": 0.18054, + "95": 0.17443, + "96": 0.18149, + "97": 0.17413, + "98": 0.17735, + "99": 0.17308, + "100": 0.18173 } } -} +} \ No newline at end of file diff --git a/tests/functional_tests/test_cases/moe/gpt3_moe_mcore_te_ep8_resume_torch_dist_dist_muon/golden_values_dev_dgx_h100.json b/tests/functional_tests/test_cases/moe/gpt3_moe_mcore_te_ep8_resume_torch_dist_dist_muon/golden_values_dev_dgx_h100.json index 78c5092db1e..4a617a6a906 100644 --- a/tests/functional_tests/test_cases/moe/gpt3_moe_mcore_te_ep8_resume_torch_dist_dist_muon/golden_values_dev_dgx_h100.json +++ b/tests/functional_tests/test_cases/moe/gpt3_moe_mcore_te_ep8_resume_torch_dist_dist_muon/golden_values_dev_dgx_h100.json @@ -6,104 +6,104 @@ "values": { "1": 10.90111, "2": 10.8957, - "3": 10.90636, - "4": 10.90606, - "5": 10.91349, - "6": 10.89747, - "7": 10.89834, - "8": 10.91366, - "9": 10.89559, - "10": 10.90375, - "11": 10.90014, - "12": 10.91887, - "13": 10.88393, - "14": 10.87584, - "15": 10.91035, - "16": 10.90325, - "17": 10.88717, - "18": 10.89416, - "19": 10.89487, - "20": 10.89092, - "21": 10.88636, - "22": 10.90341, - "23": 10.91256, - "24": 10.88547, - "25": 10.90164, - "26": 10.88848, - "27": 10.89705, - "28": 10.88168, - "29": 10.88997, - "30": 10.91097, - "31": 10.89128, - "32": 10.88933, - "33": 10.89772, - "34": 10.87339, - "35": 10.89595, - "36": 10.90751, - "37": 10.87112, - "38": 10.87839, - "39": 10.88579, - "40": 10.89247, - "41": 10.88364, - "42": 10.89496, - "43": 10.88058, - "44": 10.88492, - "45": 10.88361, - "46": 10.88582, - "47": 10.88565, - "48": 10.86424, - "49": 10.8759, - "50": 10.88259, - "51": 10.89305, - "52": 10.87123, - "53": 10.85905, - "54": 10.8702, - "55": 10.86807, - "56": 10.87324, - "57": 10.84465, - "58": 10.8563, - "59": 10.84628, - "60": 10.84343, - "61": 10.8604, - "62": 10.85675, - "63": 10.86117, - "64": 10.83852, - "65": 10.82804, - "66": 10.84584, - "67": 10.82846, - "68": 10.83007, - "69": 10.81894, - "70": 10.82797, - "71": 10.82677, - "72": 10.8088, - "73": 10.80753, - "74": 10.80659, - "75": 10.81472, - "76": 10.81086, - "77": 10.80817, - "78": 10.79255, - "79": 10.80218, - "80": 10.7879, - "81": 10.79487, - "82": 10.79556, - "83": 10.78629, - "84": 10.75658, - "85": 10.76241, - "86": 10.77699, - "87": 10.79624, - "88": 10.77534, - "89": 10.77559, - "90": 10.76414, - "91": 10.74033, - "92": 10.76012, - "93": 10.74683, - "94": 10.73435, - "95": 10.75231, - "96": 10.72344, - "97": 10.71428, - "98": 10.72631, - "99": 10.74637, - "100": 10.69592 + "3": 10.90637, + "4": 10.90594, + "5": 10.91348, + "6": 10.89741, + "7": 10.89767, + "8": 10.91394, + "9": 10.89569, + "10": 10.90411, + "11": 10.90011, + "12": 10.91923, + "13": 10.88447, + "14": 10.87534, + "15": 10.91024, + "16": 10.90314, + "17": 10.88671, + "18": 10.89362, + "19": 10.89511, + "20": 10.89108, + "21": 10.88647, + "22": 10.90362, + "23": 10.91205, + "24": 10.88572, + "25": 10.90181, + "26": 10.88836, + "27": 10.89695, + "28": 10.88143, + "29": 10.88988, + "30": 10.91064, + "31": 10.89181, + "32": 10.8893, + "33": 10.8973, + "34": 10.87376, + "35": 10.89602, + "36": 10.90778, + "37": 10.87067, + "38": 10.87869, + "39": 10.88609, + "40": 10.89314, + "41": 10.88378, + "42": 10.89495, + "43": 10.88033, + "44": 10.88452, + "45": 10.88343, + "46": 10.88649, + "47": 10.88589, + "48": 10.86401, + "49": 10.87635, + "50": 10.88222, + "51": 10.8934, + "52": 10.87115, + "53": 10.85945, + "54": 10.87099, + "55": 10.86862, + "56": 10.87284, + "57": 10.84449, + "58": 10.85626, + "59": 10.84661, + "60": 10.84288, + "61": 10.86057, + "62": 10.85689, + "63": 10.86087, + "64": 10.83867, + "65": 10.82752, + "66": 10.84604, + "67": 10.82827, + "68": 10.8301, + "69": 10.81892, + "70": 10.82842, + "71": 10.82682, + "72": 10.80818, + "73": 10.8078, + "74": 10.8061, + "75": 10.81404, + "76": 10.81141, + "77": 10.80748, + "78": 10.79208, + "79": 10.80226, + "80": 10.78869, + "81": 10.79501, + "82": 10.79581, + "83": 10.78578, + "84": 10.75661, + "85": 10.7625, + "86": 10.77652, + "87": 10.79621, + "88": 10.7752, + "89": 10.77532, + "90": 10.76478, + "91": 10.74133, + "92": 10.76027, + "93": 10.74674, + "94": 10.73396, + "95": 10.75235, + "96": 10.72346, + "97": 10.71597, + "98": 10.72639, + "99": 10.74614, + "100": 10.69587 } }, "num-zeros": { @@ -113,104 +113,104 @@ "values": { "1": 1133.0, "2": 1120.0, - "3": 1371.0, - "4": 1297.0, - "5": 1158.0, - "6": 1283.0, - "7": 1334.0, - "8": 1122.0, - "9": 1269.0, - "10": 1134.0, - "11": 1184.0, - "12": 1223.0, - "13": 1340.0, - "14": 1232.0, - "15": 1234.0, - "16": 1179.0, - "17": 1212.0, - "18": 1083.0, - "19": 1187.0, - "20": 1174.0, - "21": 1196.0, - "22": 1227.0, - "23": 1318.0, - "24": 1113.0, - "25": 1191.0, - "26": 1243.0, - "27": 1343.0, - "28": 1251.0, - "29": 1186.0, - "30": 1221.0, - "31": 1205.0, - "32": 1269.0, - "33": 1254.0, - "34": 1196.0, - "35": 1308.0, - "36": 1277.0, - "37": 1196.0, - "38": 1321.0, - "39": 1407.0, - "40": 1299.0, - "41": 1284.0, - "42": 1170.0, - "43": 1170.0, - "44": 1252.0, - "45": 1154.0, - "46": 1403.0, - "47": 1148.0, - "48": 1176.0, - "49": 1389.0, - "50": 1329.0, - "51": 1209.0, - "52": 1224.0, - "53": 1276.0, - "54": 1198.0, - "55": 1179.0, - "56": 1166.0, - "57": 1168.0, - "58": 1318.0, - "59": 1047.0, - "60": 1118.0, - "61": 1192.0, - "62": 1108.0, - "63": 1225.0, - "64": 1256.0, - "65": 1221.0, - "66": 1228.0, - "67": 1193.0, - "68": 1279.0, - "69": 1141.0, - "70": 1248.0, - "71": 1165.0, - "72": 1230.0, - "73": 1236.0, - "74": 1336.0, - "75": 1346.0, - "76": 1259.0, - "77": 1205.0, - "78": 1217.0, - "79": 1298.0, - "80": 1228.0, - "81": 1335.0, - "82": 1198.0, - "83": 1246.0, - "84": 1152.0, - "85": 1250.0, - "86": 1287.0, - "87": 1111.0, - "88": 1169.0, - "89": 1119.0, - "90": 1402.0, - "91": 1297.0, - "92": 1194.0, - "93": 1271.0, - "94": 1409.0, - "95": 980.0, - "96": 1170.0, - "97": 1212.0, - "98": 1325.0, - "99": 1125.0, - "100": 1206.0 + "3": 1370.0, + "4": 1273.0, + "5": 1168.0, + "6": 1248.0, + "7": 1339.0, + "8": 1155.0, + "9": 1178.0, + "10": 1195.0, + "11": 1286.0, + "12": 1214.0, + "13": 1366.0, + "14": 1276.0, + "15": 1265.0, + "16": 1190.0, + "17": 1215.0, + "18": 1115.0, + "19": 1230.0, + "20": 1144.0, + "21": 1170.0, + "22": 1254.0, + "23": 1197.0, + "24": 1122.0, + "25": 1279.0, + "26": 1244.0, + "27": 1410.0, + "28": 1167.0, + "29": 1160.0, + "30": 1282.0, + "31": 1186.0, + "32": 1261.0, + "33": 1219.0, + "34": 1223.0, + "35": 1263.0, + "36": 1171.0, + "37": 1198.0, + "38": 1222.0, + "39": 1377.0, + "40": 1284.0, + "41": 1322.0, + "42": 1165.0, + "43": 1167.0, + "44": 1204.0, + "45": 1123.0, + "46": 1427.0, + "47": 1227.0, + "48": 1107.0, + "49": 1439.0, + "50": 1307.0, + "51": 1250.0, + "52": 1203.0, + "53": 1254.0, + "54": 1135.0, + "55": 1164.0, + "56": 1169.0, + "57": 1178.0, + "58": 1292.0, + "59": 1034.0, + "60": 1213.0, + "61": 1247.0, + "62": 1189.0, + "63": 1277.0, + "64": 1268.0, + "65": 1213.0, + "66": 1240.0, + "67": 1249.0, + "68": 1295.0, + "69": 1175.0, + "70": 1237.0, + "71": 1218.0, + "72": 1283.0, + "73": 1255.0, + "74": 1333.0, + "75": 1313.0, + "76": 1310.0, + "77": 1225.0, + "78": 1250.0, + "79": 1320.0, + "80": 1200.0, + "81": 1283.0, + "82": 1153.0, + "83": 1232.0, + "84": 1168.0, + "85": 1254.0, + "86": 1300.0, + "87": 1105.0, + "88": 1219.0, + "89": 1172.0, + "90": 1295.0, + "91": 1270.0, + "92": 1199.0, + "93": 1333.0, + "94": 1369.0, + "95": 1022.0, + "96": 1165.0, + "97": 1283.0, + "98": 1347.0, + "99": 1177.0, + "100": 1218.0 } }, "mem-allocated-bytes": { @@ -221,103 +221,103 @@ "1": 994063360.0, "2": 994054144.0, "3": 994029056.0, - "4": 994010624.0, - "5": 994059776.0, - "6": 994042368.0, + "4": 994009600.0, + "5": 994060288.0, + "6": 994042880.0, "7": 993970176.0, - "8": 994006528.0, - "9": 994008064.0, - "10": 994039808.0, - "11": 994015232.0, - "12": 994024960.0, - "13": 994029056.0, - "14": 994030592.0, - "15": 994007040.0, + "8": 994007552.0, + "9": 994009088.0, + "10": 994041856.0, + "11": 994016256.0, + "12": 994023936.0, + "13": 994027008.0, + "14": 994031104.0, + "15": 994005504.0, "16": 993992704.0, "17": 994001408.0, - "18": 994003968.0, - "19": 994032640.0, - "20": 993997312.0, - "21": 994052608.0, + "18": 994002944.0, + "19": 994033152.0, + "20": 993996288.0, + "21": 994050048.0, "22": 994075648.0, - "23": 994049024.0, + "23": 994051072.0, "24": 993964032.0, - "25": 994045440.0, + "25": 994044928.0, "26": 994076672.0, "27": 994028544.0, - "28": 993984512.0, - "29": 994056704.0, - "30": 993990656.0, - "31": 993996800.0, - "32": 994003456.0, + "28": 993984000.0, + "29": 994055168.0, + "30": 993992192.0, + "31": 993995264.0, + "32": 994003968.0, "33": 994011136.0, - "34": 993989120.0, - "35": 993974784.0, + "34": 993989632.0, + "35": 993973760.0, "36": 994040832.0, "37": 993988608.0, "38": 994046976.0, - "39": 994054144.0, - "40": 994001408.0, - "41": 994021888.0, - "42": 993964544.0, - "43": 994032640.0, - "44": 994014208.0, - "45": 993989120.0, - "46": 994052096.0, + "39": 994051584.0, + "40": 994000384.0, + "41": 994020352.0, + "42": 993965568.0, + "43": 994033152.0, + "44": 994017280.0, + "45": 993991168.0, + "46": 994054656.0, "47": 994011136.0, - "48": 994040320.0, - "49": 993998848.0, - "50": 994044928.0, - "51": 994030592.0, - "52": 994003968.0, - "53": 993953280.0, - "54": 994051072.0, + "48": 994041344.0, + "49": 994000896.0, + "50": 994046464.0, + "51": 994032128.0, + "52": 994002944.0, + "53": 993955840.0, + "54": 994051584.0, "55": 994039808.0, - "56": 994038784.0, - "57": 994098688.0, - "58": 994035200.0, - "59": 994040832.0, - "60": 994019328.0, - "61": 993988608.0, - "62": 993980928.0, - "63": 993987584.0, - "64": 994015744.0, - "65": 994023424.0, - "66": 994037248.0, - "67": 994018816.0, - "68": 994020352.0, - "69": 993985024.0, - "70": 994067456.0, - "71": 994047488.0, - "72": 994044416.0, - "73": 993973760.0, + "56": 994038272.0, + "57": 994098176.0, + "58": 994034176.0, + "59": 994037760.0, + "60": 994018816.0, + "61": 993988096.0, + "62": 993979904.0, + "63": 993989632.0, + "64": 994015232.0, + "65": 994025472.0, + "66": 994036224.0, + "67": 994019840.0, + "68": 994020864.0, + "69": 993984512.0, + "70": 994070016.0, + "71": 994049024.0, + "72": 994043904.0, + "73": 993974784.0, "74": 994015744.0, - "75": 994060800.0, - "76": 994017792.0, - "77": 994060288.0, - "78": 994006016.0, - "79": 994072064.0, - "80": 993993728.0, + "75": 994061824.0, + "76": 994018304.0, + "77": 994061312.0, + "78": 994007552.0, + "79": 994073600.0, + "80": 993994752.0, "81": 994046464.0, - "82": 993994752.0, + "82": 993993216.0, "83": 994054656.0, "84": 994024448.0, - "85": 994032128.0, - "86": 994016256.0, + "85": 994031616.0, + "86": 994015744.0, "87": 994030080.0, - "88": 994010112.0, - "89": 994001408.0, - "90": 994024960.0, - "91": 994011136.0, - "92": 994009600.0, - "93": 994040320.0, - "94": 994015232.0, - "95": 994036736.0, - "96": 994028544.0, - "97": 993997312.0, - "98": 994035712.0, - "99": 994051584.0, - "100": 994083328.0 + "88": 994011648.0, + "89": 994003968.0, + "90": 994026496.0, + "91": 994011648.0, + "92": 994007552.0, + "93": 994039296.0, + "94": 994014208.0, + "95": 994035712.0, + "96": 994028032.0, + "97": 993996288.0, + "98": 994035200.0, + "99": 994050560.0, + "100": 994080768.0 } }, "mem-max-allocated-bytes": { @@ -329,102 +329,102 @@ "2": 3490994688.0, "3": 3490994688.0, "4": 3490994688.0, - "5": 3498078208.0, - "6": 3498078208.0, - "7": 3498078208.0, - "8": 3498078208.0, - "9": 3498078208.0, - "10": 3498078208.0, - "11": 3498078208.0, - "12": 3498078208.0, - "13": 3498078208.0, - "14": 3498078208.0, - "15": 3498078208.0, - "16": 3498078208.0, - "17": 3498078208.0, - "18": 3498078208.0, - "19": 3498078208.0, - "20": 3498078208.0, - "21": 3498078208.0, - "22": 3507077632.0, - "23": 3507077632.0, - "24": 3507077632.0, - "25": 3507077632.0, - "26": 3510859264.0, - "27": 3510859264.0, - "28": 3510859264.0, - "29": 3510859264.0, - "30": 3510859264.0, - "31": 3510859264.0, - "32": 3510859264.0, - "33": 3510859264.0, - "34": 3510859264.0, - "35": 3510859264.0, - "36": 3510859264.0, - "37": 3510859264.0, - "38": 3510859264.0, - "39": 3510859264.0, - "40": 3510859264.0, - "41": 3510859264.0, - "42": 3510859264.0, - "43": 3510859264.0, - "44": 3510859264.0, - "45": 3510859264.0, - "46": 3510859264.0, - "47": 3510859264.0, - "48": 3510859264.0, - "49": 3510859264.0, - "50": 3510859264.0, - "51": 3510859264.0, - "52": 3510859264.0, - "53": 3510859264.0, - "54": 3510859264.0, - "55": 3510859264.0, - "56": 3510859264.0, - "57": 3526288384.0, - "58": 3526288384.0, - "59": 3526288384.0, - "60": 3526288384.0, - "61": 3526288384.0, - "62": 3526288384.0, - "63": 3526288384.0, - "64": 3526288384.0, - "65": 3526288384.0, - "66": 3526288384.0, - "67": 3526288384.0, - "68": 3526288384.0, - "69": 3526288384.0, - "70": 3526288384.0, - "71": 3526288384.0, - "72": 3526288384.0, - "73": 3526288384.0, - "74": 3526288384.0, - "75": 3526288384.0, - "76": 3526288384.0, - "77": 3526288384.0, - "78": 3526288384.0, - "79": 3526288384.0, - "80": 3526288384.0, - "81": 3526288384.0, - "82": 3526288384.0, - "83": 3526288384.0, - "84": 3526288384.0, - "85": 3526288384.0, - "86": 3526288384.0, - "87": 3526288384.0, - "88": 3526288384.0, - "89": 3526288384.0, - "90": 3526288384.0, - "91": 3526288384.0, - "92": 3526288384.0, - "93": 3526288384.0, - "94": 3526288384.0, - "95": 3526288384.0, - "96": 3526288384.0, - "97": 3526288384.0, - "98": 3526288384.0, - "99": 3526288384.0, - "100": 3526288384.0 + "5": 3498184192.0, + "6": 3498184192.0, + "7": 3498184192.0, + "8": 3498184192.0, + "9": 3498184192.0, + "10": 3498184192.0, + "11": 3498184192.0, + "12": 3498184192.0, + "13": 3498184192.0, + "14": 3498184192.0, + "15": 3498184192.0, + "16": 3498184192.0, + "17": 3498184192.0, + "18": 3498184192.0, + "19": 3498184192.0, + "20": 3498184192.0, + "21": 3498184192.0, + "22": 3505588736.0, + "23": 3505588736.0, + "24": 3505588736.0, + "25": 3505588736.0, + "26": 3511026176.0, + "27": 3511026176.0, + "28": 3511026176.0, + "29": 3511026176.0, + "30": 3511026176.0, + "31": 3511026176.0, + "32": 3511026176.0, + "33": 3511026176.0, + "34": 3511026176.0, + "35": 3511026176.0, + "36": 3511026176.0, + "37": 3511026176.0, + "38": 3511026176.0, + "39": 3511026176.0, + "40": 3511026176.0, + "41": 3511026176.0, + "42": 3511026176.0, + "43": 3511026176.0, + "44": 3511026176.0, + "45": 3511026176.0, + "46": 3511026176.0, + "47": 3511026176.0, + "48": 3511026176.0, + "49": 3511026176.0, + "50": 3511026176.0, + "51": 3511026176.0, + "52": 3511026176.0, + "53": 3511026176.0, + "54": 3511026176.0, + "55": 3511026176.0, + "56": 3511026176.0, + "57": 3524790784.0, + "58": 3524790784.0, + "59": 3524790784.0, + "60": 3524790784.0, + "61": 3524790784.0, + "62": 3524790784.0, + "63": 3524790784.0, + "64": 3524790784.0, + "65": 3524790784.0, + "66": 3524790784.0, + "67": 3524790784.0, + "68": 3524790784.0, + "69": 3524790784.0, + "70": 3524790784.0, + "71": 3524790784.0, + "72": 3524790784.0, + "73": 3524790784.0, + "74": 3524790784.0, + "75": 3524790784.0, + "76": 3524790784.0, + "77": 3524790784.0, + "78": 3524790784.0, + "79": 3524790784.0, + "80": 3524790784.0, + "81": 3524790784.0, + "82": 3524790784.0, + "83": 3524790784.0, + "84": 3524790784.0, + "85": 3524790784.0, + "86": 3524790784.0, + "87": 3524790784.0, + "88": 3524790784.0, + "89": 3524790784.0, + "90": 3524790784.0, + "91": 3524790784.0, + "92": 3524790784.0, + "93": 3524790784.0, + "94": 3524790784.0, + "95": 3524790784.0, + "96": 3524790784.0, + "97": 3524790784.0, + "98": 3524790784.0, + "99": 3524790784.0, + "100": 3524790784.0 } }, "iteration-time": { @@ -433,105 +433,105 @@ "step_interval": 1, "values": { "1": "nan", - "2": 4.80343, - "3": 0.24461, - "4": 0.21905, - "5": 0.23456, - "6": 0.22151, - "7": 0.22611, - "8": 0.2404, - "9": 0.21788, - "10": 0.22319, - "11": 0.22849, - "12": 0.21807, - "13": 0.2241, - "14": 0.21328, - "15": 0.21924, - "16": 0.19766, - "17": 0.2119, - "18": 0.20845, - "19": 0.2102, - "20": 0.21676, - "21": 0.2071, - "22": 0.21066, - "23": 0.20145, - "24": 0.20861, - "25": 0.22351, - "26": 0.20212, - "27": 0.19701, - "28": 0.19726, - "29": 0.21428, - "30": 0.20896, - "31": 0.20412, - "32": 0.21173, - "33": 0.20281, - "34": 0.20432, - "35": 0.20437, - "36": 0.20527, - "37": 0.20735, - "38": 0.19734, - "39": 0.19343, - "40": 0.20184, - "41": 0.19924, - "42": 0.19429, - "43": 0.20453, - "44": 0.19924, - "45": 0.20236, - "46": 0.19519, - "47": 0.1998, - "48": 0.20296, - "49": 0.20318, - "50": 0.19638, - "51": 0.25423, - "52": 0.22224, - "53": 0.20083, - "54": 0.19378, - "55": 0.19638, - "56": 0.19433, - "57": 0.21583, - "58": 0.20227, - "59": 0.19823, - "60": 0.2012, - "61": 0.20156, - "62": 0.21045, - "63": 0.19778, - "64": 0.19188, - "65": 0.19114, - "66": 0.19429, - "67": 0.21388, - "68": 0.19833, - "69": 0.20586, - "70": 0.19598, - "71": 0.1923, - "72": 0.19411, - "73": 0.20635, - "74": 0.20087, - "75": 0.19463, - "76": 0.19501, - "77": 0.20497, - "78": 0.19948, - "79": 0.19854, - "80": 0.19801, - "81": 0.21063, - "82": 0.20059, - "83": 0.19918, - "84": 0.20511, - "85": 0.19701, - "86": 0.20162, - "87": 0.20013, - "88": 0.20349, - "89": 0.19494, - "90": 0.19192, - "91": 0.20582, - "92": 0.19576, - "93": 0.20142, - "94": 0.20002, - "95": 0.1982, - "96": 0.20164, - "97": 0.20778, - "98": 0.19874, - "99": 0.19422, - "100": 0.20224 + "2": 5.6337, + "3": 0.20761, + "4": 0.16601, + "5": 0.17758, + "6": 0.16033, + "7": 0.16776, + "8": 0.17113, + "9": 0.15346, + "10": 0.15457, + "11": 0.16483, + "12": 0.15508, + "13": 0.15225, + "14": 0.15196, + "15": 0.14958, + "16": 0.15078, + "17": 0.15458, + "18": 0.14955, + "19": 0.15431, + "20": 0.1519, + "21": 0.14861, + "22": 0.14651, + "23": 0.14479, + "24": 0.14144, + "25": 0.15196, + "26": 0.15235, + "27": 0.14741, + "28": 0.13978, + "29": 0.14317, + "30": 0.14921, + "31": 0.13637, + "32": 0.14455, + "33": 0.14591, + "34": 0.1451, + "35": 0.14294, + "36": 0.1403, + "37": 0.13387, + "38": 0.14548, + "39": 0.14062, + "40": 0.14073, + "41": 0.13594, + "42": 0.1388, + "43": 0.14951, + "44": 0.13714, + "45": 0.14626, + "46": 0.13564, + "47": 0.14163, + "48": 0.1323, + "49": 0.14165, + "50": 0.13992, + "51": 0.24023, + "52": 0.15421, + "53": 0.16392, + "54": 0.13573, + "55": 0.13475, + "56": 0.13769, + "57": 0.13657, + "58": 0.14113, + "59": 0.13818, + "60": 0.13927, + "61": 0.14266, + "62": 0.14324, + "63": 0.13522, + "64": 0.13537, + "65": 0.13339, + "66": 0.13614, + "67": 0.13038, + "68": 0.13616, + "69": 0.13313, + "70": 0.13208, + "71": 0.13396, + "72": 0.14018, + "73": 0.13913, + "74": 0.13548, + "75": 0.13183, + "76": 0.132, + "77": 0.14357, + "78": 0.13551, + "79": 0.13911, + "80": 0.14089, + "81": 0.13517, + "82": 0.13635, + "83": 0.13307, + "84": 0.13689, + "85": 0.13055, + "86": 0.14014, + "87": 0.14111, + "88": 0.1388, + "89": 0.13409, + "90": 0.13302, + "91": 0.14501, + "92": 0.14222, + "93": 0.13009, + "94": 0.14474, + "95": 0.14034, + "96": 0.13611, + "97": 0.14187, + "98": 0.13568, + "99": 0.13534, + "100": 0.14745 } } } \ No newline at end of file diff --git a/tests/functional_tests/test_cases/moe/gpt3_moe_mcore_te_ep8_resume_torch_dist_dist_muon_1node/golden_values_dev_dgx_gb200.json b/tests/functional_tests/test_cases/moe/gpt3_moe_mcore_te_ep8_resume_torch_dist_dist_muon_1node/golden_values_dev_dgx_gb200.json index 7edbab4f74a..ef2d86e16c2 100644 --- a/tests/functional_tests/test_cases/moe/gpt3_moe_mcore_te_ep8_resume_torch_dist_dist_muon_1node/golden_values_dev_dgx_gb200.json +++ b/tests/functional_tests/test_cases/moe/gpt3_moe_mcore_te_ep8_resume_torch_dist_dist_muon_1node/golden_values_dev_dgx_gb200.json @@ -6,104 +6,104 @@ "values": { "1": 10.93064, "2": 10.92264, - "3": 10.92451, - "4": 10.91716, - "5": 10.92714, - "6": 10.92064, - "7": 10.91647, - "8": 10.91788, - "9": 10.93809, - "10": 10.924, - "11": 10.91584, - "12": 10.92558, - "13": 10.9341, - "14": 10.91552, - "15": 10.92865, - "16": 10.91861, - "17": 10.92663, - "18": 10.91342, - "19": 10.93055, - "20": 10.89988, - "21": 10.94059, - "22": 10.92549, - "23": 10.92966, - "24": 10.90596, - "25": 10.91497, - "26": 10.92625, - "27": 10.92807, - "28": 10.92038, - "29": 10.92592, - "30": 10.91434, - "31": 10.91031, - "32": 10.919, - "33": 10.92014, - "34": 10.89998, - "35": 10.91944, - "36": 10.90693, - "37": 10.91287, - "38": 10.92352, - "39": 10.91196, - "40": 10.91094, - "41": 10.91323, - "42": 10.90919, - "43": 10.90589, - "44": 10.90761, - "45": 10.89741, - "46": 10.90636, - "47": 10.90203, - "48": 10.88495, - "49": 10.90015, - "50": 10.89905, - "51": 10.89814, - "52": 10.90393, - "53": 10.90158, - "54": 10.88808, - "55": 10.88643, - "56": 10.88868, - "57": 10.88748, - "58": 10.88859, - "59": 10.88263, - "60": 10.87391, - "61": 10.87367, - "62": 10.86781, - "63": 10.87607, - "64": 10.85882, - "65": 10.86026, - "66": 10.85712, - "67": 10.86113, - "68": 10.86032, - "69": 10.84979, - "70": 10.86738, - "71": 10.8492, - "72": 10.84019, - "73": 10.84868, - "74": 10.83869, - "75": 10.83621, - "76": 10.83434, - "77": 10.82938, - "78": 10.82319, - "79": 10.82846, - "80": 10.82987, - "81": 10.81761, - "82": 10.81704, - "83": 10.80437, - "84": 10.78281, - "85": 10.78319, - "86": 10.80029, - "87": 10.7956, - "88": 10.79539, - "89": 10.77706, - "90": 10.78274, - "91": 10.79015, - "92": 10.78298, - "93": 10.75507, - "94": 10.7607, - "95": 10.77273, - "96": 10.7355, - "97": 10.7421, - "98": 10.74705, - "99": 10.76264, - "100": 10.74262 + "3": 10.92453, + "4": 10.91684, + "5": 10.92748, + "6": 10.92084, + "7": 10.91663, + "8": 10.91765, + "9": 10.93835, + "10": 10.92338, + "11": 10.91561, + "12": 10.92588, + "13": 10.93465, + "14": 10.91542, + "15": 10.92828, + "16": 10.91838, + "17": 10.92685, + "18": 10.91279, + "19": 10.93034, + "20": 10.9007, + "21": 10.94026, + "22": 10.92598, + "23": 10.93005, + "24": 10.9059, + "25": 10.91516, + "26": 10.92626, + "27": 10.92772, + "28": 10.92055, + "29": 10.92569, + "30": 10.91405, + "31": 10.91037, + "32": 10.91862, + "33": 10.92051, + "34": 10.90093, + "35": 10.91967, + "36": 10.90728, + "37": 10.91249, + "38": 10.92368, + "39": 10.9118, + "40": 10.90995, + "41": 10.91309, + "42": 10.90932, + "43": 10.90607, + "44": 10.90718, + "45": 10.89743, + "46": 10.90666, + "47": 10.90179, + "48": 10.88458, + "49": 10.89985, + "50": 10.89888, + "51": 10.89767, + "52": 10.90451, + "53": 10.90156, + "54": 10.88797, + "55": 10.88694, + "56": 10.8894, + "57": 10.88717, + "58": 10.88849, + "59": 10.88234, + "60": 10.87379, + "61": 10.87315, + "62": 10.86821, + "63": 10.87594, + "64": 10.85943, + "65": 10.86096, + "66": 10.8567, + "67": 10.86183, + "68": 10.86025, + "69": 10.84972, + "70": 10.86629, + "71": 10.84873, + "72": 10.84049, + "73": 10.84902, + "74": 10.83882, + "75": 10.83635, + "76": 10.83428, + "77": 10.82965, + "78": 10.82333, + "79": 10.82875, + "80": 10.82933, + "81": 10.81781, + "82": 10.81672, + "83": 10.80466, + "84": 10.78292, + "85": 10.78303, + "86": 10.80041, + "87": 10.79499, + "88": 10.79531, + "89": 10.77724, + "90": 10.78302, + "91": 10.78962, + "92": 10.78328, + "93": 10.75473, + "94": 10.76112, + "95": 10.77316, + "96": 10.73525, + "97": 10.74204, + "98": 10.74686, + "99": 10.76311, + "100": 10.74203 } }, "num-zeros": { @@ -113,104 +113,104 @@ "values": { "1": 49180.0, "2": 50476.0, - "3": 49283.0, - "4": 48797.0, - "5": 50222.0, - "6": 50071.0, - "7": 51516.0, - "8": 49687.0, - "9": 50960.0, - "10": 52404.0, - "11": 50454.0, - "12": 48350.0, - "13": 52171.0, - "14": 49951.0, - "15": 48080.0, - "16": 49110.0, - "17": 50288.0, - "18": 52118.0, - "19": 51099.0, - "20": 52646.0, - "21": 49816.0, - "22": 49955.0, - "23": 53512.0, - "24": 51380.0, - "25": 48547.0, - "26": 52530.0, - "27": 51993.0, - "28": 48540.0, - "29": 50032.0, - "30": 51384.0, - "31": 53098.0, - "32": 49324.0, - "33": 51925.0, - "34": 50079.0, - "35": 51371.0, - "36": 51207.0, - "37": 49665.0, - "38": 50936.0, - "39": 50632.0, - "40": 50968.0, - "41": 48043.0, - "42": 50241.0, - "43": 50548.0, - "44": 47518.0, - "45": 54359.0, - "46": 47844.0, - "47": 50697.0, - "48": 50607.0, - "49": 54927.0, - "50": 51129.0, - "51": 48773.0, - "52": 50817.0, - "53": 49252.0, - "54": 48794.0, - "55": 49077.0, - "56": 48641.0, - "57": 48141.0, - "58": 52709.0, - "59": 50075.0, - "60": 49760.0, - "61": 47761.0, - "62": 50557.0, - "63": 48459.0, - "64": 56353.0, - "65": 51179.0, - "66": 48545.0, - "67": 51570.0, - "68": 49653.0, - "69": 54155.0, - "70": 49808.0, - "71": 49534.0, - "72": 52109.0, - "73": 52558.0, - "74": 49904.0, - "75": 50500.0, - "76": 52020.0, - "77": 49771.0, - "78": 49836.0, - "79": 48741.0, - "80": 52069.0, - "81": 50079.0, - "82": 46782.0, - "83": 52982.0, - "84": 48744.0, - "85": 50660.0, - "86": 50050.0, - "87": 46875.0, - "88": 48017.0, - "89": 48621.0, - "90": 51995.0, - "91": 50317.0, - "92": 50374.0, - "93": 50702.0, - "94": 51589.0, - "95": 48646.0, - "96": 50031.0, - "97": 49732.0, - "98": 48573.0, - "99": 48599.0, - "100": 47055.0 + "3": 49148.0, + "4": 49221.0, + "5": 50682.0, + "6": 49706.0, + "7": 52164.0, + "8": 49362.0, + "9": 51071.0, + "10": 52111.0, + "11": 50024.0, + "12": 48487.0, + "13": 51621.0, + "14": 50362.0, + "15": 48110.0, + "16": 49418.0, + "17": 49255.0, + "18": 51746.0, + "19": 51208.0, + "20": 52454.0, + "21": 49543.0, + "22": 50198.0, + "23": 53555.0, + "24": 51542.0, + "25": 48867.0, + "26": 51811.0, + "27": 51985.0, + "28": 48344.0, + "29": 50154.0, + "30": 50726.0, + "31": 52482.0, + "32": 49744.0, + "33": 52098.0, + "34": 50352.0, + "35": 50951.0, + "36": 51185.0, + "37": 49112.0, + "38": 50901.0, + "39": 50734.0, + "40": 51221.0, + "41": 48440.0, + "42": 50568.0, + "43": 50543.0, + "44": 47760.0, + "45": 54358.0, + "46": 47826.0, + "47": 50424.0, + "48": 50474.0, + "49": 54657.0, + "50": 50617.0, + "51": 48543.0, + "52": 50645.0, + "53": 49739.0, + "54": 48915.0, + "55": 49358.0, + "56": 48052.0, + "57": 48401.0, + "58": 52927.0, + "59": 50044.0, + "60": 49888.0, + "61": 47052.0, + "62": 50534.0, + "63": 48586.0, + "64": 56168.0, + "65": 50762.0, + "66": 49035.0, + "67": 52223.0, + "68": 49435.0, + "69": 54928.0, + "70": 50021.0, + "71": 49480.0, + "72": 52091.0, + "73": 52570.0, + "74": 49428.0, + "75": 50771.0, + "76": 51258.0, + "77": 50061.0, + "78": 50005.0, + "79": 48775.0, + "80": 52167.0, + "81": 49956.0, + "82": 47089.0, + "83": 53102.0, + "84": 48471.0, + "85": 50177.0, + "86": 49413.0, + "87": 46817.0, + "88": 48435.0, + "89": 48622.0, + "90": 52230.0, + "91": 50589.0, + "92": 50227.0, + "93": 50882.0, + "94": 51926.0, + "95": 48591.0, + "96": 50036.0, + "97": 49824.0, + "98": 48441.0, + "99": 49162.0, + "100": 47469.0 } }, "mem-allocated-bytes": { @@ -221,103 +221,103 @@ "1": 870735872.0, "2": 870754816.0, "3": 870738944.0, - "4": 870736896.0, + "4": 870737408.0, "5": 870747136.0, "6": 870734848.0, - "7": 870729728.0, + "7": 870730240.0, "8": 870758400.0, "9": 870706688.0, - "10": 870708224.0, - "11": 870691840.0, + "10": 870705664.0, + "11": 870693888.0, "12": 870722560.0, - "13": 870722048.0, + "13": 870723584.0, "14": 870712320.0, "15": 870742016.0, - "16": 870731264.0, - "17": 870742528.0, - "18": 870718976.0, - "19": 870715904.0, - "20": 870726656.0, + "16": 870732288.0, + "17": 870741504.0, + "18": 870718464.0, + "19": 870714880.0, + "20": 870730240.0, "21": 870749696.0, - "22": 870685696.0, - "23": 870727168.0, + "22": 870686208.0, + "23": 870728704.0, "24": 870740992.0, "25": 870705664.0, - "26": 870743040.0, - "27": 870745600.0, - "28": 870710272.0, - "29": 870727680.0, - "30": 870737408.0, - "31": 870700544.0, + "26": 870740992.0, + "27": 870744576.0, + "28": 870708736.0, + "29": 870727168.0, + "30": 870739456.0, + "31": 870701568.0, "32": 870722048.0, - "33": 870707712.0, + "33": 870708736.0, "34": 870750208.0, - "35": 870678016.0, - "36": 870753792.0, - "37": 870719488.0, - "38": 870744064.0, - "39": 870709760.0, - "40": 870759424.0, - "41": 870726656.0, - "42": 870705664.0, - "43": 870706176.0, - "44": 870716416.0, - "45": 870730752.0, - "46": 870681600.0, - "47": 870704640.0, - "48": 870693888.0, - "49": 870726144.0, - "50": 870706688.0, - "51": 870676480.0, - "52": 870724096.0, - "53": 870680576.0, + "35": 870679040.0, + "36": 870750720.0, + "37": 870718464.0, + "38": 870747136.0, + "39": 870708736.0, + "40": 870758912.0, + "41": 870722560.0, + "42": 870704128.0, + "43": 870708736.0, + "44": 870717952.0, + "45": 870729728.0, + "46": 870680064.0, + "47": 870703616.0, + "48": 870692864.0, + "49": 870724608.0, + "50": 870708224.0, + "51": 870676992.0, + "52": 870723072.0, + "53": 870681088.0, "54": 870741504.0, - "55": 870686208.0, - "56": 870737920.0, - "57": 870706176.0, - "58": 870713856.0, - "59": 870715904.0, - "60": 870680576.0, - "61": 870662656.0, - "62": 870695936.0, + "55": 870687744.0, + "56": 870736384.0, + "57": 870708736.0, + "58": 870713344.0, + "59": 870713344.0, + "60": 870679552.0, + "61": 870661632.0, + "62": 870694400.0, "63": 870699008.0, - "64": 870713344.0, - "65": 870717952.0, - "66": 870641152.0, - "67": 870729216.0, - "68": 870673920.0, - "69": 870702080.0, - "70": 870668288.0, - "71": 870687744.0, - "72": 870685184.0, - "73": 870686208.0, + "64": 870714368.0, + "65": 870718976.0, + "66": 870639104.0, + "67": 870726144.0, + "68": 870674432.0, + "69": 870701056.0, + "70": 870667264.0, + "71": 870687232.0, + "72": 870686720.0, + "73": 870685184.0, "74": 870650368.0, - "75": 870676992.0, - "76": 870678016.0, - "77": 870716928.0, - "78": 870710784.0, - "79": 870678528.0, + "75": 870674944.0, + "76": 870679040.0, + "77": 870714880.0, + "78": 870710272.0, + "79": 870679040.0, "80": 870626816.0, - "81": 870669312.0, - "82": 870644736.0, - "83": 870636032.0, - "84": 870677504.0, - "85": 870650880.0, - "86": 870644736.0, - "87": 870636032.0, - "88": 870648832.0, + "81": 870668800.0, + "82": 870643712.0, + "83": 870637568.0, + "84": 870681600.0, + "85": 870651392.0, + "86": 870644224.0, + "87": 870635008.0, + "88": 870651904.0, "89": 870634496.0, - "90": 870674432.0, - "91": 870637056.0, - "92": 870639104.0, - "93": 870653952.0, - "94": 870635520.0, - "95": 870631424.0, - "96": 870637056.0, - "97": 870630912.0, + "90": 870673920.0, + "91": 870638080.0, + "92": 870637568.0, + "93": 870652928.0, + "94": 870636544.0, + "95": 870632960.0, + "96": 870637568.0, + "97": 870629376.0, "98": 870680064.0, - "99": 870621184.0, - "100": 870622720.0 + "99": 870619136.0, + "100": 870623744.0 } }, "mem-max-allocated-bytes": { @@ -325,7 +325,7 @@ "end_step": 100, "step_interval": 1, "values": { - "1": 3338426368.0, + "1": 3338425856.0, "2": 3490031616.0, "3": 3490031616.0, "4": 3490031616.0, @@ -433,105 +433,105 @@ "step_interval": 1, "values": { "1": "nan", - "2": 5.31191, - "3": 0.4081, - "4": 0.75684, - "5": 0.40216, - "6": 0.79212, - "7": 0.67838, - "8": 1.14636, - "9": 0.66928, - "10": 0.39229, - "11": 1.0781, - "12": 0.3588, - "13": 0.48853, - "14": 1.01991, - "15": 0.39449, - "16": 1.00595, - "17": 0.57485, - "18": 1.08493, - "19": 0.39338, - "20": 0.52887, - "21": 0.75728, - "22": 0.76739, - "23": 0.68487, - "24": 0.3607, - "25": 0.46274, - "26": 1.16536, - "27": 0.52339, - "28": 0.61313, - "29": 0.89376, - "30": 0.35668, - "31": 0.35172, - "32": 0.98229, - "33": 0.76729, - "34": 0.45405, - "35": 0.90679, - "36": 0.62691, - "37": 0.80734, - "38": 0.52682, - "39": 0.61651, - "40": 0.72001, - "41": 0.97926, - "42": 0.57908, - "43": 0.59221, - "44": 0.65813, - "45": 0.7194, - "46": 0.76818, - "47": 0.82352, - "48": 1.01064, - "49": 0.36409, - "50": 0.3676, - "51": 0.41549, - "52": 0.39786, - "53": 0.64141, - "54": 0.81457, - "55": 0.47145, - "56": 0.75849, - "57": 0.47005, - "58": 0.54556, - "59": 0.65077, - "60": 0.67032, - "61": 0.36282, - "62": 1.07816, - "63": 0.36068, - "64": 0.81042, - "65": 0.43636, - "66": 0.65437, - "67": 0.61425, - "68": 0.36157, - "69": 0.60737, - "70": 1.00435, - "71": 0.35852, - "72": 0.49448, - "73": 0.69778, - "74": 0.4904, - "75": 1.02199, - "76": 0.43752, - "77": 0.42751, - "78": 1.80056, - "79": 0.65201, - "80": 0.64025, - "81": 0.64801, - "82": 0.44684, - "83": 0.8134, - "84": 0.73747, - "85": 0.35805, - "86": 0.39074, - "87": 0.95679, - "88": 0.91998, - "89": 0.64926, - "90": 0.48592, - "91": 0.55213, - "92": 0.83853, - "93": 0.44924, - "94": 0.63982, - "95": 0.64996, - "96": 0.50551, - "97": 0.47138, - "98": 0.72542, - "99": 1.35867, - "100": 0.35987 + "2": 6.02666, + "3": 0.43915, + "4": 0.36195, + "5": 0.35684, + "6": 0.35444, + "7": 0.36336, + "8": 0.35446, + "9": 0.34478, + "10": 0.34852, + "11": 0.34654, + "12": 0.34237, + "13": 0.34497, + "14": 0.34651, + "15": 0.34467, + "16": 0.35275, + "17": 0.34175, + "18": 0.34967, + "19": 0.34081, + "20": 0.33758, + "21": 0.34131, + "22": 0.34455, + "23": 0.35146, + "24": 0.34178, + "25": 0.34277, + "26": 0.33895, + "27": 0.3493, + "28": 0.33968, + "29": 0.34207, + "30": 0.34242, + "31": 0.34754, + "32": 0.34906, + "33": 0.34988, + "34": 0.34862, + "35": 0.34036, + "36": 0.3392, + "37": 0.34519, + "38": 0.34568, + "39": 0.34746, + "40": 0.34587, + "41": 0.34669, + "42": 0.34494, + "43": 0.34844, + "44": 0.34959, + "45": 0.34594, + "46": 0.34489, + "47": 0.34761, + "48": 0.34307, + "49": 0.33768, + "50": 0.33971, + "51": 0.4529, + "52": 0.39085, + "53": 0.34493, + "54": 0.34303, + "55": 0.34229, + "56": 0.33766, + "57": 0.33749, + "58": 0.33985, + "59": 0.34458, + "60": 0.34329, + "61": 0.34042, + "62": 0.33795, + "63": 0.3365, + "64": 0.33438, + "65": 0.33378, + "66": 0.33455, + "67": 0.3392, + "68": 0.34664, + "69": 0.34672, + "70": 0.34101, + "71": 0.33693, + "72": 0.34054, + "73": 0.34195, + "74": 0.34482, + "75": 0.33567, + "76": 0.34233, + "77": 0.34611, + "78": 0.34429, + "79": 0.34305, + "80": 0.34071, + "81": 0.34342, + "82": 0.33755, + "83": 0.33834, + "84": 0.33852, + "85": 0.3383, + "86": 0.34005, + "87": 0.34219, + "88": 0.34626, + "89": 0.34147, + "90": 0.3445, + "91": 0.34459, + "92": 0.34476, + "93": 0.34958, + "94": 0.34935, + "95": 0.77477, + "96": 0.33856, + "97": 0.34091, + "98": 0.34021, + "99": 0.34056, + "100": 0.33621 } } } \ No newline at end of file diff --git a/tests/unit_tests/test_muon_decouple_fp8_param_gather.py b/tests/unit_tests/test_muon_decouple_fp8_param_gather.py new file mode 100644 index 00000000000..b7e24ffcad3 --- /dev/null +++ b/tests/unit_tests/test_muon_decouple_fp8_param_gather.py @@ -0,0 +1,453 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +"""Bitwise ON-vs-OFF check for muon + ``--fp8-param-gather`` on the DECOUPLED +compact LayerWise layout (the default; ``--use-layer-wise-param-layout`` opts into +the padded layout). + +Construction mirrors tests/unit_tests/test_muon_mxfp8_fp8_param_gather.py +(PR #4987), but on the decouple path — the route where FP8 param-gather is +*supported* (it is rejected on the padded layout in ``validate_args``): blockwise +(Hopper) needs no ``reuse_grad_buf``; mxfp8 (Blackwell) does. Under deterministic +kernels, fp8_param_gather ON must match OFF bitwise on per-step loss, forward +output, per-param ``main_grad``, fp32 master, and (bf16) model param. Covers +``overlap_grad_reduce`` + ``overlap_param_gather`` both ON and OFF. +""" + +import gc +import os +import sys + +import pytest +import torch +from transformer_engine.pytorch.fp8 import check_fp8_support + +from megatron.core.enums import ModelType +from megatron.core.inference.utils import InferenceMode +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 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() + + +def _is_quantized(p): + return hasattr(p, 'dequantize') or hasattr(p.data, 'dequantize') + + +def _assert_equal(actual, expected, msg): + if torch.equal(actual, expected): + return + diff = (actual.float() - expected.float()).abs() + raise AssertionError( + f"{msg}: max_diff={diff.max().item()} dtype={actual.dtype}/{expected.dtype}" + ) + + +def _snapshot_masters(model): + # fp32 masters keyed by param name (stable identity across ON/OFF group layouts). + return { + n: p.main_param.detach().clone() + for n, p in model.named_parameters() + if getattr(p, 'main_param', None) is not None + } + + +def _snapshot_params(model, include_quantized=False): + # By default only non-quantized (bf16) params; fp8 weights are compared via master + grad. + # With include_quantized, fp8 params are dequantized so gathered fp8 bytes are checked directly. + out = {} + for n, p in model.named_parameters(): + if _is_quantized(p): + if include_quantized: + out[n] = p.data.dequantize().detach().clone().float() + else: + out[n] = p.detach().clone() + return out + + +def _snapshot_layerwise_grad_data(ddp): + # grad_data of the non-DistOpt LayerWise buffers (reused as the fp8 all-gather's bf16 + # receive buffer, which the param-sync finalize must re-zero). + return [ + buf.grad_data.detach().clone() + for buf in (ddp.buffers + ddp.expert_parallel_buffers) + if not buf.ddp_config.use_distributed_optimizer + ] + + +@torch.no_grad() +def _restore_initial_state(model, optimizer, params0, masters0): + # Start the ON run from the OFF run's exact init (params + fp32 masters); muon + # momentum state is empty pre-step, so it needs no restore. + for n, p in model.named_parameters(): + if n in params0: + p.data.copy_(params0[n].to(p.device)) + optimizer.reload_model_params() + for n, p in model.named_parameters(): + mp = getattr(p, 'main_param', None) + if mp is not None and n in masters0: + mp.data.copy_(masters0[n].to(mp.device)) + + +class TestMuonDecoupleFP8ParamGather: + + def setup_method(self, method): + self.seq_length = 128 + self.micro_batch_size = 1 + os.environ['CUDA_DEVICE_MAX_CONNECTIONS'] = '1' + # InferenceMode is a process-global (class-level) flag. Another test file + # in the same pytest shard can leave it active (e.g. an inference engine + # test that aborts before unset). These are training tests, so the GPT + # postprocess "Inference must always gather TP logits" assertion would + # then fire spuriously. Force training mode before each test. + InferenceMode.unset_active() + Utils.initialize_model_parallel( + tensor_model_parallel_size=1, + pipeline_model_parallel_size=1, + expert_model_parallel_size=1, + ) + + def teardown_method(self, method): + InferenceMode.unset_active() + Utils.destroy_model_parallel() + destroy_global_vars() + destroy_num_microbatches_calculator() + gc.collect() + + def model_provider(self, pre_process=True, post_process=True, **kw): + model_parallel_cuda_manual_seed(_SEED) + args = get_args() + return GPTModel( + config=core_transformer_config_from_args(args), + transformer_layer_spec=get_gpt_layer_with_transformer_engine_spec(), + vocab_size=args.vocab_size, + max_sequence_length=args.max_position_embeddings, + pre_process=pre_process, + post_process=post_process, + share_embeddings_and_output_weights=not args.untie_embeddings_and_output_weights, + position_embedding_type=args.position_embedding_type, + ) + + def _create_args( + self, fp8_param_gather, fp8_recipe, overlap, num_experts=0, expert_model_parallel_size=1 + ): + destroy_global_vars() + destroy_num_microbatches_calculator() + sys.argv = ['test_muon_decouple_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.seq_length = self.seq_length + args.micro_batch_size = self.micro_batch_size + args.create_attention_mask_in_dataloader = True + args.tensor_model_parallel_size = 1 + args.pipeline_model_parallel_size = 1 + args.context_parallel_size = 1 + args.expert_model_parallel_size = expert_model_parallel_size + args.train_iters = 10 + # Larger lr than a real run: amplifies any fp8-path discrepancy so an ON-vs-OFF + # mismatch (if one exists) shows up within a handful of steps rather than being + # lost in the low bits (per PR #5470 review). The model is tiny so this stays stable. + args.lr = 1e-3 + 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" + # muon + use_distributed_optimizer auto-routes to LayerWiseDistributedOptimizer. + 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 + # DECOUPLED compact LayerWise layout (the supported FP8 param-gather route). + args.use_layer_wise_param_layout = False + # --overlap-param-gather requires --overlap-grad-reduce (arguments.py); co-enable. + args.overlap_param_gather = overlap + args.overlap_grad_reduce = overlap + 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 # mxfp8 columnwise needs the bf16 round-trip + ) + if num_experts > 0: + # MoE variant: expert weights are 2D matrices -> Muon-managed, so they ride the + # LayerWise param path. At expt_dp == 1 (expert_model_parallel_size == world size) + # the experts are NOT all-gathered and only get the fp8 master->model copy-back; at + # expt_dp > 1 they are gathered. This covers both branches (PR #5470 review). + args.num_experts = num_experts + args.moe_router_topk = 2 + args.moe_ffn_hidden_size = args.ffn_hidden_size + args.moe_token_dispatcher_type = 'alltoall' + args.moe_grouped_gemm = False + # Deterministic routing comes from the fixed seed + deterministic_mode; drop the + # aux-loss gradient term so ON and OFF compare cleanly without router-bias drift. + args.moe_router_load_balancing_type = 'none' + args.moe_aux_loss_coeff = 0.0 + args.ddp_bucket_size = 1024 # more buckets -> exercise rs/ag overlap + validate_args(args) + set_global_variables(args, False) + return args + + def _batch(self): + d = list(range(self.seq_length)) + ids = torch.tensor(d, dtype=torch.int64).repeat((self.micro_batch_size, 1)).cuda() + labels = 1 + ids + pos = ids.clone() + 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 ids, labels, pos, mask, loss_mask + + def _build( + self, fp8_param_gather, fp8_recipe, overlap, num_experts=0, expert_model_parallel_size=1 + ): + args = self._create_args( + fp8_param_gather, + fp8_recipe, + overlap, + num_experts=num_experts, + expert_model_parallel_size=expert_model_parallel_size, + ) + set_args(args) + torch.manual_seed(_SEED) + model, optimizer, _ = setup_model_and_optimizer( + ModelType.encoder_or_decoder, self.model_provider + ) + assert len(model) == 1 + 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, model, optimizer + + def _run_steps(self, args, model, optimizer, n): + """Run ``n`` deterministic steps; return per-step loss, forward output, + per-param ``main_grad`` (pre-step), fp32 master and (bf16) param (post-step).""" + ids, labels, pos, mask, loss_mask = self._batch() + losses, outs, grads, masters, params = [], [], [], [], [] + for _ in range(n): + model[0].zero_grad_buffer() + optimizer.zero_grad() + # reuse_grad_buf aliases the bf16 staging buffer onto the just-zeroed grad buffer, so + # re-stage masters before the deferred (overlap) param all-gather. This is what + # ``force_param_sync`` / ``disable_forward_pre_hook`` do in the real training loop; + # it recurses into the sibling DistributedOptimizer (embeddings / biases / layernorm) + # that owns a byte-shard param buffer. No-op otherwise. + if args.reuse_grad_buf_for_mxfp8_param_ag and args.overlap_param_gather: + optimizer.prepare_model_params_for_param_sync() + model[0].set_is_first_microbatch() + out = model[0].forward( + input_ids=ids, + position_ids=pos, + attention_mask=mask, + labels=labels, + loss_mask=loss_mask, + ) + loss = out.mean() + loss.backward() + if args.overlap_grad_reduce: + model[0].finish_grad_sync() + grad = { + name: p.main_grad.detach().clone() + for name, p in model[0].named_parameters() + if p.main_grad is not None + } + ok, _, _ = optimizer.step() + assert ok + params.append(_snapshot_params(model[0])) + masters.append(_snapshot_masters(model[0])) + grads.append(grad) + losses.append(loss.detach().clone()) + outs.append(out.detach().clone()) + return losses, outs, grads, masters, params + + def _check_on_vs_off(self, fp8_recipe, overlap, n, num_experts=0, expert_model_parallel_size=1): + """fp8_param_gather ON must match OFF bitwise for ``n`` deterministic steps on + per-step loss / forward output / per-param main_grad / fp32 master / bf16 param.""" + with deterministic_mode(): + off_args, off_model, off_opt = self._build( + False, fp8_recipe, overlap, num_experts, expert_model_parallel_size + ) + params0, masters0 = _snapshot_params(off_model[0]), _snapshot_masters(off_model[0]) + on_args, on_model, on_opt = self._build( + True, fp8_recipe, overlap, num_experts, expert_model_parallel_size + ) + _restore_initial_state(on_model[0], on_opt, params0, masters0) + + off = [[], [], [], [], []] # loss, out, grad, master, param + on = [[], [], [], [], []] + for _ in range(n): + for dst, src in zip(off, self._run_steps(off_args, off_model, off_opt, 1)): + dst.extend(src) + for dst, src in zip(on, self._run_steps(on_args, on_model, on_opt, 1)): + dst.extend(src) + del off_model, on_model, off_opt, on_opt + gc.collect() + torch.cuda.empty_cache() + + lo, oo, go, mo, po = off + ln, on_, gn, mn, pn = on + for s in range(n): + _assert_equal(ln[s], lo[s], f"loss step {s}") + _assert_equal(on_[s], oo[s], f"output step {s}") + assert gn[s].keys() == go[s].keys(), f"grad param set mismatch step {s}" + for k in gn[s]: + _assert_equal(gn[s][k], go[s][k], f"grad step {s} {k}") + assert pn[s].keys() == po[s].keys(), f"param set mismatch step {s}" + for k in pn[s]: + _assert_equal(pn[s][k], po[s][k], f"param step {s} {k}") + common = mn[s].keys() & mo[s].keys() + assert common, f"no common masters step {s}" + for k in common: + _assert_equal(mn[s][k], mo[s][k], f"master step {s} {k}") + + @pytest.mark.parametrize("overlap", [False, True]) + @pytest.mark.parametrize("fp8_recipe", ["blockwise", "mxfp8"]) + @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") + def test_on_vs_off_bitwise_identical(self, fp8_recipe, overlap): + """fp8_param_gather ON must match OFF bitwise on the decoupled layout, for + overlap grad-reduce + param-gather both ON and OFF.""" + arch = get_device_arch_version() + if fp8_recipe == "blockwise" and arch != 9: + pytest.skip("blockwise FP8 is Hopper-only") + if fp8_recipe == "mxfp8" and arch < 10: + pytest.skip("mxfp8 requires Blackwell architecture or newer") + # 30 steps: fp8-quantization ON-vs-OFF mismatches often only surface after many + # iterations (PR #5470 review), so a handful of steps can miss a real divergence. + self._check_on_vs_off(fp8_recipe, overlap, n=30) + + @pytest.mark.parametrize("overlap", [False, True]) + @pytest.mark.parametrize("expt_dp_gt_1", [False, True]) + @pytest.mark.parametrize("fp8_recipe", ["blockwise", "mxfp8"]) + @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") + def test_moe_on_vs_off_bitwise_identical(self, fp8_recipe, overlap, expt_dp_gt_1): + """MoE variant (PR #5470 review): the 2D expert weights are Muon-managed, so they + ride the LayerWise param path. Covers both expert-data-parallel regimes: + + - ``expt_dp == 1`` (expert_model_parallel_size == world size): experts are NOT + all-gathered — they only get the fp8 master->model copy-back + (``_copy_main_params_to_model_params``'s non-gathered branch). + - ``expt_dp > 1`` (expert_model_parallel_size == 1): experts ARE gathered. + + Needs world size >= 2 to realize both regimes (at dp==1 only expt_dp==1 exists); + run with ``torchrun --nproc_per_node>=2``. + """ + world = torch.distributed.get_world_size() + if world < 2: + pytest.skip( + "MoE expt_dp coverage needs data-parallel size >= 2 (dp==1 only realizes " + "expt_dp==1); run with --nproc_per_node>=2" + ) + arch = get_device_arch_version() + if fp8_recipe == "blockwise" and arch != 9: + pytest.skip("blockwise FP8 is Hopper-only") + if fp8_recipe == "mxfp8" and arch < 10: + pytest.skip("mxfp8 requires Blackwell architecture or newer") + + # expt_dp = world_size / expert_model_parallel_size. + ep = 1 if expt_dp_gt_1 else world + # setup_method initialized model parallel with ep=1; re-init with the target EP. + Utils.destroy_model_parallel() + Utils.initialize_model_parallel( + tensor_model_parallel_size=1, + pipeline_model_parallel_size=1, + expert_model_parallel_size=ep, + ) + self._check_on_vs_off( + fp8_recipe, overlap, n=30, num_experts=8, expert_model_parallel_size=ep + ) + + @pytest.mark.parametrize("fp8_recipe", ["blockwise", "mxfp8"]) + @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") + def test_force_sync_finalizes_pending_layerwise_gather(self, fp8_recipe): + """force_sync finalize (eval/ckpt ``disable_forward_pre_hook`` path) must match the + ``finish_param_sync`` forward-pre-hook path: gathered params land in every rank's + ``param.data`` and the reused grad buffer is re-zeroed. Both asserted equal against + the reference path. Needs dp>=2 (dp==1 early-returns the all-gather); run with + ``torchrun --nproc_per_node>=2``. + """ + if torch.distributed.get_world_size() < 2: + pytest.skip( + "force_sync LayerWise gather copy-back is only exercised at data-parallel " + "size >= 2 (dp==1 skips the all-gather); run with --nproc_per_node>=2" + ) + arch = get_device_arch_version() + if fp8_recipe == "blockwise" and arch != 9: + pytest.skip("blockwise FP8 is Hopper-only") + if fp8_recipe == "mxfp8" and arch < 10: + pytest.skip("mxfp8 requires Blackwell architecture or newer") + + def _step_and_dispatch(): + args, model, opt = self._build(True, fp8_recipe, True) + self._run_steps(args, model, opt, 1) + ddp = model[0] + ddp.start_param_sync() + groups = ddp.bucket_groups + ddp.expert_parallel_bucket_groups + assert any( + g.param_gather_handle is not None for g in groups + ), "test precondition: expected a pending async param-gather handle" + return model, ddp, groups + + with deterministic_mode(): + # Reference: finish the pending gather through the forward-pre-hook path. + ref_model, ref_ddp, ref_groups = _step_and_dispatch() + for g in ref_groups: + if g.param_gather_handle is not None: + g.finish_param_sync(skip_next_bucket_dispatch=True) + ref_params = _snapshot_params(ref_model[0], include_quantized=True) + ref_grads = _snapshot_layerwise_grad_data(ref_ddp) + del ref_model, ref_ddp, ref_groups + gc.collect() + torch.cuda.empty_cache() + + # Under test: force-sync with the handle still pending. + model, ddp, groups = _step_and_dispatch() + ddp.disable_forward_pre_hook(param_sync=True) + got_params = _snapshot_params(model[0], include_quantized=True) + got_grads = _snapshot_layerwise_grad_data(ddp) + + for g in groups: + for bucket in g.buckets: + assert ( + getattr(bucket, 'layerwise_gather_list', None) is None + ), "force_sync left an unconsumed layerwise_gather_list" + + assert ref_params.keys() == got_params.keys() + for k in ref_params: + _assert_equal(got_params[k], ref_params[k], f"param after force_sync {k}") + assert len(ref_grads) == len(got_grads) and ref_grads, "expected LayerWise grad buffers" + for i, (gr, gg) in enumerate(zip(ref_grads, got_grads)): + _assert_equal(gg, gr, f"grad_data buffer {i} after force_sync") From 5e06ef969228a99e64906910a628ab886293685d Mon Sep 17 00:00:00 2001 From: Pingtian Li Date: Fri, 31 Jul 2026 16:37:35 +0800 Subject: [PATCH 03/12] Fix the dp_reshardable padding template in the decoupled ckpt save path sharded_param_state_dp_reshardable synthesizes optimizer-state padding for shards that lie entirely in a bucket's inter-param padding. That synthesis needs the state's key set and dtypes, which a rank owning no param in the bucket cannot observe locally. Derive the template per (gbuf_idx, dtype, bucket_idx) and reconcile it across the DP group with one all_gather_object. Per bucket is the axis that matters: dist_checkpointing enforces one dtype per KEY, and each bucket is its own key, so a template sampled from one bucket is not authoritative for another. The agreement assert is likewise per bucket, so buckets no rank sampled stay absent instead of aborting the save. The template cannot be derived from config: the save path runs the state through get_unscaled_state (which upcasts bf16/fp16/fp8 state to fp32 and returns int16 for store_param_remainders), and the key set comes from whatever optimizer.state holds. `step` is excluded -- it becomes a LocalNonpersistentObject and never a ShardedTensor. Device stays out of the gathered template and out of the compare. It is rank-local: under --optimizer-cpu-offload HybridDeviceOptimizer walks each rank's own shard list, so the CPU/GPU cutoff lands on a different param on every rank and a cross-rank compare would abort a legitimate save. dist_checkpointing validates dtype and shape but never device. Also reject fully_sharded_model_space up front for the decoupled compact layout. That format is independently non-functional for every DistributedOptimizer here -- it sets flattened_range on every non-factory param and ShardedTensor rejects that unconditionally -- so this assert does not fix the format; it makes the compact case fail early with an actionable message instead of dying per-param inside mapping.py. Narrow on purpose. Signed-off-by: Pingtian Li --- megatron/core/optimizer/distrib_optimizer.py | 131 +++++++++++++------ 1 file changed, 93 insertions(+), 38 deletions(-) diff --git a/megatron/core/optimizer/distrib_optimizer.py b/megatron/core/optimizer/distrib_optimizer.py index e5417e8318d..f4095ac16bf 100644 --- a/megatron/core/optimizer/distrib_optimizer.py +++ b/megatron/core/optimizer/distrib_optimizer.py @@ -1897,34 +1897,56 @@ def sharded_param_state_dp_reshardable( state = self.get_parameter_state_dp_reshardable() - # fp32 optimizer-state {key: (dtype, device)} captured before the loop below mutates - # ``state``, so an empty shard can synthesize valid padding ShardedTensors. - pad_template = None + # Optimizer-state {key: dtype} per bucket, captured before the loop below + # mutates ``state``, so a shard that is entirely padding can still synthesize valid + # ShardedTensors. Keyed per (gbuf_idx, dtype, bucket_idx) because dist_checkpointing + # enforces one dtype per KEY (validation.py, ``assert sharding.dtype == dtype``) and each + # bucket is its own key -- a template sampled from one bucket is not authoritative for + # another. ``step`` is a per-param 0-dim tensor for torch AdamW / non-TE-Apex / + # HybridDeviceOptimizer, but it becomes a LocalNonpersistentObject below and never a + # ShardedTensor, so it must not be synthesized here. + # Device stays OUT of the gathered template: it is rank-local (under + # --optimizer-cpu-offload the CPU/GPU split lands on a different param on every + # rank, because HybridDeviceOptimizer walks this rank's own shard list), and + # dist_checkpointing validates dtype/shape but never device. Comparing it across + # ranks would abort the save on a legitimate config. + pad_templates = {} + pad_device = None for _g in range(len(self.gbuf_ranges)): - for _bs_all in state[_g].values(): - for _bs in _bs_all: + for _dt_key, _bs_all in state[_g].items(): + for _b_idx, _bs in enumerate(_bs_all): if _bs: - pad_template = { - k: (v.dtype, v.device) + _key = (_g, str(_dt_key), _b_idx) + pad_templates[_key] = { + k: v.dtype for k, v in _bs[0].items() - if isinstance(v, torch.Tensor) + if isinstance(v, torch.Tensor) and k != 'step' } - break - if pad_template is not None: - break - if pad_template is not None: - break - - # A rank that owns nothing in any bucket has no local sample. The optimizer-state keys - # and their dtypes are config-determined and identical on every DP rank, so derive the - # template from config instead of gathering it across DP -- no communication needed. - if pad_template is None: - _pad_device = torch.cuda.current_device() - pad_template = { - 'param': (self.config.main_params_dtype, _pad_device), - 'exp_avg': (self.config.exp_avg_dtype, _pad_device), - 'exp_avg_sq': (self.config.exp_avg_sq_dtype, _pad_device), - } + if pad_device is None and pad_templates[_key]: + pad_device = _bs[0][next(iter(pad_templates[_key]))].device + + # A rank that owns nothing in a bucket has no local sample there, yet it may still owe + # padding coverage for it. The template cannot be derived from config: the save path runs + # the state through ``get_unscaled_state`` (which upcasts bf16/fp16/fp8 state to fp32, and + # returns int16 for ``store_param_remainders``), and the key set comes from whatever + # ``optimizer.state`` actually holds. So reconcile the observed templates across the DP + # group. The assert is per bucket -- the axis dist_checkpointing actually constrains -- + # so buckets nobody sampled simply stay absent instead of aborting the save. + if data_parallel_world_size > 1: + _gathered = [None] * data_parallel_world_size + torch.distributed.all_gather_object( + _gathered, pad_templates, group=self.data_parallel_group + ) + _merged = {} + for _candidate in _gathered: + for _k, _v in (_candidate or {}).items(): + _prev = _merged.setdefault(_k, _v) + assert _prev == _v, ( + f"padding template mismatch across DP ranks for bucket {_k}: " + f"{_prev} vs {_v}; the synthesized padding would not match the " + "real shards" + ) + pad_templates = _merged # per_bucket_numel metadata is saved separately for each TPxPP domain. for per_bucket_key in ('per_bucket_numel', 'per_bucket_numel_unpadded'): @@ -1983,18 +2005,40 @@ def sharded_param_state_dp_reshardable( pad_len = min( gbuf_local_numel, gbuf_world_numel_unpadded - world_shard_start ) - # Synthesize the padding for this shard's overlap with [0, numel_unpadded), - # using pad_template for the correct fp32 keys/dtypes. The template is - # always available (config-derived above); a None here would mean the - # coverage we must emit is being dropped -- fail loudly rather than skip. - assert ( - pad_template is not None - ), f'no padding template for {sharded_bucket_key}; coverage would be dropped' + # Synthesize this shard's overlap with [0, numel_unpadded) from the + # template observed for THIS bucket. If no rank sampled it (every shard + # of the bucket is padding), fall back to any observed template rather + # than dropping coverage; an empty map means the coverage is being + # dropped, so fail loudly. + _tpl = pad_templates.get((gbuf_idx, str(dtype), bucket_idx)) + if not _tpl: + _tpl = next(iter(pad_templates.values()), None) + assert _tpl, ( + f'no optimizer-state template available for {sharded_bucket_key}; ' + 'the padding coverage this rank owes would be dropped' + ) + # ``pad_device`` is a single rank-local device, NOT a per-bucket map. + # Deliberate: we only synthesize for buckets this rank sampled nothing + # in, so a per-bucket lookup would miss by construction. Borrowing a + # device across buckets is safe in a way borrowing a dtype is NOT -- + # dtype and shape are validated, device never is, and the padding bytes + # are dropped on load. Do NOT make this symmetric with the dtype path: + # a rank-local device inside the cross-rank compare aborts the save + # under --optimizer-cpu-offload, where the CPU/GPU cutoff differs per + # rank. bucket_state = [ { **{ - k: torch.empty(pad_len, dtype=_dt, device=_dev) - for k, (_dt, _dev) in pad_template.items() + k: torch.empty( + pad_len, + dtype=_dt, + device=( + pad_device + if pad_device is not None + else torch.cuda.current_device() + ), + ) + for k, _dt in _tpl.items() }, 'gbuf_local_start': 0, 'gbuf_local_end': pad_len, @@ -2095,6 +2139,22 @@ def sharded_param_state_fs_model_space( This will allow changing TP and PP while using DistOpt (as with other optimizers). """ + # NB: fully_sharded_model_space is independently non-functional for EVERY + # DistributedOptimizer in this tree, not just the decoupled layout -- this function sets + # ``flattened_range`` on every non-factory param below, and ShardedTensor rejects that + # unconditionally in validate_metadata_integrity. Only ShardedTensorFactory escapes, and + # only gated-MLP fc1 produces one. This assert does not fix that; it makes the compact + # LayerWise (Muon) case abort earlier and with an actionable message instead of dying + # per-param deep inside mapping.py. It is deliberately narrow: broadening it to reject + # the format outright is a separate change. + assert not ( + self.config.use_layer_wise_distributed_optimizer + and not self.config.use_layer_wise_param_layout + ), ( + "fully_sharded_model_space is not supported for the decoupled compact LayerWise " + "(Muon) optimizer layout; use dp_reshardable or fully_reshardable instead" + ) + param_to_sharded_metadata = {} model_sharded_state_dict, _ = extract_sharded_tensors_and_factories( model_sharded_state_dict @@ -2122,11 +2182,6 @@ def _get_param_state_sharded_tensors(model_param, item_slice): f" Hint: {KEEP_VARS_HINT}" ) from e - assert isinstance(sharded_metadata, ShardedTensorFactory), ( - "fully_sharded_model_space is not supported for the decoupled compact " - "LayerWise (Muon) optimizer layout" - ) - # Set DP corresponding replica_id coordinate to 0. assert ( len(sharded_metadata.replica_id) == 3 From 00639331738998b696261b141cbb468aa9ea5341 Mon Sep 17 00:00:00 2001 From: Pingtian Li Date: Fri, 31 Jul 2026 16:37:36 +0800 Subject: [PATCH 04/12] Stop inferring global emptiness from slot 0 in LayerWise param gather _allgather_helper read params_list[0][0] to pick the device and dtype for the collective. The per-rank lists come from the ping-pong assignment and the dtype split in _dispatch, either of which can leave rank 0's slot empty while other ranks own params, so this raised IndexError. `any(owned ...)` in the caller does not prevent it. Take the first non-empty list instead, mirroring the fp8 twin. The same shape appears in the expert-parallel guards, which decide "there are no expert params anywhere" from expt_dp_params_list[0] alone and then disable the expert all-gather entirely. Both sharding paths happen to place the first expert param in slot 0 (ping-pong starts at expt_dp_idx 0; the layout path gives the first param of every expert bucket shard_id 0), so the two forms are equivalent for every reachable configuration today and this fixes no observable bug. It is a defensive change: it states the intended invariant, and unlike the helper it would have failed silently rather than loudly. Signed-off-by: Pingtian Li --- megatron/core/optimizer/layer_wise_optimizer.py | 15 +++++++++++---- 1 file changed, 11 insertions(+), 4 deletions(-) diff --git a/megatron/core/optimizer/layer_wise_optimizer.py b/megatron/core/optimizer/layer_wise_optimizer.py index 624fde5fe7d..9356172e7f6 100644 --- a/megatron/core/optimizer/layer_wise_optimizer.py +++ b/megatron/core/optimizer/layer_wise_optimizer.py @@ -708,7 +708,7 @@ def _shard_params_from_layout(self, optimizers, full_param_layouts, dp_cp_size, group["params"] = local_params # Simplify when expt_dp group size is 1 or expert parallel is off. - if expt_dp_size == 1 or len(self.expt_dp_params_list[0]) == 0: + if expt_dp_size == 1 or not any(self.expt_dp_params_list): self.expt_dp_params_list = None def _build_param_sort_keys(self, model_chunks): @@ -800,7 +800,7 @@ def _shard_params_ping_pong(self, optimizers, dp_cp_size, expt_dp_size, model_ch groups["params"] = params # Simplify when expt_dp group size is 1 or expert parallel is off. - if expt_dp_size == 1 or len(self.expt_dp_params_list[0]) == 0: + if expt_dp_size == 1 or not any(self.expt_dp_params_list): self.expt_dp_params_list = None def set_bucket_layerwise_params_list(self, model_chunks): @@ -930,8 +930,15 @@ def _allgather_helper_fp8(params_list, group): # helper function to flatten local params, all-gather, # unflatten and copy to model params def _allgather_helper(params_list, group): - device = params_list[0][0].device - dtype = params_list[0][0].dtype + # Rank 0 may own zero params in this list -- the ping-pong assignment, and the + # dtype split in ``_dispatch`` below, both leave per-rank lists that can be empty. + # Mirror ``_allgather_helper_fp8``'s lookup instead of indexing rank 0 blindly. + _first = next((params[0] for params in params_list if len(params) > 0), None) + if _first is None: + # No rank owns any param in this group -> nothing to gather. + return + device = _first.device + dtype = _first.dtype rank = get_pg_rank(group) dp_size = get_pg_size(group) # Flatten this rank's params. From 873176cd91fa2ccc72eee877c5dd02ad2bcfec97 Mon Sep 17 00:00:00 2001 From: Pingtian Li Date: Fri, 31 Jul 2026 16:38:07 +0800 Subject: [PATCH 05/12] Make the fp8 param-gather test assert something, and fix its overlap sampling Three of this test's comparisons were passing for the wrong reason. _is_quantized probed for a `dequantize` attribute, but torch.Tensor.dequantize is defined on every tensor (it returns self.to(float32) for dense ones), so the predicate was always True, _snapshot_params always returned {}, and the param comparison was asserting {} == {}. Key off is_float8tensor and assert the snapshot is non-empty. _run_steps then snapshotted with include_quantized=True and fed that to the ON-vs-OFF comparison, where OFF holds plain bf16 and ON holds Q(bf16(master)): a dequantized fp32 view of the latter differs from the former by the quantization step by construction. Drop it. ON legitimately skips its fp8 params while OFF keeps them as bf16, so ON's key set is a subset; pin the dropped set at step 0 and require it to be unchanged and non-empty, rather than accepting any shrinkage. Under overlap_param_gather the param all-gather is deferred to the next forward pre-hook, so right after step() the buffer still holds pre-update values and there is no well-defined instant at which ON and OFF can be compared. Finalize before sampling -- restaging the masters first, since the LayerWise buckets alias the grad buffer that backward has just filled. Then re-arm: the forced sync leaves param_gather_dispatched=True, which would turn every subsequent forward pre-hook into a no-op and silently drop the deferred-gather coverage this parametrization exists for. Assert the pre-hook is armed so that regression cannot recur silently. Finally, the force_sync test compared its two grad_data snapshots for equality, but both paths end in grad_data.zero_(), so on the happy path that is zeros against zeros; state the post-condition directly. Its param snapshots were also asymmetric (one side included quantized storage, the other did not), which made the key-set assert pass only because both were empty. Signed-off-by: Pingtian Li --- .../test_muon_decouple_fp8_param_gather.py | 93 +++++++++++++++++-- 1 file changed, 87 insertions(+), 6 deletions(-) diff --git a/tests/unit_tests/test_muon_decouple_fp8_param_gather.py b/tests/unit_tests/test_muon_decouple_fp8_param_gather.py index b7e24ffcad3..c10a29ed39f 100644 --- a/tests/unit_tests/test_muon_decouple_fp8_param_gather.py +++ b/tests/unit_tests/test_muon_decouple_fp8_param_gather.py @@ -22,6 +22,7 @@ from transformer_engine.pytorch.fp8 import check_fp8_support from megatron.core.enums import ModelType +from megatron.core.fp8_utils import is_float8tensor from megatron.core.inference.utils import InferenceMode 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 @@ -46,7 +47,10 @@ def _is_quantized(p): - return hasattr(p, 'dequantize') or hasattr(p.data, 'dequantize') + # NB: do not probe for a ``dequantize`` attribute -- ``torch.Tensor.dequantize`` is defined + # for every tensor (it returns ``self.to(float32)`` for dense ones), so ``hasattr`` is always + # True and would classify plain bf16 params as quantized. + return is_float8tensor(p) or is_float8tensor(p.data) def _assert_equal(actual, expected, msg): @@ -68,8 +72,17 @@ def _snapshot_masters(model): def _snapshot_params(model, include_quantized=False): - # By default only non-quantized (bf16) params; fp8 weights are compared via master + grad. - # With include_quantized, fp8 params are dequantized so gathered fp8 bytes are checked directly. + """Model params for ON-vs-OFF comparison. + + fp8 params are skipped by default and covered via the fp32 master + reduced grad instead: + with fp8_param_gather OFF ``param.data`` is plain bf16, with it ON the same weight is + Float8/MXFP8 holding ``Q(bf16(master))``, so the two storages differ by the quantization + step by construction and cannot be compared directly. What remains -- layernorm, biases, + embeddings, any non-quantized weight -- is directly comparable and IS compared. + + ``include_quantized`` dequantizes instead, for the single-run comparisons (force_sync) + where both sides hold the same fp8 storage and the gathered bytes are the thing under test. + """ out = {} for n, p in model.named_parameters(): if _is_quantized(p): @@ -77,6 +90,7 @@ def _snapshot_params(model, include_quantized=False): out[n] = p.data.dequantize().detach().clone().float() else: out[n] = p.detach().clone() + assert out, "snapshot is empty -- the param comparison would assert nothing" return out @@ -265,6 +279,21 @@ def _run_steps(self, args, model, optimizer, n): # that owns a byte-shard param buffer. No-op otherwise. if args.reuse_grad_buf_for_mxfp8_param_ag and args.overlap_param_gather: optimizer.prepare_model_params_for_param_sync() + if args.overlap_param_gather: + # The deferred gather must still be pending here. A forced sync that forgets to + # re-arm leaves param_gather_dispatched=True and turns the forward pre-hook into + # a silent no-op -- the test would keep passing while losing the overlap + # coverage it exists for. + # ``all``, not ``any``: with ``any`` a single armed dense group would mask a + # refactor that drops expert_parallel_bucket_groups from the reset, leaving + # every expert group permanently dispatched. ``_groups and`` because all([]) + # is vacuously True. + _groups = model[0].bucket_groups + model[0].expert_parallel_bucket_groups + assert _groups and all(not g.param_gather_dispatched for g in _groups), ( + "forward pre-hook is not armed for every bucket group: the deferred param " + "all-gather would be skipped, so this iteration does not exercise the " + "overlap path" + ) model[0].set_is_first_microbatch() out = model[0].forward( input_ids=ids, @@ -284,6 +313,38 @@ def _run_steps(self, args, model, optimizer, n): } ok, _, _ = optimizer.step() assert ok + if args.overlap_param_gather: + # Under overlap the param all-gather is deferred to the next forward pre-hook, + # so right after ``step()`` the param buffer still holds the values gathered + # at THIS iteration's forward, i.e. pre-update -- there is no well-defined + # instant at which ON and OFF can be compared. Finalize it here. + # + # Order matters. The LayerWise (Muon) buckets are self-sufficient: their + # gather re-stages its source from ``param.main_param`` every time + # (``_stage_param_to_bf16``), so they never read stale buffer content -- + # which is why no Muon param ever diverged here. The sibling plain + # DistributedOptimizer (embeddings / biases / layernorm) instead gathers + # from a param_data shard written back at step time, and backward has since + # refilled the aliased grad buffer. Restage those masters first, otherwise + # the forced gather broadcasts gradients as parameters. This is the same + # sequence the real training loop performs around eval/checkpoint; + # ``prepare_model_params_for_param_sync`` self-guards on + # reuse_grad_buf_for_mxfp8_param_ag + overlap_param_gather. + optimizer.prepare_model_params_for_param_sync() + model[0].disable_forward_pre_hook(param_sync=True) + model[0].enable_forward_pre_hook() + # The forced sync above takes the synchronous branch and leaves + # ``param_gather_dispatched=True``, which is only cleared by finish_grad_sync + # (already past) or here. Without this reset the next forward pre-hook is a + # strict no-op, so steps 2..n would silently stop exercising the deferred + # (async dispatch -> wait -> finalize -> bucket chaining) path that the + # ``overlap`` parametrization exists to cover. Re-gathering already-correct + # values is idempotent, so the ON-vs-OFF comparison is unaffected. + model[0].reset_param_sync_dispatch_state() + # NB: include_quantized must stay False here. This snapshot feeds the ON-vs-OFF + # comparison, where OFF holds plain bf16 and ON holds Q(bf16(master)); a + # dequantized fp32 view of the latter differs from the former by the + # quantization step by construction (max_diff == 2**-9 for MXFP8). params.append(_snapshot_params(model[0])) masters.append(_snapshot_masters(model[0])) grads.append(grad) @@ -323,11 +384,25 @@ def _check_on_vs_off(self, fp8_recipe, overlap, n, num_experts=0, expert_model_p assert gn[s].keys() == go[s].keys(), f"grad param set mismatch step {s}" for k in gn[s]: _assert_equal(gn[s][k], go[s][k], f"grad step {s} {k}") - assert pn[s].keys() == po[s].keys(), f"param set mismatch step {s}" + # ON stores the fp8 weights as Float8/MXFP8 and skips them; OFF keeps the same + # weights as plain bf16, so ON's key set is a subset. What survives -- layernorms, + # biases, embeddings, the output layer -- are real bf16 param tensors compared + # bitwise. The fp8 weights are covered by the master + reduced-grad checks above. + assert pn[s].keys() <= po[s].keys(), f"param set mismatch step {s}" + assert pn[s], f"no comparable params captured step {s}" + # Pin the skipped (fp8-on-ON-only) set: a bare subset assert would silently accept + # a param dropping out of the comparison, e.g. if a copy-back regression replaced + # a bf16 ``param.data`` with quantized storage on the ON side only. + _dropped = po[s].keys() - pn[s].keys() + if s == 0: + _dropped_0 = _dropped + assert _dropped, "no fp8 params skipped -- the ON run is not using fp8 storage" + assert _dropped == _dropped_0, f"fp8-skipped param set changed at step {s}" for k in pn[s]: _assert_equal(pn[s][k], po[s][k], f"param step {s} {k}") - common = mn[s].keys() & mo[s].keys() - assert common, f"no common masters step {s}" + assert mn[s].keys() == mo[s].keys(), f"master param set mismatch step {s}" + assert mn[s], f"no masters captured step {s}" + common = mn[s].keys() for k in common: _assert_equal(mn[s][k], mo[s][k], f"master step {s} {k}") @@ -451,3 +526,9 @@ def _step_and_dispatch(): assert len(ref_grads) == len(got_grads) and ref_grads, "expected LayerWise grad buffers" for i, (gr, gg) in enumerate(zip(ref_grads, got_grads)): _assert_equal(gg, gr, f"grad_data buffer {i} after force_sync") + # Both routes end in _finalize_layerwise_param_sync -> grad_data.zero_(), so equality + # alone is zeros-vs-zeros on the happy path. State the intended post-condition + # directly: force_sync must not leave the gather payload behind. + assert ( + torch.count_nonzero(gg) == 0 + ), f"grad_data buffer {i} still holds the gather payload after force_sync" From 0604c6398f073c3b31b460a9f227d0e894d7b0d3 Mon Sep 17 00:00:00 2001 From: Pingtian Li Date: Fri, 31 Jul 2026 16:38:07 +0800 Subject: [PATCH 06/12] Cover the compact decoupled layout in the LayerWise ckpt save/load test setup_model_and_optimizer never set use_layer_wise_param_layout, so the field was only ever whatever argparse defaulted to and no test could select the layout. Thread it through both helpers, and guard it against ddp_use_layer_wise -- the flag DDP actually receives -- rather than against the optimizer-name derived use_layer_wise. The old guard admitted optimizer='dist_muon' with dist_opt=False and use_layer_wise_param_layout=True, which builds DDP without a layout while the optimizer expects the shard-aligned one. test_layer_wise_optimizer_save_load then gains a LayerWise arm that actually reaches the layout: dist_opt=True and use_param_layout=True (the DDP-level routing switch), with use_layer_wise_param_layout=False selecting the compact side. It is the only fully_reshardable coverage of that layout above tp=pp=1 and the only one with grad_reduce_in_fp32 left at its default. The padded shard-aligned layout is deliberately out of scope. Both arms must request a sharding type now that both build a real DistributedOptimizer. The default fully_sharded_model_space is unusable: the compact layout is rejected up front, and an ordinary DistOpt dies deeper in replace() with "ShardedTensor.flattened_range is not supported". dp_reshardable synthesizes padding with torch.empty, so two saves of identical state differ in the uninitialized bytes and the terminal check_equal would flake. fully_reshardable has neither problem. The plain arm is kept as an in-file control -- test_optimizer.py already covers a plain bf16 DistOpt under fully_reshardable, and harder resharding cases besides -- and is labelled as such. Signed-off-by: Pingtian Li --- .../test_layer_wise_optimizer.py | 54 ++++++++++++++----- tests/unit_tests/dist_checkpointing/utils.py | 32 +++++++++++ 2 files changed, 74 insertions(+), 12 deletions(-) diff --git a/tests/unit_tests/dist_checkpointing/test_layer_wise_optimizer.py b/tests/unit_tests/dist_checkpointing/test_layer_wise_optimizer.py index 8be5353297d..409f7120033 100644 --- a/tests/unit_tests/dist_checkpointing/test_layer_wise_optimizer.py +++ b/tests/unit_tests/dist_checkpointing/test_layer_wise_optimizer.py @@ -188,13 +188,42 @@ def test_broadcast_params(self, tp, pp): assert torch.allclose(param.data, original_params[name]) # TODO(deyuf): check bf16 False case + # Cover both sides of use_layer_wise_distributed_optimizer with a real distributed + # optimizer on each side. ``plain_distopt`` is an in-file CONTROL -- test_optimizer.py's + # resharding tests already cover a plain bf16 DistOpt under fully_reshardable, and harder + # resharding cases besides; it is here to show the same harness behaves with and without + # LayerWise. ``layerwise_compact`` is the load-bearing arm: it is the only fully_reshardable + # coverage of the compact layout above tp=pp=1, and the only one with grad_reduce_in_fp32 + # left at its default. ``use_param_layout`` is the DDP-level LayerWise routing switch + # (ddp_use_layer_wise = use_layer_wise and use_param_layout), so the LayerWise arm needs it + # True to build the decoupled layout at all; ``use_layer_wise_param_layout`` stays False, + # which selects the COMPACT layout this PR adds -- the padded shard-aligned layout is out of + # scope. The legacy ping-pong path (dist_opt=False) is covered by the other tests here. + @pytest.mark.parametrize( + 'layer_wise', + [pytest.param(False, id='plain_distopt'), pytest.param(True, id='layerwise_compact')], + ) @pytest.mark.parametrize('tp', [1, 2, 4]) @pytest.mark.parametrize('pp', [1, 2, 4]) @pytest.mark.parametrize('bf16', [True]) - def test_layer_wise_optimizer_save_load(self, tmp_path_dist_ckpt, tp, pp, bf16): + def test_layer_wise_optimizer_save_load(self, tmp_path_dist_ckpt, tp, pp, bf16, layer_wise): """Test save/load of LayerWiseDistributedOptimizer checkpoints.""" if tp * pp > 8: pytest.skip(f"TP*PP > 8 is larger than world size") + _kw = dict( + optimizer='dist_muon' if layer_wise else 'adam', + use_param_layout=layer_wise, + use_layer_wise_param_layout=False, + ) + # Both arms now build a real DistributedOptimizer, so the format has to be requested: + # - the default fully_sharded_model_space is unusable here -- the compact layout is + # rejected outright, and an ordinary DistOpt dies deeper in ``replace(...)`` with + # "ShardedTensor.flattened_range is not supported"; + # - dp_reshardable synthesizes padding with ``torch.empty``, so two saves of identical + # state differ in the uninitialized padding bytes and the check_equal below would be + # flaky (this is why test_decouple_ckpt_roundtrip_values needs its canonical view). + # fully_reshardable has neither problem and keeps the plain A-vs-B comparison valid. + _md = {'distrib_optim_sharding_type': 'fully_reshardable'} Utils.initialize_model_parallel(tp, pp) @@ -210,14 +239,14 @@ def test_layer_wise_optimizer_save_load(self, tmp_path_dist_ckpt, tp, pp, bf16): tp=tp, pp=pp, bf16=bf16, - dist_opt=False, + dist_opt=True, initialize_fn=initialize_gpt_model, - optimizer='dist_muon', + **_kw, ) # Save checkpoint A model_sharded_sd_A = model_A[0].sharded_state_dict() - optim_sd_A = optimizer_A.sharded_state_dict(model_sharded_sd_A) + optim_sd_A = optimizer_A.sharded_state_dict(model_sharded_sd_A, metadata=_md) save(optim_sd_A, ckpt_dir_A) # Create model and optimizer B with different seed @@ -226,21 +255,21 @@ def test_layer_wise_optimizer_save_load(self, tmp_path_dist_ckpt, tp, pp, bf16): tp=tp, pp=pp, bf16=bf16, - dist_opt=False, + dist_opt=True, initialize_fn=initialize_gpt_model, - optimizer='dist_muon', + **_kw, ) # Load checkpoint A into optimizer B model_sharded_sd_B = model_B[0].sharded_state_dict() load_sharded_sd = optimizer_B.sharded_state_dict( - model_sharded_sd_B, is_loading=True + model_sharded_sd_B, is_loading=True, metadata=_md ) state_dict = load(load_sharded_sd, ckpt_dir_A) optimizer_B.load_state_dict(state_dict) # Save as checkpoint B - optim_sd_B = optimizer_B.sharded_state_dict(model_sharded_sd_B) + optim_sd_B = optimizer_B.sharded_state_dict(model_sharded_sd_B, metadata=_md) save(optim_sd_B, ckpt_dir_B) Utils.destroy_model_parallel() @@ -389,10 +418,11 @@ def test_dp_reshardable_moe_synth_save(self, tmp_path_dist_ckpt): save(optim_sd, ckpt_dir) Utils.destroy_model_parallel() - # NOTE: 'fully_sharded_model_space' is intentionally NOT covered here. It is incompatible with - # the decoupled compact LayerWise (Muon) DistOpt layout regardless of precision: the sibling - # embedding / output_layer params produce a flattened_range ShardedTensor that - # dist_checkpointing rejects ('ShardedTensor.flattened_range is not supported'). This was + # NOTE: 'fully_sharded_model_space' is intentionally NOT covered here. It is non-functional + # for EVERY DistributedOptimizer in this tree, not only the decoupled compact LayerWise (Muon) + # layout: sharded_param_state_fs_model_space sets flattened_range on every non-factory param, + # and dist_checkpointing rejects that unconditionally ('ShardedTensor.flattened_range is not + # supported'); only ShardedTensorFactory (gated-MLP fc1) escapes. This was # confirmed to fail identically for bf16 and fp8 (so it is NOT fp8-specific), and it raises # non-uniformly across DP ranks, so it cannot be asserted cleanly in a distributed test. It is # a pre-existing limitation, independent of the dp_reshardable empty-bucket fix. diff --git a/tests/unit_tests/dist_checkpointing/utils.py b/tests/unit_tests/dist_checkpointing/utils.py index e7716794b67..424251be6a4 100644 --- a/tests/unit_tests/dist_checkpointing/utils.py +++ b/tests/unit_tests/dist_checkpointing/utils.py @@ -212,6 +212,7 @@ def setup_model_and_optimizer( ddp_bucket_size=None, grad_reduce_in_fp32=False, fp8=False, + use_layer_wise_param_layout=False, ): optimizer_type = optimizer use_layer_wise = False @@ -229,6 +230,19 @@ def setup_model_and_optimizer( ddp_use_dist_opt = dist_opt and not (use_layer_wise and not use_param_layout) ddp_use_layer_wise = use_layer_wise and use_param_layout + # The padded shard-aligned LayerWise layout only exists under a LayerWise optimizer, so + # ``use_layer_wise_param_layout=True`` implies ``use_layer_wise_distributed_optimizer=True``; + # the converse does not hold. Only three combinations are meaningful: + # (False, False) plain DistributedOptimizer + # (True, False) LayerWise + compact decoupled layout [the default since the flag flipped] + # (True, True) LayerWise + padded shard-aligned layout + assert not (use_layer_wise_param_layout and not ddp_use_layer_wise), ( + "use_layer_wise_param_layout=True requires LayerWise DDP routing " + "(optimizer='muon'/'dist_muon' with dist_opt=True and use_param_layout=True); " + "otherwise DDP builds untagged buffers while the optimizer expects the " + "shard-aligned layout" + ) + mock_args = parse_args(ignore_unknown_args=True) with mock.patch('megatron.training.training.get_args', new=lambda: mock_args): init_basic_mock_args(mock_args, tp, pp, bf16=bf16) @@ -250,6 +264,7 @@ def setup_model_and_optimizer( # grad_reduce_in_fp32 -> ddp_config grad_dtype=fp32 while params stay bf16, i.e. the # mixed-dtype (bf16, fp32) gradient buffer / optimizer-state bucket. mock_args.accumulate_allreduce_grads_in_fp32 = grad_reduce_in_fp32 + mock_args.use_layer_wise_param_layout = use_layer_wise_param_layout mock_args.use_distributed_optimizer = ddp_use_dist_opt mock_args.use_layer_wise_distributed_optimizer = ddp_use_layer_wise if ddp_use_layer_wise: @@ -291,6 +306,7 @@ def setup_model_and_optimizer( params_dtype=torch.bfloat16 if bf16 else torch.float, use_distributed_optimizer=ddp_use_dist_opt, use_layer_wise_distributed_optimizer=use_layer_wise, + use_layer_wise_param_layout=use_layer_wise_param_layout, optimizer=optimizer, muon_scalar_optimizer=muon_scalar_optimizer, ) @@ -386,6 +402,7 @@ def setup_moe_model_and_optimizer( use_param_layout=False, ddp_bucket_size=None, grad_reduce_in_fp32=False, + use_layer_wise_param_layout=False, ): optimizer_type = optimizer use_layer_wise = False @@ -399,6 +416,19 @@ def setup_moe_model_and_optimizer( ddp_use_dist_opt = dist_opt and not (use_layer_wise and not use_param_layout) ddp_use_layer_wise = use_layer_wise and use_param_layout + # The padded shard-aligned LayerWise layout only exists under a LayerWise optimizer, so + # ``use_layer_wise_param_layout=True`` implies ``use_layer_wise_distributed_optimizer=True``; + # the converse does not hold. Only three combinations are meaningful: + # (False, False) plain DistributedOptimizer + # (True, False) LayerWise + compact decoupled layout [the default since the flag flipped] + # (True, True) LayerWise + padded shard-aligned layout + assert not (use_layer_wise_param_layout and not ddp_use_layer_wise), ( + "use_layer_wise_param_layout=True requires LayerWise DDP routing " + "(optimizer='muon'/'dist_muon' with dist_opt=True and use_param_layout=True); " + "otherwise DDP builds untagged buffers while the optimizer expects the " + "shard-aligned layout" + ) + mock_args = parse_args(ignore_unknown_args=True) with mock.patch('megatron.training.training.get_args', new=lambda: mock_args): init_basic_mock_args(mock_args, tp, pp, bf16=bf16) @@ -413,6 +443,7 @@ def setup_moe_model_and_optimizer( # grad_reduce_in_fp32 -> ddp_config grad_dtype=fp32 while params stay bf16, i.e. the # mixed-dtype (bf16, fp32) gradient buffer / optimizer-state bucket. mock_args.accumulate_allreduce_grads_in_fp32 = grad_reduce_in_fp32 + mock_args.use_layer_wise_param_layout = use_layer_wise_param_layout mock_args.use_distributed_optimizer = ddp_use_dist_opt mock_args.use_layer_wise_distributed_optimizer = ddp_use_layer_wise if ddp_use_layer_wise: @@ -438,6 +469,7 @@ def setup_moe_model_and_optimizer( params_dtype=torch.bfloat16 if bf16 else torch.float, use_distributed_optimizer=ddp_use_dist_opt, use_layer_wise_distributed_optimizer=use_layer_wise, + use_layer_wise_param_layout=use_layer_wise_param_layout, optimizer=optimizer, ) From 68edd1cc7cf0b6258eaf16a837824b9dbc75b2a0 Mon Sep 17 00:00:00 2001 From: Pingtian Li Date: Mon, 17 Aug 2026 19:56:40 +0800 Subject: [PATCH 07/12] Tighten LayerWise FP8 gather validation Signed-off-by: Pingtian Li --- .../core/distributed/param_and_grad_buffer.py | 17 +++--- megatron/core/fp8_utils.py | 58 +++++++++++++++---- .../core/optimizer/layer_wise_optimizer.py | 16 ++--- megatron/core/optimizer/optimizer.py | 4 +- megatron/training/arguments.py | 11 +++- tests/unit_tests/test_fp8_utils.py | 34 +++++++++++ .../test_muon_decouple_fp8_param_gather.py | 16 ++++- 7 files changed, 121 insertions(+), 35 deletions(-) diff --git a/megatron/core/distributed/param_and_grad_buffer.py b/megatron/core/distributed/param_and_grad_buffer.py index a0f74649bd7..ae5df979f3f 100644 --- a/megatron/core/distributed/param_and_grad_buffer.py +++ b/megatron/core/distributed/param_and_grad_buffer.py @@ -36,6 +36,7 @@ is_grouped_mxfp8tensor, is_grouped_tensor, is_grouped_tensor_with_quantized_storage, + is_layerwise_fp8_param, is_mxfp8tensor, modify_grouped_tensor_rowwise_storage, modify_underlying_storage, @@ -491,22 +492,22 @@ def start_param_sync(self, force_sync: bool = False): local_rank = self.intra_distributed_optimizer_instance_rank group = self.intra_distributed_optimizer_instance_group layerwise_work_handles = [] - # Decouple fp8 param-gather: model params are Float8/MXFP8 but the all-gather rides bf16 - # — stage owned to bf16, gather, requantize on copy-back. Plain bf16 collapses to the - # original path. + # Decouple fp8 param-gather: supported MXFP8/blockwise model params ride BF16 — stage + # the local owner's FP32 master, gather, then requantize on copy-back. Plain BF16 + # collapses to the original path. decouple = not getattr(self.ddp_config, 'use_layer_wise_param_layout', True) for bucket in self.buckets: - # A decoupled LayerWise (Muon) bucket can MIX fp8 and bf16 params: + # A decoupled LayerWise (Muon) bucket can MIX supported fp8 and bf16 params: # merge_layerwise_fp8_grads keys fp8 Muon grads by their bf16 logical dtype so they # share ONE buffer (hence bucket) with their bf16 siblings (e.g. an MoE router / # DSA indexer / mHC weight that is not fp8-quantized). The bf16-staged path handles - # both dtypes, so a bucket holding ANY fp8 param must take it -- scanning only - # params_list[0] mis-routes a bf16-first mixed bucket into the raw + # both dtypes, so a bucket holding ANY supported fp8 param must take it -- scanning + # only params_list[0] mis-routes a bf16-first mixed bucket into the raw # _flatten_dense_tensors() path, which crashes on the MXFP8 .view(-1). bucket_is_fp8 = bool( decouple and bucket.params_list - and any(is_float8tensor(p) for p in bucket.params_list) + and any(is_layerwise_fp8_param(p) for p in bucket.params_list) ) # TODO(perf, blockwise-only): blockwise could gather the owner's fp8 rowwise data # (~2x less comm) + its small scale_inv and rebuild columnwise via transpose @@ -997,7 +998,7 @@ def group_params_for_buffers( # share ONE fp32 all_reduce buffer; a split uint8/bf16 reduction diverges ~1 ULP from OFF. if ( merge_layerwise_fp8_grads - and is_float8tensor(param) + and is_layerwise_fp8_param(param) and is_managed_by_layer_wise_optimizer ): param_dtype = param.dtype diff --git a/megatron/core/fp8_utils.py b/megatron/core/fp8_utils.py index 55df4da9111..c826725e13e 100644 --- a/megatron/core/fp8_utils.py +++ b/megatron/core/fp8_utils.py @@ -60,8 +60,17 @@ HAVE_TE_MXFP8TENSOR = True except (ImportError, ModuleNotFoundError): # MXFP8Tensor not found + MXFP8Tensor = None HAVE_TE_MXFP8TENSOR = False +try: + from transformer_engine.pytorch.tensor.float8_blockwise_tensor import Float8BlockwiseQTensor + + HAVE_TE_BLOCKWISE_FP8TENSOR = True +except (ImportError, ModuleNotFoundError): + Float8BlockwiseQTensor = None + HAVE_TE_BLOCKWISE_FP8TENSOR = False + try: from transformer_engine.pytorch.tensor.grouped_tensor import GroupedTensor @@ -136,6 +145,23 @@ def is_mxfp8tensor(tensor: torch.Tensor) -> bool: return HAVE_TE_MXFP8TENSOR and _is_instance_or_param_data(tensor, MXFP8Tensor) +def is_blockwise_float8tensor(tensor: torch.Tensor) -> bool: + """Check if a tensor is a Transformer Engine Float8BlockwiseQTensor.""" + return HAVE_TE_BLOCKWISE_FP8TENSOR and _is_instance_or_param_data( + tensor, Float8BlockwiseQTensor + ) + + +def is_layerwise_fp8_param(tensor: torch.Tensor) -> bool: + """Check if an FP8 parameter uses storage supported by LayerWise parameter gather. + + The compact LayerWise path stages whole parameters in BF16 and requantizes them on + copy-back. It supports plain MXFP8 and blockwise tensors only; generic Float8Tensor, + NVFP4, and GroupedTensor storage require different handling. + """ + return is_mxfp8tensor(tensor) or is_blockwise_float8tensor(tensor) + + def is_grouped_tensor(tensor: torch.Tensor) -> bool: """Check if a tensor is a Transformer Engine GroupedTensor.""" return HAVE_TE_GROUPED_TENSOR_CLASS and _is_instance_or_param_data(tensor, GroupedTensor) @@ -309,11 +335,22 @@ def dequantize_fp8_tensor(fp8_tensor: torch.Tensor) -> torch.Tensor: def copy_back_gathered_bf16_into_fp8_param(model_p: torch.Tensor, src_bf16: torch.Tensor) -> None: - """Requantize a gathered bf16 whole-param into an fp8 (Float8/MXFP8) model param in place. + """Copy a gathered BF16 whole-param into a compact LayerWise bucket parameter. - mxfp8 columnwise can't be derived from rowwise, so force columnwise before copy_ (TE rebuilds - both directions from the bf16); blockwise/Float8 columnwise is a lossless transpose. + Plain BF16 parameters are allowed because a LayerWise bucket can contain BF16 siblings of + supported FP8 parameters. Quantized destinations are limited to MXFP8 and blockwise tensors. + MXFP8 columnwise data cannot be derived from rowwise data, so force both usages before copy_. """ + if is_grouped_tensor_with_quantized_storage(model_p): + raise TypeError( + "LayerWise FP8 parameter gather does not support Transformer Engine GroupedTensor " + "quantized storage. Disable --moe-single-grouped-weight." + ) + if is_float8tensor(model_p) and not is_layerwise_fp8_param(model_p): + raise TypeError( + "LayerWise FP8 parameter gather supports only MXFP8Tensor and " + "Float8BlockwiseQTensor destinations." + ) if is_mxfp8tensor(model_p): quantizer = model_p.data._get_quantizer() quantizer.set_usage(rowwise=True, columnwise=True) @@ -321,16 +358,13 @@ def copy_back_gathered_bf16_into_fp8_param(model_p: torch.Tensor, src_bf16: torc def _stage_param_to_bf16(p: torch.Tensor) -> torch.Tensor: - """Stage a param to a detached bf16 whole-param for fp8 param-gather transport. - - Prefer the fp32 master (high-precision source); else dequantize an fp8 param; else copy bf16. - """ + """Stage a locally owned LayerWise parameter's FP32 master in BF16 for transport.""" main_param = getattr(p, "main_param", None) - if main_param is not None: - return main_param.detach().to(torch.bfloat16) - if is_float8tensor(p): - return dequantize_fp8_tensor(p).detach().to(torch.bfloat16) - return p.detach().to(torch.bfloat16) + if main_param is None: + raise RuntimeError( + "LayerWise FP8 parameter-gather staging requires param.main_param (FP32 master)." + ) + return main_param.detach().to(torch.bfloat16) def _resolve_callable_from_python_import_path(dotted_path: str): diff --git a/megatron/core/optimizer/layer_wise_optimizer.py b/megatron/core/optimizer/layer_wise_optimizer.py index 8a2e707d474..05e0c8aaaff 100644 --- a/megatron/core/optimizer/layer_wise_optimizer.py +++ b/megatron/core/optimizer/layer_wise_optimizer.py @@ -17,7 +17,7 @@ from ..fp8_utils import ( _stage_param_to_bf16, copy_back_gathered_bf16_into_fp8_param, - is_float8tensor, + is_layerwise_fp8_param, post_all_gather_processing, ) from .clip_grads import count_zeros_fp32, get_grad_norm_fp32 @@ -592,7 +592,7 @@ def __init__( self.shard_params(optimizers, full_param_layouts, model_chunks) # Engage FP8 param sync automatically when the decouple-managed params are actually - # quantized (fp8_param_gather on + TE Float8/MXFP8 weights). Off -> plain bf16 path. + # quantized (fp8_param_gather on + supported MXFP8/blockwise weights). Off -> plain bf16. # Also tag the gathered fp8 params: the fp8 all-gather (``_allgather_helper_fp8``) # requantizes bf16 -> each rank's fp8 ``param.data``, so the child optimizer's pre-gather # fp8 copy-back into ``param.data`` is redundant for them and is skipped. Params in these @@ -605,7 +605,7 @@ def __init__( continue for per_rank in params_list: for p in per_rank: - if is_float8tensor(p): + if is_layerwise_fp8_param(p): self.use_fp8_param_sync = True p._layer_wise_fp8_gathered = True @@ -952,7 +952,7 @@ def allgather_params(self) -> None: * **fp8** (``use_fp8_param_sync=True``): stage owned fp32 master->bf16, all-gather bf16, requantize into EVERY rank's ``param.data`` (owned included) so all hold ``Q(bf16(master))`` (== OFF/Adam). Then ``post_all_gather_processing`` rebuilds fp8 - columnwise/transpose (blockwise/Float8; mxfp8 noop since copy-back already forced it). + columnwise storage (blockwise; mxfp8 is a noop since copy-back already forced it). """ # FP8-aware variant: stage bf16, uneven all-gather bf16, requantize per rank. @@ -1003,8 +1003,8 @@ def _allgather_helper_fp8(params_list, group): copy_back_gathered_bf16_into_fp8_param(model_p, updated_bf16) # Rebuild fp8 columnwise/transpose after the gather (mirrors the overlap / DistOpt - # paths; blockwise/Float8 build it, mxfp8 is a noop). Else it'd be deferred to forward. - fp8_params = [p for params in params_list for p in params if is_float8tensor(p)] + # paths; blockwise builds it, mxfp8 is a noop). Else it'd be deferred to forward. + fp8_params = [p for params in params_list for p in params if is_layerwise_fp8_param(p)] if fp8_params: post_all_gather_processing(fp8_params) @@ -1064,11 +1064,11 @@ def _dispatch(params_list, group): # flatten is invalid. For a pure-bf16 model (no fp32 Muon params) the native group is # empty and this collapses to the original single-helper dispatch. staged = [ - [p for p in owned if is_float8tensor(p) or p.dtype != torch.float32] + [p for p in owned if is_layerwise_fp8_param(p) or p.dtype != torch.float32] for owned in params_list ] native = [ - [p for p in owned if not is_float8tensor(p) and p.dtype == torch.float32] + [p for p in owned if not is_layerwise_fp8_param(p) and p.dtype == torch.float32] for owned in params_list ] if any(owned for owned in staged): diff --git a/megatron/core/optimizer/optimizer.py b/megatron/core/optimizer/optimizer.py index 93d13007ecb..c05186171cf 100644 --- a/megatron/core/optimizer/optimizer.py +++ b/megatron/core/optimizer/optimizer.py @@ -46,7 +46,7 @@ optim_state_to_sharding_state, ) from ..dist_checkpointing.utils import add_prefix_for_sharding -from ..fp8_utils import copy_back_gathered_bf16_into_fp8_param, is_float8tensor +from ..fp8_utils import copy_back_gathered_bf16_into_fp8_param, is_layerwise_fp8_param from ..optimizer_param_scheduler import ParamGroupOverride as _ParamGroupOverride from ..transformer.module import param_is_not_shared from ..utils import log_single_rank @@ -1138,7 +1138,7 @@ def _copy_main_params_to_model_params(self): other_model_data, other_main_data = [], [] 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_float8tensor(model_param): + if is_layerwise_fp8_param(model_param): # Gathered fp8 params get ``Q(bf16(master))`` written into ``param.data`` # by the fp8 all-gather's requantize (``_allgather_helper_fp8``), which # would overwrite this copy -- so skip it for them. Non-gathered fp8 params diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py index 0e50dd06ed0..02d157f1f02 100644 --- a/megatron/training/arguments.py +++ b/megatron/training/arguments.py @@ -1699,9 +1699,8 @@ def validate_args(args, defaults={}): if not args.use_layer_wise_param_layout: # Decoupled compact LayerWise: fp8 parameter gather is supported via the FP8-aware # whole-param all-gather. Only mxfp8/blockwise (fp4 out of scope); mxfp8 needs - # reuse_grad_buf. fp4 is rejected unconditionally -- the LayerWise gather routes - # buckets by is_float8tensor, so an NVFP4 param would silently take the raw - # flatten path. + # reuse_grad_buf. fp4 is rejected unconditionally because the compact LayerWise + # gather implements only MXFP8/blockwise BF16 staging and copy-back. assert not getattr(args, 'fp4_param_gather', False), ( "Decoupled compact LayerWise DDP layout supports fp8 parameter gather only " "(mxfp8 or blockwise); fp4_param_gather is out of scope." @@ -1711,6 +1710,12 @@ def validate_args(args, defaults={}): "fp8 parameter gather on the decoupled compact LayerWise DDP layout requires " f"fp8_recipe in {{'mxfp8', 'blockwise'}}; got {args.fp8_recipe!r}." ) + assert not getattr(args, 'moe_single_grouped_weight', False), ( + "fp8 parameter gather on the decoupled compact LayerWise DDP layout does not " + "support --moe-single-grouped-weight: the LayerWise copy-back handles only " + "plain MXFP8/Float8Blockwise tensors, not Transformer Engine GroupedTensor " + "storage." + ) if args.fp8_recipe == 'mxfp8': assert args.reuse_grad_buf_for_mxfp8_param_ag, ( "mxfp8 + --fp8-param-gather on the decoupled compact LayerWise DDP layout " diff --git a/tests/unit_tests/test_fp8_utils.py b/tests/unit_tests/test_fp8_utils.py index 2fc55962ae8..42be4e25dfb 100644 --- a/tests/unit_tests/test_fp8_utils.py +++ b/tests/unit_tests/test_fp8_utils.py @@ -23,6 +23,40 @@ reason_for_no_mxfp8 = "MXFP8 requires Transformer Engine and device arch >= 10" +def test_is_layerwise_fp8_param_accepts_only_supported_storage_classes(): + class MockMXFP8Tensor: + pass + + class MockBlockwiseTensor: + pass + + class MockOtherQuantizedTensor: + pass + + with ( + patch.object(fp8_utils, 'HAVE_TE_MXFP8TENSOR', True), + patch.object(fp8_utils, 'MXFP8Tensor', MockMXFP8Tensor), + patch.object(fp8_utils, 'HAVE_TE_BLOCKWISE_FP8TENSOR', True), + patch.object(fp8_utils, 'Float8BlockwiseQTensor', MockBlockwiseTensor), + ): + assert fp8_utils.is_layerwise_fp8_param(MockMXFP8Tensor()) + assert fp8_utils.is_layerwise_fp8_param(MockBlockwiseTensor()) + assert not fp8_utils.is_layerwise_fp8_param(MockOtherQuantizedTensor()) + + +def test_stage_param_to_bf16_requires_fp32_master(): + param = nn.Parameter(torch.ones(4, dtype=torch.bfloat16)) + + with pytest.raises(RuntimeError, match=r"param\.main_param"): + fp8_utils._stage_param_to_bf16(param) + + param.main_param = torch.tensor([1.25, -2.5, 3.75, -4.0], dtype=torch.float32) + staged = fp8_utils._stage_param_to_bf16(param) + + assert staged.dtype == torch.bfloat16 + assert torch.equal(staged, param.main_param.to(torch.bfloat16)) + + @pytest.mark.skipif(not fp8_utils.HAVE_TE, reason="Transformer Engine is not installed") @pytest.mark.parametrize( ("is_init", "config_values", "te_helper"), diff --git a/tests/unit_tests/test_muon_decouple_fp8_param_gather.py b/tests/unit_tests/test_muon_decouple_fp8_param_gather.py index c10a29ed39f..cf1cf87f856 100644 --- a/tests/unit_tests/test_muon_decouple_fp8_param_gather.py +++ b/tests/unit_tests/test_muon_decouple_fp8_param_gather.py @@ -158,7 +158,13 @@ def model_provider(self, pre_process=True, post_process=True, **kw): ) def _create_args( - self, fp8_param_gather, fp8_recipe, overlap, num_experts=0, expert_model_parallel_size=1 + self, + fp8_param_gather, + fp8_recipe, + overlap, + num_experts=0, + expert_model_parallel_size=1, + single_grouped_weight=False, ): destroy_global_vars() destroy_num_microbatches_calculator() @@ -221,7 +227,9 @@ def _create_args( args.moe_router_topk = 2 args.moe_ffn_hidden_size = args.ffn_hidden_size args.moe_token_dispatcher_type = 'alltoall' - args.moe_grouped_gemm = False + args.moe_grouped_gemm = single_grouped_weight + args.moe_use_grouped_tensor = single_grouped_weight + args.moe_single_grouped_weight = single_grouped_weight # Deterministic routing comes from the fixed seed + deterministic_mode; drop the # aux-loss gradient term so ON and OFF compare cleanly without router-bias drift. args.moe_router_load_balancing_type = 'none' @@ -231,6 +239,10 @@ def _create_args( set_global_variables(args, False) return args + def test_rejects_single_grouped_weight_with_fp8_param_gather(self): + with pytest.raises(AssertionError, match="moe-single-grouped-weight"): + self._create_args(True, "blockwise", False, num_experts=8, single_grouped_weight=True) + def _batch(self): d = list(range(self.seq_length)) ids = torch.tensor(d, dtype=torch.int64).repeat((self.micro_batch_size, 1)).cuda() From f5ecccc2861cd95c623f4e5d42d50a454bb609ab Mon Sep 17 00:00:00 2001 From: Pingtian Li Date: Mon, 17 Aug 2026 20:21:09 +0800 Subject: [PATCH 08/12] Keep existing LayerWise FP8 tests unchanged Signed-off-by: Pingtian Li --- tests/unit_tests/test_fp8_utils.py | 34 ------------------- .../test_muon_decouple_fp8_param_gather.py | 16 ++------- 2 files changed, 2 insertions(+), 48 deletions(-) diff --git a/tests/unit_tests/test_fp8_utils.py b/tests/unit_tests/test_fp8_utils.py index 42be4e25dfb..2fc55962ae8 100644 --- a/tests/unit_tests/test_fp8_utils.py +++ b/tests/unit_tests/test_fp8_utils.py @@ -23,40 +23,6 @@ reason_for_no_mxfp8 = "MXFP8 requires Transformer Engine and device arch >= 10" -def test_is_layerwise_fp8_param_accepts_only_supported_storage_classes(): - class MockMXFP8Tensor: - pass - - class MockBlockwiseTensor: - pass - - class MockOtherQuantizedTensor: - pass - - with ( - patch.object(fp8_utils, 'HAVE_TE_MXFP8TENSOR', True), - patch.object(fp8_utils, 'MXFP8Tensor', MockMXFP8Tensor), - patch.object(fp8_utils, 'HAVE_TE_BLOCKWISE_FP8TENSOR', True), - patch.object(fp8_utils, 'Float8BlockwiseQTensor', MockBlockwiseTensor), - ): - assert fp8_utils.is_layerwise_fp8_param(MockMXFP8Tensor()) - assert fp8_utils.is_layerwise_fp8_param(MockBlockwiseTensor()) - assert not fp8_utils.is_layerwise_fp8_param(MockOtherQuantizedTensor()) - - -def test_stage_param_to_bf16_requires_fp32_master(): - param = nn.Parameter(torch.ones(4, dtype=torch.bfloat16)) - - with pytest.raises(RuntimeError, match=r"param\.main_param"): - fp8_utils._stage_param_to_bf16(param) - - param.main_param = torch.tensor([1.25, -2.5, 3.75, -4.0], dtype=torch.float32) - staged = fp8_utils._stage_param_to_bf16(param) - - assert staged.dtype == torch.bfloat16 - assert torch.equal(staged, param.main_param.to(torch.bfloat16)) - - @pytest.mark.skipif(not fp8_utils.HAVE_TE, reason="Transformer Engine is not installed") @pytest.mark.parametrize( ("is_init", "config_values", "te_helper"), diff --git a/tests/unit_tests/test_muon_decouple_fp8_param_gather.py b/tests/unit_tests/test_muon_decouple_fp8_param_gather.py index cf1cf87f856..c10a29ed39f 100644 --- a/tests/unit_tests/test_muon_decouple_fp8_param_gather.py +++ b/tests/unit_tests/test_muon_decouple_fp8_param_gather.py @@ -158,13 +158,7 @@ def model_provider(self, pre_process=True, post_process=True, **kw): ) def _create_args( - self, - fp8_param_gather, - fp8_recipe, - overlap, - num_experts=0, - expert_model_parallel_size=1, - single_grouped_weight=False, + self, fp8_param_gather, fp8_recipe, overlap, num_experts=0, expert_model_parallel_size=1 ): destroy_global_vars() destroy_num_microbatches_calculator() @@ -227,9 +221,7 @@ def _create_args( args.moe_router_topk = 2 args.moe_ffn_hidden_size = args.ffn_hidden_size args.moe_token_dispatcher_type = 'alltoall' - args.moe_grouped_gemm = single_grouped_weight - args.moe_use_grouped_tensor = single_grouped_weight - args.moe_single_grouped_weight = single_grouped_weight + args.moe_grouped_gemm = False # Deterministic routing comes from the fixed seed + deterministic_mode; drop the # aux-loss gradient term so ON and OFF compare cleanly without router-bias drift. args.moe_router_load_balancing_type = 'none' @@ -239,10 +231,6 @@ def _create_args( set_global_variables(args, False) return args - def test_rejects_single_grouped_weight_with_fp8_param_gather(self): - with pytest.raises(AssertionError, match="moe-single-grouped-weight"): - self._create_args(True, "blockwise", False, num_experts=8, single_grouped_weight=True) - def _batch(self): d = list(range(self.seq_length)) ids = torch.tensor(d, dtype=torch.int64).repeat((self.micro_batch_size, 1)).cuda() From 5290dd95141ee99d22822e350c02bf07cc94098b Mon Sep 17 00:00:00 2001 From: Pingtian Li Date: Mon, 17 Aug 2026 20:34:59 +0800 Subject: [PATCH 09/12] Format LayerWise optimizer argument validation Signed-off-by: Pingtian Li --- megatron/training/arguments.py | 14 ++++++++------ 1 file changed, 8 insertions(+), 6 deletions(-) diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py index 02d157f1f02..c90e6a883ff 100644 --- a/megatron/training/arguments.py +++ b/megatron/training/arguments.py @@ -41,7 +41,6 @@ from megatron.training.argument_utils import ArgumentGroupFactory, core_transformer_config_from_args # noqa: F401 # pylint: disable=unused-import - def add_megatron_arguments(parser: argparse.ArgumentParser): """"Add Megatron-LM arguments to the given parser.""" @@ -682,7 +681,7 @@ def validate_args(args, defaults={}): for elt in [args.train_data_path, args.valid_data_path, args.test_data_path]) or \ args.per_split_data_args_path is not None if use_per_split_data_path: - # Exactly one of the two has to be None if we use it. + # Exactly one of the two has to be None if we use it. assert any(elt is not None for elt in [args.train_data_path, args.valid_data_path, args.test_data_path]) is False or \ args.per_split_data_args_path is None @@ -1456,7 +1455,6 @@ def validate_args(args, defaults={}): if args.expert_model_parallel_size > 1 and 'ep_dp' not in args.high_priority_stream_groups: args.high_priority_stream_groups.append('ep_dp') - # Derive the internal gtp_weight_remat_size from the user-facing # --tensor-parallel-num-weight-shards. gtp_weight_remat_size has no CLI flag (it is excluded # from argument generation), so it is set here as a fresh attribute on args before it is @@ -1691,9 +1689,13 @@ def validate_args(args, defaults={}): args.use_distributed_optimizer = False assert not args.use_torch_fsdp2, "Emerging optimizer does not support Torch-FSDP2 for now." - assert not args.use_megatron_fsdp, "Emerging optimizer does not support Megatron-FSDP for now." - assert args.ckpt_format in ["torch", "torch_dist"], "Emerging optimizer supports torch and torch_dist checkpoint format." - + assert ( + not args.use_megatron_fsdp + ), "Emerging optimizer does not support Megatron-FSDP for now." + assert args.ckpt_format in [ + "torch", + "torch_dist", + ], "Emerging optimizer supports torch and torch_dist checkpoint format." if args.use_layer_wise_distributed_optimizer: if not args.use_layer_wise_param_layout: From 9c66a3596a97843f9a7a15f836e30c9c2e149af9 Mon Sep 17 00:00:00 2001 From: Pingtian Li Date: Mon, 17 Aug 2026 22:25:43 +0800 Subject: [PATCH 10/12] Align ModelBuilder LayerWise DDP setup Signed-off-by: Pingtian Li --- megatron/training/models/base.py | 31 ++++-- megatron/training/models/dist_utils.py | 74 ++++++++------- megatron/training/models/gpt.py | 95 ++++++++++++++----- megatron/training/models/hybrid.py | 34 ++++--- .../training/models/test_dist_utils.py | 48 +++++++++- 5 files changed, 196 insertions(+), 86 deletions(-) diff --git a/megatron/training/models/base.py b/megatron/training/models/base.py index 747c6442c38..259398b3f70 100644 --- a/megatron/training/models/base.py +++ b/megatron/training/models/base.py @@ -1,4 +1,4 @@ -# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. import abc import importlib @@ -69,12 +69,16 @@ class identified by the ``builder`` ClassVar. Must be serializable.""" # === pre-wrap and post-wrap hooks === - pre_wrap_hooks: list[Callable[[list[MegatronModule]], list[MegatronModule]]] = field(default_factory=list) + pre_wrap_hooks: list[Callable[[list[MegatronModule]], list[MegatronModule]]] = field( + default_factory=list + ) """List of functions that are executed before the model is wrapped with DDP/FSDP. Should take the model as the only argument and return a new model as the only return value. """ - post_wrap_hooks: list[Callable[[list[MegatronModule]], list[MegatronModule]]] = field(default_factory=list) + post_wrap_hooks: list[Callable[[list[MegatronModule]], list[MegatronModule]]] = field( + default_factory=list + ) """List of functions that are executed after model initialization is complete. Should take the model as the only argument and return a new model as the only return value. """ @@ -98,13 +102,15 @@ def as_dict(self) -> dict[str, Any]: """ def _as_dict(config): - result = { - "_target_": f"{config.__class__.__module__}.{config.__class__.__qualname__}", - } + result = {"_target_": f"{config.__class__.__module__}.{config.__class__.__qualname__}"} for f in dataclass_fields(config): value = getattr(config, f.name) # Skip non-serializable fields - if callable(value) or f.name.startswith("_") or f.name in ["pre_wrap_hooks", "post_wrap_hooks"]: + if ( + callable(value) + or f.name.startswith("_") + or f.name in ["pre_wrap_hooks", "post_wrap_hooks"] + ): continue if is_dataclass(value): @@ -144,7 +150,9 @@ def _from_dict(subdata): # Filter to valid fields for this class valid_fields = {f.name for f in dataclass_fields(config_cls)} - filtered_data = {k: v for k, v in subdata.items() if k in valid_fields and not k.startswith("_")} + filtered_data = { + k: v for k, v in subdata.items() if k in valid_fields and not k.startswith("_") + } # recurse on serialized nested dataclasses subconfigs = {} @@ -220,7 +228,9 @@ def build_distributed_models( use_torch_fsdp2: bool = False, wrap_with_ddp: bool = True, data_parallel_random_init: bool = False, - mixed_precision_wrapper: Callable[[Any, MegatronModule], MegatronModule] | None = Float16Module, + mixed_precision_wrapper: ( + Callable[[Any, MegatronModule], MegatronModule] | None + ) = Float16Module, model_type: ModelType = ModelType.encoder_or_decoder, use_layer_wise_distributed_optimizer: bool = False, use_layer_wise_param_layout: bool = True, @@ -239,7 +249,8 @@ def build_distributed_models( model_type: Deprecated flag, only used for backwards compatibility. use_layer_wise_distributed_optimizer: Whether the layerwise wiring runs. use_layer_wise_param_layout: When ``use_layer_wise_distributed_optimizer=True``, - controls whether to compute and supply a shard-aligned param layout to DDP. + selects the padded shard-aligned layout (``True``) or compact decoupled + layout (``False``) for LayerWise-managed buffers. Returns: List of model stages. If the model does not support virtual pipeline parallelism, diff --git a/megatron/training/models/dist_utils.py b/megatron/training/models/dist_utils.py index 30dd52f8adb..3fa498244f5 100644 --- a/megatron/training/models/dist_utils.py +++ b/megatron/training/models/dist_utils.py @@ -1,4 +1,4 @@ -# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. import logging from typing import Any, Callable @@ -31,7 +31,6 @@ from megatron.core.transformer.module import Float16Module from megatron.core.utils import get_model_config, get_pg_rank - try: from megatron.core.fp8_utils import correct_amax_history_if_needed except ImportError: @@ -86,7 +85,8 @@ def unimodal_build_distributed_models( model_type: Deprecated flag, only used for backwards compatibility. use_layer_wise_distributed_optimizer: Whether the layerwise wiring runs. use_layer_wise_param_layout: When ``use_layer_wise_distributed_optimizer=True``, - controls whether to compute and supply a shard-aligned param layout to DDP. + selects the padded shard-aligned layout (``True``) or compact decoupled + layout (``False``) for LayerWise-managed buffers. Returns: List of model stages, wrapped and ready for distributed training. @@ -98,9 +98,13 @@ def unimodal_build_distributed_models( init_model_with_meta_device = transformer_config.init_model_with_meta_device if init_model_with_meta_device: with torch.device("meta"): - model_list = build_virtual_pipeline_stages(build_model_func, pg_collection, vp_size, model_type) + model_list = build_virtual_pipeline_stages( + build_model_func, pg_collection, vp_size, model_type + ) else: - model_list = build_virtual_pipeline_stages(build_model_func, pg_collection, vp_size, model_type) + model_list = build_virtual_pipeline_stages( + build_model_func, pg_collection, vp_size, model_type + ) # Apply pre wrap hooks if pre_wrap_hook is not None: @@ -161,7 +165,8 @@ def prepare_existing_model_chunks_for_distributed_training( Pass ``None`` to skip. use_layer_wise_distributed_optimizer: Whether the layerwise wiring runs. use_layer_wise_param_layout: When ``use_layer_wise_distributed_optimizer=True``, - controls whether to compute and supply a shard-aligned param layout to DDP. + selects the padded shard-aligned layout (``True``) or compact decoupled + layout (``False``) for LayerWise-managed buffers. Returns: List of model chunks, wrapped and ready for distributed training. @@ -233,7 +238,12 @@ def _print_num_params(model: list[MegatronModule], pg_collection: ProcessGroupCo pg_collection.tp.rank(), get_pg_rank(pg_collection.gtp_remat), pg_collection.pp.rank(), - sum([sum([p.nelement() for p in model_module.parameters()]) for model_module in model]), + sum( + [ + sum([p.nelement() for p in model_module.parameters()]) + for model_module in model + ] + ), ), flush=True, ) @@ -247,7 +257,9 @@ def _wrap_with_mp_wrapper( fp16 = transformer_config.fp16 bf16 = transformer_config.bf16 if (fp16 or bf16) and mixed_precision_wrapper is not None: - model_list = [mixed_precision_wrapper(transformer_config, model_module) for model_module in model_list] + model_list = [ + mixed_precision_wrapper(transformer_config, model_module) for model_module in model_list + ] # Maintain expert bias in float32 wrapped in Float16Module for model_module in model_list: @@ -283,8 +295,8 @@ def _ddp_wrap( pg_collection: Model communication process groups. use_layer_wise_distributed_optimizer: Whether the layerwise wiring runs. use_layer_wise_param_layout: When ``use_layer_wise_distributed_optimizer=True``, - controls whether to compute and supply a shard-aligned param layout to DDP. - ``False`` keeps LayerWise on its legacy ``allgather_params`` sync path. + selects the padded shard-aligned layout (``True``) or compact decoupled + layout (``False``) for LayerWise-managed buffers. Returns: list[MegatronModule]: List of DDP/FSDP wrapped model modules @@ -292,14 +304,15 @@ def _ddp_wrap( if use_megatron_fsdp: DP = FullyShardedDataParallel if use_torch_fsdp2: - raise ValueError("Using use_megatron_fsdp and use_torch_fsdp2 at the same time is not supported.") + raise ValueError( + "Using use_megatron_fsdp and use_torch_fsdp2 at the same time is not supported." + ) elif use_torch_fsdp2: assert HAVE_FSDP2, "Torch FSDP2 requires torch>=2.4.0" DP = TorchFullyShardedDataParallel else: DP = DistributedDataParallel - if not use_torch_fsdp2: if ddp_config.num_buckets is not None: num_parameters = sum( @@ -324,12 +337,9 @@ def _ddp_wrap( # that the distributed-optimizer path provides. Mirrors wrap_model_chunks_with_ddp() in # megatron/training/training.py, which handles the non-ModelBuilder path. compute_full_param_layout = DistributedOptimizer.compute_full_param_layout - if ( - DP is DistributedDataParallel - and use_layer_wise_distributed_optimizer - and use_layer_wise_param_layout - ): + if DP is DistributedDataParallel and use_layer_wise_distributed_optimizer: ddp_config.use_distributed_optimizer = True + ddp_config.use_layer_wise_param_layout = use_layer_wise_param_layout compute_full_param_layout = LayerWiseDistributedOptimizer.compute_full_param_layout # Tag params so DDP buffer grouping routes LayerWise-managed matrices # (Muon's Newton-Schulz domain) to a shard-aligned buffer and routes @@ -354,22 +364,15 @@ def _ddp_wrap( wrapped_model = [] for model_chunk_idx, model_chunk in enumerate(model): chunk_kwargs = dict(dp_init_kwargs) - disable_bucketing = ( - (model_chunk_idx > 0) - or overlap_param_gather_with_optimizer_step - ) + disable_bucketing = (model_chunk_idx > 0) or overlap_param_gather_with_optimizer_step # Pre-compute parameter layouts for the distributed optimizer. # Only pass to DDP; FSDP variants don't accept full_param_layout. if ddp_config.use_distributed_optimizer and DP is DistributedDataParallel: - all_params = [ - p for p in model_chunk.parameters() if p.requires_grad - ] + all_params = [p for p in model_chunk.parameters() if p.requires_grad] pp_rank = pg_collection.pp.rank() effective_bucket_size = ( - None - if disable_bucketing or pp_rank > 0 - else ddp_config.bucket_size + None if disable_bucketing or pp_rank > 0 else ddp_config.bucket_size ) # Size the layout by the group the optimizer actually shards over, which is # the intra-instance group when there are several optimizer instances. Using @@ -381,9 +384,7 @@ def _ddp_wrap( all_params, effective_bucket_size, ( - intra_dp_cp_group - if intra_dp_cp_group is not None - else pg_collection.dp_cp + intra_dp_cp_group if intra_dp_cp_group is not None else pg_collection.dp_cp ).size(), ddp_config, expert_data_parallel_world_size=( @@ -443,13 +444,14 @@ def build_virtual_pipeline_stages( # Create multiple model stages for virtual pipeline model_list = [] for i in range(vp_size): - pre_process = is_vp_first_stage(vp_stage=i, vp_size=vp_size) and is_pp_first_stage(pp_group) - post_process = is_vp_last_stage(vp_stage=i, vp_size=vp_size) and is_pp_last_stage(pp_group) + pre_process = is_vp_first_stage(vp_stage=i, vp_size=vp_size) and is_pp_first_stage( + pp_group + ) + post_process = is_vp_last_stage(vp_stage=i, vp_size=vp_size) and is_pp_last_stage( + pp_group + ) model = build_model_func( - pg_collection, - pre_process=pre_process, - post_process=post_process, - vp_stage=i, + pg_collection, pre_process=pre_process, post_process=post_process, vp_stage=i ) model.model_type = model_type model_list.append(model) diff --git a/megatron/training/models/gpt.py b/megatron/training/models/gpt.py index 46dcc9b28f4..13d7286b2fc 100644 --- a/megatron/training/models/gpt.py +++ b/megatron/training/models/gpt.py @@ -1,11 +1,15 @@ -# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. import inspect import logging from typing import Any, Callable, ClassVar, Literal, override -from megatron.core.models.gpt.heterogeneous.heterogeneous_layer_specs import get_gpt_heterogeneous_layer_spec -from megatron.core.transformer.heterogeneous.heterogeneous_config import HeterogeneousTransformerConfig +from megatron.core.models.gpt.heterogeneous.heterogeneous_layer_specs import ( + get_gpt_heterogeneous_layer_spec, +) +from megatron.core.transformer.heterogeneous.heterogeneous_config import ( + HeterogeneousTransformerConfig, +) import torch from megatron.core.distributed.distributed_data_parallel_config import DistributedDataParallelConfig from megatron.core.enums import ModelType @@ -20,7 +24,9 @@ from megatron.core.process_groups_config import ProcessGroupCollection from megatron.core.transformer.spec_utils import ModuleSpec from megatron.core.transformer.module import Float16Module, MegatronModule -from megatron.core.transformer.dot_product_attention import DotProductAttention as MCoreDotProductAttention +from megatron.core.transformer.dot_product_attention import ( + DotProductAttention as MCoreDotProductAttention, +) from megatron.core.transformer.enums import AttnBackend from megatron.core.models.gpt.experimental_attention_variant_module_specs import ( get_transformer_block_with_experimental_attention_variant_spec, @@ -30,8 +36,7 @@ from megatron.training.vocab_utils import calculate_padded_vocab_size from megatron.training.models.dist_utils import unimodal_build_distributed_models -from megatron.core.transformer.transformer_config import TransformerConfig - +from megatron.core.transformer.transformer_config import TransformerConfig logger = logging.getLogger(__name__) @@ -82,7 +87,9 @@ def default_layer_spec(config: "GPTModelConfig", vp_stage: int) -> ModuleSpec: use_arbitrary_attention_mask=use_arbitrary_attention_mask, ) elif transformer_cfg.experimental_attention_variant is not None: - return get_transformer_block_with_experimental_attention_variant_spec(config=transformer_cfg, vp_stage=vp_stage) + return get_transformer_block_with_experimental_attention_variant_spec( + config=transformer_cfg, vp_stage=vp_stage + ) elif transformer_cfg.num_moe_experts is not None: return get_gpt_decoder_block_spec( transformer_cfg, @@ -96,13 +103,17 @@ def default_layer_spec(config: "GPTModelConfig", vp_stage: int) -> ModuleSpec: else: return _te_or_local_layer_spec(config, vp_stage) + def _te_or_local_layer_spec(config: "GPTModelConfig", vp_stage: int) -> ModuleSpec: """Need to be able to call just these branches for mtp transformer layer spec.""" transformer_cfg = config.transformer use_te = transformer_cfg.transformer_impl == "transformer_engine" if use_te: - if "use_te_op_fuser" in inspect.signature(get_gpt_layer_with_transformer_engine_spec).parameters: + if ( + "use_te_op_fuser" + in inspect.signature(get_gpt_layer_with_transformer_engine_spec).parameters + ): kwargs = {"use_te_op_fuser": config.use_transformer_engine_op_fuser} else: kwargs = {} @@ -135,7 +146,6 @@ def _te_or_local_layer_spec(config: "GPTModelConfig", vp_stage: int) -> ModuleSp ) - @dataclass(kw_only=True) class GPTModelConfig(ModelConfig): """Configuration for a Megatron Core GPT model. @@ -172,7 +182,9 @@ class GPTModelConfig(ModelConfig): logit_dtype: torch.dtype | None = None parallel_output: bool = True share_embeddings_and_output_weights: bool = False - position_embedding_type: Literal["learned_absolute", "rope", "mrope", "yarn", "none"] = "learned_absolute" + position_embedding_type: Literal["learned_absolute", "rope", "mrope", "yarn", "none"] = ( + "learned_absolute" + ) rotary_percent: float = 1.0 rotary_base: int = 10000 rope_scaling: bool = False @@ -196,7 +208,9 @@ def __getattr__(self, name: str, /) -> Any: raise AttributeError(f"GPTModelConfig has no attribute '{name}'") if hasattr(transformer, name): return getattr(transformer, name) - raise AttributeError(f"Neither GPTModelConfig nor TransformerConfig has any attribute '{name}'.") + raise AttributeError( + f"Neither GPTModelConfig nor TransformerConfig has any attribute '{name}'." + ) @override def __setattr__(self, name: str, value: Any, /) -> None: @@ -231,12 +245,17 @@ def finalize(self) -> None: or self.transformer.account_for_loss_in_pipeline_split ) is_pipeline_asymmetric |= ( - self.transformer.num_layers_in_first_pipeline_stage or self.transformer.num_layers_in_last_pipeline_stage + self.transformer.num_layers_in_first_pipeline_stage + or self.transformer.num_layers_in_last_pipeline_stage ) is not None - is_flexible_pp_layout = is_pipeline_asymmetric or (self.transformer.pipeline_model_parallel_layout is not None) + is_flexible_pp_layout = is_pipeline_asymmetric or ( + self.transformer.pipeline_model_parallel_layout is not None + ) if vp_size and not is_flexible_pp_layout: p_size = self.transformer.pipeline_model_parallel_size - assert (self.transformer.num_layers // p_size) % vp_size == 0, ( + assert ( + self.transformer.num_layers // p_size + ) % vp_size == 0, ( "Make sure the number of model chunks is the same across all pipeline stages." ) @@ -279,14 +298,20 @@ def build_model( transformer_layer_spec = self._model_config.transformer_layer_spec if transformer_layer_spec is None: transformer_layer_spec = default_layer_spec(self._model_config, vp_stage) - elif not isinstance(transformer_layer_spec, ModuleSpec) and callable(transformer_layer_spec): + elif not isinstance(transformer_layer_spec, ModuleSpec) and callable( + transformer_layer_spec + ): # Check if the transformer_layer_spec function accepts vp_stage parameter if "vp_stage" in inspect.signature(transformer_layer_spec).parameters: - transformer_layer_spec = transformer_layer_spec(self._model_config, vp_stage=vp_stage) + transformer_layer_spec = transformer_layer_spec( + self._model_config, vp_stage=vp_stage + ) else: transformer_layer_spec = transformer_layer_spec(self._model_config) - assert self._model_config.vocab_size is not None, "vocab_size must be configured before calling build_model()" + assert ( + self._model_config.vocab_size is not None + ), "vocab_size must be configured before calling build_model()" if self._model_config.should_pad_vocab: padded_vocab_size = calculate_padded_vocab_size( self._model_config.vocab_size, @@ -301,14 +326,20 @@ def build_model( # override spec with local backend if configured if self._model_config.attention_backend == AttnBackend.local: if hasattr(transformer_layer_spec, "submodules"): - transformer_layer_spec.submodules.self_attention.submodules.core_attention = MCoreDotProductAttention + transformer_layer_spec.submodules.self_attention.submodules.core_attention = ( + MCoreDotProductAttention + ) # Determine pre/post flags if not provided using vp + pp stage vp_size = self._model_config.virtual_pipeline_model_parallel_size if pre_process is None: - pre_process = is_vp_first_stage(vp_stage=vp_stage, vp_size=vp_size) and is_pp_first_stage(pg_collection.pp) + pre_process = is_vp_first_stage( + vp_stage=vp_stage, vp_size=vp_size + ) and is_pp_first_stage(pg_collection.pp) if post_process is None: - post_process = is_vp_last_stage(vp_stage=vp_stage, vp_size=vp_size) and is_pp_last_stage(pg_collection.pp) + post_process = is_vp_last_stage( + vp_stage=vp_stage, vp_size=vp_size + ) and is_pp_last_stage(pg_collection.pp) model = GPTModel( config=self._model_config.transformer, @@ -344,7 +375,9 @@ def build_distributed_models( use_torch_fsdp2: bool = False, wrap_with_ddp: bool = True, data_parallel_random_init: bool = True, - mixed_precision_wrapper: Callable[[Any, MegatronModule], MegatronModule] | None = Float16Module, + mixed_precision_wrapper: ( + Callable[[Any, MegatronModule], MegatronModule] | None + ) = Float16Module, model_type: ModelType = ModelType.encoder_or_decoder, use_layer_wise_distributed_optimizer: bool = False, use_layer_wise_param_layout: bool = True, @@ -364,7 +397,8 @@ def build_distributed_models( model_type: Deprecated flag, only used for backwards compatibility. use_layer_wise_distributed_optimizer: Whether the layerwise wiring runs. use_layer_wise_param_layout: When ``use_layer_wise_distributed_optimizer=True``, - controls whether to compute and supply a shard-aligned param layout to DDP. + selects the padded shard-aligned layout (``True``) or compact decoupled + layout (``False``) for LayerWise-managed buffers. Returns: List of model stages. @@ -415,14 +449,25 @@ def mtp_block_spec( if config.transformer.mtp_num_layers is not None: from megatron.core.models.gpt.gpt_layer_specs import get_gpt_mtp_block_spec - if hasattr(transformer_layer_spec, "layer_specs") and len(transformer_layer_spec.layer_specs) == 0: + if ( + hasattr(transformer_layer_spec, "layer_specs") + and len(transformer_layer_spec.layer_specs) == 0 + ): # Get the decoder layer spec explicitly if no decoder layer in the last stage, # Only happens with block spec (TransformerBlockSubmodules) when using MoE. spec = _te_or_local_layer_spec(config, vp_stage) else: - decoder_specs = get_gpt_decoder_layer_specs(transformer_cfg, use_transformer_engine=use_te, normalization=transformer_cfg.normalization, qk_l2_norm=transformer_cfg.qk_l2_norm, vp_stage=vp_stage) + decoder_specs = get_gpt_decoder_layer_specs( + transformer_cfg, + use_transformer_engine=use_te, + normalization=transformer_cfg.normalization, + qk_l2_norm=transformer_cfg.qk_l2_norm, + vp_stage=vp_stage, + ) spec = decoder_specs[-1] - return get_gpt_mtp_block_spec(transformer_cfg, spec, use_transformer_engine=use_te, vp_stage=vp_stage) + return get_gpt_mtp_block_spec( + transformer_cfg, spec, use_transformer_engine=use_te, vp_stage=vp_stage + ) else: return None diff --git a/megatron/training/models/hybrid.py b/megatron/training/models/hybrid.py index c7d1c91b118..e330899ff02 100644 --- a/megatron/training/models/hybrid.py +++ b/megatron/training/models/hybrid.py @@ -1,4 +1,4 @@ -# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. import logging from dataclasses import dataclass @@ -19,11 +19,7 @@ from megatron.core.transformer.module import Float16Module, MegatronModule from megatron.core.transformer.spec_utils import ModuleSpec from megatron.core.transformer.transformer_config import TransformerConfig -from megatron.training.models.base import ( - ModelBuilder, - ModelConfig, - compose_hooks, -) +from megatron.training.models.base import ModelBuilder, ModelConfig, compose_hooks from megatron.training.models.dist_utils import unimodal_build_distributed_models from megatron.training.vocab_utils import calculate_padded_vocab_size @@ -89,7 +85,9 @@ def __getattr__(self, name: str, /) -> Any: raise AttributeError(f"HybridModelConfig has no attribute '{name}'") if hasattr(transformer, name): return getattr(transformer, name) - raise AttributeError(f"Neither HybridModelConfig nor TransformerConfig has any attribute '{name}'.") + raise AttributeError( + f"Neither HybridModelConfig nor TransformerConfig has any attribute '{name}'." + ) @override def __setattr__(self, name: str, value: Any, /) -> None: @@ -157,13 +155,14 @@ def build_model( hybrid_stack_spec = hybrid_inference_stack_spec elif self._model_config.restore_modelopt_state: hybrid_stack_spec = get_hybrid_stack_modelopt_spec( - local_core_attention=False, - remap_te_layernorm=False, + local_core_attention=False, remap_te_layernorm=False ) else: hybrid_stack_spec = default_hybrid_stack_spec - assert self._model_config.vocab_size is not None, "vocab_size must be configured before calling build_model()" + assert ( + self._model_config.vocab_size is not None + ), "vocab_size must be configured before calling build_model()" if self._model_config.should_pad_vocab: padded_vocab_size = calculate_padded_vocab_size( self._model_config.vocab_size, @@ -173,8 +172,12 @@ def build_model( else: padded_vocab_size = self._model_config.vocab_size - pre_process = pre_process if pre_process is not None else is_pp_first_stage(pg_collection.pp) - post_process = post_process if post_process is not None else is_pp_last_stage(pg_collection.pp) + pre_process = ( + pre_process if pre_process is not None else is_pp_first_stage(pg_collection.pp) + ) + post_process = ( + post_process if post_process is not None else is_pp_last_stage(pg_collection.pp) + ) return HybridModel( config=self._model_config.transformer, hybrid_stack_spec=hybrid_stack_spec, @@ -204,7 +207,9 @@ def build_distributed_models( use_torch_fsdp2: bool = False, wrap_with_ddp: bool = True, data_parallel_random_init: bool = False, - mixed_precision_wrapper: Callable[[Any, MegatronModule], MegatronModule] | None = Float16Module, + mixed_precision_wrapper: ( + Callable[[Any, MegatronModule], MegatronModule] | None + ) = Float16Module, model_type: ModelType = ModelType.encoder_or_decoder, use_layer_wise_distributed_optimizer: bool = False, use_layer_wise_param_layout: bool = True, @@ -224,7 +229,8 @@ def build_distributed_models( model_type: Deprecated flag, only used for backwards compatibility. use_layer_wise_distributed_optimizer: Whether the layerwise wiring runs. use_layer_wise_param_layout: When ``use_layer_wise_distributed_optimizer=True``, - controls whether to compute and supply a shard-aligned param layout to DDP. + selects the padded shard-aligned layout (``True``) or compact decoupled + layout (``False``) for LayerWise-managed buffers. Returns: List of model stages. diff --git a/tests/unit_tests/training/models/test_dist_utils.py b/tests/unit_tests/training/models/test_dist_utils.py index 3d1cb6e0978..555143cee05 100644 --- a/tests/unit_tests/training/models/test_dist_utils.py +++ b/tests/unit_tests/training/models/test_dist_utils.py @@ -1,4 +1,4 @@ -# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. from unittest.mock import MagicMock, Mock, patch @@ -720,6 +720,52 @@ def test_layout_passed_to_ddp_when_distributed_optimizer( assert layout_args.args[3] is ddp_config assert layout_args.kwargs["expert_data_parallel_world_size"] == 2 + @pytest.mark.parametrize("use_layer_wise_param_layout", [False, True]) + @patch("megatron.training.models.dist_utils.tag_params_for_buffer_routing") + @patch("megatron.training.models.dist_utils.LayerWiseDistributedOptimizer") + @patch("megatron.training.models.dist_utils.DistributedDataParallel") + @patch("megatron.training.models.dist_utils.get_model_config") + @patch("torch.cuda.stream", new_callable=MagicMock) + @patch("torch.cuda.current_stream") + @patch("torch.cuda.Stream") + def test_layer_wise_layout_matches_traditional_model_entry( + self, + mock_stream, + mock_curr, + mock_ctx, + mock_cfg, + mock_ddp, + mock_layer_wise_optimizer, + mock_tag_params, + use_layer_wise_param_layout, + ): + mock_ctx.return_value.__enter__ = Mock(return_value=None) + mock_ctx.return_value.__exit__ = Mock(return_value=False) + mock_layer_wise_optimizer.compute_full_param_layout.return_value = "LAYERWISE_LAYOUT" + chunk, param = self._make_chunk_with_params() + ddp_config = self._ddp_config( + use_distributed_optimizer=False, + use_layer_wise_param_layout=not use_layer_wise_param_layout, + ) + + _ddp_wrap( + [chunk], + False, + ddp_config, + False, + pg_collection=self.pg, + use_layer_wise_distributed_optimizer=True, + use_layer_wise_param_layout=use_layer_wise_param_layout, + ) + + assert ddp_config.use_distributed_optimizer + assert ddp_config.use_layer_wise_param_layout is use_layer_wise_param_layout + mock_tag_params.assert_called_once_with([chunk]) + layout_args = mock_layer_wise_optimizer.compute_full_param_layout.call_args + assert layout_args.args[0] == [param] + assert layout_args.args[3] is ddp_config + assert mock_ddp.call_args.kwargs["full_param_layout"] == "LAYERWISE_LAYOUT" + @patch("megatron.training.models.dist_utils.DistributedDataParallel") @patch("megatron.training.models.dist_utils.get_model_config") @patch("torch.cuda.stream", new_callable=MagicMock) From ead38ad2ed305687499f97adfba39f2324c3dead Mon Sep 17 00:00:00 2001 From: Pingtian Li Date: Tue, 18 Aug 2026 14:08:30 +0800 Subject: [PATCH 11/12] Remove LayerWise checkpoint changes Signed-off-by: Pingtian Li --- .../distributed_data_parallel_config.py | 2 +- megatron/core/optimizer/distrib_optimizer.py | 179 +----------- megatron/core/optimizer/optimizer_config.py | 2 +- megatron/training/arguments.py | 2 +- megatron/training/training.py | 2 +- .../test_layer_wise_optimizer.py | 271 +----------------- tests/unit_tests/dist_checkpointing/utils.py | 117 +------- 7 files changed, 26 insertions(+), 549 deletions(-) diff --git a/megatron/core/distributed/distributed_data_parallel_config.py b/megatron/core/distributed/distributed_data_parallel_config.py index 3b86ad69e05..558b195bf9a 100644 --- a/megatron/core/distributed/distributed_data_parallel_config.py +++ b/megatron/core/distributed/distributed_data_parallel_config.py @@ -1,4 +1,4 @@ -# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved. +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. from dataclasses import dataclass from typing import Optional, Tuple diff --git a/megatron/core/optimizer/distrib_optimizer.py b/megatron/core/optimizer/distrib_optimizer.py index f4095ac16bf..54819099653 100644 --- a/megatron/core/optimizer/distrib_optimizer.py +++ b/megatron/core/optimizer/distrib_optimizer.py @@ -8,7 +8,7 @@ from collections import ChainMap from dataclasses import replace from logging import getLogger -from typing import Any, Callable, Dict, List, Optional, Set, Tuple +from typing import Any, Callable, Dict, List, Optional, Tuple import torch import torch.nn.functional @@ -255,26 +255,6 @@ def _build_model_gbuf_range(cls, param_and_grad_buffer: _ParamAndGradBuffer, buc return data - @classmethod - def _filter_gbuf_range_map( - cls, gbuf_range_map: Dict, optimizer_params: Set[torch.nn.Parameter] - ) -> Dict: - """Filter grad-buffer range maps to the params owned by this optimizer instance.""" - return { - dtype: [ - { - **range_map, - "param_map": { - param: param_range - for param, param_range in range_map["param_map"].items() - if param in optimizer_params - }, - } - for range_map in range_maps - ] - for dtype, range_maps in gbuf_range_map.items() - } - @classmethod def _build_gbuf_range_map(cls, param_and_grad_buffer: _ParamAndGradBuffer): """Builds a map between parameters and their ranges in the grad buffer. @@ -748,13 +728,6 @@ def __init__( for model_idx, buffers in self.per_model_buffers.items(): self.per_model_bucket_groups[model_idx] = partition_buckets(buffers) - optimizer_params = { - param for param_group in self.optimizer.param_groups for param in param_group['params'] - } - # Model params this optimizer owns, captured before the fp32-master swap below. Used by - # sharded_param_state_dp_reshardable / load to skip buckets it owns no param of, and to - # distinguish decoupled LayerWise setups (empty shards expected) from pure DistOpt. - self._optimizer_model_params = optimizer_params self.gbuf_ranges = [] self.per_bucket_numel = [] self.per_bucket_numel_unpadded = [] @@ -774,9 +747,7 @@ def __init__( ] } ) - self.gbuf_ranges.append( - self._filter_gbuf_range_map(self._build_gbuf_range_map(buffer), optimizer_params) - ) + self.gbuf_ranges.append(self._build_gbuf_range_map(buffer)) self.model_param_gbuf_map = self._build_model_param_gbuf_map(self.gbuf_ranges) # Add main_param field to each parameter. We will use this fp32 copy to compute @@ -1896,58 +1867,6 @@ def sharded_param_state_dp_reshardable( data_parallel_world_size = self.data_parallel_group.size() state = self.get_parameter_state_dp_reshardable() - - # Optimizer-state {key: dtype} per bucket, captured before the loop below - # mutates ``state``, so a shard that is entirely padding can still synthesize valid - # ShardedTensors. Keyed per (gbuf_idx, dtype, bucket_idx) because dist_checkpointing - # enforces one dtype per KEY (validation.py, ``assert sharding.dtype == dtype``) and each - # bucket is its own key -- a template sampled from one bucket is not authoritative for - # another. ``step`` is a per-param 0-dim tensor for torch AdamW / non-TE-Apex / - # HybridDeviceOptimizer, but it becomes a LocalNonpersistentObject below and never a - # ShardedTensor, so it must not be synthesized here. - # Device stays OUT of the gathered template: it is rank-local (under - # --optimizer-cpu-offload the CPU/GPU split lands on a different param on every - # rank, because HybridDeviceOptimizer walks this rank's own shard list), and - # dist_checkpointing validates dtype/shape but never device. Comparing it across - # ranks would abort the save on a legitimate config. - pad_templates = {} - pad_device = None - for _g in range(len(self.gbuf_ranges)): - for _dt_key, _bs_all in state[_g].items(): - for _b_idx, _bs in enumerate(_bs_all): - if _bs: - _key = (_g, str(_dt_key), _b_idx) - pad_templates[_key] = { - k: v.dtype - for k, v in _bs[0].items() - if isinstance(v, torch.Tensor) and k != 'step' - } - if pad_device is None and pad_templates[_key]: - pad_device = _bs[0][next(iter(pad_templates[_key]))].device - - # A rank that owns nothing in a bucket has no local sample there, yet it may still owe - # padding coverage for it. The template cannot be derived from config: the save path runs - # the state through ``get_unscaled_state`` (which upcasts bf16/fp16/fp8 state to fp32, and - # returns int16 for ``store_param_remainders``), and the key set comes from whatever - # ``optimizer.state`` actually holds. So reconcile the observed templates across the DP - # group. The assert is per bucket -- the axis dist_checkpointing actually constrains -- - # so buckets nobody sampled simply stay absent instead of aborting the save. - if data_parallel_world_size > 1: - _gathered = [None] * data_parallel_world_size - torch.distributed.all_gather_object( - _gathered, pad_templates, group=self.data_parallel_group - ) - _merged = {} - for _candidate in _gathered: - for _k, _v in (_candidate or {}).items(): - _prev = _merged.setdefault(_k, _v) - assert _prev == _v, ( - f"padding template mismatch across DP ranks for bucket {_k}: " - f"{_prev} vs {_v}; the synthesized padding would not match the " - "real shards" - ) - pad_templates = _merged - # per_bucket_numel metadata is saved separately for each TPxPP domain. for per_bucket_key in ('per_bucket_numel', 'per_bucket_numel_unpadded'): key = ( @@ -1979,75 +1898,9 @@ def sharded_param_state_dp_reshardable( f'.gbuf_idx_{gbuf_idx}.dtype_{dtype}.bucket_idx_{bucket_idx}' ) - # Skip buckets this optimizer owns no param of. In the decoupled compact - # LayerWise (Muon) layout our buffers also hold buckets whose params the - # LayerWise child owns; their state is saved there and every shard here is - # empty. Checking membership (vs the params[0] tag) also handles buckets that - # mix owned and LayerWise-managed params. - bucket = self.buffers[gbuf_idx].buckets[bucket_idx] - if not any(p in self._optimizer_model_params for p in bucket.params_list): - continue - - # bucket_state is built 1:1 from this rank's param_map, so an empty state is - # legitimate only when the rank owns no param range in this bucket (a small - # bucket + 64-element shard alignment can put a whole shard in padding). If it - # does own a range yet the state is empty, the state was lost -- keep the - # original strict check. - if not bucket_state: - assert not self.gbuf_ranges[gbuf_idx][dtype][bucket_idx]['param_map'], ( - f'empty dp_reshardable state for {sharded_bucket_key} but this rank ' - 'owns param ranges in the bucket (optimizer state lost)' - ) - world_shard_start = data_parallel_rank * gbuf_local_numel - if world_shard_start >= gbuf_world_numel_unpadded: - # Shard is entirely past the unpadded end (trailing padding); not saved. - continue - pad_len = min( - gbuf_local_numel, gbuf_world_numel_unpadded - world_shard_start - ) - # Synthesize this shard's overlap with [0, numel_unpadded) from the - # template observed for THIS bucket. If no rank sampled it (every shard - # of the bucket is padding), fall back to any observed template rather - # than dropping coverage; an empty map means the coverage is being - # dropped, so fail loudly. - _tpl = pad_templates.get((gbuf_idx, str(dtype), bucket_idx)) - if not _tpl: - _tpl = next(iter(pad_templates.values()), None) - assert _tpl, ( - f'no optimizer-state template available for {sharded_bucket_key}; ' - 'the padding coverage this rank owes would be dropped' - ) - # ``pad_device`` is a single rank-local device, NOT a per-bucket map. - # Deliberate: we only synthesize for buckets this rank sampled nothing - # in, so a per-bucket lookup would miss by construction. Borrowing a - # device across buckets is safe in a way borrowing a dtype is NOT -- - # dtype and shape are validated, device never is, and the padding bytes - # are dropped on load. Do NOT make this symmetric with the dtype path: - # a rank-local device inside the cross-rank compare aborts the save - # under --optimizer-cpu-offload, where the CPU/GPU cutoff differs per - # rank. - bucket_state = [ - { - **{ - k: torch.empty( - pad_len, - dtype=_dt, - device=( - pad_device - if pad_device is not None - else torch.cuda.current_device() - ), - ) - for k, _dt in _tpl.items() - }, - 'gbuf_local_start': 0, - 'gbuf_local_end': pad_len, - 'padding': True, - } - ] - # Store the synthesized shard back into ``state``: this branch rebinds - # ``bucket_state`` to a new list, so without this the shard is dropped. - gbuf_range_map_for_all_buckets[bucket_idx] = bucket_state + # The global ckpt tensors must be fully covered. + # We add extra empty padding if necessary + assert bucket_state, 'empty bucket encountered' # Insert padding between parameter tensors to ensure full coverage as needed. all_pad_tensors = {} @@ -2139,22 +1992,6 @@ def sharded_param_state_fs_model_space( This will allow changing TP and PP while using DistOpt (as with other optimizers). """ - # NB: fully_sharded_model_space is independently non-functional for EVERY - # DistributedOptimizer in this tree, not just the decoupled layout -- this function sets - # ``flattened_range`` on every non-factory param below, and ShardedTensor rejects that - # unconditionally in validate_metadata_integrity. Only ShardedTensorFactory escapes, and - # only gated-MLP fc1 produces one. This assert does not fix that; it makes the compact - # LayerWise (Muon) case abort earlier and with an actionable message instead of dying - # per-param deep inside mapping.py. It is deliberately narrow: broadening it to reject - # the format outright is a separate change. - assert not ( - self.config.use_layer_wise_distributed_optimizer - and not self.config.use_layer_wise_param_layout - ), ( - "fully_sharded_model_space is not supported for the decoupled compact LayerWise " - "(Muon) optimizer layout; use dp_reshardable or fully_reshardable instead" - ) - param_to_sharded_metadata = {} model_sharded_state_dict, _ = extract_sharded_tensors_and_factories( model_sharded_state_dict @@ -2242,12 +2079,6 @@ def load_parameter_state_from_dp_reshardable(self, state_dict): assert len(gbuf_range_maps) == 1, "single dtype supported, for now." for dtype, gbuf_range_map_for_all_buckets in gbuf_range_maps.items(): for bucket_idx, gbuf_range_map in enumerate(gbuf_range_map_for_all_buckets): - # Skip buckets this optimizer owns no param of (see the save counterpart). - if not any( - p in self._optimizer_model_params - for p in self.buffers[gbuf_idx].buckets[bucket_idx].params_list - ): - continue bucket_state = state_dict[gbuf_idx][dtype][bucket_idx] bucket_state = [ bucket_state_elem diff --git a/megatron/core/optimizer/optimizer_config.py b/megatron/core/optimizer/optimizer_config.py index 60fadd1af7b..a42a7532a09 100644 --- a/megatron/core/optimizer/optimizer_config.py +++ b/megatron/core/optimizer/optimizer_config.py @@ -1,4 +1,4 @@ -# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. import fnmatch from dataclasses import dataclass, field diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py index c90e6a883ff..2b355b94af7 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.""" diff --git a/megatron/training/training.py b/megatron/training/training.py index dbb557b7953..ac51b36f619 100644 --- a/megatron/training/training.py +++ b/megatron/training/training.py @@ -1,4 +1,4 @@ -# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. """Pretrain utilities.""" diff --git a/tests/unit_tests/dist_checkpointing/test_layer_wise_optimizer.py b/tests/unit_tests/dist_checkpointing/test_layer_wise_optimizer.py index a629b24abdd..5c20be82e04 100644 --- a/tests/unit_tests/dist_checkpointing/test_layer_wise_optimizer.py +++ b/tests/unit_tests/dist_checkpointing/test_layer_wise_optimizer.py @@ -1,4 +1,4 @@ -# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. from copy import deepcopy from functools import partial @@ -23,7 +23,6 @@ from megatron.core.utils import get_pg_size from megatron.training.arguments import parse_args from megatron.training.checkpointing import load_checkpoint, save_checkpoint -from megatron.training.utils import get_device_arch_version from tests.unit_tests.dist_checkpointing import ( TempNamedDir, init_basic_mock_args, @@ -188,42 +187,13 @@ def test_broadcast_params(self, tp, pp): assert torch.allclose(param.data, original_params[name]) # TODO(deyuf): check bf16 False case - # Cover both sides of use_layer_wise_distributed_optimizer with a real distributed - # optimizer on each side. ``plain_distopt`` is an in-file CONTROL -- test_optimizer.py's - # resharding tests already cover a plain bf16 DistOpt under fully_reshardable, and harder - # resharding cases besides; it is here to show the same harness behaves with and without - # LayerWise. ``layerwise_compact`` is the load-bearing arm: it is the only fully_reshardable - # coverage of the compact layout above tp=pp=1, and the only one with grad_reduce_in_fp32 - # left at its default. ``use_param_layout`` is the DDP-level LayerWise routing switch - # (ddp_use_layer_wise = use_layer_wise and use_param_layout), so the LayerWise arm needs it - # True to build the decoupled layout at all; ``use_layer_wise_param_layout`` stays False, - # which selects the COMPACT layout this PR adds -- the padded shard-aligned layout is out of - # scope. The legacy ping-pong path (dist_opt=False) is covered by the other tests here. - @pytest.mark.parametrize( - 'layer_wise', - [pytest.param(False, id='plain_distopt'), pytest.param(True, id='layerwise_compact')], - ) @pytest.mark.parametrize('tp', [1, 2, 4]) @pytest.mark.parametrize('pp', [1, 2, 4]) @pytest.mark.parametrize('bf16', [True]) - def test_layer_wise_optimizer_save_load(self, tmp_path_dist_ckpt, tp, pp, bf16, layer_wise): + def test_layer_wise_optimizer_save_load(self, tmp_path_dist_ckpt, tp, pp, bf16): """Test save/load of LayerWiseDistributedOptimizer checkpoints.""" if tp * pp > 8: pytest.skip(f"TP*PP > 8 is larger than world size") - _kw = dict( - optimizer='dist_muon' if layer_wise else 'adam', - use_param_layout=layer_wise, - use_layer_wise_param_layout=False, - ) - # Both arms now build a real DistributedOptimizer, so the format has to be requested: - # - the default fully_sharded_model_space is unusable here -- the compact layout is - # rejected outright, and an ordinary DistOpt dies deeper in ``replace(...)`` with - # "ShardedTensor.flattened_range is not supported"; - # - dp_reshardable synthesizes padding with ``torch.empty``, so two saves of identical - # state differ in the uninitialized padding bytes and the check_equal below would be - # flaky (this is why test_decouple_ckpt_roundtrip_values needs its canonical view). - # fully_reshardable has neither problem and keeps the plain A-vs-B comparison valid. - _md = {'distrib_optim_sharding_type': 'fully_reshardable'} Utils.initialize_model_parallel(tp, pp) @@ -239,14 +209,14 @@ def test_layer_wise_optimizer_save_load(self, tmp_path_dist_ckpt, tp, pp, bf16, tp=tp, pp=pp, bf16=bf16, - dist_opt=True, + dist_opt=False, initialize_fn=initialize_gpt_model, - **_kw, + optimizer='dist_muon', ) # Save checkpoint A model_sharded_sd_A = model_A[0].sharded_state_dict() - optim_sd_A = optimizer_A.sharded_state_dict(model_sharded_sd_A, metadata=_md) + optim_sd_A = optimizer_A.sharded_state_dict(model_sharded_sd_A) save(optim_sd_A, ckpt_dir_A) # Create model and optimizer B with different seed @@ -255,21 +225,21 @@ def test_layer_wise_optimizer_save_load(self, tmp_path_dist_ckpt, tp, pp, bf16, tp=tp, pp=pp, bf16=bf16, - dist_opt=True, + dist_opt=False, initialize_fn=initialize_gpt_model, - **_kw, + optimizer='dist_muon', ) # Load checkpoint A into optimizer B model_sharded_sd_B = model_B[0].sharded_state_dict() load_sharded_sd = optimizer_B.sharded_state_dict( - model_sharded_sd_B, is_loading=True, metadata=_md + model_sharded_sd_B, is_loading=True ) state_dict = load(load_sharded_sd, ckpt_dir_A) optimizer_B.load_state_dict(state_dict) # Save as checkpoint B - optim_sd_B = optimizer_B.sharded_state_dict(model_sharded_sd_B, metadata=_md) + optim_sd_B = optimizer_B.sharded_state_dict(model_sharded_sd_B) save(optim_sd_B, ckpt_dir_B) Utils.destroy_model_parallel() @@ -283,229 +253,6 @@ def test_layer_wise_optimizer_save_load(self, tmp_path_dist_ckpt, tp, pp, bf16, check_equal(plain_sd_A, plain_sd_B) - # grad_reduce_in_fp32=True gives the mixed-dtype (bf16 param, fp32 grad) DistOpt sibling - # buffer of the decoupled Muon layout -- the case that fails 'Failed to validate global plan' - # in sharded_param_state_dp_reshardable on the real fp8 SFT save. - @pytest.mark.parametrize('grad_reduce_in_fp32', [False, True]) - @pytest.mark.parametrize('bf16', [True]) - def test_dp_reshardable_decouple_ckpt(self, tmp_path_dist_ckpt, bf16, grad_reduce_in_fp32): - """Save/load of the decoupled compact LayerWise (Muon) optimizer in ``dp_reshardable`` - format, for both the uniform (bf16, bf16) and the mixed-dtype (bf16, fp32) sibling - DistOpt buffers. Exercises ``sharded_param_state_dp_reshardable`` on the decouple path, - including the empty-bucket and global-plan-coverage handling. - """ - Utils.initialize_model_parallel(1, 1) # tp=pp=1 -> dp = world_size - metadata = {'distrib_optim_sharding_type': 'dp_reshardable'} - - def _build(seed): - return setup_model_and_optimizer( - seed=seed, - tp=1, - pp=1, - bf16=bf16, - dist_opt=True, - initialize_fn=initialize_gpt_model, - optimizer='dist_muon', - use_param_layout=True, - grad_reduce_in_fp32=grad_reduce_in_fp32, - ) - - with TempNamedDir( - tmp_path_dist_ckpt / 'test_layer_wise_dp_reshardable', sync=True - ) as ckpt_dir: - model_A, optimizer_A = _build(2) - model_sd = model_A[0].sharded_state_dict() - optim_sd = optimizer_A.sharded_state_dict(model_sd, metadata=metadata) - save(optim_sd, ckpt_dir) - - model_B, optimizer_B = _build(3) - model_sd_B = model_B[0].sharded_state_dict() - load_sd = optimizer_B.sharded_state_dict(model_sd_B, is_loading=True, metadata=metadata) - state_dict = load(load_sd, ckpt_dir) - optimizer_B.load_state_dict(state_dict) - Utils.destroy_model_parallel() - - @pytest.mark.parametrize('ep', [2, 4]) - def test_dp_reshardable_decouple_moe_ckpt(self, tmp_path_dist_ckpt, ep): - """dp_reshardable save/load round-trip of the decoupled compact LayerWise (Muon) - optimizer on an MoE model with expert parallelism and mixed-dtype (bf16 param / - fp32 grad) sibling DistOpt buffers. Single bucket per buffer (no ``ddp_bucket_size``), - so ownership is dense and the round-trip loads cleanly. - """ - Utils.initialize_model_parallel( - tensor_model_parallel_size=1, - pipeline_model_parallel_size=1, - expert_model_parallel_size=ep, - ) - metadata = {'distrib_optim_sharding_type': 'dp_reshardable'} - - def _build(seed): - return setup_moe_model_and_optimizer( - seed=seed, - tp=1, - pp=1, - ep=ep, - bf16=True, - dist_opt=True, - optimizer='dist_muon', - use_param_layout=True, - grad_reduce_in_fp32=True, - ) - - with TempNamedDir( - tmp_path_dist_ckpt / 'test_layer_wise_dp_reshardable_moe', sync=True - ) as ckpt_dir: - model_A, optimizer_A = _build(2) - model_sd = model_A[0].sharded_state_dict() - optim_sd = optimizer_A.sharded_state_dict(model_sd, metadata=metadata) - save(optim_sd, ckpt_dir) - - model_B, optimizer_B = _build(3) - model_sd_B = model_B[0].sharded_state_dict() - load_sd = optimizer_B.sharded_state_dict(model_sd_B, is_loading=True, metadata=metadata) - state_dict = load(load_sd, ckpt_dir) - optimizer_B.load_state_dict(state_dict) - Utils.destroy_model_parallel() - - def test_dp_reshardable_moe_synth_save(self, tmp_path_dist_ckpt): - """Regression for the empty-bucket-synth coverage gap in - ``DistributedOptimizer.sharded_param_state_dp_reshardable``. - - A small ``ddp_bucket_size`` splits the sibling DistOpt buffer of the decoupled - compact LayerWise (Muon) layout into many small buckets. Params are 64-element - aligned inside a bucket, so at DP=4 (local shard = 32) some DP rank owns a shard - that lies entirely inside inter-param padding while still overlapping - ``[0, gbuf_world_numel_unpadded)`` -- the empty-bucket-synth path. Set - ``DEBUG_DP_RESHARDABLE=1`` to see the per-bucket ``synth=True`` ``[DPRESH ...]`` - lines (this config fires it ~34x/rank-bucket). - - With the fix the synthesized padding ShardedTensor is stored back into the returned - ``state`` so the per-bucket global tensor is fully covered and torch DCP ``save`` - validates the global plan. Without the store-back line those shards are dropped and - ``save`` raises ``ValueError: Failed to validate global plan`` (chunks_volume < - tensor_volume). - - Note: this asserts the *save* half only. Extending it to a load round-trip trips a - separate, pre-existing ``dp_reshardable`` multi-bucket load defect (a - ``len(bucket_state) == len(param_map)`` mismatch in - ``load_parameter_state_from_dp_reshardable``) that also fires for a non-synth - multi-bucket config, so it is out of scope for this store-back fix. - """ - Utils.initialize_model_parallel( - tensor_model_parallel_size=1, - pipeline_model_parallel_size=1, - expert_model_parallel_size=2, - ) - metadata = {'distrib_optim_sharding_type': 'dp_reshardable'} - - with TempNamedDir(tmp_path_dist_ckpt / 'moe_synth_save', sync=True) as ckpt_dir: - model, optimizer = setup_moe_model_and_optimizer( - seed=2, - tp=1, - pp=1, - ep=2, - bf16=True, - dist_opt=True, - optimizer='dist_muon', - use_param_layout=True, - grad_reduce_in_fp32=True, - ddp_bucket_size=32, - ) - model_sd = model[0].sharded_state_dict() - optim_sd = optimizer.sharded_state_dict(model_sd, metadata=metadata) - # Fails with "Failed to validate global plan" if the synth store-back line - # in sharded_param_state_dp_reshardable is removed. - save(optim_sd, ckpt_dir) - Utils.destroy_model_parallel() - - # NOTE: 'fully_sharded_model_space' is intentionally NOT covered here. It is non-functional - # for EVERY DistributedOptimizer in this tree, not only the decoupled compact LayerWise (Muon) - # layout: sharded_param_state_fs_model_space sets flattened_range on every non-factory param, - # and dist_checkpointing rejects that unconditionally ('ShardedTensor.flattened_range is not - # supported'); only ShardedTensorFactory (gated-MLP fc1) escapes. This was - # confirmed to fail identically for bf16 and fp8 (so it is NOT fp8-specific), and it raises - # non-uniformly across DP ranks, so it cannot be asserted cleanly in a distributed test. It is - # a pre-existing limitation, independent of the dp_reshardable empty-bucket fix. - @pytest.mark.parametrize('fp8', [False, True]) - @pytest.mark.parametrize('sharding_type', ['dp_reshardable', 'fully_reshardable']) - def test_decouple_ckpt_roundtrip_values(self, tmp_path_dist_ckpt, sharding_type, fp8): - """Value-level save/load round-trip of the decoupled compact LayerWise (Muon) optimizer, - for the ``dp_reshardable`` and ``fully_reshardable`` sharding formats, for both bf16 and - quantized FP8 (MXFP8) model params. - - Correctness (not just "does not crash"): save optimizer A in ``sharding_type``, load it - into a differently seeded optimizer B, then assert B's optimizer state equals A's *by - value*. If the round-trip is faithful, B's fp32 master / exp_avg / exp_avg_sq must match - A's exactly. - - The comparison is done through a padding-free *canonical view*: both A's and (post-load) - B's in-memory state are re-serialized with ``fully_reshardable`` (model-centric, no - padding, deterministic) and compared bitwise via ``load_plain_tensors`` + ``check_equal``. - This is required because ``dp_reshardable``'s own bucket-space checkpoint serializes - inter-param / empty-shard padding as uninitialized ``torch.empty`` (values discarded on - load), so two saves of identical state differ in the padding bytes and cannot be compared - directly. ``fully_reshardable`` has no such padding, so it is a faithful canonical view of - the real optimizer state for both formats under test. - - The ``fp8`` axis covers the fp8 -> fp32-master path: FP8 changes model-param storage to a - Float8Tensor while optimizer state stays fp32, so the dequantize path - (``_is_distopt_quantized_param`` in distrib_optimizer.py) must round-trip. Both bf16 and - fp8 are exercised. - """ - # setup_moe_model_and_optimizer's fp8 path uses fp8_recipe='mxfp8', which needs - # Blackwell or newer; on older archs TE raises inside dequantize. Mirrors the arch - # guard in tests/unit_tests/test_muon_decouple_fp8_param_gather.py. - if fp8 and get_device_arch_version() < 10: - pytest.skip("mxfp8 requires Blackwell architecture or newer") - - from megatron.core.dist_checkpointing import load_plain_tensors - - Utils.initialize_model_parallel(1, 1) # tp=pp=1 -> dp = world_size - metadata = {'distrib_optim_sharding_type': sharding_type} - # Padding-free canonical view used to compare optimizer state by value. - canonical = {'distrib_optim_sharding_type': 'fully_reshardable'} - - def _build(seed): - kwargs = dict( - seed=seed, - tp=1, - pp=1, - bf16=True, - dist_opt=True, - initialize_fn=initialize_gpt_model, - optimizer='dist_muon', - use_param_layout=True, - grad_reduce_in_fp32=True, - ) - if fp8: - kwargs['fp8'] = True - return setup_model_and_optimizer(**kwargs) - - tag = f'{"fp8" if fp8 else "bf16"}_{sharding_type}' - with ( - TempNamedDir(tmp_path_dist_ckpt / f'{tag}_rt', sync=True) as rt_dir, - TempNamedDir(tmp_path_dist_ckpt / f'{tag}_A', sync=True) as canon_dir_A, - TempNamedDir(tmp_path_dist_ckpt / f'{tag}_B', sync=True) as canon_dir_B, - ): - # Save A in the format under test, load it into a differently seeded B. - model_A, optimizer_A = _build(2) - model_sd_A = model_A[0].sharded_state_dict() - save(optimizer_A.sharded_state_dict(model_sd_A, metadata=metadata), rt_dir) - - model_B, optimizer_B = _build(3) - model_sd_B = model_B[0].sharded_state_dict() - load_sd = optimizer_B.sharded_state_dict(model_sd_B, is_loading=True, metadata=metadata) - optimizer_B.load_state_dict(load(load_sd, rt_dir)) - - # Compare A vs post-load B by value, through the padding-free canonical view. - save(optimizer_A.sharded_state_dict(model_sd_A, metadata=canonical), canon_dir_A) - save(optimizer_B.sharded_state_dict(model_sd_B, metadata=canonical), canon_dir_B) - Utils.destroy_model_parallel() - - Utils.initialize_model_parallel(1, 1) - check_equal(load_plain_tensors(canon_dir_A), load_plain_tensors(canon_dir_B)) - Utils.destroy_model_parallel() - @pytest.mark.parametrize('tp', [1, 2, 4]) @pytest.mark.parametrize('pp', [1, 2, 4]) def test_layer_wise_optimizer_grad_norm(self, tp, pp): diff --git a/tests/unit_tests/dist_checkpointing/utils.py b/tests/unit_tests/dist_checkpointing/utils.py index 73dcf6c22ee..ba774b34fd2 100644 --- a/tests/unit_tests/dist_checkpointing/utils.py +++ b/tests/unit_tests/dist_checkpointing/utils.py @@ -1,4 +1,4 @@ -# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. from functools import partial from typing import Any, Callable, Tuple, Union @@ -9,7 +9,6 @@ from megatron.core.dist_checkpointing.strategies.cached_metadata_filesystem_reader import ( CachedMetadataFileSystemReader, ) -from megatron.core.fp8_utils import is_float8tensor from megatron.core.models.gpt import GPTModel from megatron.core.models.gpt.gpt_layer_specs import ( get_gpt_layer_local_spec, @@ -29,7 +28,7 @@ def initialize_gpt_model( - pre_process=True, post_process=True, seed=0, use_glu=True, fp8=False, **config_kwargs + pre_process=True, post_process=True, seed=0, use_glu=True, **config_kwargs ): # These kwargs are passed through training.get_model for model construction, # but are not part of TransformerConfig; strip them before building config. @@ -46,24 +45,11 @@ def initialize_gpt_model( use_cpu_initialization=True, bf16=True, ) - if fp8: - # FP8 params only materialize with TransformerEngine layers, GPU init, and a - # hidden size that satisfies the MXFP8 32-element block-scaling alignment. The - # tiny (hidden=16, local-spec) default gives zero fp8 params. - default_config_kwargs.update( - hidden_size=128, - num_attention_heads=8, - use_cpu_initialization=False, - fp8='e4m3', - fp8_recipe='mxfp8', - fp8_param=True, - ) default_config_kwargs.update(**config_kwargs) transformer_config = TransformerConfig(**default_config_kwargs, gated_linear_unit=use_glu) - spec = get_gpt_layer_with_transformer_engine_spec() if fp8 else get_gpt_layer_local_spec() model = GPTModel( config=transformer_config, - transformer_layer_spec=spec, + transformer_layer_spec=get_gpt_layer_local_spec(), vocab_size=128, max_sequence_length=4, pre_process=pre_process, @@ -72,10 +58,6 @@ def initialize_gpt_model( with torch.no_grad(): for p in model.parameters(): - # Float8Tensor params own quantized storage; skip the plain random_ init - # (embeddings / layernorms remain plain tensors and are still randomized). - if is_float8tensor(p): - continue p.random_() return model @@ -208,10 +190,6 @@ def setup_model_and_optimizer( ep=1, etp=1, use_megatron_fsdp=False, - ddp_bucket_size=None, - grad_reduce_in_fp32=False, - fp8=False, - use_layer_wise_param_layout=False, ): optimizer_type = optimizer use_layer_wise = False @@ -229,19 +207,6 @@ def setup_model_and_optimizer( ddp_use_dist_opt = dist_opt and not (use_layer_wise and not use_param_layout) ddp_use_layer_wise = use_layer_wise and use_param_layout - # The padded shard-aligned LayerWise layout only exists under a LayerWise optimizer, so - # ``use_layer_wise_param_layout=True`` implies ``use_layer_wise_distributed_optimizer=True``; - # the converse does not hold. Only three combinations are meaningful: - # (False, False) plain DistributedOptimizer - # (True, False) LayerWise + compact decoupled layout [the default since the flag flipped] - # (True, True) LayerWise + padded shard-aligned layout - assert not (use_layer_wise_param_layout and not ddp_use_layer_wise), ( - "use_layer_wise_param_layout=True requires LayerWise DDP routing " - "(optimizer='muon'/'dist_muon' with dist_opt=True and use_param_layout=True); " - "otherwise DDP builds untagged buffers while the optimizer expects the " - "shard-aligned layout" - ) - mock_args = parse_args(ignore_unknown_args=True) with mock.patch('megatron.training.training.get_args', new=lambda: mock_args): init_basic_mock_args(mock_args, tp, pp, bf16=bf16) @@ -259,25 +224,10 @@ def setup_model_and_optimizer( mock_args.megatron_fsdp_main_grads_dtype = None mock_args.megatron_fsdp_grad_comm_dtype = None mock_args.gradient_accumulation_fusion = False - mock_args.ddp_bucket_size = ddp_bucket_size - # grad_reduce_in_fp32 -> ddp_config grad_dtype=fp32 while params stay bf16, i.e. the - # mixed-dtype (bf16, fp32) gradient buffer / optimizer-state bucket. - mock_args.accumulate_allreduce_grads_in_fp32 = grad_reduce_in_fp32 - mock_args.use_layer_wise_param_layout = use_layer_wise_param_layout mock_args.use_distributed_optimizer = ddp_use_dist_opt mock_args.use_layer_wise_distributed_optimizer = ddp_use_layer_wise if ddp_use_layer_wise: mock_args.optimizer = optimizer - # Only forward ``fp8`` to initialize_fn when enabled, so existing callers - # passing an initialize_fn without an ``fp8`` parameter keep working. - extra_init_kwargs = {} - if fp8: - # Make DDP build an fp8 param-gather buffer (MXFP8 reuses the grad buffer - # for the param all-gather). ``fp8_param`` on the TransformerConfig (set in - # ``initialize_gpt_model``) is what actually quantizes the model weights. - mock_args.fp8_param_gather = True - mock_args.reuse_grad_buf_for_mxfp8_param_ag = True - extra_init_kwargs['fp8'] = True model = get_model( partial( initialize_fn, @@ -289,23 +239,14 @@ def setup_model_and_optimizer( expert_model_parallel_size=ep, expert_tensor_parallel_size=etp, bf16=bf16, - **extra_init_kwargs, ) ) - if fp8: - # Guard against a silent config regression: if the fp8/TE path stops producing - # quantized weights the checkpoint round-trip would no longer exercise the - # fp8 -> fp32-master mapping we mean to test. - num_fp8_params = sum(1 for m in model for p in m.parameters() if is_float8tensor(p)) - assert num_fp8_params > 0, "fp8=True but no Float8Tensor params were created" - config = OptimizerConfig( bf16=bf16, params_dtype=torch.bfloat16 if bf16 else torch.float, use_distributed_optimizer=ddp_use_dist_opt, use_layer_wise_distributed_optimizer=use_layer_wise, - use_layer_wise_param_layout=use_layer_wise_param_layout, optimizer=optimizer, muon_scalar_optimizer=muon_scalar_optimizer, ) @@ -399,9 +340,6 @@ def setup_moe_model_and_optimizer( use_glu=False, optimizer='adam', use_param_layout=False, - ddp_bucket_size=None, - grad_reduce_in_fp32=False, - use_layer_wise_param_layout=False, ): optimizer_type = optimizer use_layer_wise = False @@ -415,34 +353,9 @@ def setup_moe_model_and_optimizer( ddp_use_dist_opt = dist_opt and not (use_layer_wise and not use_param_layout) ddp_use_layer_wise = use_layer_wise and use_param_layout - # The padded shard-aligned LayerWise layout only exists under a LayerWise optimizer, so - # ``use_layer_wise_param_layout=True`` implies ``use_layer_wise_distributed_optimizer=True``; - # the converse does not hold. Only three combinations are meaningful: - # (False, False) plain DistributedOptimizer - # (True, False) LayerWise + compact decoupled layout [the default since the flag flipped] - # (True, True) LayerWise + padded shard-aligned layout - assert not (use_layer_wise_param_layout and not ddp_use_layer_wise), ( - "use_layer_wise_param_layout=True requires LayerWise DDP routing " - "(optimizer='muon'/'dist_muon' with dist_opt=True and use_param_layout=True); " - "otherwise DDP builds untagged buffers while the optimizer expects the " - "shard-aligned layout" - ) - mock_args = parse_args(ignore_unknown_args=True) with mock.patch('megatron.training.training.get_args', new=lambda: mock_args): init_basic_mock_args(mock_args, tp, pp, bf16=bf16) - mock_args.ddp_bucket_size = ddp_bucket_size - # ``resolve_ddp_bucket_size`` (megatron/training/training.py) discards an explicit - # bucket_size unless overlap_grad_reduce is on -- otherwise the whole buffer is a - # single bucket. Turning overlap on (no backward is run here) lets ddp_bucket_size - # split the sibling DistOpt buffer into many small buckets so some DP rank owns a - # shard made only of inter-param alignment padding -> the empty-bucket-synth path. - if ddp_bucket_size is not None: - mock_args.overlap_grad_reduce = True - # grad_reduce_in_fp32 -> ddp_config grad_dtype=fp32 while params stay bf16, i.e. the - # mixed-dtype (bf16, fp32) gradient buffer / optimizer-state bucket. - mock_args.accumulate_allreduce_grads_in_fp32 = grad_reduce_in_fp32 - mock_args.use_layer_wise_param_layout = use_layer_wise_param_layout mock_args.use_distributed_optimizer = ddp_use_dist_opt mock_args.use_layer_wise_distributed_optimizer = ddp_use_layer_wise if ddp_use_layer_wise: @@ -468,7 +381,6 @@ def setup_moe_model_and_optimizer( params_dtype=torch.bfloat16 if bf16 else torch.float, use_distributed_optimizer=ddp_use_dist_opt, use_layer_wise_distributed_optimizer=use_layer_wise, - use_layer_wise_param_layout=use_layer_wise_param_layout, optimizer=optimizer, ) @@ -479,25 +391,12 @@ def setup_moe_model_and_optimizer( torch.manual_seed(seed + 1) model_parallel_cuda_manual_seed(seed + 1) - def _init_states(opt): - # In the decoupled compact LayerWise + DistOpt layout the top-level - # ChainedOptimizer wraps another ChainedOptimizer (LayerWise, which has no - # ``init_state_fn``) alongside a sibling DistOpt; recurse so the Muon - # Float16 sub-optimizers still get seeded, and skip optimizers without - # ``init_state_fn`` (DistOpt seeds its state elsewhere). - if isinstance(opt, ChainedOptimizer): - for child in opt.chained_optimizers: - _init_states(child) - return - if not hasattr(opt, 'init_state_fn'): - return - if not hasattr(opt, 'optimizer'): - opt.init_state_fn(opt) - else: - opt.init_state_fn(opt.optimizer) - if optimizer_type in ('muon', 'dist_muon'): - _init_states(optimizer) + for opt in optimizer.chained_optimizers: + if not hasattr(opt, 'optimizer'): + opt.init_state_fn(opt) + else: + opt.init_state_fn(opt.optimizer) else: for opt in optimizer.chained_optimizers: for group in opt.param_groups: From 098780bb9bb33775171dfc4a29d267b07c881d11 Mon Sep 17 00:00:00 2001 From: Pingtian Li Date: Tue, 18 Aug 2026 18:00:35 +0800 Subject: [PATCH 12/12] Optimize LayerWise FP8 parameter copies --- .../megatron_fsdp/param_and_grad_buffer.py | 20 ++++- .../core/distributed/param_and_grad_buffer.py | 67 ++++++++-------- megatron/core/fp8_utils.py | 78 ++++++++++++------- megatron/core/optimizer/distrib_optimizer.py | 10 ++- .../core/optimizer/layer_wise_optimizer.py | 37 ++++++--- megatron/core/optimizer/optimizer.py | 17 ++-- 6 files changed, 143 insertions(+), 86 deletions(-) diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/param_and_grad_buffer.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/param_and_grad_buffer.py index 714652dd54a..8a4bcea9c09 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/param_and_grad_buffer.py +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/param_and_grad_buffer.py @@ -1,4 +1,4 @@ -# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -66,6 +66,7 @@ from megatron.core.distributed.distributed_data_parallel_config import ( DistributedDataParallelConfig, ) + from megatron.core.fp8_utils import pop_high_precision_init_val from megatron.core.tensor_parallel import get_cuda_rng_tracker from megatron.core.utils import is_submodule @@ -77,6 +78,18 @@ from .distributed_data_parallel_config import DistributedDataParallelConfig from .utils import get_cuda_rng_tracker, is_submodule + # Standalone compatibility: MCore's shared helper is intentionally unavailable when the + # independently installable megatron_fsdp package is used without Megatron Core. + def pop_high_precision_init_val(param: torch.Tensor) -> Optional[torch.Tensor]: + """Return and clear a TE preserved high-precision initial value, if present.""" + getter = getattr(param, "get_high_precision_init_val", None) + if getter is None: + return None + + high_precision_init_val = getter() + param.clear_high_precision_init_val() + return high_precision_init_val + HAVE_MCORE = False logger.info("Megatron Core is not installed, Megatron-FSDP will run without Megatron Core.") @@ -3002,8 +3015,9 @@ def _init_each_parameter_group_buffers(self, meta_device_init_fp8_params): ) # Needed to instantiate FP8 parameters. Requires installing # TransformerEngine. - mbuf.set_item(item_id, p.get_high_precision_init_val()) - p.clear_high_precision_init_val() + high_precision_init_val = pop_high_precision_init_val(p) + assert high_precision_init_val is not None + mbuf.set_item(item_id, high_precision_init_val) else: # Insert a copy of the model weight parameter tensor into # the (high-precision) main weight buffer. diff --git a/megatron/core/distributed/param_and_grad_buffer.py b/megatron/core/distributed/param_and_grad_buffer.py index ae5df979f3f..78f5493b0a1 100644 --- a/megatron/core/distributed/param_and_grad_buffer.py +++ b/megatron/core/distributed/param_and_grad_buffer.py @@ -12,7 +12,7 @@ from typing import Dict, List, Optional, Tuple import torch -from torch._utils import _flatten_dense_tensors, _unflatten_dense_tensors +from torch._utils import _unflatten_dense_tensors from torch.distributed import _coalescing_manager import megatron.core.nccl_allocator as nccl_allocator @@ -29,8 +29,7 @@ modify_nvfp4_rowwise_storage, ) from ..fp8_utils import ( - _stage_param_to_bf16, - copy_back_gathered_bf16_into_fp8_param, + copy_back_gathered_bf16_into_fp8_params, copy_tensors_to_quantized_params, is_float8tensor, is_grouped_mxfp8tensor, @@ -188,7 +187,8 @@ def _layerwise_copy_back_gathered_params(bucket, local_rank: int, fp8_staged: bo * bf16 (``fp8_staged=False``): unflatten against the params, ``copy_`` into non-owned ``model_p.data`` (owned already hold the staged value). * fp8 (``fp8_staged=True``): the all-gather rode bf16; requantize ALL ranks (owned included) - via ``copy_back_gathered_bf16_into_fp8_param`` so every owner holds ``Q(bf16(master))``. + via ``copy_back_gathered_bf16_into_fp8_params`` so every owner holds + ``Q(bf16(master))``. no_grad: in-place copy_ on a leaf param trips autograd's in-place guard. """ @@ -198,8 +198,7 @@ def _layerwise_copy_back_gathered_params(bucket, local_rank: int, fp8_staged: bo if fp8_staged: templates = [torch.empty(p.shape, device="meta", dtype=torch.bfloat16) for p in params] updated_params = _unflatten_dense_tensors(bucket.layerwise_gather_list[idx], templates) - for updated_p, model_p in zip(updated_params, params): - copy_back_gathered_bf16_into_fp8_param(model_p, updated_p) + copy_back_gathered_bf16_into_fp8_params(params, updated_params) continue # bf16 transport: owned params already hold the staged bf16 value in their data, so only # non-owned ranks need the copy. @@ -541,35 +540,37 @@ def start_param_sync(self, force_sync: bool = False): offset += size local_slot_view = gather_list[local_rank] - # Flatten local params and copy into the local rank's slot. - # Detach from autograd since start_param_sync may be called - # during the forward pass where autograd is active. + # Copy local params directly into their views of the flat transport slot. + # This avoids allocating one BF16 tensor per parameter and then flattening it. + # Detach from autograd since start_param_sync may be called during forward. if local_size > 0: - if bucket_is_fp8: - # Decoupled layout: stage fp32 master->bf16 (high-precision source), not - # lossy dequant(fp8). Copy-back requantizes every rank, owner included. - staged = [ - _stage_param_to_bf16(p) - for p in bucket.layerwise_params_list[local_rank] - ] - flat_local_params = _flatten_dense_tensors(staged) - else: - # Padded LayerWise layout: MXFP8 params can't be flattened (view(-1) - # unsupported); gather the fp32 master (param.main_param -> bf16), which - # the receive-side copy_ re-quantizes. Non-mxfp8 params flatten as-is. - src_params = [] - for p in bucket.layerwise_params_list[local_rank]: - if is_mxfp8tensor(p): - main_param = getattr(p, "main_param", None) - assert main_param is not None, ( - "LayerWise mxfp8 param sync needs param.main_param (fp32 " - "master) to stage the all-gather source; got None." + local_offset = 0 + for param in bucket.layerwise_params_list[local_rank]: + param_numel = param.numel() + transport_view = local_slot_view[local_offset : local_offset + param_numel] + if bucket_is_fp8 or is_mxfp8tensor(param): + # Decoupled FP8 buckets stage every param from its FP32 master so + # mixed FP8/BF16 buckets all use the same high-precision source. + # Padded buckets need the same source only for MXFP8, whose tensor + # subclass cannot be flattened. + source = getattr(param, "main_param", None) + if source is None: + raise RuntimeError( + "LayerWise FP8 parameter-gather staging requires " + "param.main_param (FP32 master)." ) - src_params.append(main_param.to(param_dtype)) - else: - src_params.append(p) - flat_local_params = _flatten_dense_tensors(src_params).detach() - local_slot_view.copy_(flat_local_params) + else: + source = param + source = source.detach() + if source.numel() != param_numel: + raise RuntimeError( + "LayerWise parameter-gather staging source size mismatch: " + f"source has {source.numel()} elements, parameter has " + f"{param_numel}." + ) + transport_view.copy_(source.reshape(-1)) + local_offset += param_numel + assert local_offset == local_size bucket.layerwise_gather_list = gather_list work = torch.distributed.all_gather( diff --git a/megatron/core/fp8_utils.py b/megatron/core/fp8_utils.py index c826725e13e..63d0eccbc36 100644 --- a/megatron/core/fp8_utils.py +++ b/megatron/core/fp8_utils.py @@ -176,6 +176,21 @@ def is_grouped_tensor_with_quantized_storage(tensor: torch.Tensor) -> bool: return rowwise_data is not None and rowwise_data.dtype == torch.uint8 +def pop_high_precision_init_val(param: torch.Tensor) -> Optional[torch.Tensor]: + """Return and clear a Transformer Engine preserved high-precision initial value. + + The returned tensor is left unmodified so each optimizer path can preserve its + existing slicing, cloning, device placement, and dtype conversion behavior. + """ + getter = getattr(param, "get_high_precision_init_val", None) + if getter is None: + return None + + high_precision_init_val = getter() + param.clear_high_precision_init_val() + return high_precision_init_val + + def _get_grouped_quantized_recipe(tensor: torch.Tensor): """Return TE recipe for grouped quantized storage, or None if unavailable.""" tensor = _unwrap_parameter_data(tensor) @@ -334,37 +349,46 @@ def dequantize_fp8_tensor(fp8_tensor: torch.Tensor) -> torch.Tensor: return fp8_tensor.from_float8() -def copy_back_gathered_bf16_into_fp8_param(model_p: torch.Tensor, src_bf16: torch.Tensor) -> None: - """Copy a gathered BF16 whole-param into a compact LayerWise bucket parameter. +def copy_back_gathered_bf16_into_fp8_params( + model_params: List[torch.Tensor], srcs_bf16: List[torch.Tensor] +) -> None: + """Copy gathered BF16 whole-params into compact LayerWise bucket parameters. Plain BF16 parameters are allowed because a LayerWise bucket can contain BF16 siblings of supported FP8 parameters. Quantized destinations are limited to MXFP8 and blockwise tensors. - MXFP8 columnwise data cannot be derived from rowwise data, so force both usages before copy_. + MXFP8 columnwise data cannot be derived from rowwise data, so force both usages before the + batched quantized copy. Validate the entire batch before mutating any quantizer usage. """ - if is_grouped_tensor_with_quantized_storage(model_p): - raise TypeError( - "LayerWise FP8 parameter gather does not support Transformer Engine GroupedTensor " - "quantized storage. Disable --moe-single-grouped-weight." - ) - if is_float8tensor(model_p) and not is_layerwise_fp8_param(model_p): - raise TypeError( - "LayerWise FP8 parameter gather supports only MXFP8Tensor and " - "Float8BlockwiseQTensor destinations." + if len(model_params) != len(srcs_bf16): + raise ValueError( + "LayerWise FP8 parameter gather copy-back requires one source per parameter: " + f"got {len(model_params)} parameters and {len(srcs_bf16)} sources." ) - if is_mxfp8tensor(model_p): - quantizer = model_p.data._get_quantizer() + + mxfp8_quantizers = [] + for model_p in model_params: + if is_grouped_tensor_with_quantized_storage(model_p): + raise TypeError( + "LayerWise FP8 parameter gather does not support Transformer Engine " + "GroupedTensor quantized storage. Disable --moe-single-grouped-weight." + ) + if is_float8tensor(model_p) and not is_layerwise_fp8_param(model_p): + raise TypeError( + "LayerWise FP8 parameter gather supports only MXFP8Tensor and " + "Float8BlockwiseQTensor destinations." + ) + if is_mxfp8tensor(model_p): + mxfp8_quantizers.append(model_p.data._get_quantizer()) + + for quantizer in mxfp8_quantizers: quantizer.set_usage(rowwise=True, columnwise=True) - model_p.data.copy_(src_bf16) + copy_tensors_to_quantized_params(model_params, srcs_bf16) -def _stage_param_to_bf16(p: torch.Tensor) -> torch.Tensor: - """Stage a locally owned LayerWise parameter's FP32 master in BF16 for transport.""" - main_param = getattr(p, "main_param", None) - if main_param is None: - raise RuntimeError( - "LayerWise FP8 parameter-gather staging requires param.main_param (FP32 master)." - ) - return main_param.detach().to(torch.bfloat16) + +def copy_back_gathered_bf16_into_fp8_param(model_p: torch.Tensor, src_bf16: torch.Tensor) -> None: + """Single-parameter compatibility wrapper for LayerWise BF16 copy-back.""" + copy_back_gathered_bf16_into_fp8_params([model_p], [src_bf16]) def _resolve_callable_from_python_import_path(dotted_path: str): @@ -406,11 +430,9 @@ def _get_custom_recipe(quantizer_factory_python_path: str) -> Union[Fp8Recipe, F try: custom_recipe = transformer_engine.common.recipe.CustomRecipe(qfactory=quantizer_factory) except AttributeError: - raise ValueError( - """CustomRecipe recipe is not available in this version of - Transformer Engine. Please make sure you are using TE version - >= 2.9.0.dev0.""" - ) + raise ValueError("""CustomRecipe recipe is not available in this version of + Transformer Engine. Please make sure you are using TE version + >= 2.9.0.dev0.""") return custom_recipe diff --git a/megatron/core/optimizer/distrib_optimizer.py b/megatron/core/optimizer/distrib_optimizer.py index 54819099653..83351acc39a 100644 --- a/megatron/core/optimizer/distrib_optimizer.py +++ b/megatron/core/optimizer/distrib_optimizer.py @@ -58,6 +58,7 @@ get_grouped_quantized_members, is_float8tensor, is_grouped_tensor_with_quantized_storage, + pop_high_precision_init_val, quantize_param_shard, ) from ..transformer.fsdp_dtensor_checkpoint import handle_experts_in_state_dict @@ -428,15 +429,16 @@ def _build_model_and_main_param_groups( if is_nvfp4tensor(model_param) or cls._is_distopt_quantized_param( model_param ): - if hasattr(model_param, 'get_high_precision_init_val'): + high_precision_init_val = pop_high_precision_init_val(model_param) + if high_precision_init_val is not None: shard_main_param = ( - model_param.get_high_precision_init_val() - .view(-1)[param_range.start : param_range.end] + high_precision_init_val.view(-1)[ + param_range.start : param_range.end + ] .clone() .to(model_param.device) .float() ) - model_param.clear_high_precision_init_val() else: shard_main_param = model_param.float().view(-1)[ param_range.start : param_range.end diff --git a/megatron/core/optimizer/layer_wise_optimizer.py b/megatron/core/optimizer/layer_wise_optimizer.py index 05e0c8aaaff..e8887aac372 100644 --- a/megatron/core/optimizer/layer_wise_optimizer.py +++ b/megatron/core/optimizer/layer_wise_optimizer.py @@ -15,8 +15,7 @@ from megatron.core.utils import get_pg_rank, get_pg_size, log_single_rank from ..fp8_utils import ( - _stage_param_to_bf16, - copy_back_gathered_bf16_into_fp8_param, + copy_back_gathered_bf16_into_fp8_params, is_layerwise_fp8_param, post_all_gather_processing, ) @@ -968,17 +967,34 @@ def _allgather_helper_fp8(params_list, group): # No rank owns any param in this buffer -> nothing to gather. return - # Stage fp32 master->bf16 (high-precision source), not lossy dequant(fp8). - owned = params_list[rank] - src = ( - _flatten_dense_tensors([_stage_param_to_bf16(p) for p in owned]) - if len(owned) > 0 - else torch.empty(0, device=device, dtype=torch.bfloat16) - ) flat_sizes = [sum(p.numel() for p in params) for params in params_list] if max(flat_sizes) == 0: return + # Stage FP32 masters directly into one flat BF16 transport tensor rather than + # allocating a temporary BF16 tensor per parameter and flattening those copies. + owned = params_list[rank] + src = torch.empty(flat_sizes[rank], device=device, dtype=torch.bfloat16) + offset = 0 + for param in owned: + main_param = getattr(param, "main_param", None) + if main_param is None: + raise RuntimeError( + "LayerWise FP8 parameter-gather staging requires " + "param.main_param (FP32 master)." + ) + param_numel = param.numel() + main_param = main_param.detach() + if main_param.numel() != param_numel: + raise RuntimeError( + "LayerWise parameter-gather staging source size mismatch: " + f"source has {main_param.numel()} elements, parameter has " + f"{param_numel}." + ) + src[offset : offset + param_numel].copy_(main_param.reshape(-1)) + offset += param_numel + assert offset == flat_sizes[rank] + gather_list = [] for i in range(dp_size): if i == rank: @@ -999,8 +1015,7 @@ def _allgather_helper_fp8(params_list, group): torch.empty(p.shape, device="meta", dtype=torch.bfloat16) for p in params ] updated_params = _unflatten_dense_tensors(gather_list[idx], templates) - for updated_bf16, model_p in zip(updated_params, params): - copy_back_gathered_bf16_into_fp8_param(model_p, updated_bf16) + copy_back_gathered_bf16_into_fp8_params(params, updated_params) # Rebuild fp8 columnwise/transpose after the gather (mirrors the overlap / DistOpt # paths; blockwise builds it, mxfp8 is a noop). Else it'd be deferred to forward. diff --git a/megatron/core/optimizer/optimizer.py b/megatron/core/optimizer/optimizer.py index c05186171cf..618eaf13dc3 100644 --- a/megatron/core/optimizer/optimizer.py +++ b/megatron/core/optimizer/optimizer.py @@ -46,7 +46,11 @@ optim_state_to_sharding_state, ) from ..dist_checkpointing.utils import add_prefix_for_sharding -from ..fp8_utils import copy_back_gathered_bf16_into_fp8_param, is_layerwise_fp8_param +from ..fp8_utils import ( + copy_back_gathered_bf16_into_fp8_params, + is_layerwise_fp8_param, + pop_high_precision_init_val, +) from ..optimizer_param_scheduler import ParamGroupOverride as _ParamGroupOverride from ..transformer.module import param_is_not_shared from ..utils import log_single_rank @@ -1020,15 +1024,14 @@ def __init__( # Seed the fp32 master from the high-precision pre-quantization init # for fp8 params (not the lossy fp8 dequant), matching DistOpt so # fp8_param_gather ON/OFF hold an identical master at iter 0. - if hasattr(param, 'get_high_precision_init_val'): + high_precision_init_val = pop_high_precision_init_val(param) + if high_precision_init_val is not None: main_param = ( - param.get_high_precision_init_val() - .detach() + high_precision_init_val.detach() .clone() .to(param.device) .float() ) - param.clear_high_precision_init_val() else: main_param = param.detach().clone().float() # Copy tensor model parallel attributes. @@ -1145,8 +1148,8 @@ def _copy_main_params_to_model_params(self): # (e.g. MoE experts at expt_dp == 1, which the all-gather skips) are not # tagged and still get their ``Q(bf16(master))`` written here. if not getattr(model_param, '_layer_wise_fp8_gathered', False): - copy_back_gathered_bf16_into_fp8_param( - model_param, main_param.detach().to(torch.bfloat16) + copy_back_gathered_bf16_into_fp8_params( + [model_param], [main_param.detach().to(torch.bfloat16)] ) else: other_model_data.append(model_param.data)