diff --git a/megatron/core/distributed/fsdp/mcore_fsdp_adapter.py b/megatron/core/distributed/fsdp/mcore_fsdp_adapter.py index a7c0d5802ab..7432a7f9a36 100644 --- a/megatron/core/distributed/fsdp/mcore_fsdp_adapter.py +++ b/megatron/core/distributed/fsdp/mcore_fsdp_adapter.py @@ -13,6 +13,7 @@ # limitations under the License. import logging +import random from typing import List, Optional try: @@ -22,6 +23,7 @@ except ImportError: HAVE_EINOPS = False +import numpy as np import torch import torch.distributed as dist @@ -32,10 +34,11 @@ except ImportError: HAVE_DTENSOR = False -from megatron.core import parallel_state +from megatron.core import parallel_state, tensor_parallel from megatron.core.config_logger import has_config_logger_enabled, log_config_to_disk from megatron.core.distributed.data_parallel_base import _BaseDataParallel from megatron.core.distributed.distributed_data_parallel_config import DistributedDataParallelConfig +from megatron.core.extensions.transformer_engine import TELinear from megatron.core.process_groups_config import ProcessGroupCollection from megatron.core.transformer.transformer_config import TransformerConfig from megatron.core.transformer.transformer_layer import TransformerLayer @@ -95,6 +98,8 @@ def __init__( else: self.fsdp_unit_modules = [] + self._fix_tensor_parallel_attributes(module) + super().__init__( config=config, module=MegatronFSDP( @@ -119,6 +124,8 @@ def __init__( self.module.state_dict_for_save_checkpoint = self.module.state_dict self.state_dict_for_save_checkpoint = self.state_dict + self.sync_rng_states_across_tp_group() + def load_state_dict(self, state_dict, strict=True): """ Load the state dictionary into the module. @@ -141,6 +148,44 @@ def load_state_dict(self, state_dict, strict=True): self.module.load_state_dict(custom_state_dict, strict=strict) + def _fix_tensor_parallel_attributes(self, module): + is_expert_param = lambda n, p: ".experts." in n + is_router_param = lambda n, p: ".router.weight" in n + + if parallel_state.get_tensor_model_parallel_group(): + tp_size = parallel_state.get_tensor_model_parallel_group().size() + else: + tp_size = 1 + + if parallel_state.get_expert_tensor_parallel_group(): + expt_tp_size = parallel_state.get_expert_tensor_parallel_group().size() + else: + expt_tp_size = 1 + + param_to_direct_module = {} + for name, m in module.named_modules(): + for p in m.parameters(recurse=False): + param_to_direct_module[p] = (name, m) + + for name, param in module.named_parameters(): + if is_expert_param(name, param) and expt_tp_size > 1: + setattr(param, "_mcore_tp", True) + if "linear_fc1.weight" in name: + setattr(param, "_tp_partition_dim", 0) + elif "linear_fc2.weight" in name: + setattr(param, "_tp_partition_dim", 1) + + if not is_expert_param(name, param) and tp_size > 1: + m_name, direct_module = param_to_direct_module[param] + if isinstance(direct_module, (TELinear,)): + parallel_mode = getattr(direct_module, "parallel_mode", None) + if parallel_mode is None: + setattr(param, "_mcore_tp", True) + setattr(param, "_tp_duplicated", True) + elif is_router_param(name, param): + setattr(param, "_mcore_tp", True) + setattr(param, "_tp_duplicated", True) + def _init_dist_index(self, pg_collection): """ Initialize the distributed index for the module. @@ -154,6 +199,7 @@ def _init_dist_index(self, pg_collection): enable_hsdp = self.ddp_config.num_distributed_optimizer_instances > 1 if pg_collection is None: tp_group = parallel_state.get_tensor_model_parallel_group() + expt_tp_group = parallel_state.get_expert_tensor_parallel_group() if enable_hsdp: dp_cp_group = parallel_state.get_data_parallel_group( with_context_parallel=True, partial_data_parallel=True @@ -168,8 +214,11 @@ def _init_dist_index(self, pg_collection): ) outer_fsdp_group = None hybrid_fsdp_group = None + expt_dp_group = parallel_state.get_expert_data_parallel_group() + ep_group = parallel_state.get_expert_model_parallel_group() else: tp_group = getattr(pg_collection, 'tp', None) + expt_tp_group = getattr(pg_collection, 'expt_tp', None) if enable_hsdp: dp_cp_group = pg_collection.intra_dp_cp outer_fsdp_group = pg_collection.inter_dist_opt @@ -178,11 +227,17 @@ def _init_dist_index(self, pg_collection): dp_cp_group = pg_collection.dp_cp outer_fsdp_group = None hybrid_fsdp_group = None + expt_dp_group = getattr(pg_collection, 'expt_dp', None) + ep_group = getattr(pg_collection, 'ep', None) if tp_group is None: single_rank_group = dist.new_group(ranks=[dist.get_rank()]) tp_group = single_rank_group + if expt_tp_group is None: + single_rank_group = dist.new_group(ranks=[dist.get_rank()]) + expt_tp_group = single_rank_group + if enable_hsdp: mesh = _get_hsdp_tp_mesh(outer_fsdp_group, dp_cp_group, tp_group) dist_index = FSDPDistributedIndex( @@ -199,6 +254,17 @@ def _init_dist_index(self, pg_collection): hybrid_fsdp_group=hybrid_fsdp_group, ) else: + if ep_group is not None: + expt_mesh = _get_dp_tp_mesh(expt_dp_group, expt_tp_group, ep_size=ep_group.size()) + expt_device_mesh = DeviceMesh.from_group( + [expt_dp_group, expt_tp_group], + device_type="cuda", + mesh=expt_mesh.tolist(), + mesh_dim_names=["dp_cp", "tp"], + ) + else: + expt_device_mesh = None + mesh = _get_dp_tp_mesh(dp_cp_group, tp_group) dist_index = FSDPDistributedIndex( device_mesh=DeviceMesh.from_group( @@ -209,8 +275,11 @@ def _init_dist_index(self, pg_collection): ), dp_shard_dim="dp_cp", tp_dim="tp", + expt_device_mesh=expt_device_mesh, ) + self.tp_group = tp_group + return dist_index def stop_communication(self): @@ -220,6 +289,20 @@ def stop_communication(self): self.module.synchronize_gradient_reduce() self.module.synchronize_param_gather() + def sync_rng_states_across_tp_group(self): + """ + Synchronize the tensor parallel random number generator states. + """ + if self.tp_group.size() <= 1: + return + + if self.tp_group.rank() == 0: + broadcast_list = [_get_rng_state_dict()] + else: + broadcast_list = [None] + torch.distributed.broadcast_object_list(broadcast_list, group=self.tp_group, group_src=0) + _load_rng_state_dict(broadcast_list[0]) + def _get_hsdp_tp_mesh(outer_fsdp_dp_group, dp_cp_group, tp_group): assert HAVE_EINOPS, "einops is not installed. Please install it with `pip install einops`." @@ -273,29 +356,46 @@ def _get_hsdp_tp_mesh(outer_fsdp_dp_group, dp_cp_group, tp_group): return mesh -def _get_dp_tp_mesh(dp_cp_group, tp_group): +def _get_dp_tp_mesh(dp_cp_group, tp_group, ep_size=1): assert HAVE_EINOPS, "einops is not installed. Please install it with `pip install einops`." world_size = dist.get_world_size() tp_size = dist.get_world_size(tp_group) if tp_group is not None else 1 - # TODO: Supports configurable (dp, cp, tp) order. - mesh = einops.rearrange(torch.arange(world_size), "(dp_cp tp) -> dp_cp tp", tp=tp_size) + # TODO: Supports configurable (dp, cp, ep, tp) order. + mesh = einops.rearrange( + torch.arange(world_size), + "(dp_cp ep tp) -> ep dp_cp tp", + dp_cp=dp_cp_group.size(), + tp=tp_size, + ep=ep_size, + ) - mesh_dp_ranks = einops.rearrange(mesh, 'dp_cp tp -> tp dp_cp', tp=tp_size) + mesh_dp_ranks = einops.rearrange(mesh, 'ep dp_cp tp -> (ep tp) dp_cp', dp_cp=dp_cp_group.size()) dp_cp_group_ranks = dist.get_process_group_ranks(dp_cp_group) assert _check_mesh_ranks_and_group_ranks_are_consistent(mesh_dp_ranks, dp_cp_group_ranks), ( f"[Megatron-FSDP] Data Parallel ranks in the mesh {mesh_dp_ranks} " f"do not match the ranks in the DP group {dp_cp_group_ranks}." ) - mesh_tp_ranks = einops.rearrange(mesh, 'dp_cp tp -> (dp_cp) tp', tp=tp_size) + mesh_tp_ranks = einops.rearrange(mesh, 'ep dp_cp tp -> (dp_cp ep) tp', tp=tp_size) tp_group_ranks = dist.get_process_group_ranks(tp_group) assert _check_mesh_ranks_and_group_ranks_are_consistent(mesh_tp_ranks, tp_group_ranks), ( f"[Megatron-FSDP] Tensor Parallel ranks in the mesh {mesh_tp_ranks} " f"do not match the ranks in the TP group {tp_group_ranks}." ) - return mesh + # Exclude the expert parallel dimension + rank = dist.get_rank() + dp_tp_meshes = [per_ep_mesh for per_ep_mesh in mesh if rank in per_ep_mesh.reshape(-1).tolist()] + assert ( + len(dp_tp_meshes) == 1 + ), f"[Megatron-FSDP] Current rank {rank} is not unique in the mesh ranks {mesh.tolist()}." + assert len(dp_tp_meshes[0].reshape(-1).tolist()) == dp_cp_group.size() * tp_group.size(), ( + f"[Megatron-FSDP] DP-TP mesh size {len(dp_tp_meshes[0].reshape(-1).tolist())} " + f"does not match expected size {dp_cp_group.size() * tp_group.size()}." + ) + + return dp_tp_meshes[0] def _check_mesh_ranks_and_group_ranks_are_consistent(mesh_ranks, group_ranks): @@ -310,3 +410,22 @@ def _check_mesh_ranks_and_group_ranks_are_consistent(mesh_ranks, group_ranks): f"{mesh_ranks.tolist()} does not match the group ranks {group_ranks}." ) return sorted(current_ranks[0]) == sorted(group_ranks) + + +def _get_rng_state_dict(): + rng_state_dict = { + 'random_rng_state': random.getstate(), + 'np_rng_state': np.random.get_state(), + 'torch_rng_state': torch.get_rng_state(), + 'cuda_rng_state': torch.cuda.get_rng_state(), + 'rng_tracker_states': tensor_parallel.get_cuda_rng_tracker().get_states(), + } + return rng_state_dict + + +def _load_rng_state_dict(rng_state_dict): + random.setstate(rng_state_dict['random_rng_state']) + np.random.set_state(rng_state_dict['np_rng_state']) + torch.set_rng_state(rng_state_dict['torch_rng_state']) + torch.cuda.set_rng_state(rng_state_dict['cuda_rng_state']) + tensor_parallel.get_cuda_rng_tracker().set_states(rng_state_dict['rng_tracker_states']) diff --git a/megatron/core/distributed/fsdp/src/README.md b/megatron/core/distributed/fsdp/src/README.md index d879c6c26f8..9e036f22f67 100644 --- a/megatron/core/distributed/fsdp/src/README.md +++ b/megatron/core/distributed/fsdp/src/README.md @@ -127,6 +127,12 @@ device_mesh[("dp_shard", "cp")]._flatten("dp_shard_cp") # Only required if using HSDP. Otherwise, don't pass hybrid_fsdp_group. device_mesh[("dp_outer", "dp_shard", "cp")]._flatten("hsdp") hsdp_group = device_mesh["hsdp"].get_group() +# Initialize DeviceMesh for expert parallel (EP) modules when using FSDP + EP. +expert_device_mesh = torch.distributed.device_mesh.init_device_mesh( + "cuda", + mesh_shape=(expt_dp_shard_size, expt_tp_size), + mesh_dim_names=("dp_shard", "tp"), +) # Fully-shards your model and distributes your optimizer. model, optimizer = fully_shard( @@ -145,6 +151,8 @@ model, optimizer = fully_shard( tp_dim="tp", # Only required when using HSDP. Otherwise, set this to None. hybrid_fsdp_group=hsdp_group, + # Only required for FSDP + EP. Otherwise, set this to None. + expt_device_mesh=expt_device_mesh, # FSDP Sharding Strategy: no_shard (0) / optim (1) / optim_grads (2) / optim_grads_params (3) zero_dp_strategy=3, outer_dp_sharding_strategy=1, @@ -192,6 +200,9 @@ optimizer.load_state_dict(ckpt_state_dict["optimizer"]) - `tp_dim` is the name of the sub-mesh used for tensor parallelism (TP), which is required for `(FSDP, TP)`-strided sharding when using Megatron-LM or Torch-native `DTensor` TP. - For more information about tensor parallelism, refer to: [Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism](https://arxiv.org/abs/1909.08053). - `hybrid_fsdp_group` is the `ProcessGroup` which contains all ranks in the flattened `dp_shard_dim` and `dp_outer_dim` sub-meshes utilized to specify the `(DP-Outer, DP-Shard)` sharded coordinate system for the weight and gradient buffers. Required for HSDP. +- `expt_device_mesh` is another [`torch.distributed.DeviceMesh`](https://docs.pytorch.org/docs/stable/distributed.html#devicemesh) tailored for the expert parallel (EP) modules in `MegatronFSDP`. + - `dp_shard_dim` is the name of the sub-mesh required for FSDP sharding of the EP modules, enabling expert data parallelism (EDP). + - `tp_dim` is the name of the sub-mesh used for expert tensor parallelism (ETP), which is required for `(FSDP, ETP)`-strided sharding when using Megatron-LM or Torch-native `DTensor` ETP. - `init_model_with_meta_device` has `MegatronFSDP` initialize your `meta`-device model in shards on every CUDA device to avoid OOM when initializing extremely large models that cannot fit on a single device. Users can initialize their model on a [`meta`-device](https://docs.pytorch.org/docs/stable/meta.html) (`with torch.device('meta'): ...`), and ``MegatronFSDP`` will further shard and initialize the model parameters layer-by-layer adhering to the customizable `module.reset_parameters` method, which prevents the entire model from being allocated in memory at any point during runtime. - Defaults to `False`. - Note that the `device` argument which installs your model on a specific device or rank will be deactivated when `init_model_with_meta_device=True`. diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/fully_shard.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/fully_shard.py index 24e86cede72..e98362a1a03 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/fully_shard.py +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/fully_shard.py @@ -64,6 +64,7 @@ def fully_shard_model( dp_outer_dim: Optional[str] = None, tp_dim: Optional[str] = None, hybrid_fsdp_group: Optional[torch.distributed.ProcessGroup] = None, + expt_device_mesh: Optional[DeviceMesh] = None, fsdp_unit_modules: Optional[Sequence[Type[torch.nn.Module]] | Sequence[str]] = None, zero_dp_strategy: str | int = 3, outer_dp_sharding_strategy: str | int = 0, @@ -183,8 +184,10 @@ def fully_shard_model( tp_dim=tp_dim, # Only required for HSDP. hybrid_fsdp_group=hybrid_fsdp_group, - # Access to flattened DP rank assignments for HFSDP. + # Access to flattened DP rank assignments for HSDP. hsdp_outer_dp_shard=_outer_fsdp_sharding, + # Only required for Megatron-FSDP + EP. + expt_device_mesh=expt_device_mesh, ) # Wrap model in Megatron FSDP. @@ -330,6 +333,7 @@ def fully_shard( dp_outer_dim: Optional[str] = None, tp_dim: Optional[str] = None, hybrid_fsdp_group: Optional[torch.distributed.ProcessGroup] = None, + expt_device_mesh: Optional[DeviceMesh] = None, fsdp_unit_modules: Optional[Sequence[Type[torch.nn.Module]] | Sequence[str]] = None, zero_dp_strategy: str | int = 3, outer_dp_sharding_strategy: str | int = 0, @@ -391,6 +395,9 @@ def fully_shard( by flattening the outer-FSDP (dp_outer_dim) and FSDP (dp_shard_dim) process groups or sub-meshes. Defaults to None. Required for HSDP, i.e. if dp_outer_dim is not None. + expt_device_mesh (Optional[DeviceMesh]): + Expert parallel device mesh object defining the topology for MoE distributed training. + fsdp_unit_modules (Optional[Sequence[Type[torch.nn.Module]] | Sequence[str]]): List of (sub-)module classes or (sub-)module class import paths that are "units", which are torch.nn.Module(s) that are sharded and scheduled by Megatron-FSDP. @@ -503,6 +510,7 @@ def fully_shard( dp_outer_dim=dp_outer_dim, tp_dim=tp_dim, hybrid_fsdp_group=hybrid_fsdp_group, + expt_device_mesh=expt_device_mesh, fsdp_unit_modules=fsdp_unit_modules, zero_dp_strategy=zero_dp_strategy, outer_dp_sharding_strategy=outer_dp_sharding_strategy, diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/megatron_fsdp.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/megatron_fsdp.py index 10a8ae14d65..d6ef5f6210e 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/megatron_fsdp.py +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/megatron_fsdp.py @@ -235,7 +235,10 @@ def __init__( self.dist_index = dist_index # If Megatron Expert Parallelism is enabled, you need to provide an expt_dp_group. - if has_expert_parameters and self.dist_index.get_expert_dp_group() is None: + if ( + has_expert_parameters + and self.dist_index.get_fsdp_group(is_expert_parallel=True) is None + ): raise ValueError( "[Megatron-FSDP] Megatron Expert Parallelism is enabled, but no expt_dp_group is" "provided." @@ -353,9 +356,7 @@ def _init_fsdp_param_and_grad_buffer(self): ) # Set the suggested communication unit size for reduce-scatter and all-gather pipelines. - suggested_communication_unit_size = ( - self.ddp_config.suggested_communication_unit_size or 1_000_000_000 - ) + suggested_communication_unit_size = self.ddp_config.suggested_communication_unit_size if suggested_communication_unit_size is None: if self.data_parallel_sharding_strategy == "optim_grads_params": total_param_elements = 0 @@ -370,6 +371,8 @@ def _init_fsdp_param_and_grad_buffer(self): suggested_communication_unit_size = total_param_elements // total_fsdp_module * 2 elif self.bucket_size is not None: suggested_communication_unit_size = self.bucket_size + else: + suggested_communication_unit_size = 1_000_000_000 # Cap to 1B elements. suggested_communication_unit_size = max( 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 c8116150d52..bdf480d867b 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 @@ -34,7 +34,14 @@ from torch.distributed.tensor.device_mesh import _mesh_resources from .uneven_dtensor import update_uneven_dtensor_chunk_metadata, validate_uneven_dtensor -from .utils import _MODEL_PARALLEL_RNG_TRACKER_NAME, FSDPDistributedIndex, get_global_memory_buffer +from .utils import ( + _MODEL_PARALLEL_RNG_TRACKER_NAME, + FSDPDistributedIndex, + get_global_memory_buffer, + get_mcore_tensor_parallel_partition_dim, + is_mcore_tensor_model_parallel, + is_mcore_tensor_parallel_duplicated, +) logger = logging.getLogger(__name__) @@ -1299,7 +1306,7 @@ def _does_param_require_new_bucket(param): and policy.data_parallel_sharding_strategy != "no_shard" ) - is_expert_parameter = lambda p: not getattr(p, "allreduce", True) + is_expert_parameter = lambda n, p: ".experts." in n # Step 1: Group the parameters according to their execution order and attributes. # FSDP unit module parameters are split into multiple parameter sub-groups. @@ -1313,7 +1320,7 @@ def _does_param_require_new_bucket(param): if is_float8tensor(param) or meta_device_init_fp8_params.get(name, False) else param.dtype ), - is_expert_param=is_expert_parameter(param), + is_expert_param=is_expert_parameter(name, param), requires_grad=param.requires_grad, fsdp_unit_id=None, ) @@ -2257,6 +2264,10 @@ def _reset_parameters(self, old_params, new_params): self.param_to_direct_module[new_param] = self.param_to_direct_module[old_param] del self.param_to_direct_module[old_param] + for tp_attr in ["_mcore_tp", "_tp_partition_dim", "_tp_duplicated"]: + if getattr(old_param, tp_attr, None) is not None: + setattr(new_param, tp_attr, getattr(old_param, tp_attr)) + for item_id, p in enumerate(self.params): if p in param_map: new_p = param_map[p] @@ -2340,6 +2351,7 @@ def _init_distributed_params(self): is_expert_param=pg.is_expert_param, run_check=True, update_uneven_dtensor_chunk_meta=True, + force_sync_tp_duplicated_param=True, ) dist_main_weight[param_name] = dist_param elif wbuf: @@ -2351,6 +2363,7 @@ def _init_distributed_params(self): is_expert_param=pg.is_expert_param, run_check=True, update_uneven_dtensor_chunk_meta=True, + force_sync_tp_duplicated_param=True, ) dist_main_weight[param_name] = dist_param else: @@ -2365,6 +2378,7 @@ def _init_distributed_params(self): is_expert_param=pg.is_expert_param, run_check=True, update_uneven_dtensor_chunk_meta=False, + force_sync_tp_duplicated_param=True, ) dist_main_weight[param_name] = dist_param @@ -2399,6 +2413,9 @@ def set_param_attribute(): "partition_dim", "partition_stride", "is_embedding_or_output_parameter", + "_mcore_tp", + "_tp_duplicated", + "_tp_partition_dim", ]: if hasattr(orig_param, attr_name): setattr(param, attr_name, getattr(orig_param, attr_name)) @@ -3546,7 +3563,9 @@ def to_local_if_dtensor(tensor): return tensor -def _get_fsdp_tensor_spec(param, dist_index: FSDPDistributedIndex, is_sharded_param): +def _get_fsdp_tensor_spec( + param, dist_index: FSDPDistributedIndex, is_sharded_param, is_expert_param +): """ Get the DeviceMesh for the parameter and modify the placement for Megatron-FSDP. """ @@ -3557,7 +3576,7 @@ def _get_fsdp_tensor_spec(param, dist_index: FSDPDistributedIndex, is_sharded_pa dtensor_mesh = getattr(dtensor_spec, "mesh", None) # Validate that the DTensor root mesh is identical to the Megatron-FSDP device mesh. - megatron_fsdp_global_mesh = dist_index.get_root_mesh() + megatron_fsdp_global_mesh = dist_index.get_root_mesh(is_expert_parallel=is_expert_param) dtensor_global_mesh = _mesh_resources.get_root_mesh(dtensor_mesh) # FIXME(boxiangw): add or megatron_fsdp_global_mesh != dtensor_global_mesh: # _mesh_resources.get_root_mesh(dtensor_mesh) is not getting the correct root mesh @@ -3602,7 +3621,7 @@ def _get_fsdp_tensor_spec(param, dist_index: FSDPDistributedIndex, is_sharded_pa placements = [Shard(0), dtensor_placement] shard_order = [1, 0] - device_mesh = dist_index.get_submesh(mesh_dim_names) + device_mesh = dist_index.get_submesh(mesh_dim_names, is_expert_parallel=is_expert_param) if shard_order is not None: setattr(device_mesh, "_shard_order", shard_order) @@ -3627,7 +3646,7 @@ def _get_fsdp_tensor_spec(param, dist_index: FSDPDistributedIndex, is_sharded_pa else: placements = [Shard(0)] - device_mesh = dist_index.get_submesh(mesh_dim_names) + device_mesh = dist_index.get_submesh(mesh_dim_names, is_expert_parallel=is_expert_param) if shard_order is not None: setattr(device_mesh, "_shard_order", shard_order) @@ -3642,6 +3661,7 @@ def make_fsdp_dtensor( is_expert_param: bool = False, run_check: bool = False, update_uneven_dtensor_chunk_meta: bool = False, + force_sync_tp_duplicated_param: bool = False, ): """ Creates a distributed tensor (DTensor) from a local tensor with support for @@ -3720,38 +3740,39 @@ def make_fsdp_dtensor( orig_param = param # Handle tensor model parallel specific logic - if getattr(param, "tensor_model_parallel", False): + if is_mcore_tensor_model_parallel(param): # Ensure parameter is not already a DTensor assert not isinstance(param, DTensor), ( - "[Megatron-FSDP] Parameter is already a DTensor, yet tensor_model_parallel " - "is True. Check usage." + "[Megatron-FSDP] Parameter is already a DTensor, yet tensor_model_parallel " "is True." ) - # Validate M-Core TP attributes - assert hasattr( - param, "partition_dim" - ), "[Megatron-FSDP] tensor_model_parallel param missing 'partition_dim'." - assert hasattr( - param, "partition_stride" - ), "[Megatron-FSDP] tensor_model_parallel param missing 'partition_stride'." - assert ( - param.partition_stride == 1 - ), "[Megatron-FSDP] Only partition_stride=1 is currently supported for " - "tensor_model_parallel." - - tp_dim = param.partition_dim - tp_mesh = dist_index.get_submesh(dist_index.tp_dim) - - # Adjust shape for global dimension + tp_mesh = dist_index.get_submesh(dist_index.tp_dim, is_expert_parallel=is_expert_param) + global_shape = list(param.shape) if tp_mesh.mesh.numel() > 1: - global_shape = list(param.shape) - global_shape[tp_dim] *= tp_mesh.mesh.numel() + if is_mcore_tensor_parallel_duplicated(param): + placements = [Replicate()] + if force_sync_tp_duplicated_param: + if local_tensor.numel() > 0: + torch.distributed.broadcast( + local_tensor, group=tp_mesh.get_group(), group_src=0 + ) + elif run_check: + # TODO: Implement consistency check for duplicated TP parameters + pass + else: + tp_dim = get_mcore_tensor_parallel_partition_dim(param) + assert tp_dim is not None, ( + "[Megatron-FSDP] Parameter is not tensor model parallel, " + "yet tensor_model_parallel is True." + ) + placements = [Shard(tp_dim)] + global_shape[tp_dim] *= tp_mesh.mesh.numel() # Construct TP-sharded DTensor using Megatron-style placement param = DTensor.from_local( - local_tensor=param, + local_tensor=local_tensor, device_mesh=tp_mesh, - placements=[Shard(tp_dim)], + placements=placements, run_check=run_check, shape=global_shape, stride=torch.empty(global_shape).stride(), @@ -3759,7 +3780,7 @@ def make_fsdp_dtensor( # Get FSDP-configured mesh and placements from provided param device_mesh, placements = _get_fsdp_tensor_spec( - param, dist_index, is_sharded_param=is_sharded_param + param, dist_index, is_sharded_param=is_sharded_param, is_expert_param=is_expert_param ) # Reshape local tensor for sharded layouts beyond 1D diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/uneven_dtensor.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/uneven_dtensor.py index 523d8fae333..490d80c0f21 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/uneven_dtensor.py +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/uneven_dtensor.py @@ -365,7 +365,9 @@ def _assemble_full_tensor_from_uneven_chunks( # Wrap into a replicated DTensor and return return DTensor.from_local( - full_tensor, placements=[Replicate()], device_mesh=dtensor.device_mesh + full_tensor, + placements=[Replicate()] * len(dtensor.placements), + device_mesh=dtensor.device_mesh, ) diff --git a/megatron/core/distributed/fsdp/src/megatron_fsdp/utils.py b/megatron/core/distributed/fsdp/src/megatron_fsdp/utils.py index 1dfe08b90f4..b94a332bb0d 100644 --- a/megatron/core/distributed/fsdp/src/megatron_fsdp/utils.py +++ b/megatron/core/distributed/fsdp/src/megatron_fsdp/utils.py @@ -675,6 +675,7 @@ def __init__( tp_dim: Optional[str] = None, hybrid_fsdp_group: Optional[torch.distributed.ProcessGroup] = None, hsdp_outer_dp_shard: bool = False, + expt_device_mesh: Optional[DeviceMesh] = None, ): """ Args: @@ -691,6 +692,8 @@ def __init__( in hybrid FSDP. Specifying outer sharding will lift the bucket sharding coordinate system to flattened ranks of (dp_shard, dp_outer) instead of just sharding across dp_shard ranks and replicating across dp_outer ranks. + expt_device_mesh (Optional[DeviceMesh]): The expert parallel device mesh + to use for the DistributedIndex. """ # Device mesh arguments. self.device_mesh = device_mesh @@ -701,6 +704,11 @@ def __init__( self.use_hybrid_fsdp = dp_outer_dim is not None # Helper flag to denote if we are outer-sharding in hybrid FSDP. self.hsdp_outer_dp_shard = hsdp_outer_dp_shard + self.expt_device_mesh = expt_device_mesh + + # Handling the situation where M-Core MoE EP=1 + if self.expt_device_mesh is None: + self.expt_device_mesh = device_mesh # Hybrid FSDP Process Groups # Retrieve the FSDP process group from the DeviceMesh. @@ -719,6 +727,14 @@ def __init__( # combination of the outer-FSDP and FSDP process groups. self.hybrid_fsdp_group = hybrid_fsdp_group + # Retrieve the expert parallel process groups from the DeviceMesh. + self.expt_fsdp_group = ( + self.expt_device_mesh[self.dp_shard_dim].get_group() + if self.expt_device_mesh is not None + and contains_submesh(self.expt_device_mesh, self.dp_shard_dim) + else None + ) + """ Store a persistent reference to the core device meshes that back Megatron-FSDP. This is necessary because _MeshEnv (_mesh_resources) may not persist: @@ -732,26 +748,33 @@ def __init__( FIXME(@cspades): Identify the root cause of this behavior. """ self.mesh_library = {} - # TP Mesh + + def register_submesh(device_mesh, submesh, is_expert_parallel): + """Register a submesh with identifier: (*submesh, is_expert_parallel) + in the mesh library.""" + if contains_submesh(device_mesh, submesh): + submesh_identifier = tuple(list(submesh) + [is_expert_parallel]) + self.mesh_library[submesh_identifier] = device_mesh[submesh] + + # Define common submesh patterns tp_submesh = (self.tp_dim,) - if contains_submesh(self.device_mesh, tp_submesh): - self.mesh_library[tp_submesh] = self.device_mesh[tp_submesh] - # HSDP-TP Mesh hsdp_tp_submesh = (self.dp_outer_dim, self.dp_shard_dim, self.tp_dim) - if contains_submesh(self.device_mesh, hsdp_tp_submesh): - self.mesh_library[hsdp_tp_submesh] = self.device_mesh[hsdp_tp_submesh] - # FSDP-TP Mesh fsdp_tp_submesh = (self.dp_shard_dim, self.tp_dim) - if contains_submesh(self.device_mesh, fsdp_tp_submesh): - self.mesh_library[fsdp_tp_submesh] = self.device_mesh[fsdp_tp_submesh] - # HSDP Mesh hsdp_submesh = (self.dp_outer_dim, self.dp_shard_dim) - if contains_submesh(self.device_mesh, hsdp_submesh): - self.mesh_library[hsdp_submesh] = self.device_mesh[hsdp_submesh] - # FSDP Mesh fsdp_submesh = (self.dp_shard_dim,) - if contains_submesh(self.device_mesh, fsdp_submesh): - self.mesh_library[fsdp_submesh] = self.device_mesh[fsdp_submesh] + + # Register non-EP submeshes + register_submesh(self.device_mesh, tp_submesh, False) + register_submesh(self.device_mesh, hsdp_tp_submesh, False) + register_submesh(self.device_mesh, fsdp_tp_submesh, False) + register_submesh(self.device_mesh, hsdp_submesh, False) + register_submesh(self.device_mesh, fsdp_submesh, False) + + # Register EP submeshes + if self.expt_device_mesh is not None: + register_submesh(self.expt_device_mesh, tp_submesh, True) + register_submesh(self.expt_device_mesh, fsdp_tp_submesh, True) + register_submesh(self.expt_device_mesh, fsdp_submesh, True) # Validate FSDP arguments. if self.fsdp_group is None: @@ -776,36 +799,54 @@ def __init__( "process groups or sub-meshes." ) - def get_submesh(self, mesh_dim_names: str | Sequence[str]) -> DeviceMesh: + def get_submesh( + self, mesh_dim_names: str | Sequence[str], is_expert_parallel: bool = False + ) -> DeviceMesh: """ - Retrieve an Megatron-FSDP-registered sub-mesh by name(s). + Retrieve an Megatron-FSDP-registered submesh by name(s). """ if isinstance(mesh_dim_names, str): mesh_dim_names = (mesh_dim_names,) - # Search for the sub-mesh in the mesh library. - device_submesh = self.mesh_library.get(tuple(mesh_dim_names), None) + + # Construct submesh identifier: (*mesh_dim_names, is_expert_parallel) + submesh_identifier = tuple(list(mesh_dim_names) + [is_expert_parallel]) + + # Retrieve the submesh from the mesh library + device_submesh = self.mesh_library.get(submesh_identifier, None) + if device_submesh is None: - if self.tp_dim is None: - # Warn about not specifying tp_dim for - # layers or frameworks that depend on this. + # Warn about not specifying tp_dim for layers or frameworks that depend on this. + if self.tp_dim is None and not is_expert_parallel: logger.warning( - "[FSDPDistributedIndex] Note: For TransformerEngine, or other machine learning " - "frameworks like Megatron that assume TP=1, you must specify tp_dim to use " - "Megatron-FSDP. Create a trivial TP dimension by setting the TP dimension size " + "[FSDPDistributedIndex] Note: For TransformerEngine, or " + "other machine learning frameworks like Megatron that assume " + "TP=1, you must specify tp_dim to use Megatron-FSDP. " + "Create a trivial TP dimension by setting the TP dimension size " "to 1 in the DeviceMesh.\n" f"DeviceMesh: {self.device_mesh}" ) + elif self.tp_dim is None and is_expert_parallel: + logger.warning( + "[FSDPDistributedIndex] Note: For TransformerEngine, or " + "other machine learning frameworks like Megatron that assume " + "ETP=1, you must specify tp_dim to use Megatron-FSDP. " + "Create a trivial ETP dimension by setting the ETP dimension size " + "to 1 in the DeviceMesh.\n" + f"DeviceMesh: {self.expt_device_mesh}" + ) + raise ValueError( - f"[FSDPDistributedIndex][get_submesh] No sub-mesh with " - f"mesh_dim_names={mesh_dim_names} has been registered with Megatron-FSDP." + f"[FSDPDistributedIndex][get_submesh] No submesh with " + f"mesh_dim_names={mesh_dim_names}, is_expert_parallel={is_expert_parallel} " + f"has been registered with Megatron-FSDP." ) + return device_submesh def get_dp_group(self, is_expert_parallel: bool = False) -> ProcessGroup: """Get the data parallel process group.""" if is_expert_parallel: - # Expert parallel is not supported - return None + return self.expt_fsdp_group if self.use_hybrid_fsdp: return self.hybrid_fsdp_group return self.fsdp_group @@ -813,8 +854,7 @@ def get_dp_group(self, is_expert_parallel: bool = False) -> ProcessGroup: def get_fsdp_group(self, is_expert_parallel: bool = False) -> ProcessGroup: """Get the FSDP process group.""" if is_expert_parallel: - # Expert parallel is not supported - return None + return self.expt_fsdp_group return self.fsdp_group def get_outer_fsdp_group(self) -> ProcessGroup: @@ -826,7 +866,7 @@ def get_outer_fsdp_group(self) -> ProcessGroup: def get_root_mesh(self, is_expert_parallel: bool = False) -> DeviceMesh: """Get the device mesh.""" if is_expert_parallel: - raise NotImplementedError("Expert parallel is not supported in Megatron-FSDP.") + return self.expt_device_mesh return self.device_mesh def get_logical_hybrid_fsdp_rank(self): @@ -924,3 +964,29 @@ def create_updated_function_signature(original_function, **extended_kwargs: dict # Return the updated function signature. return inspect.Signature(params) + + +def is_mcore_tensor_model_parallel(param: torch.Tensor) -> bool: + """ + Check if the given parameter is Megatron-Core tensor model parallel. + """ + return getattr(param, "_mcore_tp", False) or getattr(param, "tensor_model_parallel", False) + + +def is_mcore_tensor_parallel_duplicated(param: torch.Tensor) -> bool: + """ + Check if the given parameter is Megatron-Core tensor model parallel and duplicated. + """ + return getattr(param, "_tp_duplicated", False) + + +def get_mcore_tensor_parallel_partition_dim(param: torch.Tensor) -> Optional[int]: + """ + Get the partition dimension for a Megatron-Core tensor model parallel parameter. + """ + if is_mcore_tensor_model_parallel(param): + if hasattr(param, "_tp_partition_dim"): + return param._tp_partition_dim + else: + return param.partition_dim + return None diff --git a/megatron/core/models/common/embeddings/yarn_rotary_pos_embedding.py b/megatron/core/models/common/embeddings/yarn_rotary_pos_embedding.py index 507472f789f..455a7757d28 100644 --- a/megatron/core/models/common/embeddings/yarn_rotary_pos_embedding.py +++ b/megatron/core/models/common/embeddings/yarn_rotary_pos_embedding.py @@ -130,9 +130,9 @@ def forward(self, max_seq_len: int, offset: int = 0, packed_seq: bool = False) - self.original_max_position_embeddings, self.correction_range_round_to_int, ) - inv_freq_mask = 1.0 - _yarn_linear_ramp_mask(low, high, self.dim // 2).to( - device=self.inv_freq_extra.device, dtype=torch.float32 - ) + inv_freq_mask = 1.0 - _yarn_linear_ramp_mask( + low, high, self.dim // 2, device=self.inv_freq_extra.device + ).to(dtype=torch.float32) inv_freq = self.inv_freq_inter * (1 - inv_freq_mask) + self.inv_freq_extra * inv_freq_mask seq = ( @@ -211,11 +211,11 @@ def _yarn_find_correction_range( return max(low, 0), min(high, dim - 1) # Clamp values just in case -def _yarn_linear_ramp_mask(min: float, max: float, dim: int) -> Tensor: +def _yarn_linear_ramp_mask(min: float, max: float, dim: int, device: torch.device) -> Tensor: if min == max: max += 0.001 # Prevent singularity - linear_func = (torch.arange(dim, dtype=torch.float32) - min) / (max - min) + linear_func = (torch.arange(dim, dtype=torch.float32, device=device) - min) / (max - min) ramp_func = torch.clamp(linear_func, 0, 1) return ramp_func diff --git a/megatron/core/optimizer/__init__.py b/megatron/core/optimizer/__init__.py index 307538fad22..c254b2f6882 100644 --- a/megatron/core/optimizer/__init__.py +++ b/megatron/core/optimizer/__init__.py @@ -34,6 +34,7 @@ from megatron.core import parallel_state from megatron.core.optimizer.cpu_offloading.hybrid_optimizer import HybridDeviceOptimizer from megatron.core.process_groups_config import ProcessGroupCollection +from megatron.core.transformer.fsdp_dtensor_checkpoint import get_global_unique_param_name from ..distributed.param_and_grad_buffer import _ParamAndGradBuffer from ..transformer.module import MegatronModule @@ -481,6 +482,7 @@ def get_megatron_optimizer( use_gloo_process_groups: bool = True, default_skip_embedding_weight_decay: bool = False, pg_collection: Optional[ProcessGroupCollection] = None, + dump_param_to_param_group_map: Optional[str] = None, ) -> MegatronOptimizer: """Retrieve the Megatron optimizer for model chunks. @@ -502,6 +504,7 @@ def get_megatron_optimizer( This is useful if you do not want embeddings to shrink to zero in training as recommended in https://arxiv.org/abs/2312.16903 pg_collection: Optional unified process group for distributed training. + dump_param_to_param_group_map (Optional[str]): path to dump parameter to param group map. Returns: Instance of MegatronOptimizer. @@ -579,6 +582,9 @@ def get_megatron_optimizer( return ChainedOptimizer(optimizers) + if dump_param_to_param_group_map is not None: + param_to_param_group = {} + param_group_id = 0 for dense_model_chunks, overlap_param_gather_with_optimizer_step in zip( all_dense_model_chunks, overlap_param_gather_with_optimizer_step_flags ): @@ -597,6 +603,12 @@ def get_megatron_optimizer( model_chunk.overlap_param_gather_with_optimizer_step = ( overlap_param_gather_with_optimizer_step ) + if dump_param_to_param_group_map is not None: + for param_group in param_groups: + for param in param_group["params"]: + param_name = get_global_unique_param_name(model_chunks, param) + param_to_param_group[param_name] = param_group_id + param_group_id += 1 # Pass Gloo process groups into optimizer only if needed. optimizers.append( @@ -626,6 +638,12 @@ def get_megatron_optimizer( buffer_name='expert_parallel_buffers', default_skip_embedding_weight_decay=default_skip_embedding_weight_decay, ) + if dump_param_to_param_group_map is not None: + for param_group in moe_param_groups: + for param in param_group["params"]: + param_name = get_global_unique_param_name(model_chunks, param) + param_to_param_group[param_name] = param_group_id + param_group_id += 1 if len(moe_param_groups) > 0: expt_model_parallel_rank = get_pg_rank(expt_tp_pp_group) # Pass Gloo process groups into optimizer only if needed. @@ -648,4 +666,9 @@ def get_megatron_optimizer( ) ) + if dump_param_to_param_group_map is not None: + torch.distributed.checkpoint.save( + state_dict=param_to_param_group, checkpoint_id=dump_param_to_param_group_map + ) + return ChainedOptimizer(optimizers) diff --git a/megatron/core/optimizer/distrib_optimizer.py b/megatron/core/optimizer/distrib_optimizer.py index 2925edcce60..8b4740516e2 100644 --- a/megatron/core/optimizer/distrib_optimizer.py +++ b/megatron/core/optimizer/distrib_optimizer.py @@ -47,6 +47,7 @@ from ..dist_checkpointing.utils import extract_sharded_tensors_and_factories from ..distributed.param_and_grad_buffer import _ParamAndGradBuffer, partition_buckets from ..fp8_utils import dequantize_fp8_tensor, is_float8tensor, quantize_param_shard +from ..transformer.fsdp_dtensor_checkpoint import handle_experts_in_state_dict from ..transformer.module import MegatronModule from .grad_scaler import MegatronGradScaler from .optimizer import MixedPrecisionOptimizer, _zero_grad_group_helper, param_group_identifier_keys @@ -1152,6 +1153,7 @@ def _param_name(self, param: torch.nn.Parameter) -> str: "Ensure that each model chunk has unique parameter names." ) name_to_param.update(_name_to_param) + name_to_param = handle_experts_in_state_dict(name_to_param) self.param_to_name = {param: name for name, param in name_to_param.items()} assert ( param in self.param_to_name diff --git a/megatron/core/transformer/fsdp_dtensor_checkpoint.py b/megatron/core/transformer/fsdp_dtensor_checkpoint.py index dad1947a183..9ef3f1f1b82 100644 --- a/megatron/core/transformer/fsdp_dtensor_checkpoint.py +++ b/megatron/core/transformer/fsdp_dtensor_checkpoint.py @@ -12,18 +12,160 @@ # See the License for the specific language governing permissions and # limitations under the License. +import logging +import re + import torch +import torch.distributed as dist +from torch.distributed.checkpoint import default_planner + +logger = logging.getLogger(__name__) try: + from torch.distributed import DeviceMesh + from torch.distributed._tensor import DTensor + from torch.distributed.checkpoint.metadata import TensorStorageMetadata + from torch.distributed.tensor.placement_types import Replicate, Shard + from megatron.core.distributed.fsdp.src.megatron_fsdp.param_and_grad_buffer import ( make_fsdp_dtensor, ) + from megatron.core.distributed.fsdp.src.megatron_fsdp.uneven_dtensor import ( + gather_uneven_dtensor_to_full_tensor, + ) + from megatron.core.distributed.fsdp.src.megatron_fsdp.utils import ( + get_mcore_tensor_parallel_partition_dim, + is_mcore_tensor_model_parallel, + ) HAVE_MEGATRON_FSDP = True except ImportError: HAVE_MEGATRON_FSDP = False +from megatron.core import parallel_state from megatron.core.tensor_parallel.layers import copy_tensor_model_parallel_attributes +from megatron.core.transformer.transformer_layer import TransformerLayer + + +def get_ep_layer_offset(): + """ + Get the expert layer offset for the current model. + """ + from megatron.training.global_vars import get_args + + args = get_args() + ep_size = parallel_state.get_expert_model_parallel_world_size() + ep_rank = parallel_state.get_expert_model_parallel_rank() + num_local_experts = args.num_experts // ep_size if args.num_experts else 0 + local_expert_offset = ep_rank * num_local_experts + + return local_expert_offset + + +def get_total_num_experts(): + """ + Get the total number of experts for the current model. + """ + from megatron.training.global_vars import get_args + + args = get_args() + return args.num_experts if args.num_experts else 0 + + +def get_expert_index_from_key(key): + """Extract expert index from various expert key formats. + + Supported formats: + - GroupedMLP: 'mlp.experts.linear_fc1.weight0', 'mlp.experts.linear_fc2.weight0' + - SequentialMLP: 'mlp.experts.local_experts.0.linear_fc1.weight', + 'mlp.experts.local_experts.0.linear_fc2.weight' + + Returns: + int: Expert index if found, None otherwise. + """ + # GroupedMLP: index is at the end after 'weight' + if 'mlp.experts.linear_fc1.weight' in key or 'mlp.experts.linear_fc2.weight' in key: + m = re.search(r'^.*\.mlp\.experts\.linear_fc\d\.weight(\d+)', key) + assert m, f"Failed to parse expert index from key: {key}" + return int(m.group(1)) + # SequentialMLP: index is between 'local_experts.' and next '.' + elif 'mlp.experts.local_experts' in key: + m = re.search(r'^.*\.mlp\.experts\.local_experts\.(\d+)', key) + assert m, f"Failed to parse expert index from key: {key}" + return int(m.group(1)) + return None + + +def handle_experts_in_state_dict(state_dict): + """ + Rewrite expert keys in state dict. + """ + local_expert_start = get_ep_layer_offset() + local_expert_end = get_total_num_experts() + + def should_keep_expert_key(expert_index): + """Determine if this rank should keep this expert key based on expert index""" + if expert_index is None: + # If we can't determine expert index, keep the key (non-expert weights) + return True + + # Check if this expert belongs to this rank + return local_expert_start <= expert_index < local_expert_end + + def replace_expert_index_in_key(key, expert_index, state_dict): + """Replace expert index in key with new index corresponding to the current rank""" + new_expert_index = expert_index + local_expert_start + # GroupedMLP: 'mlp.experts.linear_fc1.weight0', 'mlp.experts.linear_fc2.weight0' + if 'mlp.experts.linear_fc1.weight' in key or 'mlp.experts.linear_fc2.weight' in key: + # Handle SwiGLU weight{idx}_w and weight{idx}_v format + if key.endswith('_w') or key.endswith('_v'): + suffix = key[-2:] # '_w' or '_v' + new_key = key.replace( + f'weight{expert_index}{suffix}', f'weight{new_expert_index}{suffix}' + ) + # Handle regular weight{idx} format + else: + new_key = key.replace(f'weight{expert_index}', f'weight{new_expert_index}') + # SequentialMLP: index is between 'local_experts.' and next '.' + elif 'mlp.experts.local_experts' in key: + new_key = key.replace( + f'local_experts.{expert_index}.', f'local_experts.{new_expert_index}.' + ) + else: + raise ValueError(f"Unexpected expert key format: {key}") + + state_dict[new_key] = state_dict[key] + del state_dict[key] + + # Process model state dict + state_dict = state_dict.copy() + for key in list(state_dict.keys()): + expert_index = get_expert_index_from_key(key) + if not should_keep_expert_key(expert_index): + replace_expert_index_in_key(key, expert_index, state_dict) + + return state_dict + + +def expert_param_local_key(key): + """Get the module parameter corresponding to the key.""" + local_expert_offset = get_ep_layer_offset() + expert_index = get_expert_index_from_key(key) + if expert_index is not None: + new_expert_index = expert_index - local_expert_offset + # GroupedMLP: 'mlp.experts.linear_fc1.weight0', 'mlp.experts.linear_fc2.weight0' + if 'mlp.experts.linear_fc1.weight' in key or 'mlp.experts.linear_fc2.weight' in key: + new_key = key.replace(f'weight{expert_index}', f'weight{new_expert_index}') + # SequentialMLP: index is between 'local_experts.' and next '.' + elif 'mlp.experts.local_experts' in key: + new_key = key.replace( + f'local_experts.{expert_index}.', f'local_experts.{new_expert_index}.' + ) + else: + raise ValueError(f"Unexpected expert key format: {key}") + key = new_key + + return key def handle_swiglu_in_state_dict(model, model_state_dict, optimizer_state_dict): @@ -43,7 +185,29 @@ def intersection(s1, s2): def offset_slice(s, offset): return slice(s.start + offset, s.stop + offset) - def split_swiglu_linear_fc1(data, dist_param, swiglu_shard_axis): + def is_swiglu_key(key): + """ + Check if this key should be handled as SwiGLU linear_fc1 weight or bias. + """ + # Non-expert MLP: 'mlp.linear_fc1.weight', 'mlp.linear_fc1.bias' + # GroupedMLP: 'mlp.experts.linear_fc1.weight0', 'mlp.experts.linear_fc1.bias0' + # SequentialMLP: 'mlp.experts.local_experts.0.linear_fc1.weight', + # 'mlp.experts.local_experts.0.linear_fc1.bias' + return any( + re.search(pat, key) + for pat in [ + r"(.*)\.mlp\.linear_fc1\.weight$", + r"(.*)\.mlp\.linear_fc1\.bias$", + r"(.*)\.mlp\.experts\.linear_fc1\.weight(\d+)$", + r"(.*)\.mlp\.experts\.linear_fc1\.bias(\d+)$", + r"(.*)\.mlp\.experts\.local_experts\.(\d+)\.linear_fc1\.weight$", + r"(.*)\.mlp\.experts\.local_experts\.(\d+)\.linear_fc1\.bias$", + r"(.*)\.mlp\.shared_experts\.linear_fc1\.weight$", + r"(.*)\.mlp\.shared_experts\.linear_fc1\.bias$", + ] + ) + + def split_swiglu_linear_fc1(data, dist_param, swiglu_shard_axis, is_expert_param): """ Split the SWiGLU linear_fc1 parameter into two parts: weight_w and weight_v. """ @@ -55,7 +219,9 @@ def split_swiglu_linear_fc1(data, dist_param, swiglu_shard_axis): fsdp_slice = dist_param.megatron_fsdp_slice megatron_fsdp_dist_index = dist_param.megatron_fsdp_dist_index - tp_mesh = megatron_fsdp_dist_index.get_submesh([megatron_fsdp_dist_index.tp_dim]) + tp_mesh = megatron_fsdp_dist_index.get_submesh( + [megatron_fsdp_dist_index.tp_dim], is_expert_parallel=is_expert_param + ) data_size = data.numel() // tp_mesh.mesh.numel() w_slice = slice(0, data_size // 2) v_slice = slice(data_size // 2, data_size) @@ -75,8 +241,9 @@ def split_swiglu_linear_fc1(data, dist_param, swiglu_shard_axis): # Fake parameters w and v are used to provide the correct parameter # shape and Tensor-Parallelism information. per_tp_rank_shape = list(data.shape) - if getattr(dist_param, "tensor_model_parallel", False): - tp_dim = dist_param.partition_dim + if is_mcore_tensor_model_parallel(dist_param): + tp_dim = get_mcore_tensor_parallel_partition_dim(dist_param) + assert tp_dim is not None, "Tensor model parallel dimension not found" per_tp_rank_shape[tp_dim] //= tp_mesh.mesh.numel() linear_fc1_meta = torch.empty(*per_tp_rank_shape, device="meta") w_meta, v_meta = torch.chunk(linear_fc1_meta, 2, dim=swiglu_shard_axis) @@ -87,6 +254,7 @@ def split_swiglu_linear_fc1(data, dist_param, swiglu_shard_axis): weight_w.data, w_meta, dist_index=megatron_fsdp_dist_index, + is_expert_param=is_expert_param, run_check=True, update_uneven_dtensor_chunk_meta=True, ) @@ -94,16 +262,21 @@ def split_swiglu_linear_fc1(data, dist_param, swiglu_shard_axis): weight_v.data, v_meta, dist_index=megatron_fsdp_dist_index, + is_expert_param=is_expert_param, run_check=True, update_uneven_dtensor_chunk_meta=True, ) return weight_w, weight_v + model_state_dict = model_state_dict.copy() for key in list(model_state_dict.keys()): - if key.endswith('mlp.linear_fc1.weight') or key.endswith('mlp.linear_fc1.bias'): + if is_swiglu_key(key): dist_param = model.get_parameter(f"module.{key}") weight_w, weight_v = split_swiglu_linear_fc1( - model_state_dict[key], dist_param, swiglu_shard_axis=0 + model_state_dict[key], + dist_param, + swiglu_shard_axis=0, + is_expert_param='mlp.experts' in key, ) # Update the model state dict with the new keys @@ -111,26 +284,32 @@ def split_swiglu_linear_fc1(data, dist_param, swiglu_shard_axis): model_state_dict[f"{key}_v"] = weight_v del model_state_dict[key] - try: - optimizer_state_dict = optimizer_state_dict["state"] - except KeyError: - optimizer_state_dict = {} + if optimizer_state_dict is not None: + optimizer_state_dict = optimizer_state_dict.copy() + if len(optimizer_state_dict["state"]) != 0: + opt_state_dict = optimizer_state_dict["state"] + new_opt_state_dict = {} + for key in list(opt_state_dict.keys()): + # Only process SWIGLU keys + if not is_swiglu_key(key): + new_opt_state_dict[key] = opt_state_dict[key] + continue + new_opt_state_dict[f"{key}_w"] = opt_state_dict[key].copy() + new_opt_state_dict[f"{key}_v"] = opt_state_dict[key].copy() + for subkey in ["exp_avg", "exp_avg_sq"]: + dist_param = model.get_parameter(expert_param_local_key(key[len("module.") :])) + weight_w, weight_v = split_swiglu_linear_fc1( + opt_state_dict[key][subkey], + dist_param, + swiglu_shard_axis=0, + is_expert_param="mlp.experts" in key, + ) + # Update the optimizer state dict with the new keys + new_opt_state_dict[f"{key}_w"][subkey] = weight_w + new_opt_state_dict[f"{key}_v"][subkey] = weight_v + optimizer_state_dict["state"] = new_opt_state_dict - if len(optimizer_state_dict) != 0: - for key in list(optimizer_state_dict.keys()): - if not (key.endswith('mlp.linear_fc1.weight') or key.endswith('mlp.linear_fc1.bias')): - continue - optimizer_state_dict[f"{key}_w"] = optimizer_state_dict[key].copy() - optimizer_state_dict[f"{key}_v"] = optimizer_state_dict[key].copy() - for subkey in ["exp_avg", "exp_avg_sq"]: - dist_param = model.get_parameter(key[len("module.") :]) - weight_w, weight_v = split_swiglu_linear_fc1( - optimizer_state_dict[key][subkey], dist_param, swiglu_shard_axis=0 - ) - # Update the optimizer state dict with the new keys - optimizer_state_dict[f"{key}_w"][subkey] = weight_w - optimizer_state_dict[f"{key}_v"][subkey] = weight_v - del optimizer_state_dict[key] + return model_state_dict, optimizer_state_dict def handle_fp8_extra_state_case(model_state_dict): @@ -162,7 +341,7 @@ def flatten_state_dict(obj, parent_key="", sep="."): return items -def print_diff_in_state_dicts(state_dict_metadata, load_state_dict): +def print_diff_in_state_dicts(state_dict_metadata, load_state_dict, limit=100): """ Print the differences between two state dicts: metadata state dict and load state dict. This function compares the keys and shapes of the tensors in both dicts. @@ -172,24 +351,105 @@ def print_diff_in_state_dicts(state_dict_metadata, load_state_dict): meta_keys = set(state_dict_metadata.keys()) load_keys = set(load_state_dict.keys()) - only_in_meta = meta_keys - load_keys - only_in_load = load_keys - meta_keys - in_both = meta_keys & load_keys + only_in_meta = list(meta_keys - load_keys) + only_in_load = list(load_keys - meta_keys) + in_both = list(meta_keys & load_keys) - print("Keys only in checkpoint metadata_state_dict:") - for k in sorted(only_in_meta): - print(f" {k}") + logger.info(f"Keys only in checkpoint metadata_state_dict(first {limit}):") + for k in sorted(only_in_meta[:limit]): + logger.info(f" {k}") - print("\nKeys only in load_state_dict:") - for k in sorted(only_in_load): - print(f" {k}") + logger.info(f"\nKeys only in load_state_dict(first {limit}):") + for k in sorted(only_in_load[:limit]): + logger.info(f" {k}") - print("\nKeys in both but with different shapes:") - for k in sorted(in_both): + logger.info(f"\nKeys in both but with different shapes(first {limit}):") + for k in sorted(in_both[:limit]): v_meta = state_dict_metadata[k] v_load = load_state_dict[k] # If tensors, compare shape; else, compare type/values meta_shape = v_meta.size if hasattr(v_meta, "size") else type(v_meta) load_shape = v_load.shape if hasattr(v_load, "shape") else type(v_load) if meta_shape != load_shape: - print(f" {k}: meta shape={meta_shape}, load shape={load_shape}") + logger.info(f" {k}: meta shape={meta_shape}, load shape={load_shape}") + + +def validate_loaded_state_dict(state_dict, checkpoint_path): + """ + Validate the loaded state dict against the expected structure and types. + """ + assert HAVE_MEGATRON_FSDP, "This function requires Megatron-FSDP to be installed." + + # Initialize reader + reader = torch.distributed.checkpoint.FileSystemReader(checkpoint_path) + metadata = reader.read_metadata() + flat_state_dict = flatten_state_dict(state_dict) + + for key, value in flat_state_dict.items(): + tensor_metadata = metadata.state_dict_metadata[key] + + if not isinstance(tensor_metadata, TensorStorageMetadata): + continue + if not isinstance(value, DTensor): + load_item_dict = {key: torch.empty_like(value)} + else: + load_item_dict = { + key: torch.distributed.tensor.empty( + tensor_metadata.size, + dtype=tensor_metadata.properties.dtype, + device_mesh=DeviceMesh.from_group( + group=dist.group.WORLD, + device_type="cuda", + mesh=torch.arange(dist.get_world_size()), + mesh_dim_names=("world",), + ), + placements=[Shard(0)], + ) + } + torch.distributed.checkpoint.load( + load_item_dict, storage_reader=reader, planner=default_planner.DefaultLoadPlanner() + ) + if isinstance(value, DTensor): + full_value = gather_uneven_dtensor_to_full_tensor(value) + loaded_tensor = load_item_dict[key].redistribute( + placements=[Replicate()] * len(value.placements) + ) + assert torch.allclose( + loaded_tensor._local_tensor, full_value._local_tensor, atol=1e-8, rtol=1e-5 + ), f"key: {key}; {loaded_tensor} {full_value}" + else: + assert torch.allclose( + value, load_item_dict[key] + ), f"key: {key}; {value} {load_item_dict[key]}" + + +def get_global_unique_param_name(model_chunks, param): + """ + Get the global unique parameter name for a given model and parameter. + """ + param_name = None + for model in model_chunks: + for name, p in model.named_parameters(): + if p is param: + param_name = name + break + if param_name is None: + raise ValueError("Parameter not found in model chunks") + + # Get PP unique parameter name + if re.search(r"layers\.(\d+)", param_name) and "mtp" not in param_name: + tf_layer_number = -1 + for module in model.modules(): + if not isinstance(module, TransformerLayer): + continue + for p in module.parameters(): + if p is param: + tf_layer_number = module.layer_number + break + if tf_layer_number != -1: + param_name = re.sub(r"layers\.(\d+)", f"layers.{tf_layer_number - 1}", param_name) + + # Get EP unique parameter name + param_name = list(handle_experts_in_state_dict({param_name: None}).keys())[0] + + return param_name diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py index 8e5f343b73c..cd1de6a5118 100644 --- a/megatron/training/arguments.py +++ b/megatron/training/arguments.py @@ -2271,6 +2271,10 @@ def _add_training_args(parser): help="Use torch.optim.Optimizer instead of Megatron's optimizer in optimizer cpu offload mode.") group.add_argument('--overlap-cpu-optimizer-d2h-h2d', action='store_true', default=False, help='Overlap CPU optimizer step, gradients D2H and updated parameters H2D.') + group.add_argument('--dump-param-to-param-group-map', type=str, default=None, + help="Path to a file containing parameter-to-parameter-group mapping. " + "Provide a JSON file that specifies which parameters belong to which " + "parameter group for global coordination.") group.add_argument('--no-pin-cpu-grads', action='store_false', dest='pin_cpu_grads', help='Disable pinning of CPU memory for gradients.') group.add_argument('--no-pin-cpu-params', action='store_false', dest='pin_cpu_params', diff --git a/megatron/training/checkpointing.py b/megatron/training/checkpointing.py index 71b9cd97021..93c23255f4c 100644 --- a/megatron/training/checkpointing.py +++ b/megatron/training/checkpointing.py @@ -42,9 +42,10 @@ try: from megatron.core.distributed.fsdp.src.megatron_fsdp.uneven_dtensor import preprocess_state_dict_for_uneven_dtensor from megatron.core.transformer.fsdp_dtensor_checkpoint import ( + print_diff_in_state_dicts, handle_fp8_extra_state_case, handle_swiglu_in_state_dict, - print_diff_in_state_dicts, + handle_experts_in_state_dict, ) HAVE_MEGATRON_FSDP = True except ImportError: @@ -561,6 +562,9 @@ def save_checkpoint(iteration, model, optimizer, opt_param_scheduler, num_floati # TODO Handle non-empty directories (e.g., after a crash during saving). ensure_directory_exists(checkpoint_name, check_parent=False) + if ckpt_format == "fsdp_dtensor": + state_dict = preprocess_fsdp_dtensor_state_dict(args, state_dict, model[0]) + fs_storage_writer = torch.distributed.checkpoint.FileSystemWriter(checkpoint_name) torch.distributed.checkpoint.save( state_dict=state_dict, @@ -784,9 +788,17 @@ def maybe_save_dataloader_state(train_iterator, iteration, dataloader_save_path) torch.save(dataloader_save_dict, data_state_save_path) -def generate_state_dict(args, model, optimizer, opt_param_scheduler, - rng_state, iteration=None, - optim_sd_kwargs=None, model_sd_kwargs=None, rerun_state=None): +def generate_state_dict( + args, + model, + optimizer, + opt_param_scheduler, + rng_state, + iteration=None, + optim_sd_kwargs=None, + model_sd_kwargs=None, + rerun_state=None, +): """Generate a state dict from given model, optimizer, scheduler, rng state and others. """ # Arguments, iteration, and model. @@ -839,16 +851,27 @@ def generate_state_dict(args, model, optimizer, opt_param_scheduler, if not args.no_save_rng and rng_state: state_dict["rng_state"] = rng_state - # fsdp_dtensor ckpt specific state dict preprocessing - if args.ckpt_format == "fsdp_dtensor": - assert HAVE_MEGATRON_FSDP, "Megatron FSDP is enabled but Megatron-FSDP is not available." - assert len(model) == 1, "FSDP DTensor checkpoints are not supported for multiple models." - if args.swiglu: - state_dict = state_dict.copy() - handle_swiglu_in_state_dict( - model[0], state_dict["model"], state_dict["optimizer"]) - handle_fp8_extra_state_case(state_dict["model"]) - preprocess_state_dict_for_uneven_dtensor(state_dict) + return state_dict + + +def preprocess_fsdp_dtensor_state_dict(args, raw_state_dict, model): + state_dict = raw_state_dict.copy() + handle_fp8_extra_state_case(state_dict["model"]) + if args.swiglu: + if "optimizer" in state_dict: + model_state_dict, optimizer_state_dict = handle_swiglu_in_state_dict( + model, state_dict["model"], state_dict["optimizer"] + ) + state_dict["model"] = model_state_dict + state_dict["optimizer"] = optimizer_state_dict + else: + model_state_dict, _ = handle_swiglu_in_state_dict( + model, state_dict["model"], None + ) + state_dict["model"] = model_state_dict + if args.num_experts: + state_dict["model"] = handle_experts_in_state_dict(state_dict["model"]) + preprocess_state_dict_for_uneven_dtensor(state_dict) return state_dict @@ -1169,6 +1192,12 @@ def _load_base_checkpoint( if rank0: return {}, checkpoint_name, release, CheckpointType.FSDP_DTENSOR + state_dict = sharded_state_dict + raw_optimizer_state_dict = state_dict["optimizer"].copy() if "optimizer" in state_dict else None + raw_model_state_dict = state_dict["model"].copy() if "model" in state_dict else None + model = state_dict.pop("_model") + state_dict = preprocess_fsdp_dtensor_state_dict(args, state_dict, model[0]) + ckpt_type = CheckpointType.FSDP_DTENSOR fs_storage_reader = torch.distributed.checkpoint.FileSystemReader(checkpoint_name) allow_partial_load = not getattr(args, 'strict_fsdp_dtensor_load', False) @@ -1177,15 +1206,20 @@ def _load_base_checkpoint( rank = torch.distributed.get_rank() import time as _time _time.sleep(rank * 0.001) # Make that logs of different ranks do not overlap - print_diff_in_state_dicts(state_dict_metadata, sharded_state_dict) + print_diff_in_state_dicts(state_dict_metadata, state_dict) planner = default_planner.DefaultLoadPlanner(allow_partial_load=allow_partial_load) torch.distributed.checkpoint.load_state_dict( - state_dict=sharded_state_dict, + state_dict=state_dict, storage_reader=fs_storage_reader, planner=planner, ) - state_dict = sharded_state_dict + + if raw_optimizer_state_dict is not None: + state_dict["optimizer"] = raw_optimizer_state_dict + + if raw_model_state_dict is not None: + state_dict["model"] = raw_model_state_dict else: raise NotImplementedError(f"checkpoint format {ckpt_format} not supported") @@ -1520,7 +1554,7 @@ def load_checkpoint(ddp_model, optimizer, opt_param_scheduler, load_arg='load', except FileNotFoundError: state_dict_metadata = {} - gen_sd_rerun_state = None + gen_sd_rerun_state = {} gen_sd_opt_param_scheduler = None gen_sd_rng_state = None gen_sd_optim = None @@ -1537,7 +1571,7 @@ def load_checkpoint(ddp_model, optimizer, opt_param_scheduler, load_arg='load', optim_sd_kwargs = dict(metadata=_build_sharded_state_dict_metadata(args), is_loading=True) - load_kwargs["sharded_state_dict"] = generate_state_dict( + state_dict = generate_state_dict( args, model=model, optimizer=gen_sd_optim, @@ -1547,6 +1581,8 @@ def load_checkpoint(ddp_model, optimizer, opt_param_scheduler, load_arg='load', rerun_state=gen_sd_rerun_state, iteration=1, ) + state_dict["_model"] = model + load_kwargs["sharded_state_dict"] = state_dict state_dict, checkpoint_name, release, ckpt_type = _load_base_checkpoint( load_dir, args, rank0=False, checkpointing_context=checkpointing_context, diff --git a/megatron/training/training.py b/megatron/training/training.py index f805dab0f15..bda9e42dc82 100644 --- a/megatron/training/training.py +++ b/megatron/training/training.py @@ -1210,6 +1210,7 @@ def setup_model_and_optimizer( # If the user is asking for a non-zero embedding init std, skip weight decay for embeddings # to avoid embeddings from shrinking to zero as recommended in https://arxiv.org/abs/2312.16903 default_skip_embedding_weight_decay=args.embedding_init_method_std is not None, + dump_param_to_param_group_map=args.dump_param_to_param_group_map, ) else: optimizer = get_megatron_muon_optimizer( diff --git a/tests/functional_tests/test_cases/moe/gpt3_mcore_te_tp2_pp2_ep4_etp1_memory_speed/golden_values_dev_dgxh100_coreweave.json b/tests/functional_tests/test_cases/moe/gpt3_mcore_te_tp2_pp2_ep4_etp1_memory_speed/golden_values_dev_dgxh100_coreweave.json index 0f2637a9511..717ae3f5fa6 100644 --- a/tests/functional_tests/test_cases/moe/gpt3_mcore_te_tp2_pp2_ep4_etp1_memory_speed/golden_values_dev_dgxh100_coreweave.json +++ b/tests/functional_tests/test_cases/moe/gpt3_mcore_te_tp2_pp2_ep4_etp1_memory_speed/golden_values_dev_dgxh100_coreweave.json @@ -4,56 +4,56 @@ "end_step": 50, "step_interval": 1, "values": { - "1": 11.04748, - "2": 11.03561, - "3": 9.58774, - "4": 9.25819, - "5": 9.53583, - "6": 9.8804, - "7": 9.48247, - "8": 8.93575, - "9": 8.65813, - "10": 9.0567, - "11": 8.49445, - "12": 8.52444, - "13": 8.45239, - "14": 7.97323, - "15": 8.0476, - "16": 8.07971, - "17": 8.09081, - "18": 7.76437, - "19": 8.14892, - "20": 7.89868, - "21": 7.59371, - "22": 7.54743, - "23": 7.43222, - "24": 7.4302, - "25": 7.67579, - "26": 7.06929, - "27": 7.62041, - "28": 7.32495, - "29": 7.49042, - "30": 7.64391, - "31": 7.39435, - "32": 7.58789, - "33": 7.64037, - "34": 7.69778, - "35": 7.20998, - "36": 7.08538, - "37": 7.42584, - "38": 7.18804, - "39": 7.55054, - "40": 7.54446, - "41": 7.49287, - "42": 7.24937, - "43": 7.23587, - "44": 7.41595, - "45": 7.18755, - "46": 6.89949, - "47": 7.29966, - "48": 7.14134, - "49": 7.58963, - "50": 7.03602 + "1": 11.04722, + "2": 11.03572, + "3": 9.58802, + "4": 9.25807, + "5": 9.46595, + "6": 9.99646, + "7": 9.50952, + "8": 8.97596, + "9": 8.64768, + "10": 9.40103, + "11": 8.86556, + "12": 8.63563, + "13": 8.52125, + "14": 8.08824, + "15": 8.1958, + "16": 8.22112, + "17": 8.14098, + "18": 7.8386, + "19": 8.23438, + "20": 7.95361, + "21": 7.62549, + "22": 7.60352, + "23": 7.47957, + "24": 7.46573, + "25": 7.70343, + "26": 7.10719, + "27": 7.64313, + "28": 7.34582, + "29": 7.5169, + "30": 7.67511, + "31": 7.41799, + "32": 7.61213, + "33": 7.66582, + "34": 7.73101, + "35": 7.23081, + "36": 7.10765, + "37": 7.4476, + "38": 7.21053, + "39": 7.57508, + "40": 7.5662, + "41": 7.51605, + "42": 7.27243, + "43": 7.25706, + "44": 7.44, + "45": 7.21244, + "46": 6.92421, + "47": 7.32604, + "48": 7.17147, + "49": 7.62154, + "50": 7.0624 } }, "num-zeros": { @@ -62,55 +62,55 @@ "step_interval": 1, "values": { "1": 38802612.0, - "2": 38543592.0, - "3": 38739528.0, - "4": 279937824.0, - "5": 259189728.0, - "6": 271446400.0, - "7": 604773504.0, - "8": 768892544.0, - "9": 645824128.0, - "10": 744257088.0, - "11": 718888576.0, - "12": 746732544.0, - "13": 871990976.0, - "14": 821645632.0, - "15": 724250816.0, - "16": 932241472.0, - "17": 648958912.0, - "18": 649120000.0, - "19": 925992960.0, - "20": 989207936.0, - "21": 819324096.0, - "22": 736955072.0, - "23": 910497792.0, - "24": 876716672.0, - "25": 843170688.0, - "26": 809573824.0, - "27": 854086912.0, - "28": 802857664.0, - "29": 805523328.0, - "30": 775645184.0, - "31": 771754624.0, - "32": 749733696.0, - "33": 718385216.0, - "34": 724771200.0, - "35": 737655104.0, - "36": 690419968.0, - "37": 673203456.0, - "38": 627239552.0, - "39": 614047168.0, - "40": 607288512.0, - "41": 582590592.0, - "42": 548211200.0, - "43": 532740640.0, - "44": 554239168.0, - "45": 514790528.0, - "46": 350258560.0, - "47": 472420128.0, - "48": 453788736.0, - "49": 440597216.0, - "50": 303063296.0 + "2": 38543656.0, + "3": 38739356.0, + "4": 273649600.0, + "5": 252887040.0, + "6": 255692384.0, + "7": 598483264.0, + "8": 787737984.0, + "9": 696133120.0, + "10": 505146368.0, + "11": 718888640.0, + "12": 872597184.0, + "13": 947495104.0, + "14": 1076398976.0, + "15": 856390592.0, + "16": 1048635648.0, + "17": 831370688.0, + "18": 963679552.0, + "19": 970018240.0, + "20": 935737344.0, + "21": 904189312.0, + "22": 887937280.0, + "23": 894777856.0, + "24": 703744192.0, + "25": 909232512.0, + "26": 875633216.0, + "27": 894981376.0, + "28": 919242816.0, + "29": 931351552.0, + "30": 929784768.0, + "31": 941621376.0, + "32": 885000768.0, + "33": 828484096.0, + "34": 822284800.0, + "35": 832032128.0, + "36": 787939392.0, + "37": 770719808.0, + "38": 561204672.0, + "39": 617201536.0, + "40": 695374592.0, + "41": 698978816.0, + "42": 692913728.0, + "43": 668003776.0, + "44": 673780992.0, + "45": 631182912.0, + "46": 444613312.0, + "47": 591957824.0, + "48": 617363968.0, + "49": 585295808.0, + "50": 570423872.0 } }, "mem-allocated-bytes": { @@ -118,56 +118,56 @@ "end_step": 50, "step_interval": 1, "values": { - "1": 6637267456.0, - "2": 6637269504.0, - "3": 6637269504.0, - "4": 6637269504.0, - "5": 6637269504.0, - "6": 6637269504.0, - "7": 6637269504.0, - "8": 6637269504.0, - "9": 6637269504.0, - "10": 6637269504.0, - "11": 6637269504.0, - "12": 6637269504.0, - "13": 6637269504.0, - "14": 6637269504.0, - "15": 6637269504.0, - "16": 6637269504.0, - "17": 6637269504.0, - "18": 6637269504.0, - "19": 6637269504.0, - "20": 6637269504.0, - "21": 6637269504.0, - "22": 6637269504.0, - "23": 6637269504.0, - "24": 6637269504.0, - "25": 6637269504.0, - "26": 6637269504.0, - "27": 6637269504.0, - "28": 6637269504.0, - "29": 6637269504.0, - "30": 6637269504.0, - "31": 6637269504.0, - "32": 6637269504.0, - "33": 6637269504.0, - "34": 6637269504.0, - "35": 6637269504.0, - "36": 6637269504.0, - "37": 6637269504.0, - "38": 6637269504.0, - "39": 6637269504.0, - "40": 6637269504.0, - "41": 6637269504.0, - "42": 6637269504.0, - "43": 6637269504.0, - "44": 6637269504.0, - "45": 6637269504.0, - "46": 6637269504.0, - "47": 6637269504.0, - "48": 6637269504.0, - "49": 6637269504.0, - "50": 6637269504.0 + "1": 6637272576.0, + "2": 6637274624.0, + "3": 6637274624.0, + "4": 6637274624.0, + "5": 6637274624.0, + "6": 6637274624.0, + "7": 6637274624.0, + "8": 6637274624.0, + "9": 6637274624.0, + "10": 6637274624.0, + "11": 6637274624.0, + "12": 6637274624.0, + "13": 6637274624.0, + "14": 6637274624.0, + "15": 6637274624.0, + "16": 6637274624.0, + "17": 6637274624.0, + "18": 6637274624.0, + "19": 6637274624.0, + "20": 6637274624.0, + "21": 6637274624.0, + "22": 6637274624.0, + "23": 6637274624.0, + "24": 6637274624.0, + "25": 6637274624.0, + "26": 6637274624.0, + "27": 6637274624.0, + "28": 6637274624.0, + "29": 6637274624.0, + "30": 6637274624.0, + "31": 6637274624.0, + "32": 6637274624.0, + "33": 6637274624.0, + "34": 6637274624.0, + "35": 6637274624.0, + "36": 6637274624.0, + "37": 6637274624.0, + "38": 6637274624.0, + "39": 6637274624.0, + "40": 6637274624.0, + "41": 6637274624.0, + "42": 6637274624.0, + "43": 6637274624.0, + "44": 6637274624.0, + "45": 6637274624.0, + "46": 6637274624.0, + "47": 6637274624.0, + "48": 6637274624.0, + "49": 6637274624.0, + "50": 6637274624.0 } }, "mem-max-allocated-bytes": { @@ -175,56 +175,56 @@ "end_step": 50, "step_interval": 1, "values": { - "1": 55055331328.0, - "2": 57809321984.0, - "3": 57918455808.0, - "4": 57918455808.0, - "5": 57918455808.0, - "6": 57918455808.0, - "7": 57918455808.0, - "8": 57918455808.0, - "9": 57918455808.0, - "10": 57918455808.0, - "11": 57918455808.0, - "12": 57918455808.0, - "13": 57931390976.0, - "14": 57931390976.0, - "15": 57931390976.0, - "16": 57931390976.0, - "17": 57931390976.0, - "18": 57931390976.0, - "19": 57931390976.0, - "20": 57931390976.0, - "21": 57931390976.0, - "22": 57931390976.0, - "23": 57931390976.0, - "24": 57931390976.0, - "25": 57931390976.0, - "26": 57931390976.0, - "27": 57931390976.0, - "28": 57931390976.0, - "29": 57931390976.0, - "30": 57931390976.0, - "31": 57931390976.0, - "32": 58003226624.0, - "33": 58003226624.0, - "34": 58003226624.0, - "35": 58003226624.0, - "36": 58003226624.0, - "37": 58003226624.0, - "38": 58003226624.0, - "39": 58003226624.0, - "40": 58003226624.0, - "41": 58003226624.0, - "42": 58003226624.0, - "43": 58003226624.0, - "44": 58183614464.0, - "45": 58234208256.0, - "46": 58555555840.0, - "47": 58555555840.0, - "48": 58555555840.0, - "49": 58555555840.0, - "50": 58780934144.0 + "1": 55056003072.0, + "2": 57810763776.0, + "3": 57920647168.0, + "4": 57920647168.0, + "5": 57920647168.0, + "6": 57920647168.0, + "7": 57920647168.0, + "8": 57920647168.0, + "9": 57920647168.0, + "10": 57920647168.0, + "11": 57920647168.0, + "12": 57920647168.0, + "13": 57920647168.0, + "14": 57920647168.0, + "15": 57920647168.0, + "16": 57920647168.0, + "17": 57920647168.0, + "18": 57920647168.0, + "19": 57920647168.0, + "20": 57920647168.0, + "21": 57920647168.0, + "22": 57920647168.0, + "23": 57920647168.0, + "24": 57920647168.0, + "25": 57920647168.0, + "26": 57920647168.0, + "27": 57920647168.0, + "28": 57920647168.0, + "29": 57920647168.0, + "30": 57920647168.0, + "31": 57920647168.0, + "32": 57920647168.0, + "33": 57920647168.0, + "34": 57961472000.0, + "35": 57961472000.0, + "36": 57961472000.0, + "37": 57961472000.0, + "38": 57961472000.0, + "39": 57961472000.0, + "40": 57961472000.0, + "41": 57961472000.0, + "42": 57961472000.0, + "43": 57961472000.0, + "44": 57961472000.0, + "45": 57961472000.0, + "46": 57961472000.0, + "47": 57961472000.0, + "48": 57961472000.0, + "49": 57961472000.0, + "50": 57961472000.0 } }, "mtp_1 loss": { @@ -232,56 +232,56 @@ "end_step": 50, "step_interval": 1, "values": { - "1": 11.07654, - "2": 11.07406, - "3": 10.53881, - "4": 10.09803, - "5": 9.81154, - "6": 10.06236, - "7": 9.79762, - "8": 9.07117, - "9": 8.87049, - "10": 9.127, - "11": 8.49853, - "12": 8.53046, - "13": 8.42444, - "14": 7.847, - "15": 7.99077, - "16": 8.05015, - "17": 8.00064, - "18": 7.73104, - "19": 8.11087, - "20": 7.82933, - "21": 7.52501, - "22": 7.49916, - "23": 7.36982, - "24": 7.37235, - "25": 7.61578, - "26": 7.02029, - "27": 7.56014, - "28": 7.2681, - "29": 7.44399, - "30": 7.58618, - "31": 7.32468, - "32": 7.50596, - "33": 7.5715, - "34": 7.63581, - "35": 7.15224, - "36": 7.01784, - "37": 7.35163, - "38": 7.12551, - "39": 7.48656, - "40": 7.47408, - "41": 7.42096, - "42": 7.17595, - "43": 7.16059, - "44": 7.34289, - "45": 7.11969, - "46": 6.82753, - "47": 7.23525, - "48": 7.08042, - "49": 7.51043, - "50": 6.9735 + "1": 11.07648, + "2": 11.07404, + "3": 10.53854, + "4": 10.09813, + "5": 9.81166, + "6": 10.09741, + "7": 9.79481, + "8": 9.0642, + "9": 8.86016, + "10": 9.34039, + "11": 8.51318, + "12": 8.59467, + "13": 8.5292, + "14": 7.95757, + "15": 8.06962, + "16": 8.11802, + "17": 8.06993, + "18": 7.80587, + "19": 8.19192, + "20": 7.8906, + "21": 7.57063, + "22": 7.55091, + "23": 7.41606, + "24": 7.42454, + "25": 7.65274, + "26": 7.05583, + "27": 7.59747, + "28": 7.29984, + "29": 7.472, + "30": 7.61908, + "31": 7.35179, + "32": 7.52979, + "33": 7.59161, + "34": 7.66287, + "35": 7.17383, + "36": 7.04133, + "37": 7.37081, + "38": 7.1443, + "39": 7.50879, + "40": 7.48921, + "41": 7.43802, + "42": 7.19405, + "43": 7.17581, + "44": 7.35785, + "45": 7.13985, + "46": 6.84014, + "47": 7.25094, + "48": 7.09407, + "49": 7.52321, + "50": 6.98987 } }, "iteration-time": { @@ -289,56 +289,56 @@ "end_step": 50, "step_interval": 1, "values": { - "1": 69.29797, - "2": 1.7261, - "3": 1.40981, - "4": 2.16562, - "5": 1.7862, - "6": 1.7469, - "7": 1.96688, - "8": 1.97301, - "9": 1.74665, - "10": 1.69613, - "11": 1.02979, - "12": 1.02408, - "13": 1.03261, - "14": 1.02432, - "15": 1.0529, - "16": 1.04491, - "17": 1.03693, - "18": 1.03399, - "19": 1.03627, - "20": 1.02284, - "21": 1.01667, - "22": 1.02932, - "23": 1.03591, - "24": 1.03466, - "25": 1.03149, - "26": 1.03165, - "27": 1.02342, - "28": 1.03777, - "29": 1.04061, - "30": 1.05641, - "31": 1.02382, - "32": 1.01775, - "33": 1.03039, - "34": 1.03693, - "35": 1.03153, - "36": 1.02699, - "37": 1.02756, - "38": 1.02919, - "39": 1.01773, - "40": 1.03491, - "41": 1.03152, - "42": 1.03035, - "43": 1.0221, - "44": 1.05201, - "45": 1.02579, - "46": 1.02798, - "47": 1.03857, - "48": 1.02772, - "49": 1.0408, - "50": 1.03745 + "1": 93.39829, + "2": 1.82958, + "3": 1.3241, + "4": 2.19661, + "5": 2.13156, + "6": 1.75452, + "7": 2.08539, + "8": 1.58016, + "9": 1.60816, + "10": 1.03407, + "11": 1.01797, + "12": 1.0168, + "13": 1.01666, + "14": 1.0748, + "15": 1.04137, + "16": 1.05864, + "17": 1.05961, + "18": 1.03233, + "19": 1.02728, + "20": 1.02917, + "21": 1.04313, + "22": 1.03054, + "23": 1.0313, + "24": 1.03789, + "25": 1.04414, + "26": 1.05561, + "27": 1.03361, + "28": 1.03142, + "29": 1.02437, + "30": 1.02195, + "31": 1.0172, + "32": 1.03318, + "33": 1.03742, + "34": 1.03628, + "35": 1.03575, + "36": 1.05127, + "37": 1.03273, + "38": 1.03381, + "39": 1.02923, + "40": 1.02986, + "41": 1.03249, + "42": 1.033, + "43": 1.03169, + "44": 1.03818, + "45": 1.02736, + "46": 1.02698, + "47": 1.03158, + "48": 1.02471, + "49": 1.03674, + "50": 1.0291 } } } \ No newline at end of file diff --git a/tests/functional_tests/test_cases/moe/gpt3_mcore_te_tp2_pp2_ep4_etp1_mtp_resume_torch_dist_fp8/golden_values_dev_dgxh100_coreweave.json b/tests/functional_tests/test_cases/moe/gpt3_mcore_te_tp2_pp2_ep4_etp1_mtp_resume_torch_dist_fp8/golden_values_dev_dgxh100_coreweave.json index 0af1bff480e..adec1b3bd58 100644 --- a/tests/functional_tests/test_cases/moe/gpt3_mcore_te_tp2_pp2_ep4_etp1_mtp_resume_torch_dist_fp8/golden_values_dev_dgxh100_coreweave.json +++ b/tests/functional_tests/test_cases/moe/gpt3_mcore_te_tp2_pp2_ep4_etp1_mtp_resume_torch_dist_fp8/golden_values_dev_dgxh100_coreweave.json @@ -4,56 +4,56 @@ "end_step": 50, "step_interval": 1, "values": { - "1": 11.04624, - "2": 11.03476, - "3": 9.59903, - "4": 9.26301, - "5": 9.36373, - "6": 9.59608, - "7": 9.45214, - "8": 8.95198, - "9": 8.65952, - "10": 9.17778, - "11": 9.21306, - "12": 8.68184, - "13": 8.6038, - "14": 8.01576, - "15": 8.13595, - "16": 8.20124, - "17": 8.13602, - "18": 7.83369, - "19": 8.22974, - "20": 7.9452, - "21": 7.62338, - "22": 7.60791, - "23": 7.48374, - "24": 7.46559, - "25": 7.71274, - "26": 7.12081, - "27": 7.64626, - "28": 7.35234, - "29": 7.52084, - "30": 7.67784, - "31": 7.42246, - "32": 7.6137, - "33": 7.66159, - "34": 7.72817, - "35": 7.23134, - "36": 7.10612, - "37": 7.44953, - "38": 7.20946, - "39": 7.57073, - "40": 7.56124, - "41": 7.51119, - "42": 7.27048, - "43": 7.25633, - "44": 7.43634, - "45": 7.21132, - "46": 6.91913, - "47": 7.32211, - "48": 7.16551, - "49": 7.6155, - "50": 7.05648 + "1": 11.04577, + "2": 11.03578, + "3": 9.5968, + "4": 9.26068, + "5": 9.09365, + "6": 8.97825, + "7": 9.18096, + "8": 8.70673, + "9": 8.55632, + "10": 8.85377, + "11": 8.31245, + "12": 8.35862, + "13": 8.28114, + "14": 7.73951, + "15": 7.91242, + "16": 7.94944, + "17": 7.89918, + "18": 7.64375, + "19": 8.02647, + "20": 7.73813, + "21": 7.44557, + "22": 7.43367, + "23": 7.31291, + "24": 7.30268, + "25": 7.57549, + "26": 6.98093, + "27": 7.50005, + "28": 7.241, + "29": 7.40369, + "30": 7.51839, + "31": 7.29514, + "32": 7.47818, + "33": 7.52568, + "34": 7.57647, + "35": 7.12091, + "36": 6.97439, + "37": 7.30929, + "38": 7.09349, + "39": 7.43659, + "40": 7.45122, + "41": 7.37904, + "42": 7.14627, + "43": 7.13408, + "44": 7.30886, + "45": 7.08523, + "46": 6.8067, + "47": 7.21159, + "48": 7.0245, + "49": 7.50096, + "50": 6.92687 } }, "num-zeros": { @@ -61,56 +61,56 @@ "end_step": 50, "step_interval": 1, "values": { - "1": 38802568, - "2": 38543544, - "3": 41886704, - "4": 264367872, - "5": 224737792, - "6": 302994528, - "7": 645808768, - "8": 775291136, - "9": 765475328, - "10": 675259904, - "11": 615098624, - "12": 702764352, - "13": 934951360, - "14": 1060699008, - "15": 802967296, - "16": 1026771392, - "17": 756706880, - "18": 715253696, - "19": 929126208, - "20": 875969472, - "21": 665188032, - "22": 903854976, - "23": 747044352, - "24": 920777856, - "25": 733230528, - "26": 863183104, - "27": 879318336, - "28": 916219136, - "29": 909384256, - "30": 879622720, - "31": 866425152, - "32": 819074560, - "33": 589493056, - "34": 772011648, - "35": 778655488, - "36": 759651584, - "37": 761302144, - "38": 463804224, - "39": 543038400, - "40": 497278720, - "41": 658241792, - "42": 661600512, - "43": 495713632, - "44": 673788672, - "45": 470873536, - "46": 614455040, - "47": 554219584, - "48": 570200064, - "49": 557109312, - "50": 347212736 + "1": 38802664.0, + "2": 38543552.0, + "3": 38740472.0, + "4": 273766176.0, + "5": 196515488.0, + "6": 432153600.0, + "7": 715038528.0, + "8": 797328960.0, + "9": 696279488.0, + "10": 668928192.0, + "11": 583742720.0, + "12": 595799040.0, + "13": 695916288.0, + "14": 617245056.0, + "15": 629936832.0, + "16": 639940800.0, + "17": 642766016.0, + "18": 664898112.0, + "19": 671247104.0, + "20": 602545216.0, + "21": 542607872.0, + "22": 551419008.0, + "23": 533094816.0, + "24": 527647904.0, + "25": 570717824.0, + "26": 510874176.0, + "27": 498748096.0, + "28": 510353632.0, + "29": 506802112.0, + "30": 486336928.0, + "31": 410143360.0, + "32": 372280800.0, + "33": 369351776.0, + "34": 353666688.0, + "35": 344549376.0, + "36": 278456576.0, + "37": 289517152.0, + "38": 274950816.0, + "39": 242921776.0, + "40": 223597264.0, + "41": 186386944.0, + "42": 180387488.0, + "43": 224573440.0, + "44": 217714800.0, + "45": 143723568.0, + "46": 161525888.0, + "47": 120124336.0, + "48": 183368272.0, + "49": 154411968.0, + "50": 167778288.0 } }, "mem-allocated-bytes": { @@ -118,56 +118,56 @@ "end_step": 50, "step_interval": 1, "values": { - "1": 7321308672, - "2": 7321310720, - "3": 7321310720, - "4": 7321310720, - "5": 7321310720, - "6": 7321310720, - "7": 7321310720, - "8": 7321310720, - "9": 7321310720, - "10": 7321310720, - "11": 7321310720, - "12": 7321310720, - "13": 7321310720, - "14": 7321310720, - "15": 7321310720, - "16": 7321310720, - "17": 7321310720, - "18": 7321310720, - "19": 7321310720, - "20": 7321310720, - "21": 7321310720, - "22": 7321310720, - "23": 7321310720, - "24": 7321310720, - "25": 7321310720, - "26": 7321310720, - "27": 7321310720, - "28": 7321310720, - "29": 7321310720, - "30": 7321310720, - "31": 7321310720, - "32": 7321310720, - "33": 7321310720, - "34": 7321310720, - "35": 7321310720, - "36": 7321310720, - "37": 7321310720, - "38": 7321310720, - "39": 7321310720, - "40": 7321310720, - "41": 7321310720, - "42": 7321310720, - "43": 7321310720, - "44": 7321310720, - "45": 7321310720, - "46": 7321310720, - "47": 7321310720, - "48": 7321310720, - "49": 7321310720, - "50": 7321310720 + "1": 7321336320.0, + "2": 7321338368.0, + "3": 7321338368.0, + "4": 7321338368.0, + "5": 7321338368.0, + "6": 7321338368.0, + "7": 7321338368.0, + "8": 7321338368.0, + "9": 7321338368.0, + "10": 7321338368.0, + "11": 7321338368.0, + "12": 7321338368.0, + "13": 7321338368.0, + "14": 7321338368.0, + "15": 7321338368.0, + "16": 7321338368.0, + "17": 7321338368.0, + "18": 7321338368.0, + "19": 7321338368.0, + "20": 7321338368.0, + "21": 7321338368.0, + "22": 7321338368.0, + "23": 7321338368.0, + "24": 7321338368.0, + "25": 7321338368.0, + "26": 7321338368.0, + "27": 7321338368.0, + "28": 7321338368.0, + "29": 7321338368.0, + "30": 7321338368.0, + "31": 7321338368.0, + "32": 7321338368.0, + "33": 7321338368.0, + "34": 7321338368.0, + "35": 7321338368.0, + "36": 7321338368.0, + "37": 7321338368.0, + "38": 7321338368.0, + "39": 7321338368.0, + "40": 7321338368.0, + "41": 7321338368.0, + "42": 7321338368.0, + "43": 7321338368.0, + "44": 7321338368.0, + "45": 7321338368.0, + "46": 7321338368.0, + "47": 7321338368.0, + "48": 7321338368.0, + "49": 7321338368.0, + "50": 7321338368.0 } }, "mem-max-allocated-bytes": { @@ -175,56 +175,56 @@ "end_step": 50, "step_interval": 1, "values": { - "1": 54396813312, - "2": 57149165568, - "3": 57165475840, - "4": 57165475840, - "5": 57165475840, - "6": 57165475840, - "7": 57165475840, - "8": 57165475840, - "9": 57165475840, - "10": 57165475840, - "11": 57165475840, - "12": 57165475840, - "13": 57165475840, - "14": 57165475840, - "15": 57165475840, - "16": 57165475840, - "17": 57165475840, - "18": 57165475840, - "19": 57165475840, - "20": 57165475840, - "21": 57165475840, - "22": 57165475840, - "23": 57165475840, - "24": 57165475840, - "25": 57165475840, - "26": 57165475840, - "27": 57165475840, - "28": 57165475840, - "29": 57165475840, - "30": 57165475840, - "31": 57165475840, - "32": 57165475840, - "33": 57165475840, - "34": 57165475840, - "35": 57165475840, - "36": 57165475840, - "37": 57165475840, - "38": 57165475840, - "39": 57165475840, - "40": 57295986688, - "41": 57295986688, - "42": 57331482624, - "43": 57360437248, - "44": 57561960448, - "45": 57561960448, - "46": 57561960448, - "47": 57585307648, - "48": 57602347008, - "49": 57823961088, - "50": 57823961088 + "1": 54402162688.0, + "2": 57150373888.0, + "3": 57150373888.0, + "4": 57150373888.0, + "5": 57150373888.0, + "6": 57150373888.0, + "7": 57150373888.0, + "8": 57150373888.0, + "9": 57150373888.0, + "10": 57150373888.0, + "11": 57150373888.0, + "12": 57150373888.0, + "13": 57150373888.0, + "14": 57150373888.0, + "15": 57150373888.0, + "16": 57150373888.0, + "17": 57150373888.0, + "18": 57150373888.0, + "19": 57150373888.0, + "20": 57150373888.0, + "21": 57150373888.0, + "22": 57150373888.0, + "23": 57150373888.0, + "24": 57150373888.0, + "25": 57150373888.0, + "26": 57150373888.0, + "27": 57150373888.0, + "28": 57150373888.0, + "29": 57150373888.0, + "30": 57150373888.0, + "31": 57150373888.0, + "32": 57150373888.0, + "33": 57150373888.0, + "34": 57150373888.0, + "35": 57152438272.0, + "36": 57344114688.0, + "37": 57344114688.0, + "38": 57449279488.0, + "39": 57449279488.0, + "40": 57449279488.0, + "41": 57449279488.0, + "42": 57449279488.0, + "43": 57449279488.0, + "44": 57449279488.0, + "45": 57470353408.0, + "46": 57470353408.0, + "47": 57470353408.0, + "48": 57470353408.0, + "49": 57470353408.0, + "50": 57470353408.0 } }, "mtp_1 loss": { @@ -232,56 +232,56 @@ "end_step": 50, "step_interval": 1, "values": { - "1": 11.07779, - "2": 11.07564, - "3": 10.52904, - "4": 10.08924, - "5": 9.81101, - "6": 9.88786, - "7": 9.72987, - "8": 9.02044, - "9": 8.8145, - "10": 9.09362, - "11": 8.77612, - "12": 8.56714, - "13": 8.54777, - "14": 8.04338, - "15": 8.10946, - "16": 8.13231, - "17": 8.0853, - "18": 7.83475, - "19": 8.21923, - "20": 7.91097, - "21": 7.58489, - "22": 7.56231, - "23": 7.44204, - "24": 7.44303, - "25": 7.67594, - "26": 7.07138, - "27": 7.60696, - "28": 7.30925, - "29": 7.48219, - "30": 7.62699, - "31": 7.3655, - "32": 7.54203, - "33": 7.60199, - "34": 7.66716, - "35": 7.18385, - "36": 7.05252, - "37": 7.38377, - "38": 7.15521, - "39": 7.51639, - "40": 7.4929, - "41": 7.44762, - "42": 7.20298, - "43": 7.18681, - "44": 7.36683, - "45": 7.15506, - "46": 6.85064, - "47": 7.26072, - "48": 7.10489, - "49": 7.53477, - "50": 6.99715 + "1": 11.07769, + "2": 11.07625, + "3": 10.52909, + "4": 10.08687, + "5": 9.82013, + "6": 9.48246, + "7": 9.54169, + "8": 8.83661, + "9": 8.64933, + "10": 8.95821, + "11": 8.32934, + "12": 8.36033, + "13": 8.26936, + "14": 7.73441, + "15": 7.87122, + "16": 7.9153, + "17": 7.86923, + "18": 7.61191, + "19": 7.99919, + "20": 7.72174, + "21": 7.4147, + "22": 7.40336, + "23": 7.27676, + "24": 7.28557, + "25": 7.53782, + "26": 6.94933, + "27": 7.48504, + "28": 7.20219, + "29": 7.38696, + "30": 7.51152, + "31": 7.26613, + "32": 7.45631, + "33": 7.51482, + "34": 7.57527, + "35": 7.10374, + "36": 6.97224, + "37": 7.31053, + "38": 7.08607, + "39": 7.44371, + "40": 7.43612, + "41": 7.37848, + "42": 7.13561, + "43": 7.11558, + "44": 7.30254, + "45": 7.08147, + "46": 6.78911, + "47": 7.21791, + "48": 7.03066, + "49": 7.46668, + "50": 6.93251 } }, "iteration-time": { @@ -289,56 +289,56 @@ "end_step": 50, "step_interval": 1, "values": { - "1": 98.46571, - "2": 1.63304, - "3": 1.32772, - "4": 1.63453, - "5": 1.11673, - "6": 1.14377, - "7": 1.33213, - "8": 1.32699, - "9": 1.07499, - "10": 1.12938, - "11": 1.07438, - "12": 1.11078, - "13": 1.06958, - "14": 1.08718, - "15": 1.10547, - "16": 1.07557, - "17": 1.08606, - "18": 1.0832, - "19": 1.08226, - "20": 1.126, - "21": 1.08645, - "22": 1.07978, - "23": 1.07859, - "24": 1.08221, - "25": 1.08192, - "26": 1.09185, - "27": 1.0923, - "28": 1.09562, - "29": 1.10486, - "30": 1.10038, - "31": 1.09094, - "32": 1.08693, - "33": 1.0883, - "34": 1.08169, - "35": 1.08611, - "36": 1.07758, - "37": 1.07933, - "38": 1.08289, - "39": 1.07885, - "40": 1.08075, - "41": 1.0781, - "42": 1.08028, - "43": 1.08035, - "44": 1.08973, - "45": 1.08944, - "46": 1.07483, - "47": 1.08306, - "48": 1.07701, - "49": 1.0768, - "50": 1.07022 + "1": 92.7075, + "2": 1.62502, + "3": 1.31213, + "4": 1.71707, + "5": 1.11852, + "6": 1.39151, + "7": 1.37049, + "8": 1.22293, + "9": 1.10694, + "10": 1.11053, + "11": 1.10169, + "12": 1.14642, + "13": 1.11639, + "14": 1.12927, + "15": 1.12868, + "16": 1.11899, + "17": 1.10545, + "18": 1.11542, + "19": 1.11417, + "20": 1.11349, + "21": 1.11071, + "22": 1.11032, + "23": 1.11836, + "24": 1.11402, + "25": 1.11546, + "26": 1.10471, + "27": 1.10368, + "28": 1.09929, + "29": 1.10324, + "30": 1.10507, + "31": 1.10255, + "32": 1.10727, + "33": 1.1043, + "34": 1.10476, + "35": 1.10252, + "36": 1.10053, + "37": 1.1068, + "38": 1.09229, + "39": 1.08165, + "40": 1.07889, + "41": 1.07583, + "42": 1.07174, + "43": 1.07738, + "44": 1.08604, + "45": 1.09529, + "46": 1.08309, + "47": 1.08896, + "48": 1.08318, + "49": 1.08597, + "50": 1.08649 } } } \ No newline at end of file diff --git a/tests/functional_tests/test_cases/moe/gpt3_mcore_te_tp2_pp2_ep4_etp1_mtp_resume_torch_dist_fp8/golden_values_dev_dgxh100_eos.json b/tests/functional_tests/test_cases/moe/gpt3_mcore_te_tp2_pp2_ep4_etp1_mtp_resume_torch_dist_fp8/golden_values_dev_dgxh100_eos.json index 585139e83c9..b7df693e1f7 100644 --- a/tests/functional_tests/test_cases/moe/gpt3_mcore_te_tp2_pp2_ep4_etp1_mtp_resume_torch_dist_fp8/golden_values_dev_dgxh100_eos.json +++ b/tests/functional_tests/test_cases/moe/gpt3_mcore_te_tp2_pp2_ep4_etp1_mtp_resume_torch_dist_fp8/golden_values_dev_dgxh100_eos.json @@ -4,56 +4,56 @@ "end_step": 50, "step_interval": 1, "values": { - "1": 11.04624, - "2": 11.03476, - "3": 9.59903, - "4": 9.26301, - "5": 9.36373, - "6": 9.59608, - "7": 9.45214, - "8": 8.95198, - "9": 8.65952, - "10": 9.17778, - "11": 9.21306, - "12": 8.68184, - "13": 8.6038, - "14": 8.01576, - "15": 8.13595, - "16": 8.20124, - "17": 8.13602, - "18": 7.83369, - "19": 8.22974, - "20": 7.9452, - "21": 7.62338, - "22": 7.60791, - "23": 7.48374, - "24": 7.46559, - "25": 7.71274, - "26": 7.12081, - "27": 7.64626, - "28": 7.35234, - "29": 7.52084, - "30": 7.67784, - "31": 7.42246, - "32": 7.6137, - "33": 7.66159, - "34": 7.72817, - "35": 7.23134, - "36": 7.10612, - "37": 7.44953, - "38": 7.20946, - "39": 7.57073, - "40": 7.56124, - "41": 7.51119, - "42": 7.27048, - "43": 7.25633, - "44": 7.43634, - "45": 7.21132, - "46": 6.91913, - "47": 7.32211, - "48": 7.16551, - "49": 7.6155, - "50": 7.05648 + "1": 11.04577, + "2": 11.03578, + "3": 9.5968, + "4": 9.26068, + "5": 9.09365, + "6": 8.97825, + "7": 9.18096, + "8": 8.70673, + "9": 8.55632, + "10": 8.85377, + "11": 8.31245, + "12": 8.35862, + "13": 8.28114, + "14": 7.73951, + "15": 7.91242, + "16": 7.94944, + "17": 7.89918, + "18": 7.64375, + "19": 8.02647, + "20": 7.73813, + "21": 7.44557, + "22": 7.43367, + "23": 7.31291, + "24": 7.30268, + "25": 7.57549, + "26": 6.98093, + "27": 7.50005, + "28": 7.241, + "29": 7.40369, + "30": 7.51839, + "31": 7.29514, + "32": 7.47818, + "33": 7.52568, + "34": 7.57647, + "35": 7.12091, + "36": 6.97439, + "37": 7.30929, + "38": 7.09349, + "39": 7.43659, + "40": 7.45122, + "41": 7.37904, + "42": 7.14627, + "43": 7.13408, + "44": 7.30886, + "45": 7.08523, + "46": 6.8067, + "47": 7.21159, + "48": 7.0245, + "49": 7.50096, + "50": 6.92687 } }, "num-zeros": { @@ -61,56 +61,56 @@ "end_step": 50, "step_interval": 1, "values": { - "1": 38802568, - "2": 38543544, - "3": 41886704, - "4": 264367872, - "5": 224737792, - "6": 302994528, - "7": 645808768, - "8": 775291136, - "9": 765475328, - "10": 675259904, - "11": 615098624, - "12": 702764352, - "13": 934951360, - "14": 1060699008, - "15": 802967296, - "16": 1026771392, - "17": 756706880, - "18": 715253696, - "19": 929126208, - "20": 875969472, - "21": 665188032, - "22": 903854976, - "23": 747044352, - "24": 920777856, - "25": 733230528, - "26": 863183104, - "27": 879318336, - "28": 916219136, - "29": 909384256, - "30": 879622720, - "31": 866425152, - "32": 819074560, - "33": 589493056, - "34": 772011648, - "35": 778655488, - "36": 759651584, - "37": 761302144, - "38": 463804224, - "39": 543038400, - "40": 497278720, - "41": 658241792, - "42": 661600512, - "43": 495713632, - "44": 673788672, - "45": 470873536, - "46": 614455040, - "47": 554219584, - "48": 570200064, - "49": 557109312, - "50": 347212736 + "1": 38802664.0, + "2": 38543552.0, + "3": 38740472.0, + "4": 273766176.0, + "5": 196515488.0, + "6": 432153600.0, + "7": 715038528.0, + "8": 797328960.0, + "9": 696279488.0, + "10": 668928192.0, + "11": 583742720.0, + "12": 595799040.0, + "13": 695916288.0, + "14": 617245056.0, + "15": 629936832.0, + "16": 639940800.0, + "17": 642766016.0, + "18": 664898112.0, + "19": 671247104.0, + "20": 602545216.0, + "21": 542607872.0, + "22": 551419008.0, + "23": 533094816.0, + "24": 527647904.0, + "25": 570717824.0, + "26": 510874176.0, + "27": 498748096.0, + "28": 510353632.0, + "29": 506802112.0, + "30": 486336928.0, + "31": 410143360.0, + "32": 372280800.0, + "33": 369351776.0, + "34": 353666688.0, + "35": 344549376.0, + "36": 278456576.0, + "37": 289517152.0, + "38": 274950816.0, + "39": 242921776.0, + "40": 223597264.0, + "41": 186386944.0, + "42": 180387488.0, + "43": 224573440.0, + "44": 217714800.0, + "45": 143723568.0, + "46": 161525888.0, + "47": 120124336.0, + "48": 183368272.0, + "49": 154411968.0, + "50": 167778288.0 } }, "mem-allocated-bytes": { @@ -118,56 +118,56 @@ "end_step": 50, "step_interval": 1, "values": { - "1": 7321308672, - "2": 7321310720, - "3": 7321310720, - "4": 7321310720, - "5": 7321310720, - "6": 7321310720, - "7": 7321310720, - "8": 7321310720, - "9": 7321310720, - "10": 7321310720, - "11": 7321310720, - "12": 7321310720, - "13": 7321310720, - "14": 7321310720, - "15": 7321310720, - "16": 7321310720, - "17": 7321310720, - "18": 7321310720, - "19": 7321310720, - "20": 7321310720, - "21": 7321310720, - "22": 7321310720, - "23": 7321310720, - "24": 7321310720, - "25": 7321310720, - "26": 7321310720, - "27": 7321310720, - "28": 7321310720, - "29": 7321310720, - "30": 7321310720, - "31": 7321310720, - "32": 7321310720, - "33": 7321310720, - "34": 7321310720, - "35": 7321310720, - "36": 7321310720, - "37": 7321310720, - "38": 7321310720, - "39": 7321310720, - "40": 7321310720, - "41": 7321310720, - "42": 7321310720, - "43": 7321310720, - "44": 7321310720, - "45": 7321310720, - "46": 7321310720, - "47": 7321310720, - "48": 7321310720, - "49": 7321310720, - "50": 7321310720 + "1": 7321336320.0, + "2": 7321338368.0, + "3": 7321338368.0, + "4": 7321338368.0, + "5": 7321338368.0, + "6": 7321338368.0, + "7": 7321338368.0, + "8": 7321338368.0, + "9": 7321338368.0, + "10": 7321338368.0, + "11": 7321338368.0, + "12": 7321338368.0, + "13": 7321338368.0, + "14": 7321338368.0, + "15": 7321338368.0, + "16": 7321338368.0, + "17": 7321338368.0, + "18": 7321338368.0, + "19": 7321338368.0, + "20": 7321338368.0, + "21": 7321338368.0, + "22": 7321338368.0, + "23": 7321338368.0, + "24": 7321338368.0, + "25": 7321338368.0, + "26": 7321338368.0, + "27": 7321338368.0, + "28": 7321338368.0, + "29": 7321338368.0, + "30": 7321338368.0, + "31": 7321338368.0, + "32": 7321338368.0, + "33": 7321338368.0, + "34": 7321338368.0, + "35": 7321338368.0, + "36": 7321338368.0, + "37": 7321338368.0, + "38": 7321338368.0, + "39": 7321338368.0, + "40": 7321338368.0, + "41": 7321338368.0, + "42": 7321338368.0, + "43": 7321338368.0, + "44": 7321338368.0, + "45": 7321338368.0, + "46": 7321338368.0, + "47": 7321338368.0, + "48": 7321338368.0, + "49": 7321338368.0, + "50": 7321338368.0 } }, "mem-max-allocated-bytes": { @@ -175,56 +175,56 @@ "end_step": 50, "step_interval": 1, "values": { - "1": 54396813312, - "2": 57149165568, - "3": 57165475840, - "4": 57165475840, - "5": 57165475840, - "6": 57165475840, - "7": 57165475840, - "8": 57165475840, - "9": 57165475840, - "10": 57165475840, - "11": 57165475840, - "12": 57165475840, - "13": 57165475840, - "14": 57165475840, - "15": 57165475840, - "16": 57165475840, - "17": 57165475840, - "18": 57165475840, - "19": 57165475840, - "20": 57165475840, - "21": 57165475840, - "22": 57165475840, - "23": 57165475840, - "24": 57165475840, - "25": 57165475840, - "26": 57165475840, - "27": 57165475840, - "28": 57165475840, - "29": 57165475840, - "30": 57165475840, - "31": 57165475840, - "32": 57165475840, - "33": 57165475840, - "34": 57165475840, - "35": 57165475840, - "36": 57165475840, - "37": 57165475840, - "38": 57165475840, - "39": 57165475840, - "40": 57295986688, - "41": 57295986688, - "42": 57331482624, - "43": 57360437248, - "44": 57561960448, - "45": 57561960448, - "46": 57561960448, - "47": 57585307648, - "48": 57602347008, - "49": 57823961088, - "50": 57823961088 + "1": 54402162688.0, + "2": 57150373888.0, + "3": 57150373888.0, + "4": 57150373888.0, + "5": 57150373888.0, + "6": 57150373888.0, + "7": 57150373888.0, + "8": 57150373888.0, + "9": 57150373888.0, + "10": 57150373888.0, + "11": 57150373888.0, + "12": 57150373888.0, + "13": 57150373888.0, + "14": 57150373888.0, + "15": 57150373888.0, + "16": 57150373888.0, + "17": 57150373888.0, + "18": 57150373888.0, + "19": 57150373888.0, + "20": 57150373888.0, + "21": 57150373888.0, + "22": 57150373888.0, + "23": 57150373888.0, + "24": 57150373888.0, + "25": 57150373888.0, + "26": 57150373888.0, + "27": 57150373888.0, + "28": 57150373888.0, + "29": 57150373888.0, + "30": 57150373888.0, + "31": 57150373888.0, + "32": 57150373888.0, + "33": 57150373888.0, + "34": 57150373888.0, + "35": 57152438272.0, + "36": 57344114688.0, + "37": 57344114688.0, + "38": 57449279488.0, + "39": 57449279488.0, + "40": 57449279488.0, + "41": 57449279488.0, + "42": 57449279488.0, + "43": 57449279488.0, + "44": 57449279488.0, + "45": 57470353408.0, + "46": 57470353408.0, + "47": 57470353408.0, + "48": 57470353408.0, + "49": 57470353408.0, + "50": 57470353408.0 } }, "mtp_1 loss": { @@ -232,56 +232,56 @@ "end_step": 50, "step_interval": 1, "values": { - "1": 11.07779, - "2": 11.07564, - "3": 10.52904, - "4": 10.08924, - "5": 9.81101, - "6": 9.88786, - "7": 9.72987, - "8": 9.02044, - "9": 8.8145, - "10": 9.09362, - "11": 8.77612, - "12": 8.56714, - "13": 8.54777, - "14": 8.04338, - "15": 8.10946, - "16": 8.13231, - "17": 8.0853, - "18": 7.83475, - "19": 8.21923, - "20": 7.91097, - "21": 7.58489, - "22": 7.56231, - "23": 7.44204, - "24": 7.44303, - "25": 7.67594, - "26": 7.07138, - "27": 7.60696, - "28": 7.30925, - "29": 7.48219, - "30": 7.62699, - "31": 7.3655, - "32": 7.54203, - "33": 7.60199, - "34": 7.66716, - "35": 7.18385, - "36": 7.05252, - "37": 7.38377, - "38": 7.15521, - "39": 7.51639, - "40": 7.4929, - "41": 7.44762, - "42": 7.20298, - "43": 7.18681, - "44": 7.36683, - "45": 7.15506, - "46": 6.85064, - "47": 7.26072, - "48": 7.10489, - "49": 7.53477, - "50": 6.99715 + "1": 11.07769, + "2": 11.07625, + "3": 10.52909, + "4": 10.08687, + "5": 9.82013, + "6": 9.48246, + "7": 9.54169, + "8": 8.83661, + "9": 8.64933, + "10": 8.95821, + "11": 8.32934, + "12": 8.36033, + "13": 8.26936, + "14": 7.73441, + "15": 7.87122, + "16": 7.9153, + "17": 7.86923, + "18": 7.61191, + "19": 7.99919, + "20": 7.72174, + "21": 7.4147, + "22": 7.40336, + "23": 7.27676, + "24": 7.28557, + "25": 7.53782, + "26": 6.94933, + "27": 7.48504, + "28": 7.20219, + "29": 7.38696, + "30": 7.51152, + "31": 7.26613, + "32": 7.45631, + "33": 7.51482, + "34": 7.57527, + "35": 7.10374, + "36": 6.97224, + "37": 7.31053, + "38": 7.08607, + "39": 7.44371, + "40": 7.43612, + "41": 7.37848, + "42": 7.13561, + "43": 7.11558, + "44": 7.30254, + "45": 7.08147, + "46": 6.78911, + "47": 7.21791, + "48": 7.03066, + "49": 7.46668, + "50": 6.93251 } }, "iteration-time": { @@ -289,56 +289,56 @@ "end_step": 50, "step_interval": 1, "values": { - "1": 89.12995, - "2": 1.33749, - "3": 1.24205, - "4": 1.63759, - "5": 1.13139, - "6": 1.12938, - "7": 1.37914, - "8": 1.3886, - "9": 1.10046, - "10": 1.11649, - "11": 1.11259, - "12": 1.10822, - "13": 1.10532, - "14": 1.11189, - "15": 1.1132, - "16": 1.10539, - "17": 1.11434, - "18": 1.11836, - "19": 1.11073, - "20": 1.11278, - "21": 1.11212, - "22": 1.10671, - "23": 1.11034, - "24": 1.11107, - "25": 1.11085, - "26": 1.10756, - "27": 1.10109, - "28": 1.1069, - "29": 1.11354, - "30": 1.11254, - "31": 1.10893, - "32": 1.11311, - "33": 1.10722, - "34": 1.10243, - "35": 1.10358, - "36": 1.09746, - "37": 1.09875, - "38": 1.10151, - "39": 1.10188, - "40": 1.10069, - "41": 1.10545, - "42": 1.10709, - "43": 1.1028, - "44": 1.10723, - "45": 1.10614, - "46": 1.09997, - "47": 1.1053, - "48": 1.10274, - "49": 1.09986, - "50": 1.10191 + "1": 95.02242, + "2": 1.29728, + "3": 1.24413, + "4": 1.67309, + "5": 1.12527, + "6": 1.39226, + "7": 1.33351, + "8": 1.19614, + "9": 1.10737, + "10": 1.09796, + "11": 1.10736, + "12": 1.10105, + "13": 1.10552, + "14": 1.11007, + "15": 1.09853, + "16": 1.10142, + "17": 1.09718, + "18": 1.10103, + "19": 1.10339, + "20": 1.1069, + "21": 1.10541, + "22": 1.10374, + "23": 1.1028, + "24": 1.1, + "25": 1.09935, + "26": 1.09318, + "27": 1.09779, + "28": 1.09457, + "29": 1.09, + "30": 1.09267, + "31": 1.08899, + "32": 1.09268, + "33": 1.08757, + "34": 1.08991, + "35": 1.09705, + "36": 1.09429, + "37": 1.09459, + "38": 1.08857, + "39": 1.09547, + "40": 1.09224, + "41": 1.089, + "42": 1.08879, + "43": 1.0834, + "44": 1.08212, + "45": 1.08363, + "46": 1.08596, + "47": 1.07798, + "48": 1.07329, + "49": 1.07678, + "50": 1.07483 } } } \ No newline at end of file diff --git a/tests/functional_tests/test_cases/moe/gpt3_mcore_te_tp2_pp2_ep4_etp1_resume_torch_dist_attn_cudagraph/golden_values_dev_dgxh100_coreweave.json b/tests/functional_tests/test_cases/moe/gpt3_mcore_te_tp2_pp2_ep4_etp1_resume_torch_dist_attn_cudagraph/golden_values_dev_dgxh100_coreweave.json index 58eb3fc16cd..8cea616921e 100644 --- a/tests/functional_tests/test_cases/moe/gpt3_mcore_te_tp2_pp2_ep4_etp1_resume_torch_dist_attn_cudagraph/golden_values_dev_dgxh100_coreweave.json +++ b/tests/functional_tests/test_cases/moe/gpt3_mcore_te_tp2_pp2_ep4_etp1_resume_torch_dist_attn_cudagraph/golden_values_dev_dgxh100_coreweave.json @@ -4,56 +4,56 @@ "end_step": 50, "step_interval": 1, "values": { - "1": 10.95004, - "2": 10.9521, - "3": 10.5115, - "4": 9.96454, - "5": 9.93941, - "6": 9.67273, - "7": 10.20975, - "8": 9.49716, - "9": 9.55902, - "10": 9.79742, - "11": 9.30109, - "12": 9.40483, - "13": 9.39546, - "14": 8.84681, - "15": 9.02444, - "16": 9.07121, - "17": 9.04574, - "18": 8.75678, - "19": 9.18159, - "20": 8.8595, - "21": 8.53503, - "22": 8.55182, - "23": 8.42441, - "24": 8.37608, - "25": 8.64304, - "26": 7.97393, - "27": 8.56806, - "28": 8.19764, - "29": 8.3928, - "30": 8.67283, - "31": 8.289, - "32": 8.43572, - "33": 8.5568, - "34": 8.66018, - "35": 8.07934, - "36": 7.94976, - "37": 8.29565, - "38": 7.98044, - "39": 8.39201, - "40": 8.35513, - "41": 8.31876, - "42": 8.0583, - "43": 8.03283, - "44": 8.24243, - "45": 8.10277, - "46": 7.61696, - "47": 8.15273, - "48": 8.00569, - "49": 8.38688, - "50": 7.81491 + "1": 10.94971, + "2": 10.95163, + "3": 10.51641, + "4": 9.9652, + "5": 9.94116, + "6": 9.67394, + "7": 10.19887, + "8": 9.50035, + "9": 9.54982, + "10": 9.79667, + "11": 9.30128, + "12": 9.40566, + "13": 9.39438, + "14": 8.84572, + "15": 9.02231, + "16": 9.06973, + "17": 9.04712, + "18": 8.75662, + "19": 9.18074, + "20": 8.86175, + "21": 8.53558, + "22": 8.55288, + "23": 8.42513, + "24": 8.37683, + "25": 8.64426, + "26": 7.9756, + "27": 8.57026, + "28": 8.1987, + "29": 8.39406, + "30": 8.67631, + "31": 8.29096, + "32": 8.43692, + "33": 8.55897, + "34": 8.66123, + "35": 8.08, + "36": 7.95214, + "37": 8.2979, + "38": 7.98177, + "39": 8.39281, + "40": 8.35852, + "41": 8.32006, + "42": 8.05954, + "43": 8.03381, + "44": 8.24236, + "45": 8.1025, + "46": 7.61814, + "47": 8.15364, + "48": 8.00693, + "49": 8.38704, + "50": 7.81592 } }, "num-zeros": { @@ -61,56 +61,56 @@ "end_step": 50, "step_interval": 1, "values": { - "1": 19403624.0, - "2": 19274194.0, - "3": 19372760.0, - "4": 86525248.0, - "5": 148575568.0, - "6": 145226704.0, - "7": 171879984.0, - "8": 195785248.0, - "9": 164124752.0, - "10": 167684736.0, - "11": 221077344.0, - "12": 200384224.0, - "13": 248872528.0, - "14": 211169424.0, - "15": 214304608.0, - "16": 216075632.0, - "17": 267845984.0, - "18": 170470336.0, - "19": 176865072.0, - "20": 187955392.0, - "21": 225750704.0, - "22": 247396816.0, - "23": 211643856.0, - "24": 205638464.0, - "25": 277022272.0, - "26": 291562304.0, - "27": 225789840.0, - "28": 288202368.0, - "29": 198390384.0, - "30": 213302208.0, - "31": 227204752.0, - "32": 271112416.0, - "33": 231840432.0, - "34": 203575536.0, - "35": 191152368.0, - "36": 222566928.0, - "37": 177810112.0, - "38": 228708544.0, - "39": 211168784.0, - "40": 215603968.0, - "41": 200089440.0, - "42": 228529888.0, - "43": 198782848.0, - "44": 141902272.0, - "45": 181922816.0, - "46": 115369856.0, - "47": 170214176.0, - "48": 137292832.0, - "49": 97654936.0, - "50": 160979632.0 + "1": 19403704.0, + "2": 19274216.0, + "3": 22517470.0, + "4": 83429816.0, + "5": 139167728.0, + "6": 138921280.0, + "7": 173470304.0, + "8": 200511856.0, + "9": 165696320.0, + "10": 166120112.0, + "11": 213254416.0, + "12": 187847360.0, + "13": 231586656.0, + "14": 226879072.0, + "15": 219025920.0, + "16": 205179664.0, + "17": 280450432.0, + "18": 181477792.0, + "19": 191026096.0, + "20": 186395632.0, + "21": 233632576.0, + "22": 231696832.0, + "23": 216390688.0, + "24": 215133760.0, + "25": 233079504.0, + "26": 244437920.0, + "27": 222637584.0, + "28": 278773952.0, + "29": 253409264.0, + "30": 240036736.0, + "31": 236599008.0, + "32": 205066624.0, + "33": 263303312.0, + "34": 200444544.0, + "35": 199033824.0, + "36": 243001216.0, + "37": 151181872.0, + "38": 175301280.0, + "39": 219001024.0, + "40": 220307936.0, + "41": 217385856.0, + "42": 230074176.0, + "43": 208226784.0, + "44": 148172720.0, + "45": 141103744.0, + "46": 132664976.0, + "47": 179619392.0, + "48": 118381144.0, + "49": 86643984.0, + "50": 113798320.0 } }, "mem-allocated-bytes": { @@ -118,56 +118,56 @@ "end_step": 50, "step_interval": 1, "values": { - "1": 4883602432.0, - "2": 4885017088.0, - "3": 4882657792.0, - "4": 4883046912.0, - "5": 4883725824.0, - "6": 4883713536.0, - "7": 4883040768.0, - "8": 4883273216.0, - "9": 4882952704.0, - "10": 4885949952.0, - "11": 4883990016.0, - "12": 4887679488.0, - "13": 4884011520.0, - "14": 4882899456.0, - "15": 4883515904.0, - "16": 4883990016.0, - "17": 4883410432.0, - "18": 4883673600.0, - "19": 4882903552.0, - "20": 4884541952.0, - "21": 4883138048.0, - "22": 4883247616.0, - "23": 4883839488.0, - "24": 4885058048.0, - "25": 4882676224.0, - "26": 4884058624.0, - "27": 4884724224.0, - "28": 4884874752.0, - "29": 4883127808.0, - "30": 4883252736.0, - "31": 4882955776.0, - "32": 4885190144.0, - "33": 4883845632.0, - "34": 4884392448.0, - "35": 4883083776.0, - "36": 4883851776.0, - "37": 4885246464.0, - "38": 4882680320.0, - "39": 4884296192.0, - "40": 4884689408.0, - "41": 4882836992.0, - "42": 4883972608.0, - "43": 4884519424.0, - "44": 4883354112.0, - "45": 4883495424.0, - "46": 4882788864.0, - "47": 4883144192.0, - "48": 4883688960.0, - "49": 4884182528.0, - "50": 4885279232.0 + "1": 4883287040.0, + "2": 4883441152.0, + "3": 4881697280.0, + "4": 4883730944.0, + "5": 4882556416.0, + "6": 4882616832.0, + "7": 4883438080.0, + "8": 4881568256.0, + "9": 4883173888.0, + "10": 4882272768.0, + "11": 4883676672.0, + "12": 4881393152.0, + "13": 4883141120.0, + "14": 4883697152.0, + "15": 4882622976.0, + "16": 4881830400.0, + "17": 4881658368.0, + "18": 4881863168.0, + "19": 4883804672.0, + "20": 4881795584.0, + "21": 4883333632.0, + "22": 4882194944.0, + "23": 4882084352.0, + "24": 4884065792.0, + "25": 4881804800.0, + "26": 4883596800.0, + "27": 4883047936.0, + "28": 4882476544.0, + "29": 4883087872.0, + "30": 4882151936.0, + "31": 4882625024.0, + "32": 4883104256.0, + "33": 4882526720.0, + "34": 4882292224.0, + "35": 4882485760.0, + "36": 4882867712.0, + "37": 4882634240.0, + "38": 4882610688.0, + "39": 4881474048.0, + "40": 4881961472.0, + "41": 4882663936.0, + "42": 4881860096.0, + "43": 4881499648.0, + "44": 4883392000.0, + "45": 4882392576.0, + "46": 4882815488.0, + "47": 4883113472.0, + "48": 4882158080.0, + "49": 4881207808.0, + "50": 4881588736.0 } }, "mem-max-allocated-bytes": { @@ -175,56 +175,56 @@ "end_step": 50, "step_interval": 1, "values": { - "1": 41210470400.0, - "2": 41210470400.0, - "3": 41210470400.0, - "4": 41210470400.0, - "5": 41210470400.0, - "6": 41210470400.0, - "7": 41210470400.0, - "8": 41210470400.0, - "9": 41210470400.0, - "10": 41210470400.0, - "11": 41210470400.0, - "12": 41210470400.0, - "13": 41210470400.0, - "14": 41210470400.0, - "15": 41210470400.0, - "16": 41210470400.0, - "17": 41210470400.0, - "18": 41210470400.0, - "19": 41210470400.0, - "20": 41210470400.0, - "21": 41210470400.0, - "22": 41210470400.0, - "23": 41210470400.0, - "24": 41210470400.0, - "25": 41210470400.0, - "26": 41210470400.0, - "27": 41210470400.0, - "28": 41210470400.0, - "29": 41210470400.0, - "30": 41210470400.0, - "31": 41210470400.0, - "32": 41210470400.0, - "33": 41210470400.0, - "34": 41210470400.0, - "35": 41210470400.0, - "36": 41210470400.0, - "37": 41210470400.0, - "38": 41210470400.0, - "39": 41210470400.0, - "40": 41210470400.0, - "41": 41210470400.0, - "42": 41210470400.0, - "43": 41210470400.0, - "44": 41210470400.0, - "45": 41210470400.0, - "46": 41210470400.0, - "47": 41210470400.0, - "48": 41210470400.0, - "49": 41210470400.0, - "50": 41210470400.0 + "1": 41208348672.0, + "2": 41208348672.0, + "3": 41208348672.0, + "4": 41208348672.0, + "5": 41208348672.0, + "6": 41208348672.0, + "7": 41208348672.0, + "8": 41208348672.0, + "9": 41208348672.0, + "10": 41208348672.0, + "11": 41208348672.0, + "12": 41208348672.0, + "13": 41208348672.0, + "14": 41208348672.0, + "15": 41208348672.0, + "16": 41208348672.0, + "17": 41208348672.0, + "18": 41208348672.0, + "19": 41208348672.0, + "20": 41208348672.0, + "21": 41208348672.0, + "22": 41208348672.0, + "23": 41208348672.0, + "24": 41208348672.0, + "25": 41208348672.0, + "26": 41208348672.0, + "27": 41208348672.0, + "28": 41208348672.0, + "29": 41208348672.0, + "30": 41208348672.0, + "31": 41208348672.0, + "32": 41208348672.0, + "33": 41208348672.0, + "34": 41208348672.0, + "35": 41208348672.0, + "36": 41208348672.0, + "37": 41208348672.0, + "38": 41208348672.0, + "39": 41208348672.0, + "40": 41208348672.0, + "41": 41208348672.0, + "42": 41208348672.0, + "43": 41208348672.0, + "44": 41208348672.0, + "45": 41208348672.0, + "46": 41208348672.0, + "47": 41208348672.0, + "48": 41208348672.0, + "49": 41208348672.0, + "50": 41208348672.0 } }, "iteration-time": { @@ -232,56 +232,56 @@ "end_step": 50, "step_interval": 1, "values": { - "1": 86.8085, - "2": 1.10913, - "3": 0.99097, - "4": 0.89412, - "5": 1.25997, - "6": 0.98162, - "7": 0.98318, - "8": 1.13296, - "9": 0.88126, - "10": 0.8633, - "11": 2.2744, - "12": 4.5393, - "13": 3.22763, - "14": 1.64923, - "15": 0.86595, - "16": 0.86575, - "17": 0.85272, - "18": 0.85454, - "19": 0.85281, - "20": 0.87018, - "21": 0.84654, - "22": 0.8494, - "23": 0.84882, - "24": 0.84482, - "25": 0.85311, - "26": 0.84678, - "27": 0.84096, - "28": 0.8412, - "29": 0.84156, - "30": 0.84475, - "31": 0.84747, - "32": 0.85058, - "33": 0.84977, - "34": 0.8479, - "35": 0.85234, - "36": 0.85012, - "37": 0.85087, - "38": 0.84594, - "39": 0.84558, - "40": 0.84807, - "41": 0.84183, - "42": 0.8439, - "43": 0.84221, - "44": 0.84248, - "45": 0.84257, - "46": 0.83922, - "47": 0.84311, - "48": 0.84159, - "49": 0.84011, - "50": 0.8353 + "1": 89.10928, + "2": 1.08143, + "3": 0.94222, + "4": 0.89675, + "5": 1.34524, + "6": 1.06972, + "7": 1.00314, + "8": 1.04961, + "9": 0.86611, + "10": 0.86248, + "11": 0.98739, + "12": 0.86057, + "13": 0.86777, + "14": 0.85834, + "15": 0.8559, + "16": 0.85522, + "17": 0.84644, + "18": 0.85748, + "19": 0.85218, + "20": 0.85342, + "21": 0.84029, + "22": 0.84342, + "23": 0.84297, + "24": 0.83925, + "25": 0.8439, + "26": 0.85696, + "27": 0.83981, + "28": 0.84643, + "29": 0.8433, + "30": 0.86234, + "31": 0.85636, + "32": 0.84184, + "33": 0.84501, + "34": 0.84316, + "35": 0.83806, + "36": 0.84143, + "37": 0.84447, + "38": 0.84137, + "39": 0.84133, + "40": 0.84321, + "41": 0.84019, + "42": 0.84164, + "43": 0.83741, + "44": 0.84203, + "45": 0.83966, + "46": 0.84109, + "47": 0.83945, + "48": 0.84001, + "49": 0.84194, + "50": 0.83578 } } } \ No newline at end of file diff --git a/tests/functional_tests/test_cases/moe/gpt3_mcore_te_tp2_zp_z3_resume_torch_dist_te_8experts2parallel_top2router/golden_values_dev_dgx_h100.json b/tests/functional_tests/test_cases/moe/gpt3_mcore_te_tp2_zp_z3_resume_torch_dist_te_8experts2parallel_top2router/golden_values_dev_dgx_h100.json index 1ba051f4889..0835e95b926 100644 --- a/tests/functional_tests/test_cases/moe/gpt3_mcore_te_tp2_zp_z3_resume_torch_dist_te_8experts2parallel_top2router/golden_values_dev_dgx_h100.json +++ b/tests/functional_tests/test_cases/moe/gpt3_mcore_te_tp2_zp_z3_resume_torch_dist_te_8experts2parallel_top2router/golden_values_dev_dgx_h100.json @@ -1 +1,142 @@ -{"lm loss": {"start_step": 1, "end_step": 100, "step_interval": 5, "values": {"1": 10.83281, "5": 10.85975, "10": 10.79613, "15": 10.80527, "20": 10.72502, "25": 10.53599, "30": 10.3571, "35": 10.24605, "40": 10.05992, "45": 9.7836, "50": 9.8722, "55": 9.83189, "60": 9.45075, "65": 8.89679, "70": 9.71414, "75": 9.39795, "80": 9.38169, "85": 9.58585, "90": 9.7999, "95": 9.50528, "100": 9.37224}}, "num-zeros": {"start_step": 1, "end_step": 100, "step_interval": 5, "values": {"1": 27013.0, "5": 31736.0, "10": 25785.0, "15": 30383.0, "20": 28435.0, "25": 27493.0, "30": 30329.0, "35": 31750.0, "40": 34279.0, "45": 34634.0, "50": 38531.0, "55": 37465.0, "60": 40172.0, "65": 40624.0, "70": 44852.0, "75": 39231.0, "80": 130535.0, "85": 123250.0, "90": 47793.0, "95": 167340.0, "100": 163328.0}}, "mem-allocated-bytes": {"start_step": 1, "end_step": 100, "step_interval": 5, "values": {"1": 814390272.0, "5": 814420480.0, "10": 814376448.0, "15": 814376960.0, "20": 814373376.0, "25": 814321152.0, "30": 814306304.0, "35": 814292992.0, "40": 814288896.0, "45": 814272000.0, "50": 814262272.0, "55": 814258688.0, "60": 814268416.0, "65": 814220800.0, "70": 814266880.0, "75": 814318080.0, "80": 814285312.0, "85": 814289408.0, "90": 814315520.0, "95": 814320128.0, "100": 814311424.0}}, "mem-max-allocated-bytes": {"start_step": 1, "end_step": 100, "step_interval": 5, "values": {"1": 2111314944.0, "5": 2370209280.0, "10": 2370209280.0, "15": 2370209280.0, "20": 2370209280.0, "25": 2370209280.0, "30": 2370209280.0, "35": 2370209280.0, "40": 2370209280.0, "45": 2370209280.0, "50": 2370209280.0, "55": 2370209280.0, "60": 2370209280.0, "65": 2370209280.0, "70": 2370209280.0, "75": 2370209280.0, "80": 2370209280.0, "85": 2370209280.0, "90": 2370209280.0, "95": 2370209280.0, "100": 2370209280.0}}, "iteration-time": {"start_step": 1, "end_step": 100, "step_interval": 5, "values": {"1": 20.98318, "5": 0.79797, "10": 0.74028, "15": 0.67279, "20": 0.62948, "25": 0.61132, "30": 0.61547, "35": 0.6152, "40": 0.60421, "45": 0.59124, "50": 0.5891, "55": 0.57048, "60": 0.54799, "65": 0.52185, "70": 0.51195, "75": 0.50105, "80": 0.4628, "85": 0.45992, "90": 0.46498, "95": 0.4599, "100": 0.42568}}} \ No newline at end of file +{ + "lm loss": { + "start_step": 1, + "end_step": 100, + "step_interval": 5, + "values": { + "1": 10.82922, + "5": 10.85652, + "10": 10.79298, + "15": 10.8067, + "20": 10.72654, + "25": 10.53282, + "30": 10.35802, + "35": 10.24483, + "40": 10.05533, + "45": 9.77951, + "50": 9.86874, + "55": 9.82995, + "60": 9.449, + "65": 8.89366, + "70": 9.71127, + "75": 9.39451, + "80": 9.38198, + "85": 9.58333, + "90": 9.79944, + "95": 9.50213, + "100": 9.37131 + } + }, + "num-zeros": { + "start_step": 1, + "end_step": 100, + "step_interval": 5, + "values": { + "1": 27245.0, + "5": 31369.0, + "10": 25870.0, + "15": 29830.0, + "20": 28243.0, + "25": 27636.0, + "30": 30387.0, + "35": 31488.0, + "40": 34779.0, + "45": 35158.0, + "50": 38234.0, + "55": 37133.0, + "60": 40450.0, + "65": 40947.0, + "70": 43436.0, + "75": 39925.0, + "80": 51863.0, + "85": 2145177.0, + "90": 51330.0, + "95": 45247.0, + "100": 163741.0 + } + }, + "mem-allocated-bytes": { + "start_step": 1, + "end_step": 100, + "step_interval": 5, + "values": { + "1": 787511296.0, + "5": 787542016.0, + "10": 787500032.0, + "15": 787499008.0, + "20": 787500032.0, + "25": 787446272.0, + "30": 787429888.0, + "35": 787413504.0, + "40": 787409920.0, + "45": 787394560.0, + "50": 787384320.0, + "55": 787383808.0, + "60": 787389952.0, + "65": 787346432.0, + "70": 787387904.0, + "75": 787437568.0, + "80": 787405312.0, + "85": 787407360.0, + "90": 787441664.0, + "95": 787445248.0, + "100": 787433472.0 + } + }, + "mem-max-allocated-bytes": { + "start_step": 1, + "end_step": 100, + "step_interval": 5, + "values": { + "1": 2465793024.0, + "5": 2492764160.0, + "10": 2492764160.0, + "15": 2492764160.0, + "20": 2492764160.0, + "25": 2492764160.0, + "30": 2492764160.0, + "35": 2492764160.0, + "40": 2492764160.0, + "45": 2492764160.0, + "50": 2492764160.0, + "55": 2492764160.0, + "60": 2492764160.0, + "65": 2492764160.0, + "70": 2492764160.0, + "75": 2492764160.0, + "80": 2492764160.0, + "85": 2492764160.0, + "90": 2492764160.0, + "95": 2492764160.0, + "100": 2492764160.0 + } + }, + "iteration-time": { + "start_step": 1, + "end_step": 100, + "step_interval": 5, + "values": { + "1": 9.68104, + "5": 0.32859, + "10": 0.30772, + "15": 0.31234, + "20": 0.29254, + "25": 0.29296, + "30": 0.31344, + "35": 0.31026, + "40": 0.30514, + "45": 0.30481, + "50": 0.30324, + "55": 0.29929, + "60": 0.30103, + "65": 0.32008, + "70": 0.31307, + "75": 0.2933, + "80": 0.29351, + "85": 0.29283, + "90": 0.29375, + "95": 0.29458, + "100": 0.29103 + } + } +} \ No newline at end of file diff --git a/tests/functional_tests/test_cases/moe/gpt3_mcore_te_tp2_zp_z3_resume_torch_dist_te_8experts2parallel_top2router/golden_values_dev_dgxh100_coreweave.json b/tests/functional_tests/test_cases/moe/gpt3_mcore_te_tp2_zp_z3_resume_torch_dist_te_8experts2parallel_top2router/golden_values_dev_dgxh100_coreweave.json new file mode 100644 index 00000000000..7e299df5257 --- /dev/null +++ b/tests/functional_tests/test_cases/moe/gpt3_mcore_te_tp2_zp_z3_resume_torch_dist_te_8experts2parallel_top2router/golden_values_dev_dgxh100_coreweave.json @@ -0,0 +1,537 @@ +{ + "lm loss": { + "start_step": 1, + "end_step": 100, + "step_interval": 1, + "values": { + "1": 10.82922, + "2": 10.84163, + "3": 10.84245, + "4": 10.82, + "5": 10.85652, + "6": 10.86906, + "7": 10.83778, + "8": 10.84312, + "9": 10.84423, + "10": 10.79298, + "11": 10.86697, + "12": 10.86875, + "13": 10.86207, + "14": 10.86919, + "15": 10.8067, + "16": 10.8057, + "17": 10.77686, + "18": 10.79541, + "19": 10.78384, + "20": 10.72654, + "21": 10.69491, + "22": 10.54462, + "23": 10.6993, + "24": 10.58151, + "25": 10.53282, + "26": 10.58817, + "27": 10.601, + "28": 10.57563, + "29": 10.58022, + "30": 10.35802, + "31": 10.08769, + "32": 10.44466, + "33": 10.4477, + "34": 10.18704, + "35": 10.24483, + "36": 10.19713, + "37": 10.32294, + "38": 10.17101, + "39": 10.37026, + "40": 10.05533, + "41": 10.09491, + "42": 10.17971, + "43": 9.78263, + "44": 9.91346, + "45": 9.77951, + "46": 9.75648, + "47": 10.09647, + "48": 9.80391, + "49": 9.46649, + "50": 9.86874, + "51": 9.79428, + "52": 9.68303, + "53": 10.03314, + "54": 9.9113, + "55": 9.82995, + "56": 9.57839, + "57": 9.42377, + "58": 9.80549, + "59": 9.53292, + "60": 9.449, + "61": 9.65293, + "62": 9.95672, + "63": 9.33775, + "64": 9.74194, + "65": 8.89366, + "66": 9.67317, + "67": 9.33002, + "68": 9.76517, + "69": 9.76336, + "70": 9.71127, + "71": 9.59511, + "72": 9.54797, + "73": 9.47124, + "74": 8.89297, + "75": 9.39451, + "76": 9.04721, + "77": 10.04318, + "78": 9.70313, + "79": 9.35169, + "80": 9.38198, + "81": 9.45146, + "82": 9.67546, + "83": 9.27658, + "84": 9.39241, + "85": 9.58333, + "86": 9.04518, + "87": 9.56487, + "88": 9.72459, + "89": 9.57019, + "90": 9.79944, + "91": 9.30737, + "92": 9.3313, + "93": 9.04109, + "94": 8.80259, + "95": 9.50213, + "96": 9.5021, + "97": 9.28183, + "98": 9.64883, + "99": 8.8594, + "100": 9.37131 + } + }, + "num-zeros": { + "start_step": 1, + "end_step": 100, + "step_interval": 1, + "values": { + "1": 27245.0, + "2": 28958.0, + "3": 29464.0, + "4": 28046.0, + "5": 31369.0, + "6": 33287.0, + "7": 31200.0, + "8": 26921.0, + "9": 30008.0, + "10": 25870.0, + "11": 33681.0, + "12": 30344.0, + "13": 32737.0, + "14": 33315.0, + "15": 29830.0, + "16": 32475.0, + "17": 30747.0, + "18": 30381.0, + "19": 31032.0, + "20": 28243.0, + "21": 29224.0, + "22": 27340.0, + "23": 34119.0, + "24": 29049.0, + "25": 27636.0, + "26": 30662.0, + "27": 32009.0, + "28": 33355.0, + "29": 34714.0, + "30": 30387.0, + "31": 28212.0, + "32": 33411.0, + "33": 34696.0, + "34": 30053.0, + "35": 31488.0, + "36": 32943.0, + "37": 35829.0, + "38": 33740.0, + "39": 37632.0, + "40": 34779.0, + "41": 33958.0, + "42": 36396.0, + "43": 34088.0, + "44": 34090.0, + "45": 35158.0, + "46": 36174.0, + "47": 39772.0, + "48": 36516.0, + "49": 36733.0, + "50": 38234.0, + "51": 38608.0, + "52": 37030.0, + "53": 42442.0, + "54": 40944.0, + "55": 37133.0, + "56": 41001.0, + "57": 37524.0, + "58": 42317.0, + "59": 40804.0, + "60": 40450.0, + "61": 41478.0, + "62": 39766.0, + "63": 37941.0, + "64": 42197.0, + "65": 40947.0, + "66": 44094.0, + "67": 41958.0, + "68": 40060.0, + "69": 42189.0, + "70": 43436.0, + "71": 42748.0, + "72": 44280.0, + "73": 47478.0, + "74": 41456.0, + "75": 39925.0, + "76": 43490.0, + "77": 45636.0, + "78": 2141470.0, + "79": 46055.0, + "80": 51863.0, + "81": 151341.0, + "82": 49835.0, + "83": 143360.0, + "84": 2141546.0, + "85": 2145177.0, + "86": 132114.0, + "87": 2147022.0, + "88": 59899.0, + "89": 162883.0, + "90": 51330.0, + "91": 2141901.0, + "92": 44946.0, + "93": 138194.0, + "94": 2145772.0, + "95": 45247.0, + "96": 135045.0, + "97": 53170.0, + "98": 168576.0, + "99": 2141797.0, + "100": 163741.0 + } + }, + "mem-allocated-bytes": { + "start_step": 1, + "end_step": 100, + "step_interval": 1, + "values": { + "1": 787516416.0, + "2": 787540992.0, + "3": 787524096.0, + "4": 787512320.0, + "5": 787547136.0, + "6": 787537920.0, + "7": 787512832.0, + "8": 787524608.0, + "9": 787528192.0, + "10": 787505152.0, + "11": 787522048.0, + "12": 787520000.0, + "13": 787529728.0, + "14": 787529216.0, + "15": 787504128.0, + "16": 787513344.0, + "17": 787503104.0, + "18": 787489280.0, + "19": 787514880.0, + "20": 787505152.0, + "21": 787479552.0, + "22": 787486208.0, + "23": 787478528.0, + "24": 787486208.0, + "25": 787451392.0, + "26": 787482112.0, + "27": 787470848.0, + "28": 787450368.0, + "29": 787458048.0, + "30": 787435008.0, + "31": 787406848.0, + "32": 787424256.0, + "33": 787435520.0, + "34": 787426304.0, + "35": 787418624.0, + "36": 787436544.0, + "37": 787428352.0, + "38": 787436544.0, + "39": 787417600.0, + "40": 787415040.0, + "41": 787405824.0, + "42": 787415040.0, + "43": 787367936.0, + "44": 787392512.0, + "45": 787399680.0, + "46": 787355136.0, + "47": 787411456.0, + "48": 787354112.0, + "49": 787374080.0, + "50": 787389440.0, + "51": 787375616.0, + "52": 787383808.0, + "53": 787379712.0, + "54": 787384832.0, + "55": 787388928.0, + "56": 787388928.0, + "57": 787351040.0, + "58": 787382784.0, + "59": 787374080.0, + "60": 787395072.0, + "61": 787405312.0, + "62": 787405824.0, + "63": 787373056.0, + "64": 787388928.0, + "65": 787351552.0, + "66": 787386880.0, + "67": 787392000.0, + "68": 787399168.0, + "69": 787383296.0, + "70": 787393024.0, + "71": 787406848.0, + "72": 787400704.0, + "73": 787401216.0, + "74": 787403264.0, + "75": 787442688.0, + "76": 787444736.0, + "77": 787445760.0, + "78": 787395072.0, + "79": 787430400.0, + "80": 787410432.0, + "81": 787412992.0, + "82": 787427840.0, + "83": 787428864.0, + "84": 787412480.0, + "85": 787412480.0, + "86": 787394560.0, + "87": 787452928.0, + "88": 787414528.0, + "89": 787404800.0, + "90": 787446784.0, + "91": 787446272.0, + "92": 787446784.0, + "93": 787430400.0, + "94": 787440128.0, + "95": 787450368.0, + "96": 787454976.0, + "97": 787427328.0, + "98": 787475968.0, + "99": 787419136.0, + "100": 787438592.0 + } + }, + "mem-max-allocated-bytes": { + "start_step": 1, + "end_step": 100, + "step_interval": 1, + "values": { + "1": 2479493120.0, + "2": 2485449728.0, + "3": 2487249408.0, + "4": 2487249408.0, + "5": 2495991808.0, + "6": 2495991808.0, + "7": 2495991808.0, + "8": 2495991808.0, + "9": 2495991808.0, + "10": 2495991808.0, + "11": 2495991808.0, + "12": 2495991808.0, + "13": 2495991808.0, + "14": 2495991808.0, + "15": 2495991808.0, + "16": 2495991808.0, + "17": 2495991808.0, + "18": 2495991808.0, + "19": 2495991808.0, + "20": 2495991808.0, + "21": 2495991808.0, + "22": 2495991808.0, + "23": 2495991808.0, + "24": 2495991808.0, + "25": 2495991808.0, + "26": 2495991808.0, + "27": 2495991808.0, + "28": 2495991808.0, + "29": 2495991808.0, + "30": 2495991808.0, + "31": 2495991808.0, + "32": 2495991808.0, + "33": 2495991808.0, + "34": 2495991808.0, + "35": 2495991808.0, + "36": 2495991808.0, + "37": 2495991808.0, + "38": 2495991808.0, + "39": 2495991808.0, + "40": 2495991808.0, + "41": 2495991808.0, + "42": 2495991808.0, + "43": 2495991808.0, + "44": 2495991808.0, + "45": 2495991808.0, + "46": 2495991808.0, + "47": 2495991808.0, + "48": 2495991808.0, + "49": 2495991808.0, + "50": 2495991808.0, + "51": 2495991808.0, + "52": 2495991808.0, + "53": 2495991808.0, + "54": 2495991808.0, + "55": 2495991808.0, + "56": 2495991808.0, + "57": 2495991808.0, + "58": 2495991808.0, + "59": 2495991808.0, + "60": 2495991808.0, + "61": 2495991808.0, + "62": 2495991808.0, + "63": 2495991808.0, + "64": 2495991808.0, + "65": 2495991808.0, + "66": 2495991808.0, + "67": 2495991808.0, + "68": 2495991808.0, + "69": 2495991808.0, + "70": 2495991808.0, + "71": 2495991808.0, + "72": 2495991808.0, + "73": 2495991808.0, + "74": 2495991808.0, + "75": 2495991808.0, + "76": 2495991808.0, + "77": 2495991808.0, + "78": 2495991808.0, + "79": 2495991808.0, + "80": 2495991808.0, + "81": 2495991808.0, + "82": 2495991808.0, + "83": 2495991808.0, + "84": 2495991808.0, + "85": 2495991808.0, + "86": 2495991808.0, + "87": 2495991808.0, + "88": 2495991808.0, + "89": 2495991808.0, + "90": 2495991808.0, + "91": 2495991808.0, + "92": 2495991808.0, + "93": 2495991808.0, + "94": 2495991808.0, + "95": 2495991808.0, + "96": 2495991808.0, + "97": 2495991808.0, + "98": 2495991808.0, + "99": 2495991808.0, + "100": 2495991808.0 + } + }, + "iteration-time": { + "start_step": 1, + "end_step": 100, + "step_interval": 1, + "values": { + "1": 12.11313, + "2": 0.4805, + "3": 0.36965, + "4": 0.36695, + "5": 0.31705, + "6": 0.31275, + "7": 0.31299, + "8": 0.29866, + "9": 0.28961, + "10": 0.28859, + "11": 0.29067, + "12": 0.29044, + "13": 0.29806, + "14": 0.29287, + "15": 0.29391, + "16": 0.3175, + "17": 0.28363, + "18": 0.2818, + "19": 0.29347, + "20": 0.28931, + "21": 0.29103, + "22": 0.28444, + "23": 0.28907, + "24": 0.27608, + "25": 0.28277, + "26": 0.28656, + "27": 0.28921, + "28": 0.30243, + "29": 0.30435, + "30": 0.31231, + "31": 0.30439, + "32": 0.31412, + "33": 0.28887, + "34": 0.29613, + "35": 0.29738, + "36": 0.29754, + "37": 0.3019, + "38": 0.2933, + "39": 0.2944, + "40": 0.29283, + "41": 0.29592, + "42": 0.29673, + "43": 0.29319, + "44": 0.30127, + "45": 0.29921, + "46": 0.29904, + "47": 0.28795, + "48": 0.29918, + "49": 0.28711, + "50": 0.29645, + "51": 0.28777, + "52": 0.29536, + "53": 0.2847, + "54": 0.28286, + "55": 0.2874, + "56": 0.28699, + "57": 0.28614, + "58": 0.29825, + "59": 0.28363, + "60": 0.29423, + "61": 0.29226, + "62": 0.2896, + "63": 0.28065, + "64": 0.29533, + "65": 0.29842, + "66": 0.28487, + "67": 0.28419, + "68": 0.29474, + "69": 0.28383, + "70": 0.28417, + "71": 0.29253, + "72": 0.28737, + "73": 0.27923, + "74": 0.28728, + "75": 0.29383, + "76": 0.28157, + "77": 0.64771, + "78": 0.29148, + "79": 0.28742, + "80": 0.29245, + "81": 0.28827, + "82": 0.28368, + "83": 0.28963, + "84": 0.29234, + "85": 0.28183, + "86": 0.28337, + "87": 0.27879, + "88": 0.28388, + "89": 0.28309, + "90": 0.28852, + "91": 0.28254, + "92": 0.28375, + "93": 0.28633, + "94": 0.28567, + "95": 0.28235, + "96": 0.28513, + "97": 0.27951, + "98": 0.27851, + "99": 0.28336, + "100": 0.27744 + } + } +} \ No newline at end of file diff --git a/tests/functional_tests/test_cases/moe/gpt3_mcore_te_tp2_zp_z3_resume_torch_dist_te_8experts2parallel_top2router/model_config.yaml b/tests/functional_tests/test_cases/moe/gpt3_mcore_te_tp2_zp_z3_resume_torch_dist_te_8experts2parallel_top2router/model_config.yaml index 3ecd68b9841..8874f9cf045 100644 --- a/tests/functional_tests/test_cases/moe/gpt3_mcore_te_tp2_zp_z3_resume_torch_dist_te_8experts2parallel_top2router/model_config.yaml +++ b/tests/functional_tests/test_cases/moe/gpt3_mcore_te_tp2_zp_z3_resume_torch_dist_te_8experts2parallel_top2router/model_config.yaml @@ -56,7 +56,7 @@ MODEL_ARGS: --attention-softmax-in-fp32: true --use-checkpoint-opt_param-scheduler: true --use-mcore-models: true - --ckpt-format: torch_dist + --ckpt-format: fsdp_dtensor --dist-ckpt-optim-fully-reshardable: true --dist-ckpt-strictness: log_all # backward compatibility for TE changes --data-cache-path: ${DATA_CACHE_PATH} diff --git a/tests/functional_tests/test_cases/moe/gpt3_moe_mcore_te_tp2_pp2_ep4_etp1_fine_grained_offloading/golden_values_dev_dgxh100_coreweave.json b/tests/functional_tests/test_cases/moe/gpt3_moe_mcore_te_tp2_pp2_ep4_etp1_fine_grained_offloading/golden_values_dev_dgxh100_coreweave.json index b3f192ba287..73fb00c9231 100644 --- a/tests/functional_tests/test_cases/moe/gpt3_moe_mcore_te_tp2_pp2_ep4_etp1_fine_grained_offloading/golden_values_dev_dgxh100_coreweave.json +++ b/tests/functional_tests/test_cases/moe/gpt3_moe_mcore_te_tp2_pp2_ep4_etp1_fine_grained_offloading/golden_values_dev_dgxh100_coreweave.json @@ -4,56 +4,56 @@ "end_step": 50, "step_interval": 1, "values": { - "1": 11.07546, - "2": 11.03837, - "3": 9.66011, - "4": 9.91381, - "5": 9.32909, - "6": 9.13922, - "7": 9.13574, - "8": 8.65508, - "9": 8.51394, - "10": 8.8409, - "11": 8.29149, - "12": 8.34581, - "13": 8.25518, - "14": 7.73711, - "15": 7.86249, - "16": 7.9371, - "17": 7.89319, - "18": 7.63123, - "19": 7.99731, - "20": 7.74538, - "21": 7.44348, - "22": 7.42249, - "23": 7.29714, - "24": 7.27462, - "25": 7.54574, - "26": 6.96838, - "27": 7.50556, - "28": 7.22743, - "29": 7.36588, - "30": 7.52622, - "31": 7.27026, - "32": 7.45521, - "33": 7.50954, - "34": 7.55686, - "35": 7.10177, - "36": 6.96431, - "37": 7.28463, - "38": 7.0808, - "39": 7.40923, - "40": 7.43338, - "41": 7.38496, - "42": 7.15749, - "43": 7.15858, - "44": 7.28852, - "45": 7.16793, - "46": 6.78468, - "47": 7.4114, - "48": 7.0027, - "49": 7.46249, - "50": 6.92151 + "1": 11.07559, + "2": 11.03834, + "3": 9.66022, + "4": 9.91367, + "5": 9.3291, + "6": 9.13927, + "7": 9.13591, + "8": 8.65527, + "9": 8.51396, + "10": 8.84095, + "11": 8.29144, + "12": 8.34584, + "13": 8.25509, + "14": 7.73685, + "15": 7.86273, + "16": 7.93699, + "17": 7.89257, + "18": 7.63116, + "19": 7.99719, + "20": 7.7453, + "21": 7.44298, + "22": 7.42242, + "23": 7.29721, + "24": 7.27467, + "25": 7.54562, + "26": 6.96839, + "27": 7.50569, + "28": 7.22761, + "29": 7.36579, + "30": 7.52635, + "31": 7.27036, + "32": 7.45548, + "33": 7.50952, + "34": 7.55694, + "35": 7.10212, + "36": 6.96414, + "37": 7.28438, + "38": 7.08049, + "39": 7.40908, + "40": 7.4335, + "41": 7.38491, + "42": 7.15766, + "43": 7.15867, + "44": 7.28831, + "45": 7.16729, + "46": 6.78429, + "47": 7.40937, + "48": 7.00259, + "49": 7.46241, + "50": 6.92143 } }, "num-zeros": { @@ -63,54 +63,54 @@ "values": { "1": 911219392.0, "2": 910960384.0, - "3": 911156352.0, - "4": 912204800.0, - "5": 920796544.0, - "6": 940387968.0, - "7": 990599872.0, - "8": 976457728.0, - "9": 998097664.0, - "10": 995852672.0, - "11": 994583680.0, - "12": 977344896.0, - "13": 1028141824.0, - "14": 1007166208.0, - "15": 987423616.0, - "16": 993054784.0, - "17": 982319168.0, - "18": 998261760.0, - "19": 984696320.0, - "20": 982914752.0, - "21": 979667456.0, - "22": 953988864.0, - "23": 972353984.0, - "24": 964792064.0, - "25": 958512192.0, - "26": 946928512.0, + "3": 911156288.0, + "4": 913253376.0, + "5": 921845056.0, + "6": 941436672.0, + "7": 993745472.0, + "8": 974360512.0, + "9": 999146112.0, + "10": 992706944.0, + "11": 991438144.0, + "12": 979442048.0, + "13": 1029190272.0, + "14": 1008214656.0, + "15": 988472000.0, + "16": 988861120.0, + "17": 979173312.0, + "18": 996164608.0, + "19": 979453440.0, + "20": 982914688.0, + "21": 975473344.0, + "22": 955037568.0, + "23": 969208128.0, + "24": 965840832.0, + "25": 953269440.0, + "26": 949025536.0, "27": 948458304.0, - "28": 949643968.0, - "29": 942877440.0, + "28": 951741184.0, + "29": 943926272.0, "30": 935020160.0, - "31": 935327616.0, - "32": 934281088.0, - "33": 921805568.0, - "34": 928189312.0, - "35": 922202496.0, - "36": 924246656.0, - "37": 920661248.0, + "31": 933230336.0, + "32": 930086848.0, + "33": 922853952.0, + "34": 927140800.0, + "35": 925348224.0, + "36": 925295168.0, + "37": 922758272.0, "38": 922930752.0, - "39": 922322816.0, - "40": 921856512.0, - "41": 920227968.0, + "39": 922322880.0, + "40": 921856640.0, + "41": 920227776.0, "42": 918353664.0, - "43": 918607040.0, - "44": 914948032.0, - "45": 914295232.0, + "43": 919655616.0, + "44": 914948224.0, + "45": 916392512.0, "46": 914344448.0, "47": 911769536.0, - "48": 912013312.0, - "49": 910349440.0, - "50": 914351552.0 + "48": 912013248.0, + "49": 910349376.0, + "50": 914351616.0 } }, "mem-allocated-bytes": { @@ -175,56 +175,56 @@ "end_step": 50, "step_interval": 1, "values": { - "1": 41739952128.0, - "2": 43687571456.0, - "3": 43687571456.0, - "4": 43983216640.0, - "5": 43983216640.0, - "6": 43983216640.0, - "7": 43983216640.0, - "8": 44024635392.0, - "9": 44041216000.0, - "10": 44041216000.0, - "11": 44041216000.0, - "12": 44041216000.0, - "13": 44041216000.0, - "14": 44041216000.0, - "15": 44041216000.0, - "16": 44041216000.0, - "17": 44041216000.0, - "18": 44041216000.0, - "19": 44041216000.0, - "20": 44041216000.0, - "21": 44041216000.0, - "22": 44041216000.0, - "23": 44041216000.0, - "24": 44041216000.0, - "25": 44041216000.0, - "26": 44041216000.0, - "27": 44041216000.0, - "28": 44041216000.0, - "29": 44041326592.0, - "30": 44162326528.0, - "31": 44220485632.0, - "32": 44270411776.0, - "33": 44293799936.0, - "34": 44293799936.0, - "35": 44293799936.0, - "36": 44293799936.0, - "37": 44293799936.0, - "38": 44293799936.0, - "39": 44293799936.0, - "40": 44293799936.0, - "41": 44293799936.0, - "42": 44293799936.0, - "43": 44293799936.0, - "44": 44293799936.0, - "45": 44293799936.0, - "46": 44293799936.0, - "47": 44293799936.0, - "48": 44293799936.0, - "49": 44293799936.0, - "50": 44293799936.0 + "1": 41740259328.0, + "2": 43687292928.0, + "3": 43687292928.0, + "4": 43984064512.0, + "5": 43984064512.0, + "6": 43984064512.0, + "7": 43984064512.0, + "8": 44026380288.0, + "9": 44041506816.0, + "10": 44041506816.0, + "11": 44041506816.0, + "12": 44041506816.0, + "13": 44041506816.0, + "14": 44041506816.0, + "15": 44041506816.0, + "16": 44041506816.0, + "17": 44041506816.0, + "18": 44041506816.0, + "19": 44041506816.0, + "20": 44041506816.0, + "21": 44041506816.0, + "22": 44041506816.0, + "23": 44041506816.0, + "24": 44041506816.0, + "25": 44041506816.0, + "26": 44041506816.0, + "27": 44041506816.0, + "28": 44041506816.0, + "29": 44044173312.0, + "30": 44164231168.0, + "31": 44221079552.0, + "32": 44271415296.0, + "33": 44290232320.0, + "34": 44290232320.0, + "35": 44290232320.0, + "36": 44290232320.0, + "37": 44290232320.0, + "38": 44290232320.0, + "39": 44290232320.0, + "40": 44290232320.0, + "41": 44290232320.0, + "42": 44290232320.0, + "43": 44290232320.0, + "44": 44290232320.0, + "45": 44290232320.0, + "46": 44290232320.0, + "47": 44290232320.0, + "48": 44290232320.0, + "49": 44290232320.0, + "50": 44290232320.0 } }, "mtp_1 loss": { @@ -232,56 +232,56 @@ "end_step": 50, "step_interval": 1, "values": { - "1": 11.08617, - "2": 11.10475, - "3": 10.48001, - "4": 10.13466, - "5": 9.79047, - "6": 9.50601, - "7": 9.5113, - "8": 8.85336, - "9": 8.66683, - "10": 8.95866, - "11": 8.29315, - "12": 8.36982, - "13": 8.25544, - "14": 7.73322, + "1": 11.08623, + "2": 11.1047, + "3": 10.47999, + "4": 10.13471, + "5": 9.79045, + "6": 9.50607, + "7": 9.51139, + "8": 8.85331, + "9": 8.66688, + "10": 8.95867, + "11": 8.29318, + "12": 8.36986, + "13": 8.25545, + "14": 7.73323, "15": 7.86639, - "16": 7.92442, - "17": 7.86278, - "18": 7.61012, - "19": 8.00269, - "20": 7.73019, - "21": 7.4165, - "22": 7.41478, - "23": 7.28671, - "24": 7.27903, - "25": 7.54456, - "26": 6.96542, - "27": 7.50538, - "28": 7.20607, - "29": 7.377, - "30": 7.52777, - "31": 7.27094, - "32": 7.4604, + "16": 7.92438, + "17": 7.86276, + "18": 7.61004, + "19": 8.00261, + "20": 7.73004, + "21": 7.41636, + "22": 7.41466, + "23": 7.28656, + "24": 7.27882, + "25": 7.54458, + "26": 6.96533, + "27": 7.5053, + "28": 7.20603, + "29": 7.37687, + "30": 7.52783, + "31": 7.27097, + "32": 7.46043, "33": 7.51419, - "34": 7.56867, - "35": 7.09252, - "36": 6.96015, - "37": 7.29846, - "38": 7.0742, - "39": 7.43347, - "40": 7.43116, - "41": 7.40919, + "34": 7.56879, + "35": 7.09276, + "36": 6.96019, + "37": 7.29843, + "38": 7.07417, + "39": 7.43338, + "40": 7.43134, + "41": 7.40946, "42": 7.15527, - "43": 7.15652, - "44": 7.30441, - "45": 7.1893, - "46": 6.77296, - "47": 7.45045, - "48": 7.02403, - "49": 7.45719, - "50": 6.92656 + "43": 7.15684, + "44": 7.30429, + "45": 7.18917, + "46": 6.77286, + "47": 7.44985, + "48": 7.02383, + "49": 7.4572, + "50": 6.92645 } }, "iteration-time": { @@ -289,56 +289,56 @@ "end_step": 50, "step_interval": 1, "values": { - "1": 64.40054, - "2": 2.16564, - "3": 3.72378, - "4": 1.63174, - "5": 2.30947, - "6": 1.7246, - "7": 1.5089, - "8": 1.60943, - "9": 1.48606, - "10": 1.47162, - "11": 1.05608, - "12": 1.3309, - "13": 1.06824, - "14": 1.41914, - "15": 1.10033, - "16": 1.15759, - "17": 1.23897, - "18": 1.10439, - "19": 1.11869, - "20": 1.09363, - "21": 1.23622, - "22": 1.14797, - "23": 1.23037, - "24": 1.03991, - "25": 1.07795, - "26": 1.04416, - "27": 1.03654, - "28": 1.04098, - "29": 1.03502, - "30": 1.02909, - "31": 1.17935, - "32": 1.14717, - "33": 1.05403, - "34": 1.13894, - "35": 1.04538, - "36": 1.04367, - "37": 1.0843, - "38": 1.04631, - "39": 1.06131, - "40": 1.06988, - "41": 1.09756, - "42": 1.04759, - "43": 1.09649, - "44": 1.05666, - "45": 1.05249, - "46": 1.04539, - "47": 1.04041, - "48": 1.04904, - "49": 1.04777, - "50": 1.06237 + "1": 89.89187, + "2": 2.19484, + "3": 3.80506, + "4": 1.63188, + "5": 2.52939, + "6": 2.46374, + "7": 1.5097, + "8": 1.75664, + "9": 1.62191, + "10": 1.35808, + "11": 1.04295, + "12": 1.35317, + "13": 1.07545, + "14": 1.42301, + "15": 1.10347, + "16": 1.28287, + "17": 1.22104, + "18": 1.07676, + "19": 1.08763, + "20": 1.12221, + "21": 1.25145, + "22": 1.04596, + "23": 1.22539, + "24": 1.06194, + "25": 1.11205, + "26": 1.05389, + "27": 1.03357, + "28": 1.0291, + "29": 1.04027, + "30": 1.06631, + "31": 1.18617, + "32": 1.142, + "33": 1.03842, + "34": 1.12457, + "35": 1.04164, + "36": 1.04698, + "37": 1.07674, + "38": 1.03833, + "39": 1.03043, + "40": 1.02697, + "41": 1.11388, + "42": 1.04538, + "43": 1.03328, + "44": 1.04873, + "45": 1.03241, + "46": 1.03847, + "47": 1.04164, + "48": 1.04077, + "49": 1.03715, + "50": 1.02734 } } } \ No newline at end of file diff --git a/tests/functional_tests/test_cases/moe/gpt3_moe_mcore_te_tp2_pp2_ep4_etp1_fine_grained_offloading/golden_values_dev_dgxh100_eos.json b/tests/functional_tests/test_cases/moe/gpt3_moe_mcore_te_tp2_pp2_ep4_etp1_fine_grained_offloading/golden_values_dev_dgxh100_eos.json index d7372742ca7..0a6724a3e95 100644 --- a/tests/functional_tests/test_cases/moe/gpt3_moe_mcore_te_tp2_pp2_ep4_etp1_fine_grained_offloading/golden_values_dev_dgxh100_eos.json +++ b/tests/functional_tests/test_cases/moe/gpt3_moe_mcore_te_tp2_pp2_ep4_etp1_fine_grained_offloading/golden_values_dev_dgxh100_eos.json @@ -4,56 +4,56 @@ "end_step": 50, "step_interval": 1, "values": { - "1": 11.07546, - "2": 11.03837, - "3": 9.66011, - "4": 9.91381, - "5": 9.32909, - "6": 9.13922, - "7": 9.13574, - "8": 8.65508, - "9": 8.51394, - "10": 8.8409, - "11": 8.29149, - "12": 8.34581, - "13": 8.25518, - "14": 7.73711, - "15": 7.86249, - "16": 7.9371, - "17": 7.89319, - "18": 7.63123, - "19": 7.99731, - "20": 7.74538, - "21": 7.44348, - "22": 7.42249, - "23": 7.29714, - "24": 7.27462, - "25": 7.54574, - "26": 6.96838, - "27": 7.50556, - "28": 7.22743, - "29": 7.36588, - "30": 7.52622, - "31": 7.27026, - "32": 7.45521, - "33": 7.50954, - "34": 7.55686, - "35": 7.10177, - "36": 6.96431, - "37": 7.28463, - "38": 7.0808, - "39": 7.40923, - "40": 7.43338, - "41": 7.38496, - "42": 7.15749, - "43": 7.15858, - "44": 7.28852, - "45": 7.16793, - "46": 6.78468, - "47": 7.4114, - "48": 7.0027, - "49": 7.46249, - "50": 6.92151 + "1": 11.07559, + "2": 11.03834, + "3": 9.66022, + "4": 9.91367, + "5": 9.3291, + "6": 9.13927, + "7": 9.13591, + "8": 8.65527, + "9": 8.51396, + "10": 8.84095, + "11": 8.29144, + "12": 8.34584, + "13": 8.25509, + "14": 7.73685, + "15": 7.86273, + "16": 7.93699, + "17": 7.89257, + "18": 7.63116, + "19": 7.99719, + "20": 7.7453, + "21": 7.44298, + "22": 7.42242, + "23": 7.29721, + "24": 7.27467, + "25": 7.54562, + "26": 6.96839, + "27": 7.50569, + "28": 7.22761, + "29": 7.36579, + "30": 7.52635, + "31": 7.27036, + "32": 7.45548, + "33": 7.50952, + "34": 7.55694, + "35": 7.10212, + "36": 6.96414, + "37": 7.28438, + "38": 7.08049, + "39": 7.40908, + "40": 7.4335, + "41": 7.38491, + "42": 7.15766, + "43": 7.15867, + "44": 7.28831, + "45": 7.16729, + "46": 6.78429, + "47": 7.40937, + "48": 7.00259, + "49": 7.46241, + "50": 6.92143 } }, "num-zeros": { @@ -63,54 +63,54 @@ "values": { "1": 911219392.0, "2": 910960384.0, - "3": 911156352.0, - "4": 912204800.0, - "5": 920796544.0, - "6": 940387968.0, - "7": 990599872.0, - "8": 976457728.0, - "9": 998097664.0, - "10": 995852672.0, - "11": 994583680.0, - "12": 977344896.0, - "13": 1028141824.0, - "14": 1007166208.0, - "15": 987423616.0, - "16": 993054784.0, - "17": 982319168.0, - "18": 998261760.0, - "19": 984696320.0, - "20": 982914752.0, - "21": 979667456.0, - "22": 953988864.0, - "23": 972353984.0, - "24": 964792064.0, - "25": 958512192.0, - "26": 946928512.0, + "3": 911156288.0, + "4": 913253376.0, + "5": 921845056.0, + "6": 941436672.0, + "7": 993745472.0, + "8": 974360512.0, + "9": 999146112.0, + "10": 992706944.0, + "11": 991438144.0, + "12": 979442048.0, + "13": 1029190272.0, + "14": 1008214656.0, + "15": 988472000.0, + "16": 988861120.0, + "17": 979173312.0, + "18": 996164608.0, + "19": 979453440.0, + "20": 982914688.0, + "21": 975473344.0, + "22": 955037568.0, + "23": 969208128.0, + "24": 965840832.0, + "25": 953269440.0, + "26": 949025536.0, "27": 948458304.0, - "28": 949643968.0, - "29": 942877440.0, + "28": 951741184.0, + "29": 943926272.0, "30": 935020160.0, - "31": 935327616.0, - "32": 934281088.0, - "33": 921805568.0, - "34": 928189312.0, - "35": 922202496.0, - "36": 924246656.0, - "37": 920661248.0, + "31": 933230336.0, + "32": 930086848.0, + "33": 922853952.0, + "34": 927140800.0, + "35": 925348224.0, + "36": 925295168.0, + "37": 922758272.0, "38": 922930752.0, - "39": 922322816.0, - "40": 921856512.0, - "41": 920227968.0, + "39": 922322880.0, + "40": 921856640.0, + "41": 920227776.0, "42": 918353664.0, - "43": 918607040.0, - "44": 914948032.0, - "45": 914295232.0, + "43": 919655616.0, + "44": 914948224.0, + "45": 916392512.0, "46": 914344448.0, "47": 911769536.0, - "48": 912013312.0, - "49": 910349440.0, - "50": 914351552.0 + "48": 912013248.0, + "49": 910349376.0, + "50": 914351616.0 } }, "mem-allocated-bytes": { @@ -175,56 +175,56 @@ "end_step": 50, "step_interval": 1, "values": { - "1": 41739952128.0, - "2": 43687571456.0, - "3": 43687571456.0, - "4": 43983216640.0, - "5": 43983216640.0, - "6": 43983216640.0, - "7": 43983216640.0, - "8": 44024635392.0, - "9": 44041216000.0, - "10": 44041216000.0, - "11": 44041216000.0, - "12": 44041216000.0, - "13": 44041216000.0, - "14": 44041216000.0, - "15": 44041216000.0, - "16": 44041216000.0, - "17": 44041216000.0, - "18": 44041216000.0, - "19": 44041216000.0, - "20": 44041216000.0, - "21": 44041216000.0, - "22": 44041216000.0, - "23": 44041216000.0, - "24": 44041216000.0, - "25": 44041216000.0, - "26": 44041216000.0, - "27": 44041216000.0, - "28": 44041216000.0, - "29": 44041326592.0, - "30": 44162326528.0, - "31": 44220485632.0, - "32": 44270411776.0, - "33": 44293799936.0, - "34": 44293799936.0, - "35": 44293799936.0, - "36": 44293799936.0, - "37": 44293799936.0, - "38": 44293799936.0, - "39": 44293799936.0, - "40": 44293799936.0, - "41": 44293799936.0, - "42": 44293799936.0, - "43": 44293799936.0, - "44": 44293799936.0, - "45": 44293799936.0, - "46": 44293799936.0, - "47": 44293799936.0, - "48": 44293799936.0, - "49": 44293799936.0, - "50": 44293799936.0 + "1": 41740259328.0, + "2": 43687292928.0, + "3": 43687292928.0, + "4": 43984064512.0, + "5": 43984064512.0, + "6": 43984064512.0, + "7": 43984064512.0, + "8": 44026380288.0, + "9": 44041506816.0, + "10": 44041506816.0, + "11": 44041506816.0, + "12": 44041506816.0, + "13": 44041506816.0, + "14": 44041506816.0, + "15": 44041506816.0, + "16": 44041506816.0, + "17": 44041506816.0, + "18": 44041506816.0, + "19": 44041506816.0, + "20": 44041506816.0, + "21": 44041506816.0, + "22": 44041506816.0, + "23": 44041506816.0, + "24": 44041506816.0, + "25": 44041506816.0, + "26": 44041506816.0, + "27": 44041506816.0, + "28": 44041506816.0, + "29": 44044173312.0, + "30": 44164231168.0, + "31": 44221079552.0, + "32": 44271415296.0, + "33": 44290232320.0, + "34": 44290232320.0, + "35": 44290232320.0, + "36": 44290232320.0, + "37": 44290232320.0, + "38": 44290232320.0, + "39": 44290232320.0, + "40": 44290232320.0, + "41": 44290232320.0, + "42": 44290232320.0, + "43": 44290232320.0, + "44": 44290232320.0, + "45": 44290232320.0, + "46": 44290232320.0, + "47": 44290232320.0, + "48": 44290232320.0, + "49": 44290232320.0, + "50": 44290232320.0 } }, "mtp_1 loss": { @@ -232,56 +232,56 @@ "end_step": 50, "step_interval": 1, "values": { - "1": 11.08617, - "2": 11.10475, - "3": 10.48001, - "4": 10.13466, - "5": 9.79047, - "6": 9.50601, - "7": 9.5113, - "8": 8.85336, - "9": 8.66683, - "10": 8.95866, - "11": 8.29315, - "12": 8.36982, - "13": 8.25544, - "14": 7.73322, + "1": 11.08623, + "2": 11.1047, + "3": 10.47999, + "4": 10.13471, + "5": 9.79045, + "6": 9.50607, + "7": 9.51139, + "8": 8.85331, + "9": 8.66688, + "10": 8.95867, + "11": 8.29318, + "12": 8.36986, + "13": 8.25545, + "14": 7.73323, "15": 7.86639, - "16": 7.92442, - "17": 7.86278, - "18": 7.61012, - "19": 8.00269, - "20": 7.73019, - "21": 7.4165, - "22": 7.41478, - "23": 7.28671, - "24": 7.27903, - "25": 7.54456, - "26": 6.96542, - "27": 7.50538, - "28": 7.20607, - "29": 7.377, - "30": 7.52777, - "31": 7.27094, - "32": 7.4604, + "16": 7.92438, + "17": 7.86276, + "18": 7.61004, + "19": 8.00261, + "20": 7.73004, + "21": 7.41636, + "22": 7.41466, + "23": 7.28656, + "24": 7.27882, + "25": 7.54458, + "26": 6.96533, + "27": 7.5053, + "28": 7.20603, + "29": 7.37687, + "30": 7.52783, + "31": 7.27097, + "32": 7.46043, "33": 7.51419, - "34": 7.56867, - "35": 7.09252, - "36": 6.96015, - "37": 7.29846, - "38": 7.0742, - "39": 7.43347, - "40": 7.43116, - "41": 7.40919, + "34": 7.56879, + "35": 7.09276, + "36": 6.96019, + "37": 7.29843, + "38": 7.07417, + "39": 7.43338, + "40": 7.43134, + "41": 7.40946, "42": 7.15527, - "43": 7.15652, - "44": 7.30441, - "45": 7.1893, - "46": 6.77296, - "47": 7.45045, - "48": 7.02403, - "49": 7.45719, - "50": 6.92656 + "43": 7.15684, + "44": 7.30429, + "45": 7.18917, + "46": 6.77286, + "47": 7.44985, + "48": 7.02383, + "49": 7.4572, + "50": 6.92645 } }, "iteration-time": { @@ -289,56 +289,56 @@ "end_step": 50, "step_interval": 1, "values": { - "1": 87.63934, - "2": 1.98402, - "3": 3.95877, - "4": 1.64812, - "5": 2.312, - "6": 2.02902, - "7": 1.56333, - "8": 1.66703, - "9": 1.6393, - "10": 1.40472, - "11": 1.086, - "12": 1.34921, - "13": 1.0854, - "14": 1.4242, - "15": 1.09539, - "16": 1.79766, - "17": 1.2562, - "18": 1.08887, - "19": 1.08371, - "20": 1.10071, - "21": 1.25979, - "22": 1.3212, - "23": 1.25044, - "24": 1.05384, - "25": 1.11356, - "26": 1.0605, - "27": 1.03418, - "28": 1.0405, - "29": 1.05174, - "30": 1.04166, - "31": 1.20036, - "32": 1.12936, - "33": 1.02917, - "34": 1.13473, - "35": 1.02829, - "36": 1.04352, - "37": 1.0843, - "38": 1.03714, - "39": 1.04534, - "40": 1.07031, - "41": 1.07618, - "42": 1.03008, - "43": 1.06043, - "44": 1.04049, - "45": 1.02875, - "46": 1.03669, - "47": 1.03128, - "48": 1.02808, - "49": 1.03038, - "50": 1.04621 + "1": 85.92313, + "2": 1.99152, + "3": 3.91366, + "4": 1.68454, + "5": 2.53883, + "6": 2.55539, + "7": 1.60104, + "8": 1.70562, + "9": 1.72325, + "10": 1.4332, + "11": 1.07958, + "12": 1.399, + "13": 1.10259, + "14": 1.43922, + "15": 1.12046, + "16": 1.33695, + "17": 1.24765, + "18": 1.11257, + "19": 1.10335, + "20": 1.12919, + "21": 1.27711, + "22": 1.09482, + "23": 1.27635, + "24": 1.112, + "25": 1.17791, + "26": 1.10426, + "27": 1.09103, + "28": 1.08338, + "29": 1.07904, + "30": 1.08709, + "31": 1.2237, + "32": 1.18059, + "33": 1.07913, + "34": 1.17232, + "35": 1.09059, + "36": 1.09648, + "37": 1.12683, + "38": 1.10153, + "39": 1.09557, + "40": 1.07747, + "41": 1.12905, + "42": 1.09275, + "43": 1.08609, + "44": 1.08042, + "45": 1.08321, + "46": 1.0732, + "47": 1.08666, + "48": 1.08865, + "49": 1.08808, + "50": 1.08086 } } } \ No newline at end of file diff --git a/tests/test_utils/recipes/moe.yaml b/tests/test_utils/recipes/moe.yaml index 649da3ba518..53047ff4a3b 100644 --- a/tests/test_utils/recipes/moe.yaml +++ b/tests/test_utils/recipes/moe.yaml @@ -106,14 +106,13 @@ products: - environment: [dev] scope: [mr, mr-github] platforms: [dgx_h100] - # TODO: The migration of custom fsdp causes EP + FSDP to be temporarily unavailable, which will be fixed in a subsequent MR. - # - test_case: [gpt3_mcore_te_tp2_zp_z3_resume_torch_dist_te_8experts2parallel_top2router] - # products: - # - environment: [dev] - # scope: [mr] - # platforms: [dgx_h100] - # - environment: [lts] - # scope: [nightly] + - test_case: [gpt3_mcore_te_tp2_zp_z3_resume_torch_dist_te_8experts2parallel_top2router] + products: + - environment: [dev] + scope: [mr] + platforms: [dgx_h100] + - environment: [lts] + scope: [nightly] - test_case: [gpt3_mcore_te_tp2_pp1_te_8experts2parallel_ddp_average_in_collective] products: - environment: [dev] diff --git a/tools/checkpoint/checkpoint_inspector.py b/tools/checkpoint/checkpoint_inspector.py index 34afa27755f..c62f0ca7417 100644 --- a/tools/checkpoint/checkpoint_inspector.py +++ b/tools/checkpoint/checkpoint_inspector.py @@ -8,6 +8,8 @@ import time import re import shutil +from typing import Optional +import tempfile import click import torch @@ -19,6 +21,7 @@ FileSystemReader, FileSystemWriter, ) +from torch.distributed.checkpoint.format_utils import dcp_to_torch_save from torch.distributed.checkpoint.metadata import ( BytesStorageMetadata, TensorStorageMetadata, @@ -64,7 +67,8 @@ def cli(): @cli.command() @click.argument("checkpoint_dir", type=click.Path(exists=True)) @click.option("--enable-msc", is_flag=True, help="Enable MultiStorageClient feature.") -def inspect(checkpoint_dir, enable_msc): +@click.option("--not-ignore-param-to-group-meta", is_flag=True, help="Ignore parameter-to-group metadata.") +def inspect(checkpoint_dir, enable_msc, not_ignore_param_to_group_meta): """Inspect a Megatron Core Distributed Checkpoint""" ckpt_path = Path(checkpoint_dir) @@ -138,6 +142,8 @@ def inspect(checkpoint_dir, enable_msc): ] click.echo(" | ".join(stats) + "\n") + ignore_param_to_group_meta = not not_ignore_param_to_group_meta + ignore_param_to_group_meta_count = 0 for key, value in metadata.state_dict_metadata.items(): bullet = click.style("►", fg="blue") key_styled = click.style(key, fg="green") @@ -147,11 +153,18 @@ def inspect(checkpoint_dir, enable_msc): shape = click.style(f"{tuple(value.size)}", fg="magenta") click.echo(f" {bullet} {key_styled} [{dtype}, shape={shape}]") elif isinstance(value, BytesStorageMetadata): + if ignore_param_to_group_meta and key.startswith("optimizer.param_to_group_meta."): + ignore_param_to_group_meta_count += 1 + continue click.echo(f" {bullet} {key_styled} {click.style('[BYTES]', fg='yellow')}") else: click.echo( f" {bullet} {key_styled} {click.style('[UNKNOWN TYPE]', fg='red')}" ) + if ignore_param_to_group_meta: + click.echo( + click.style(f"Ignored parameter-to-group metadata: {ignore_param_to_group_meta_count}", fg="yellow") + ) # MCore data section try: @@ -323,8 +336,10 @@ def convert_checkpoint( output_dir, swiglu, process_group, + optimizer_param_to_group_prefix="optimizer.param_to_group_meta.module.module.module", optimizer_state_prefix="optimizer.state.module.module.module", model_weight_prefix="model.module", + param_to_param_group_map={}, ): """Convert a Megatron Core Distributed Checkpoint from torch_dist to standard fsdp_dtensor format.""" device_mesh = DeviceMesh.from_group(process_group, device_type="cuda") @@ -371,6 +386,104 @@ def _free_up_some_gpu_memory(): gc.collect() torch.cuda.empty_cache() + def split_layers( + key: str, + value: torch.Tensor, + orig_shape: Optional[torch.Size] = None, + ) -> dict[str, torch.Tensor]: + """ + Split layers into separate tensors. + """ + _free_up_some_gpu_memory() + layers = {} + for i, v in enumerate(split_dtensor(value, 1, dim=0)): + v = gather_uneven_dtensor_to_full_tensor(v).reshape( + orig_shape[1:] if orig_shape else value.shape[1:] + ).redistribute(placements=[Shard(0)]) + + layer_key = key.replace(".layers.", f".layers.{i}.") + layers[layer_key] = v + + return layers + + def split_expert_weights( + key: str, + value: torch.Tensor, + orig_shape: Optional[torch.Size] = None, + ) -> dict[str, torch.Tensor]: + """ + Split expert weights into separate tensors for each expert. + """ + experts = {} + layer_key = key.replace(".experts.experts.", ".experts.") + expert_weights = split_dtensor(value, 1, dim=0) + for expert_idx, expert_weight in enumerate(expert_weights): + layer_key_parts = layer_key.split(".weight", 1) + if len(layer_key_parts) == 1: + expert_key = f"{layer_key}{expert_idx}" + elif len(layer_key_parts) == 2: + expert_key = f"{layer_key_parts[0]}.weight{expert_idx}{layer_key_parts[1]}" + else: + raise ValueError(f"Unexpected expert layer key: {layer_key}") + + expert_weight = gather_uneven_dtensor_to_full_tensor(expert_weight) + expert_shape = orig_shape[1:] if orig_shape else value.shape[1:] + # Handle optimizer states for expert linear_fc2 when ETP is enabled + if ( + layer_key.startswith("optimizer.state.") + and "linear_fc2" in layer_key + and expert_weight.shape[-2] > 1 + ): + tp_size = expert_weight.shape[-2] + rows, cols = expert_shape + # Reshape to split column dimension by tp_size + expert_weight = expert_weight.reshape( + *expert_weight.shape[:-1], rows, cols // tp_size + ) + dims = list(range(expert_weight.ndim)) + dims[-3], dims[-2] = dims[-2], dims[-3] + expert_weight = ( + expert_weight.permute(*dims) + .reshape(expert_shape) + .redistribute(placements=[Shard(0)]) + ) + else: + expert_weight = expert_weight.reshape(expert_shape).redistribute( + placements=[Shard(0)] + ) + experts[expert_key] = expert_weight + return experts + + def is_swiglu_key(key): + return any(re.search(pat, key) for pat in [ + r"(.*)\.mlp\.linear_fc1\.weight", + r"(.*)\.mlp\.linear_fc1\.bias", + r"(.*)\.mlp\.experts\.linear_fc1\.weight(\d+)", + r"(.*)\.mlp\.experts\.linear_fc1\.bias(\d+)", + r"(.*)\.mlp\.experts\.local_experts\.(\d+)\.linear_fc1\.weight", + r"(.*)\.mlp\.experts\.local_experts\.(\d+)\.linear_fc1\.bias", + r"(.*)\.mlp\.shared_experts\.linear_fc1\.weight", + r"(.*)\.mlp\.shared_experts\.linear_fc1\.bias", + ]) + + def split_swiglu_weight(key: str, value: torch.Tensor) -> dict[str, torch.Tensor]: + """ + Split SwiGLU weights into separate tensors. + """ + value = gather_uneven_dtensor_to_full_tensor(value) + swiglu_w_and_v = {} + w, v = torch.chunk(value, 2, dim=0) + w = w.redistribute(placements=[Shard(0)]) + v = v.redistribute(placements=[Shard(0)]) + w_key = re.sub(r'(weight\d*)(.*)', r'\1_w\2', key) + v_key = re.sub(r'(weight\d*)(.*)', r'\1_v\2', key) + swiglu_w_and_v[w_key] = w + swiglu_w_and_v[v_key] = v + return swiglu_w_and_v + + def has_layer_index(key: str) -> bool: + return bool(re.search(r"layers\.(\d+)\.", key)) + while state_dict: key, value = state_dict.popitem() if torch.distributed.get_rank() == 0: @@ -387,9 +500,11 @@ def _free_up_some_gpu_memory(): # Special handling for optimizer state key_list = key.split(".") new_key = f"{optimizer_state_prefix}.{'.'.join(key_list[3:])}.{key_list[2]}" + is_param = False else: # Special handling for module parameters new_key = f"{model_weight_prefix}.{key}" + is_param = True # Handle dist-opt flatten tensors if ( @@ -406,68 +521,47 @@ def _free_up_some_gpu_memory(): else: orig_shape = None - # Handle multi-layer tensors - if ".layers." in new_key: - n_layer = value.shape[0] - - _free_up_some_gpu_memory() - per_layer_values = [ - gather_uneven_dtensor_to_full_tensor(v).redistribute( - placements=[Shard(len(v.shape) - 1)] - ) - for v in split_dtensor(value, 1, dim=0) - ] - for i in range(n_layer): - if orig_shape is not None: - layer_shape = orig_shape[1:] - else: - layer_shape = value.shape[1:] - - per_layer_values[i] = ( - per_layer_values[i] - .reshape(layer_shape) - .redistribute(placements=[Shard(0)]) - ) - for i in range(0, n_layer): - layer_key = new_key.replace(".layers.", f".layers.{i}.") - if swiglu and "mlp.linear_fc1.weight" in layer_key: - # Special case for SwiGLU - w, v = torch.chunk(per_layer_values[i], 2, dim=0) - w = w.redistribute(placements=[Shard(0)]) - v = v.redistribute(placements=[Shard(0)]) - w_key = layer_key.replace( - "mlp.linear_fc1.weight", "mlp.linear_fc1.weight_w" - ) - v_key = layer_key.replace( - "mlp.linear_fc1.weight", "mlp.linear_fc1.weight_v" - ) - # Store both w and v in the state_dict - fsdp_dtensor_state_dict[w_key] = w - fsdp_dtensor_state_dict[v_key] = v - elif ( - "experts.experts.linear_fc1.weight" in layer_key - or "experts.experts.linear_fc2.weight" in layer_key + # Handle multi-layer / experts tensors + split_tensors = {} + if ".layers." in new_key and not has_layer_index(new_key): + split_tensors = split_layers(new_key, value, orig_shape) + elif ".experts.experts." in new_key: + split_tensors = split_expert_weights(new_key, value, orig_shape) + else: + if orig_shape: + value = gather_uneven_dtensor_to_full_tensor(value) + # Handle optimizer states with partition_dim=1 when TP is enabled + if ( + new_key.startswith("optimizer.state.") + and value.ndim > 2 + and value.shape[-2] > 1 ): - # Special case for MoE - layer_key = layer_key.replace(".experts.experts.", ".experts.") - expert_weights = torch.split(per_layer_values[i], 1, dim=0) - for expert_idx, expert_weight in enumerate(expert_weights): - expert_key = f"{layer_key}{expert_idx}" - fsdp_dtensor_state_dict[expert_key] = expert_weight.squeeze( - 0 - ) + tp_size = value.shape[-2] + rows, cols = orig_shape + # Reshape to split column dimension by tp_size + value = value.reshape(*value.shape[:-1], rows, cols // tp_size) + dims = list(range(value.ndim)) + dims[-3], dims[-2] = dims[-2], dims[-3] + value = ( + value.permute(*dims) + .reshape(orig_shape) + .redistribute(placements=[Shard(0)]) + ) else: - # General case - fsdp_dtensor_state_dict[layer_key] = per_layer_values[i] - else: - if orig_shape is not None: - _free_up_some_gpu_memory() - value = ( - value.redistribute(placements=[Replicate()]) - .reshape(orig_shape) - .redistribute(placements=[Shard(0)]) - ) - fsdp_dtensor_state_dict[new_key] = value + value = value.reshape(orig_shape).redistribute(placements=[Shard(0)]) + split_tensors = {new_key: value} + + # Handle SWiGLU weights + for key, value in list(split_tensors.items()): + if swiglu and is_swiglu_key(key): + swiglu_w_and_v = split_swiglu_weight(key, value) + split_tensors.update(swiglu_w_and_v) + del split_tensors[key] + + fsdp_dtensor_state_dict.update(split_tensors) + if is_param and key in param_to_param_group_map: + for new_key in split_tensors.keys(): + param_to_param_group_map[new_key] = param_to_param_group_map[key] elif key.startswith("rng_state"): # Skip RNG states continue @@ -530,6 +624,15 @@ def _free_up_some_gpu_memory(): ) ) common_state = common_strategy.load_common(input_dir) + try: + if "param_groups" in common_state["optimizer"]: + ckpt_param_groups = common_state["optimizer"]["param_groups"] + else: + ckpt_param_groups = [] + for opt_state_dict in common_state["optimizer"].values(): + ckpt_param_groups.extend(opt_state_dict["optimizer"]["param_groups"]) + except: + ckpt_param_groups = None common_state = flatten(common_state) for key, value in common_state.items(): if key.startswith("optimizer.optimizer.param_groups."): @@ -541,12 +644,29 @@ def _free_up_some_gpu_memory(): ) fsdp_dtensor_state_dict[key] = value + # set up per-parameter param_groups + if param_to_param_group_map and ckpt_param_groups is not None: + for name in list(fsdp_dtensor_state_dict.keys()): + if not name.startswith(model_weight_prefix) or name.endswith(".expert_bias"): + continue + + assert name in param_to_param_group_map, f"Missing param group for {name}" + param_group_id = param_to_param_group_map[name] + assert param_group_id < len(ckpt_param_groups), f"Invalid param group id {param_group_id} for {name}" + name_without_prefix = name[len(model_weight_prefix):] + fsdp_dtensor_state_dict[ + f"{optimizer_param_to_group_prefix}.{name_without_prefix}" + ] = ckpt_param_groups[param_group_id] + if "checkpoint_version" not in fsdp_dtensor_state_dict: fsdp_dtensor_state_dict["checkpoint_version"] = 3.0 # Save modified checkpoint save_checkpoint_with_pickle_protocol(fsdp_dtensor_state_dict, output_dir) + dist.barrier() # Synchronize all ranks + dist.destroy_process_group() + @cli.command() @click.argument("input_dir", type=click.Path(exists=True)) @@ -560,12 +680,6 @@ def _free_up_some_gpu_memory(): "--oom-traceback", is_flag=True, help="Enable OOM traceback for debugging." ) @click.option("--enable-msc", is_flag=True, help="Enable MultiStorageClient feature.") -@click.option( - "--distributed-timeout-minutes", - default=10, - type=int, - help="Timeout for distributed operations in minutes.", -) @click.option( "--output-optimizer-state-prefix", default="optimizer.state.module.module.module", @@ -576,15 +690,21 @@ def _free_up_some_gpu_memory(): default="model.module", help="Prefix for model weight keys in the checkpoint.", ) +@click.option( + "--param-to-param-group-map-json", + type=str, + default="{}", + help="JSON string representing the param to parameter group map." +) def convert_torch_dist_to_fsdp_dtensor( input_dir, output_dir, swiglu, oom_traceback, enable_msc, - distributed_timeout_minutes, output_optimizer_state_prefix, output_model_weight_prefix, + param_to_param_group_map_json, ): """Convert a Megatron Core Distributed Checkpoint from torch_dist to fsdp_dtensor format.""" if not enable_msc: @@ -624,10 +744,13 @@ def oom_observer(device, alloc, device_alloc, device_free): ckpt_path = Path(input_dir) output_dir = Path(output_dir) + with open(param_to_param_group_map_json, "r") as f: + param_to_param_group_map = json.load(f) convert_checkpoint( ckpt_path, output_dir, swiglu, process_group=dist.group.WORLD, optimizer_state_prefix=output_optimizer_state_prefix, model_weight_prefix=output_model_weight_prefix, + param_to_param_group_map=param_to_param_group_map, ) click.echo( @@ -742,6 +865,109 @@ def modify_state_dict(input_dir, output_dir, op, enable_msc): ) +def _compare_two_checkpoint(checkpoint_1, checkpoint_2): + reader_1 = FileSystemReader(checkpoint_1) + metadata_1 = reader_1.read_metadata() + + reader_2 = FileSystemReader(checkpoint_2) + metadata_2 = reader_2.read_metadata() + + keys_1 = set(metadata_1.state_dict_metadata.keys()) + keys_2 = set(metadata_2.state_dict_metadata.keys()) + + click.echo(click.style("Comparing checkpoints...", fg="blue")) + + # Compare keys + missing_in_1 = keys_2 - keys_1 + missing_in_2 = keys_1 - keys_2 + common_keys = keys_1 & keys_2 + + click.echo(click.style("Keys missing in checkpoint 1:", fg="red")) + for key in missing_in_1: + click.echo(click.style(f" - {key}", fg="red")) + + click.echo(click.style("Keys missing in checkpoint 2:", fg="red")) + for key in missing_in_2: + click.echo(click.style(f" - {key}", fg="red")) + + # Compare common keys + click.echo(click.style("Common keys in both checkpoints:", fg="green")) + for key in common_keys: + meta_1 = metadata_1.state_dict_metadata[key] + meta_2 = metadata_2.state_dict_metadata[key] + + if not isinstance(meta_1, TensorStorageMetadata): + continue + + if meta_1.size != meta_2.size or meta_1.properties.dtype != meta_2.properties.dtype: + click.echo(click.style(f" - {key} (metadata differ) meta_1: {meta_1}, meta_2: {meta_2}", fg="red")) + else: + value_1 = torch.empty(meta_1.size, dtype=meta_1.properties.dtype) + value_2 = value_1.clone() + + dcp.load({key: value_1}, storage_reader=reader_1, planner=DefaultLoadPlanner()) + dcp.load({key: value_2}, storage_reader=reader_2, planner=DefaultLoadPlanner()) + + if not torch.allclose( + value_1, value_2, atol=1e-8, rtol=1e-5 + ): + click.echo(click.style(f" - {key} (values differ) value_1: {value_1}, value_2: {value_2}", fg="red")) + + +@cli.command() +@click.argument("checkpoint_1", type=click.Path(exists=True)) +@click.argument("checkpoint_2", type=click.Path(exists=True)) +@click.option("--enable-msc", is_flag=True, help="Enable MultiStorageClient feature.") +def compare_two_checkpoint(checkpoint_1, checkpoint_2, enable_msc): + """ + Compare two checkpoints. + """ + init_process_group(f"compare_two_checkpoint from {checkpoint_1} to {checkpoint_2}") + + if not enable_msc: + MultiStorageClientFeature.disable() + + _compare_two_checkpoint( + Path(checkpoint_1), + Path(checkpoint_2), + ) + + click.echo( + click.style( + f"Comparison between {checkpoint_1} and {checkpoint_2} completed.", fg="green", bold=True + ) + ) + + +@cli.command() +@click.argument("torch_dcp_dir", type=click.Path(exists=True)) +def print_torch_dcp_in_json(torch_dcp_dir, model_weight_prefix="model.module"): + # Use a temporary file context + with tempfile.NamedTemporaryFile(suffix=".pth") as tmp_file: + # Convert distributed checkpoint directory to a single-file checkpoint + dcp_to_torch_save(torch_dcp_dir, tmp_file.name) + + # Load the state dict from the temporary file + state_dict = torch.load(tmp_file.name, map_location="cpu") + + click.echo(f"torch dcp content: {json.dumps(state_dict)}") + + # Replace all "module.module." with model_weight_prefix in dict keys + new_state_dict = {} + for key, value in state_dict.items(): + new_key = key.replace("module.module", model_weight_prefix) + new_state_dict[new_key] = value + + # Convert state dict to JSON-serializable format + serializable_dict = {k: v.tolist() if hasattr(v, "tolist") else v for k, v in new_state_dict.items()} + + # Save to a JSON file + json_file_path = os.path.join(torch_dcp_dir, "param_to_param_group_map.json") + with open(json_file_path, "w") as json_file: + json.dump(serializable_dict, json_file, indent=2) + click.echo(f"Saved converted param_to_param_group_map to: {json_file_path}") + + def init_process_group(message): rank = int(os.getenv("RANK", "0")) world_size = int(os.getenv("WORLD_SIZE", "1"))