diff --git a/megatron/core/distributed/distributed_data_parallel.py b/megatron/core/distributed/distributed_data_parallel.py index 763368a9814..ebbd876c258 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 import weakref @@ -182,8 +182,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/distributed_data_parallel_config.py b/megatron/core/distributed/distributed_data_parallel_config.py index cf8d963f3a1..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 @@ -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/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 5761f0f3100..78f5493b0a1 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. +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +import dataclasses import fnmatch import functools import logging @@ -11,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 @@ -28,11 +29,13 @@ modify_nvfp4_rowwise_storage, ) from ..fp8_utils import ( + copy_back_gathered_bf16_into_fp8_params, copy_tensors_to_quantized_params, is_float8tensor, 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, @@ -176,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_params`` 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) + 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. + 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. @@ -298,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 quantized_params = [] @@ -400,6 +460,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: @@ -427,10 +491,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: 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: - # 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 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 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_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 + # (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 @@ -453,26 +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: - # 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)) + 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)." + ) else: - src_params.append(p) - flat_local_params = _flatten_dense_tensors(src_params).detach() - local_slot_view.copy_(flat_local_params) + 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( @@ -485,23 +583,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. @@ -572,29 +654,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): @@ -842,6 +902,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 @@ -886,7 +950,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. @@ -905,6 +971,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). @@ -927,6 +995,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_layerwise_fp8_param(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 ) @@ -1069,6 +1146,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. @@ -1119,6 +1207,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: @@ -1688,6 +1797,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. @@ -1709,22 +1822,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 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) + + 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 @@ -1743,11 +1906,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: @@ -1761,17 +1925,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, ) @@ -1779,16 +1943,22 @@ 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, ) ) + return bucket_groups diff --git a/megatron/core/fp8_utils.py b/megatron/core/fp8_utils.py index fc79cc9a7db..63d0eccbc36 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""" @@ -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) @@ -150,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) @@ -308,6 +349,48 @@ def dequantize_fp8_tensor(fp8_tensor: torch.Tensor) -> torch.Tensor: return fp8_tensor.from_float8() +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 the + batched quantized copy. Validate the entire batch before mutating any quantizer usage. + """ + 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." + ) + + 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) + + copy_tensors_to_quantized_params(model_params, srcs_bf16) + + +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): """Resolve a Python import path like 'pkg.mod.func' to a callable. @@ -347,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 782d3263570..83351acc39a 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.""" @@ -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 @@ -613,7 +615,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 a791d262954..e8887aac372 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,11 @@ 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 ( + copy_back_gathered_bf16_into_fp8_params, + is_layerwise_fp8_param, + post_all_gather_processing, +) from .clip_grads import count_zeros_fp32, get_grad_norm_fp32 from .optimizer import ( ChainedOptimizer, @@ -71,6 +76,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``. @@ -470,9 +485,22 @@ 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 - buffer_groups = group_params_for_buffers(params, ddp_config.grad_reduce_in_fp32) + # 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 + + # 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: @@ -487,6 +515,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: + # 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 + ) + 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 @@ -518,6 +555,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 @@ -541,20 +579,42 @@ 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 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 + 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 + # 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_layerwise_fp8_param(p): + self.use_fp8_param_sync = True + p._layer_wise_fp8_gathered = True - # 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. @@ -587,6 +647,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() self.tp_group = self.pg_collection.tp self.expert_tp_group = getattr(self.pg_collection, 'expt_tp', self.tp_group) @@ -611,7 +692,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` @@ -640,7 +721,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.""" @@ -707,19 +788,47 @@ 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 _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. @@ -732,12 +841,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,7 +880,7 @@ def _shard_params_ping_pong(self, optimizers, dp_cp_size, expt_dp_size): 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): @@ -821,13 +943,98 @@ 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 storage (blockwise; mxfp8 is a 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 + + 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: + 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) + 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. + 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) + # 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. @@ -862,10 +1069,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_layerwise_fp8_param(p) or p.dtype != torch.float32] + for owned in params_list + ] + native = [ + [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): + 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 5252f6b86eb..618eaf13dc3 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,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_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 @@ -678,6 +683,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 +803,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 +814,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 +1021,19 @@ 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. + high_precision_init_val = pop_high_precision_init_val(param) + if high_precision_init_val is not None: + main_param = ( + high_precision_init_val.detach() + .clone() + .to(param.device) + .float() + ) + 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 +1134,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_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 + # (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_params( + [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/core/optimizer/optimizer_config.py b/megatron/core/optimizer/optimizer_config.py index 20045108c89..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 @@ -330,6 +330,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 0e0c6664860..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.""" @@ -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,54 @@ 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: + # 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 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." + ) + 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}." + ) + 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 " + "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 " + "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." + ) + 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: @@ -3071,17 +3114,6 @@ def _add_distributed_args(parser): help='Use distributed optimizer.') group.add_argument('--megatron-fsdp-version', type=int, default=1, choices=[1, 2], help='Megatron-FSDP implementation version. Defaults to 1.') - 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-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,' @@ -3136,6 +3168,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/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/megatron/training/training.py b/megatron/training/training.py index 11a9b946525..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.""" @@ -1650,7 +1650,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, @@ -1661,12 +1660,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 @@ -1682,11 +1682,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, @@ -1709,13 +1707,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: + # 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 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 @@ -1981,9 +1984,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/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/model_config.yaml b/tests/functional_tests/test_cases/moe/gpt3_moe_mcore_te_ep8_resume_torch_dist_dist_muon/model_config.yaml index 316b155d685..bc3bd5484f1 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,6 +66,5 @@ 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 LAUNCHER: ft_launcher 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/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 da53a3e7e5f..d184ee46673 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,6 +66,5 @@ 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 LAUNCHER: ft_launcher 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 bc3bd5484f1..815cdc93cb6 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,5 +66,6 @@ MODEL_ARGS: --async-save: true --use-persistent-ckpt-worker: true --use-distributed-optimizer: true + --use-layer-wise-param-layout: true TEST_TYPE: ckpt-resume LAUNCHER: ft_launcher 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 d184ee46673..3c8af4f29a2 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,5 +66,6 @@ MODEL_ARGS: --async-save: true --use-persistent-ckpt-worker: true --use-distributed-optimizer: true + --use-layer-wise-param-layout: true TEST_TYPE: ckpt-resume LAUNCHER: ft_launcher 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..c10a29ed39f --- /dev/null +++ b/tests/unit_tests/test_muon_decouple_fp8_param_gather.py @@ -0,0 +1,534 @@ +# 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.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 +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): + # 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): + 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): + """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): + if include_quantized: + 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 + + +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() + 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, + 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 + 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) + 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}") + # 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}") + 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}") + + @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") + # 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" 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)