diff --git a/megatron/core/extensions/transformer_engine.py b/megatron/core/extensions/transformer_engine.py index 085109c60eb..8a0e1f6b853 100644 --- a/megatron/core/extensions/transformer_engine.py +++ b/megatron/core/extensions/transformer_engine.py @@ -1707,11 +1707,11 @@ def __init__( self.disable_parameter_transpose_cache = self.config.disable_parameter_transpose_cache extra_kwargs = _get_extra_te_kwargs(config) + self.delay_wgrad_compute = ( self.config.delay_wgrad_compute or self.config.overlap_dispatch_backward_with_experts_wgrad ) - if self.delay_wgrad_compute: if is_te_min_version("2.3.0"): extra_kwargs["delay_wgrad_compute"] = True @@ -1755,6 +1755,14 @@ def __init__( tp_size = 1 tp_group_for_te = None + if is_te_min_version("2.14.0"): + extra_kwargs["single_grouped_weight"] = getattr( + config, "moe_single_grouped_weight", False + ) + extra_kwargs["single_grouped_bias"] = getattr( + config, "moe_single_grouped_bias", False + ) + super().__init__( num_gemms=num_gemms, in_features=input_size, @@ -1791,6 +1799,10 @@ def __init__( setattr(weight, "partition_dim", part_dim) setattr(weight, "partition_stride", 1) + self._register_load_state_dict_pre_hook( + type(self)._normalize_grouped_parameter_keys, with_module=True + ) + def merge_extra_states( self, state_dict, @@ -1881,6 +1893,76 @@ def merge_extra_states( self._register_load_state_dict_pre_hook(merge_extra_states, with_module=True) + def _normalize_grouped_parameter_keys( + self, + state_dict, + prefix, + local_metadata, + strict, + missing_keys, + unexpected_keys, + error_msgs, + ): + """Make grouped checkpoint keys compatible across parameter layouts. + + Registered as a load_state_dict pre-hook to bridge checkpoints saved + in one layout (single grouped tensor vs per-GEMM indexed tensors) + and a model expecting the other. + """ + + def maybe_remap_param(param_name: str, single_grouped: bool) -> None: + grouped_key = f"{prefix}{param_name}" + indexed_keys = [ + f"{prefix}{param_name}{gemm_idx}" for gemm_idx in range(self.num_gemms) + ] + has_grouped_key = grouped_key in state_dict + has_any_indexed_key = any(key in state_dict for key in indexed_keys) + has_all_indexed_keys = all(key in state_dict for key in indexed_keys) + + if single_grouped: + if has_grouped_key or not has_all_indexed_keys: + return + state_dict[grouped_key] = torch.stack( + [state_dict.pop(key) for key in indexed_keys], dim=0 + ) + else: + if has_any_indexed_key or not has_grouped_key: + return + split_tensors = self._split_grouped_checkpoint_tensor( + state_dict.pop(grouped_key), grouped_key + ) + for gemm_idx, tensor in enumerate(split_tensors): + state_dict[f"{prefix}{param_name}{gemm_idx}"] = tensor + + maybe_remap_param("weight", getattr(self, "single_grouped_weight", False)) + if self.use_bias: + maybe_remap_param("bias", getattr(self, "single_grouped_bias", False)) + + def _split_grouped_checkpoint_tensor( + self, tensor: torch.Tensor, checkpoint_key: str + ) -> list[torch.Tensor]: + """Split grouped checkpoint tensor into one tensor per GEMM.""" + if hasattr(tensor, "split_into_quantized_tensors") and callable( + tensor.split_into_quantized_tensors + ): + grouped_tensors = getattr(tensor, "quantized_tensors", None) + if grouped_tensors is None: + grouped_tensors = tensor.split_into_quantized_tensors() + if len(grouped_tensors) != self.num_gemms: + raise RuntimeError( + f"Grouped checkpoint tensor {checkpoint_key} has {len(grouped_tensors)} " + f"groups, expected {self.num_gemms}." + ) + return list(grouped_tensors) + if tensor.ndim > 0 and tensor.shape[0] == self.num_gemms: + return list(tensor.unbind(dim=0)) + if tensor.ndim > 0 and tensor.shape[0] % self.num_gemms == 0: + return list(torch.chunk(tensor, self.num_gemms, dim=0)) + raise RuntimeError( + f"Cannot split checkpoint tensor {checkpoint_key} with shape {tuple(tensor.shape)} " + f"into {self.num_gemms} GEMM shards." + ) + def finish_init(self, quantization_config: QuantizationConfig): """Post-init of quantization override""" if quantization_config is None: @@ -1987,6 +2069,21 @@ def _sharded_state_dict_grouped( singleton_local_shards = (metadata or {}).get('singleton_local_shards', False) sharded_state_dict = {} full_state_dict = self.state_dict(prefix="", keep_vars=True) + grouped_split_cache = {} + + def get_gemm_tensor(param_name: str, gemm_idx: int) -> torch.Tensor: + indexed_name = f"{param_name}{gemm_idx}" + if indexed_name in full_state_dict: + return full_state_dict[indexed_name] + if param_name not in full_state_dict: + raise KeyError(indexed_name) + if param_name not in grouped_split_cache: + grouped_split_cache[param_name] = self._split_grouped_checkpoint_tensor( + full_state_dict[param_name], param_name + ) + grouped_splits = grouped_split_cache[param_name] + return grouped_splits[gemm_idx] + num_global_experts = get_pg_size(self._pg_collection.ep) * self.num_gemms local_expert_indices_offset = get_pg_rank(self._pg_collection.ep) * self.num_gemms ep_axis = len(sharded_offsets) @@ -1994,11 +2091,11 @@ def _sharded_state_dict_grouped( for gemm_idx in range(self.num_gemms): global_expert_idx = local_expert_indices_offset + gemm_idx state_dict = { - f"{gemm_idx}.weight": full_state_dict[f"weight{gemm_idx}"], + f"{gemm_idx}.weight": get_gemm_tensor("weight", gemm_idx), f"{gemm_idx}._extra_state": extra_states[gemm_idx], } if self.use_bias: - state_dict[f"{gemm_idx}.bias"] = full_state_dict[f"bias{gemm_idx}"] + state_dict[f"{gemm_idx}.bias"] = get_gemm_tensor("bias", gemm_idx) if singleton_local_shards: expert_prefix = f"{global_expert_idx}.{prefix}" new_sharded_offsets = sharded_offsets diff --git a/megatron/core/optimizer/distrib_optimizer.py b/megatron/core/optimizer/distrib_optimizer.py index 95e8a0c407b..e0d3c2a54ac 100644 --- a/megatron/core/optimizer/distrib_optimizer.py +++ b/megatron/core/optimizer/distrib_optimizer.py @@ -379,9 +379,11 @@ def _build_model_and_main_param_groups( # Generate sharded model param. if ( - is_float8tensor(model_param) and config.fp8_recipe != "delayed" + cls._is_distopt_quantized_param(model_param) + and config.fp8_recipe != "delayed" ) or is_nvfp4tensor(model_param): - # MXFP8Tensor, BlockwiseQTensor, and NVFP4Tensor don't support view(-1) + # MXFP8Tensor, BlockwiseQTensor, grouped quantized tensors, and NVFP4Tensor + # don't support view(-1). shard_model_param = None else: shard_model_param = model_param.detach().view(-1)[ @@ -400,7 +402,9 @@ def _build_model_and_main_param_groups( # precision at the beginning of training (this problem will not occur if the # training is long enough or if the main params are loaded from a # checkpoint). - if is_nvfp4tensor(model_param) or is_float8tensor(model_param): + if is_nvfp4tensor(model_param) or cls._is_distopt_quantized_param( + model_param + ): if hasattr(model_param, 'get_high_precision_init_val'): shard_main_param = ( model_param.get_high_precision_init_val() @@ -1062,6 +1066,70 @@ def _get_main_param_and_optimizer_states(self, model_param): tensors[k] = v return tensors + @staticmethod + def _is_grouped_quantized_tensor(tensor: torch.Tensor) -> bool: + """Check if tensor is a TE GroupedTensor using quantized storage.""" + return ( + hasattr(tensor, "split_into_quantized_tensors") + and callable(tensor.split_into_quantized_tensors) + and getattr(tensor, "quantizer", None) is not None + ) + + @classmethod + def _is_distopt_quantized_param(cls, tensor: torch.Tensor) -> bool: + """Check if tensor should follow quantized parameter path in dist optimizer.""" + return is_float8tensor(tensor) or cls._is_grouped_quantized_tensor(tensor) + + def _expand_quantized_param_shard_for_cast( + self, + model_param: torch.Tensor, + shard_main_param: Optional[torch.Tensor], + start_offset: Optional[int], + ): + """Expand one quantized model param to cast-ready entries. + + For grouped quantized tensors, split into member quantized tensors and map the sharded + master slice to per-member offset ranges, while preserving deterministic ordering across + DP ranks. + """ + if not self._is_grouped_quantized_tensor(model_param): + return [model_param], [shard_main_param], [start_offset] + + quantized_members = model_param.quantized_tensors + if quantized_members is None: + quantized_members = model_param.split_into_quantized_tensors() + + shard_start = 0 if start_offset is None else start_offset + shard_size = 0 if shard_main_param is None else shard_main_param.numel() + shard_end = shard_start + shard_size + shard_flat = None if shard_main_param is None else shard_main_param.view(-1) + + expanded_model_params = [] + expanded_shard_main_params = [] + expanded_start_offsets = [] + member_offset = 0 + for member in quantized_members: + member_numel = member.numel() + member_start = member_offset + member_end = member_start + member_numel + overlap_start = max(member_start, shard_start) + overlap_end = min(member_end, shard_end) + + member_master = None + member_start_offset = None + if overlap_start < overlap_end: + local_start = overlap_start - shard_start + local_end = overlap_end - shard_start + member_master = shard_flat[local_start:local_end] + member_start_offset = overlap_start - member_start + + expanded_model_params.append(member) + expanded_shard_main_params.append(member_master) + expanded_start_offsets.append(member_start_offset) + member_offset = member_end + + return expanded_model_params, expanded_shard_main_params, expanded_start_offsets + def _set_main_param_and_optimizer_states(self, model_param, tensors): """Set the main param and optimizer states corresponding to the input model_param. @@ -2298,7 +2366,7 @@ def split_state_dict_if_needed(self, state_dict): fp8_gbuf_indices = [] for gbuf_idx, gbuf_range_maps in enumerate(self.gbuf_ranges): for dtype, _ in gbuf_range_maps.items(): - if is_float8tensor(self.buffers[gbuf_idx].params[0]): + if self._is_distopt_quantized_param(self.buffers[gbuf_idx].params[0]): fp8_gbuf_indices.append(gbuf_idx) if len(fp8_gbuf_indices) == 0: return @@ -2320,7 +2388,7 @@ def split_state_dict_if_needed(self, state_dict): new_state_dict = {'buckets_coalesced': state_dict['buckets_coalesced']} for gbuf_idx, gbuf_range_maps in enumerate(self.gbuf_ranges): for dtype, _ in gbuf_range_maps.items(): - if not is_float8tensor(self.buffers[gbuf_idx].params[0]): + if not self._is_distopt_quantized_param(self.buffers[gbuf_idx].params[0]): new_state_dict[gbuf_idx] = state_dict[dtype_to_gbuf_idx[dtype]] for fp8_gbuf_idx in fp8_gbuf_indices: @@ -2522,7 +2590,7 @@ def _get_fp8_params_and_shard_fp32_from_fp8(self): idx = 0 for buffer in buffers: for param in buffer.params: - if is_float8tensor(param): + if self._is_distopt_quantized_param(param): fp8_params.append(param) shard_fp32_from_fp8.append(None) shard_offsets_in_fp8.append(None) @@ -2537,7 +2605,7 @@ def get_shard_fp32_from_fp8(shard_main_groups, model_groups): """ for shard_main_group, model_group in zip(shard_main_groups, model_groups): for shard_main_param, model_param in zip(shard_main_group, model_group): - if is_float8tensor(model_param): + if self._is_distopt_quantized_param(model_param): param_range_map = self._get_model_param_range_map(model_param) param_range = param_range_map["param"] assert param_range.size == shard_main_param.nelement() @@ -2660,8 +2728,32 @@ def _copy_main_params_to_model_params(self): return if self.ddp_config.fp8_param_gather: + # Grouped quantized tensors expose one logical parameter backed by multiple TE + # quantized members. Expand them so quantize_param_shard receives member-aligned + # master shards instead of a shard over the grouped wrapper. + fp8_params, shard_fp32_from_fp8, shard_offsets_in_fp8 = ( + self._get_fp8_params_and_shard_fp32_from_fp8() + ) + expanded_fp8_params = [] + expanded_shard_fp32_from_fp8 = [] + expanded_shard_offsets_in_fp8 = [] + for model_param, shard_main_param, start_offset in zip( + fp8_params, shard_fp32_from_fp8, shard_offsets_in_fp8 + ): + sub_model_params, sub_shard_main_params, sub_start_offsets = ( + self._expand_quantized_param_shard_for_cast( + model_param, shard_main_param, start_offset + ) + ) + expanded_fp8_params.extend(sub_model_params) + expanded_shard_fp32_from_fp8.extend(sub_shard_main_params) + expanded_shard_offsets_in_fp8.extend(sub_start_offsets) + quantize_param_shard( - *self._get_fp8_params_and_shard_fp32_from_fp8(), self.data_parallel_group + expanded_fp8_params, + expanded_shard_fp32_from_fp8, + expanded_shard_offsets_in_fp8, + self.data_parallel_group, ) elif self.ddp_config.fp4_param_gather: # Quantize FP32 master shards back to NVFP4 model params (rowwise only) @@ -2688,7 +2780,7 @@ def copy_group_params(shard_main_groups, model_groups): world_range.start : world_range.end ] - if is_float8tensor(model_param): + if self._is_distopt_quantized_param(model_param): # FP8 params are quantized in the above "quantize_param_shard" function. continue elif is_nvfp4tensor(model_param): @@ -2730,6 +2822,66 @@ def _copy_main_params_to_param_buffer(self): shard_param_buffer.copy_(shard_main_param) + @staticmethod + def _normalize_state_dict_for_grouped_params(state_dict_flat, model_chunk): + """Normalize state dict keys when grouped/indexed parameter formats differ. + + TEGroupedLinear modules can store weights as either a single grouped tensor + (single_grouped_weight=True, key "weight") or individually indexed tensors + (single_grouped_weight=False, keys "weight0", "weight1", ...). When the + checkpoint was saved with a different setting, the keys won't match the + current model's parameters. This mirrors the + normalize_grouped_parameter_keys hook registered in TEGroupedLinear.__init__ + which only fires during load_state_dict and therefore does not cover the + optimizer reload_model_params path. + """ + for module_name, module in model_chunk.named_modules(): + if not ( + hasattr(module, 'num_gemms') and hasattr(module, '_split_grouped_checkpoint_tensor') + ): + continue + + clean_name = module_name + while clean_name.startswith("module."): + clean_name = clean_name[len("module.") :] + + num_gemms = module.num_gemms + + params_to_check = [("weight", getattr(module, "single_grouped_weight", False))] + if getattr(module, "use_bias", False): + params_to_check.append(("bias", getattr(module, "single_grouped_bias", False))) + + for param_name, single_grouped in params_to_check: + grouped_suffix = f"{clean_name}.{param_name}" if clean_name else param_name + indexed_suffixes = [f"{grouped_suffix}{i}" for i in range(num_gemms)] + + grouped_matches = [k for k in state_dict_flat if k.endswith(grouped_suffix)] + indexed_match_map = {} + for i, idx_suffix in enumerate(indexed_suffixes): + matches = [k for k in state_dict_flat if k.endswith(idx_suffix)] + if len(matches) == 1: + indexed_match_map[i] = matches[0] + + if single_grouped: + if grouped_matches or len(indexed_match_map) != num_gemms: + continue + indexed_keys = [indexed_match_map[i] for i in range(num_gemms)] + grouped_tensor = torch.stack( + [state_dict_flat.pop(k) for k in indexed_keys], dim=0 + ) + key_prefix = indexed_keys[0][: len(indexed_keys[0]) - len(indexed_suffixes[0])] + state_dict_flat[f"{key_prefix}{grouped_suffix}"] = grouped_tensor + else: + if indexed_match_map or len(grouped_matches) != 1: + continue + grouped_key = grouped_matches[0] + split_tensors = module._split_grouped_checkpoint_tensor( + state_dict_flat.pop(grouped_key), grouped_key + ) + key_prefix = grouped_key[: len(grouped_key) - len(grouped_suffix)] + for i, tensor in enumerate(split_tensors): + state_dict_flat[f"{key_prefix}{indexed_suffixes[i]}"] = tensor + def _build_model_param_to_state_dict_param_map(self, state_dict): """Create a map from model params to tensors in state_dict based on their names.""" state_dict_list = [] @@ -2751,6 +2903,7 @@ def _build_model_param_to_state_dict_param_map(self, state_dict): model_param_to_state_dict_param_map = {} for chunk_idx, model_chunk in enumerate(self.model_chunks): + self._normalize_state_dict_for_grouped_params(state_dict_list[chunk_idx], model_chunk) names_in_state_dict = set(state_dict_list[chunk_idx].keys()) for name, model_param in model_chunk.named_parameters(): while name.startswith("module."): @@ -2810,8 +2963,12 @@ def copy_group_params(model_groups, shard_main_groups): # Use param from state_dict to initialize main_param model_param = model_param_to_state_dict_param_map[model_param] - if is_float8tensor(model_param): - shard_model_param = dequantize_fp8_tensor(model_param).view(-1)[ + if self._is_distopt_quantized_param(model_param): + if self._is_grouped_quantized_tensor(model_param): + dequantized_model_param = model_param.float() + else: + dequantized_model_param = dequantize_fp8_tensor(model_param) + shard_model_param = dequantized_model_param.view(-1)[ param_range.start : param_range.end ] else: diff --git a/megatron/core/transformer/moe/experts.py b/megatron/core/transformer/moe/experts.py index 2828d086e5d..209e54cfe08 100644 --- a/megatron/core/transformer/moe/experts.py +++ b/megatron/core/transformer/moe/experts.py @@ -5,6 +5,7 @@ from collections.abc import Callable from copy import deepcopy from dataclasses import dataclass +from itertools import chain from math import ceil from typing import Optional, Protocol, Tuple @@ -33,6 +34,7 @@ from megatron.core.transformer.moe.moe_utils import ( ProcessGroupCollection, get_align_size_for_quantization, + skip_routed_expert_padding, ) from megatron.core.transformer.moe.token_dispatcher_inference import ( InferenceAllGatherDispatcherBase, @@ -44,10 +46,14 @@ sharded_state_dict_default, ) from megatron.core.typed_torch import apply_module, not_none +from megatron.core.utils import is_te_min_version if HAVE_TE: + import transformer_engine as te + from megatron.core.extensions.transformer_engine import Fp8Padding, Fp8Unpadding else: + te = None # type: ignore[assignment, misc] Fp8Padding, Fp8Unpadding = None, None try: @@ -242,10 +248,32 @@ def __init__( set_save_original_input(self.linear_fc1) + # Fused implementation with Transformer Engine op fuser API + if self.config.use_transformer_engine_op_fuser: + assert ( + self._is_fused_impl_supported() + ), "Fused GroupedMLP is not supported for this configuration." + self._with_fused_impl: bool = self.config.use_transformer_engine_op_fuser + self._fused_ops: Optional[Tuple[torch.nn.Module]] = None + if ( + self.config.gated_linear_unit + and self.config.moe_mlp_glu_interleave_size is not None + and not self._with_fused_impl + ): + logger.warning( + "`moe_mlp_glu_interleave_size=%s` is enabled, but fused MoE MLP implementation " + "is not supported for this configuration. The non-fused path may incur extra " + "tensor reordering/copy overhead each forward pass.", + self.config.moe_mlp_glu_interleave_size, + ) + if self.config.fp8 or self.config.fp4: assert HAVE_TE, "FP8 and FP4 requires TE." - self.quantization_padding = Fp8Padding(self.num_local_experts) - self.quantization_unpadding = Fp8Unpadding(self.num_local_experts) + align_size = 256 if self._with_fused_impl else None + self.quantization_padding = Fp8Padding(self.num_local_experts, align_size=align_size) + self.quantization_unpadding = Fp8Unpadding( + self.num_local_experts, align_size=align_size + ) @staticmethod def _apply_bias(intermediate_parallel, bias_parallel, tokens_per_expert, permuted_probs): @@ -267,62 +295,238 @@ def _apply_bias(intermediate_parallel, bias_parallel, tokens_per_expert, permute .to(intermediate_parallel.dtype) ) - def bias_act_func(self, intermediate_parallel, bias_parallel, permuted_probs): - """ - Applies bias and activation function to the output of linear_fc1. - """ - if self.config.use_te_activation_func: - if bias_parallel is not None: - intermediate_parallel = intermediate_parallel + bias_parallel - intermediate_parallel = self.activation_func(intermediate_parallel) - if permuted_probs is not None: - original_dtype = intermediate_parallel.dtype - intermediate_parallel = intermediate_parallel * permuted_probs - intermediate_parallel = intermediate_parallel.to(original_dtype) - elif self.config.bias_activation_fusion: - if self.activation_func == F.silu and self.config.gated_linear_unit: - # dtype is handled inside the fused kernel - intermediate_parallel = weighted_bias_swiglu_impl( - intermediate_parallel, - bias_parallel, - permuted_probs, - self.config.activation_func_fp8_input_store, - ) - elif self.activation_func == quick_gelu and self.config.gated_linear_unit: - intermediate_parallel = weighted_bias_quick_geglu_impl( - intermediate_parallel, - bias_parallel, - permuted_probs, - self.config.activation_func_fp8_input_store, - self.config.glu_linear_offset, - self.config.activation_func_clamp_value, + def _is_fused_impl_supported(self) -> bool: + """Check if the TE op fuser supports implementing this module.""" + + # Check Transformer Engine installation + if not HAVE_TE: + return False # Transformer Engine is not available + try: + from transformer_engine.pytorch.ops import GroupedLinear, ScaledSwiGLU + except ImportError: + return False # Transformer Engine version is too old + + if not is_te_min_version("2.14.0"): + return False + + # Check for unsupported features + if self.tp_group.size() > 1: + return False # Tensor parallelism is not supported + if self.offload_expert_fc1 or self.offload_moe_act: + return False # Fine-grained activation offloading is not supported + if self.config.moe_apply_probs_on_input: + return False # Pre-multiplying probs is not supported + + # Check grouped linear modules + if not isinstance(self.linear_fc1, te.pytorch.GroupedLinear): + return False + if not isinstance(self.linear_fc2, te.pytorch.GroupedLinear): + return False + + # Check activation: SwiGLU (ScaledSwiGLU) or quick GEGLU (ScaledClampedQGeGLU, TE >= 2.15) + if self.config.gated_linear_unit: + if self.activation_func == F.silu: + return True + if self.activation_func == quick_gelu: + try: + from transformer_engine.pytorch.ops import ScaledClampedQGeGLU # noqa: F401 + except ImportError: + return False + return True + + return False + + def _make_fused_ops(self) -> torch.nn.Module: + """Construct fused module for FC1, activation, and FC2.""" + + assert HAVE_TE, "_make_fused_ops requires Transformer Engine." + + # Container for fusible ops + ops = te.pytorch.ops.Sequential() + + # Check if there are 1 or "num_gemms" params in the GroupedLinear module. + fc1_single_grouped_weight = self.linear_fc1.single_grouped_weight + fc1_weight_dtype = ( + self.linear_fc1.weight.dtype + if fc1_single_grouped_weight + else self.linear_fc1.weight0.dtype + ) + fc2_single_grouped_weight = self.linear_fc2.single_grouped_weight + fc2_weight_dtype = ( + self.linear_fc2.weight.dtype + if fc2_single_grouped_weight + else self.linear_fc2.weight0.dtype + ) + fc1_single_grouped_bias = self.linear_fc1.single_grouped_bias + fc2_single_grouped_bias = self.linear_fc2.single_grouped_bias + # Mirror the wrapper's combined delay-wgrad mode (config.delay_wgrad_compute OR + # config.overlap_dispatch_backward_with_experts_wgrad) — see TEGroupedLinear.__init__. + # Using config.delay_wgrad_compute alone would silently drop the overlap optimization + # for runs that enable it via overlap_dispatch_backward_with_experts_wgrad. + fc1_delay_wgrad_compute = self.linear_fc1.delay_wgrad_compute + fc2_delay_wgrad_compute = self.linear_fc2.delay_wgrad_compute + + # Create a parameterless op shell and then attach the existing GroupedLinear weights below. + # Using meta avoids allocating duplicate weights for the fused wrapper. + op = te.pytorch.ops.GroupedLinear( + self.linear_fc1.num_gemms, + self.linear_fc1.in_features, + self.linear_fc1.out_features, + bias=self.linear_fc1.use_bias, + device="meta", + dtype=fc1_weight_dtype, + accumulate_into_main_grad=self.linear_fc1.fuse_wgrad_accumulation, + single_grouped_weight=fc1_single_grouped_weight, + single_grouped_bias=fc1_single_grouped_bias, + delay_wgrad_compute=fc1_delay_wgrad_compute, + ) + + # Copy the weights from GroupedLinear module to GroupedLinear op. + if fc1_single_grouped_weight: + setattr(op, "weight", getattr(self.linear_fc1, "weight")) + + for idx in range(self.linear_fc1.num_gemms): + if not fc1_single_grouped_weight: + setattr(op, f"weight{idx}", getattr(self.linear_fc1, f"weight{idx}")) + if self.linear_fc1.use_bias and not fc1_single_grouped_bias: + setattr(op, f"bias{idx}", getattr(self.linear_fc1, f"bias{idx}")) + if self.linear_fc1.use_bias and fc1_single_grouped_bias: + setattr(op, "bias", getattr(self.linear_fc1, "bias")) + ops.append(op) + + # Activation and post-multiply probs (SwiGLU or clamped quick-GeGLU) + glu_interleave = self.config.moe_mlp_glu_interleave_size + if self.activation_func == F.silu and self.config.gated_linear_unit: + op = te.pytorch.ops.ScaledSwiGLU(glu_interleave_size=glu_interleave) + elif self.activation_func == quick_gelu and self.config.gated_linear_unit: + clamp = self.config.activation_func_clamp_value + if clamp is not None: + op = te.pytorch.ops.ScaledClampedQGeGLU( + glu_interleave_size=glu_interleave, limit=clamp ) else: - raise ValueError("Only support fusion of swiglu and quick_gelu in TEGroupedMLP.") - elif self.activation_func == squared_relu and self.config.use_fused_weighted_squared_relu: - assert bias_parallel is None, "Bias is not supported with fused weighted squared relu." - intermediate_parallel = weighted_squared_relu_impl( - intermediate_parallel, permuted_probs - ) + op = te.pytorch.ops.ScaledClampedQGeGLU(glu_interleave_size=glu_interleave) else: - if self.config.gated_linear_unit: - - def glu(x): - x_glu, x_linear = torch.chunk(x, 2, dim=-1) - if (val := self.config.activation_func_clamp_value) is not None: - x_glu = x_glu.clamp(min=None, max=val) - x_linear = x_linear.clamp(min=-val, max=val) - return self.config.activation_func(x_glu) * ( - x_linear + self.config.glu_linear_offset - ) + raise RuntimeError( + "_make_fused_ops expected SwiGLU or quick_gelu with gated_linear_unit; " + "call _is_fused_impl_supported() before constructing fused ops." + ) + ops.append(op) + + # FC2 + op = te.pytorch.ops.GroupedLinear( + self.linear_fc2.num_gemms, + self.linear_fc2.in_features, + self.linear_fc2.out_features, + bias=self.linear_fc2.use_bias, + device="meta", + dtype=fc2_weight_dtype, + accumulate_into_main_grad=self.linear_fc2.fuse_wgrad_accumulation, + single_grouped_weight=fc2_single_grouped_weight, + single_grouped_bias=fc2_single_grouped_bias, + delay_wgrad_compute=fc2_delay_wgrad_compute, + ) - intermediate_parallel = glu(intermediate_parallel) - else: - intermediate_parallel = self.activation_func(intermediate_parallel) - original_dtype = intermediate_parallel.dtype - intermediate_parallel = intermediate_parallel * permuted_probs - intermediate_parallel = intermediate_parallel.to(original_dtype) - return intermediate_parallel + # Copy the weights from GroupedLinear module to GroupedLinear op. + if fc2_single_grouped_weight: + setattr(op, "weight", getattr(self.linear_fc2, "weight")) + + for idx in range(self.linear_fc2.num_gemms): + if not fc2_single_grouped_weight: + setattr(op, f"weight{idx}", getattr(self.linear_fc2, f"weight{idx}")) + if self.linear_fc2.use_bias and not fc2_single_grouped_bias: + setattr(op, f"bias{idx}", getattr(self.linear_fc2, f"bias{idx}")) + if self.linear_fc2.use_bias and fc2_single_grouped_bias: + setattr(op, "bias", getattr(self.linear_fc2, "bias")) + ops.append(op) + + # Emulate submodule pre-forward hooks + ops.register_forward_pre_hook(self._make_fused_impl_pre_forward_hook()) + + return ops + + def _make_fused_impl_pre_forward_hook(self) -> Callable: + """Make function that calls submodule pre-forward callback hooks. + + This is intended for compatibility with + DistributedDataParallel hooks that trigger parameter + all-gathers. It does not support general pre-forward hooks + since they may manipulate intermediate tensors that are never + instantiated by the fused implementation. + + """ + + def forward_pre_hook(module, *_) -> None: + for submodule in chain(self.linear_fc1.modules(), self.linear_fc2.modules()): + for hook in submodule._forward_pre_hooks.values(): + # Assume that hook does not interact with input + ret = hook(submodule, None) + if ret is not None: + raise RuntimeError( + f"Applying a fused implementation for {self.__class__.__name__}, " + f"but a {submodule.__class__.__name__} submodule " + "has a pre-forward hook that modifies the input tensor." + ) + + return forward_pre_hook + + def _fused_forward( + self, + permuted_local_hidden_states: torch.Tensor, + tokens_per_expert: torch.Tensor, + permuted_probs: torch.Tensor, + ) -> torch.Tensor: + """Forward pass using Transformer Engine operation fuser API.""" + + # Construct fused impl if needed + # Note: We initialize during the first forward pass in case + # the params are modified after the constructor. + # Note: The fused impl is stored in a tuple to avoid + # registering submodules. + if self._fused_ops is None: + self._fused_ops = (self._make_fused_ops(),) + (ops,) = self._fused_ops + + # Apply padding if needed + unpadded_tokens_per_expert = None + if skip_routed_expert_padding(self.config): + pass + elif self.config.fp8 or self.config.fp4: + tokens_per_expert = tokens_per_expert.tolist() + unpadded_tokens_per_expert = tokens_per_expert + permuted_local_hidden_states, tokens_per_expert = self.quantization_padding( + permuted_local_hidden_states, tokens_per_expert + ) + permuted_probs, _ = self.quantization_padding( + permuted_probs.unsqueeze(-1), unpadded_tokens_per_expert + ) + permuted_probs = permuted_probs.squeeze(-1) + tokens_per_expert = torch.tensor( + tokens_per_expert, dtype=torch.int, device=permuted_probs.device + ) + + # Call fused impl + output = ops( + permuted_local_hidden_states, + tokens_per_expert, # FC1 + permuted_probs, # Scaled SwiGLU + tokens_per_expert, # FC2 + ) + + # Remove padding if needed + if unpadded_tokens_per_expert is not None: + output = self.quantization_unpadding(output, unpadded_tokens_per_expert) + + return output + + @staticmethod + def _remove_glu_interleaving(x: torch.Tensor, interleave_size: int) -> torch.Tensor: + """Reorder interleaved GLU blocks so gate and linear halves are contiguous.""" + shape = x.size() + x = x.reshape(-1, shape[-1] // (2 * interleave_size), 2, interleave_size) + x = x.transpose(1, 2).contiguous() + x = x.view(shape) + return x def forward( self, @@ -341,17 +545,29 @@ def forward( Return: output (torch.Tensor): The output of the local experts. """ + + # Call fused impl if enabled + if self._with_fused_impl: + output = self._fused_forward( + permuted_local_hidden_states, tokens_per_expert, permuted_probs + ) + output_bias = None + return output, output_bias + + # Apply padding if needed + unpadded_tokens_per_expert = None tokens_per_expert: list[int] = tokens_per_expert.tolist() - if self.config.fp8 or self.config.fp4: - actual_tokens_per_expert = tokens_per_expert + permuted_probs = permuted_probs.unsqueeze(-1) + if skip_routed_expert_padding(self.config): + pass + elif self.config.fp8 or self.config.fp4: + unpadded_tokens_per_expert = tokens_per_expert permuted_local_hidden_states, tokens_per_expert = self.quantization_padding( permuted_local_hidden_states, tokens_per_expert ) permuted_probs, _ = self.quantization_padding( - permuted_probs.unsqueeze(-1), actual_tokens_per_expert + permuted_probs, unpadded_tokens_per_expert ) - else: - permuted_probs = permuted_probs.unsqueeze(-1) if self.config.moe_apply_probs_on_input: assert ( @@ -376,15 +592,90 @@ def forward( forced_released_tensors=[permuted_local_hidden_states], ) + def bias_act_func(intermediate_parallel, bias_parallel, permuted_probs): + + # Whether activation function is interleaved GLU + with_glu_interleaving = ( + self.config.gated_linear_unit + and self.config.moe_mlp_glu_interleave_size is not None + ) + + if self.config.use_te_activation_func: + if bias_parallel is not None: + intermediate_parallel = intermediate_parallel + bias_parallel + if with_glu_interleaving: + intermediate_parallel = self._remove_glu_interleaving( + intermediate_parallel, self.config.moe_mlp_glu_interleave_size + ) + intermediate_parallel = self.activation_func(intermediate_parallel) + if permuted_probs is not None: + original_dtype = intermediate_parallel.dtype + intermediate_parallel = intermediate_parallel * permuted_probs + intermediate_parallel = intermediate_parallel.to(original_dtype) + elif self.config.bias_activation_fusion and not with_glu_interleaving: + if self.activation_func == F.silu and self.config.gated_linear_unit: + # dtype is handled inside the fused kernel + intermediate_parallel = weighted_bias_swiglu_impl( + intermediate_parallel, + bias_parallel, + permuted_probs, + self.config.activation_func_fp8_input_store, + ) + elif self.activation_func == quick_gelu and self.config.gated_linear_unit: + intermediate_parallel = weighted_bias_quick_geglu_impl( + intermediate_parallel, + bias_parallel, + permuted_probs, + self.config.activation_func_fp8_input_store, + self.config.glu_linear_offset, + self.config.activation_func_clamp_value, + ) + else: + raise ValueError( + "Only support fusion of swiglu and quick_gelu in TEGroupedMLP." + ) + elif ( + self.activation_func == squared_relu and self.config.use_fused_weighted_squared_relu + ): + assert ( + bias_parallel is None + ), "Bias is not supported with fused weighted squared relu." + intermediate_parallel = weighted_squared_relu_impl( + intermediate_parallel, permuted_probs + ) + else: + if self.config.gated_linear_unit: + + def glu(x): + if with_glu_interleaving: + x = self._remove_glu_interleaving( + x, self.config.moe_mlp_glu_interleave_size + ) + x_glu, x_linear = torch.chunk(x, 2, dim=-1) + if (val := self.config.activation_func_clamp_value) is not None: + x_glu = x_glu.clamp(min=None, max=val) + x_linear = x_linear.clamp(min=-val, max=val) + return self.config.activation_func(x_glu) * ( + x_linear + self.config.glu_linear_offset + ) + + intermediate_parallel = glu(intermediate_parallel) + else: + intermediate_parallel = self.activation_func(intermediate_parallel) + original_dtype = intermediate_parallel.dtype + intermediate_parallel = intermediate_parallel * permuted_probs + intermediate_parallel = intermediate_parallel.to(original_dtype) + return intermediate_parallel + if self.activation_recompute: self.activation_checkpoint = tensor_parallel.CheckpointWithoutOutput() with off_interface(self.offload_moe_act, fc1_output, "moe_act") as fc1_output: bias_act_output = self.activation_checkpoint.checkpoint( - self.bias_act_func, fc1_output, bias_parallel, permuted_probs + bias_act_func, fc1_output, bias_parallel, permuted_probs ) else: with off_interface(self.offload_moe_act, fc1_output, "moe_act") as fc1_output: - bias_act_output = self.bias_act_func(fc1_output, bias_parallel, permuted_probs) + bias_act_output = bias_act_func(fc1_output, bias_parallel, permuted_probs) output, output_bias = apply_module(self.linear_fc2)(bias_act_output, tokens_per_expert) if self.activation_recompute: self.activation_checkpoint.discard_output_and_register_recompute(output) @@ -398,8 +689,8 @@ def forward( output = self._apply_bias(output, output_bias, tokens_per_expert, permuted_probs) # upad and concat the output - if self.config.fp8 or self.config.fp4: - output = self.quantization_unpadding(output, actual_tokens_per_expert) + if unpadded_tokens_per_expert is not None: + output = self.quantization_unpadding(output, unpadded_tokens_per_expert) output_bias = None @@ -453,6 +744,31 @@ def backward_dw(self): If an error occurs during execution, it is caught and re-raised with a descriptive message. """ + # Match the wrapper's combined delay-wgrad mode used in _make_fused_ops so that + # `overlap_dispatch_backward_with_experts_wgrad`-driven runs invoke the deferred + # wgrad pass through the fused children instead of falling through to no-op + # backward_dw() on linear_fc{1,2} (whose forward never ran in the fused path). + if self._with_fused_impl and self.linear_fc1.delay_wgrad_compute: + if self._fused_ops is not None: + (seq,) = self._fused_ops + fused_children = list(seq.children()) + assert len(fused_children) >= 3, "expected FC1, activation, FC2 in fused TE ops" + fused_children[2].backward_dw() + fused_children[0].backward_dw() + # DDP registers wgrad hooks on the original linear_fc1/fc2 module objects + # (those are in the nn.Module tree), but backward_dw() is called on the + # NEW GroupedLinear instances created by _make_fused_ops(). We must + # explicitly fire the hooks on the originals so DDP can zero param.grad + # and trigger reduce-scatter – otherwise param.grad is never cleared and + # AccumulateGrad performs a spurious add_ into main_grad. + # TODO: find a better place to invoke _trigger_wgrad_accumulation_and_reduce_hooks. + # The wgrad hook registration lives in TE while the trigger is issued here + # in MCore, so the hook lifecycle is split across both codebases. Consolidate + # ownership on one side (either register+trigger entirely in TE, or expose + # the fused backward_dw through MCore) to remove this fragmentation. + self.linear_fc2._trigger_wgrad_accumulation_and_reduce_hooks() + self.linear_fc1._trigger_wgrad_accumulation_and_reduce_hooks() + return self.linear_fc2.backward_dw() self.linear_fc1.backward_dw() diff --git a/megatron/core/transformer/moe/moe_utils.py b/megatron/core/transformer/moe/moe_utils.py index f258f3474ae..231dade4379 100644 --- a/megatron/core/transformer/moe/moe_utils.py +++ b/megatron/core/transformer/moe/moe_utils.py @@ -1429,14 +1429,33 @@ def get_align_size_for_quantization(config: TransformerConfig) -> int: Returns: int: The alignment size for quantization. """ + # CUTLASS kernel for grouped GEMM assumes 256 alignment. + if config.use_transformer_engine_op_fuser: + return 256 if config.fp8: return get_fp8_align_size(config.fp8_recipe) - elif config.fp4: + if config.fp4: return get_fp4_align_size(config.fp4_recipe) # Only FP8 or FP4 requires padding. Defaults to 0. return 0 +def skip_routed_expert_padding(config: TransformerConfig) -> bool: + """Whether the expert module should skip quantization padding. + + Returns True when padding is already applied by the router or the + HybridEP dispatcher. + """ + if config.moe_router_padding_for_quantization: + return True + if ( + config.moe_token_dispatcher_type == "flex" + and config.moe_flex_dispatcher_backend == "hybridep" + ): + return True + return False + + # TODO(Hepteract): delete the usage of the global parallel_state. # Initialize process groups with the global parallel_state. def get_default_pg_collection() -> ProcessGroupCollection: diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py index bb044787b9c..7e3b5e56a13 100644 --- a/megatron/core/transformer/transformer_config.py +++ b/megatron/core/transformer/transformer_config.py @@ -450,6 +450,10 @@ class TransformerConfig(ModelParallelConfig): fused_residual_rmsnorm: bool = False """If True, fuses residual connection and RMSNorm backward pass when TE is used.""" + use_transformer_engine_op_fuser: bool = False + """If True, submodules may use Transformer Engine's operation fuser + API to enable advanced fusions.""" + #################### # activation recomputation #################### @@ -732,6 +736,16 @@ class TransformerConfig(ModelParallelConfig): GEMM feature introduced since CUTLASS 2.8 (https://github.com/fanshiqing/grouped_gemm). """ + moe_single_grouped_weight: bool = False + """When using TE GroupedLinear for MoE experts, store expert weights as a single grouped + parameter via Transformer Engine's `GroupedTensor`. Requires ``moe_grouped_gemm=True``. + """ + + moe_single_grouped_bias: bool = False + """When using TE GroupedLinear for MoE experts, store expert biases as a single grouped + parameter via Transformer Engine's `GroupedTensor`. Requires ``moe_grouped_gemm=True`` + and ``add_bias_linear=True``.""" + moe_aux_loss_coeff: Union[float, List[float]] = 0.0 """Scaling coefficient for the aux loss. A starting value of 1e-2 is recommended. If a list of load balancing types is provided for `moe_router_load_balancing_type`, @@ -821,6 +835,15 @@ class TransformerConfig(ModelParallelConfig): When permute_fusion_into_hybridep is True, this sets the number of SMs for the unpermute part (only 1 block per SM).""" + moe_mlp_glu_interleave_size: Optional[int] = None + """When set, GLU activations in the MoE grouped MLP layer will use a + block interleaved format. Instead of interpreting the input tensor + as a concatenation of gates and linear units, it will be + interpreted as alternating blocks of gates and linear units. + + This data format is experimental and primarily intended to enable + advanced fused kernels.""" + ################## # Context Parallel ################## @@ -1274,6 +1297,32 @@ def __post_init__(self): "Please set num_moe_experts or remove moe_ffn_hidden_size." ) + if self.moe_single_grouped_weight or self.moe_single_grouped_bias: + if not self.moe_grouped_gemm: + raise ValueError( + "moe_single_grouped_weight and moe_single_grouped_bias require " + "moe_grouped_gemm=True." + ) + if not is_te_min_version("2.14.0"): + raise ValueError( + "moe_single_grouped_weight and moe_single_grouped_bias require " + f"transformer-engine>=2.14.0, but your version is {get_te_version()}." + ) + if self.moe_single_grouped_weight: + # The dist-optimizer's quantized-param shard path on the single-grouped-weight + # storage is only validated for fp8 mode with the mxfp8 recipe today; other + # combinations have a known numerical issue tracked in upstream PR + # NVIDIA/Megatron-LM#4621. Reject at construction time so users don't silently + # train on a broken numerical path. (moe_single_grouped_bias is not gated: + # biases aren't quantized, so they don't enter the buggy code path.) + if self.fp4 or not self.fp8 or self.fp8_recipe != Fp8Recipe.mxfp8: + raise ValueError( + "moe_single_grouped_weight is currently supported only with fp8 mode " + "and fp8_recipe='mxfp8'." + ) + if self.moe_single_grouped_bias and not self.add_bias_linear: + raise ValueError("moe_single_grouped_bias requires add_bias_linear=True.") + if self.moe_enable_deepep: if self.moe_token_dispatcher_type != "flex": raise ValueError("DeepEP backend is only supported with flex token dispatcher.") diff --git a/megatron/training/checkpointing.py b/megatron/training/checkpointing.py index 43a0eb9d12a..fac717557a7 100644 --- a/megatron/training/checkpointing.py +++ b/megatron/training/checkpointing.py @@ -1554,6 +1554,8 @@ def _set_arg(arg_name, old_arg_name=None, force=False): _set_arg('moe_token_dispatcher_type', force=False) _set_arg('moe_router_pre_softmax', force=True) _set_arg('moe_grouped_gemm', force=True) + _set_arg('moe_single_grouped_weight', force=True) + _set_arg('moe_single_grouped_bias', force=True) _set_arg('moe_shared_expert_intermediate_size', force=True) _set_arg('moe_router_score_function', force=True) _set_arg('moe_router_enable_expert_bias', force=True) diff --git a/tests/functional_tests/test_cases/moe/gpt3_mcore_te_tp1_pp1_te_4experts_groupedGEMM_op_fuser/golden_values_dev_dgx_h100.json b/tests/functional_tests/test_cases/moe/gpt3_mcore_te_tp1_pp1_te_4experts_groupedGEMM_op_fuser/golden_values_dev_dgx_h100.json new file mode 100644 index 00000000000..8d900a14ce5 --- /dev/null +++ b/tests/functional_tests/test_cases/moe/gpt3_mcore_te_tp1_pp1_te_4experts_groupedGEMM_op_fuser/golden_values_dev_dgx_h100.json @@ -0,0 +1,287 @@ +{ + "lm loss": { + "start_step": 1, + "end_step": 50, + "step_interval": 1, + "values": { + "1": 10.87658, + "2": 10.87401, + "3": 10.89527, + "4": 10.88157, + "5": 10.89275, + "6": 10.87883, + "7": 10.89477, + "8": 10.87687, + "9": 10.89461, + "10": 10.89105, + "11": 10.89671, + "12": 10.88599, + "13": 10.87443, + "14": 10.86612, + "15": 10.89764, + "16": 10.86465, + "17": 10.87144, + "18": 10.86503, + "19": 10.86598, + "20": 10.84921, + "21": 10.87341, + "22": 10.84988, + "23": 10.86536, + "24": 10.84142, + "25": 10.87489, + "26": 10.83702, + "27": 10.85594, + "28": 10.80589, + "29": 10.81923, + "30": 10.80475, + "31": 10.79606, + "32": 10.80516, + "33": 10.80757, + "34": 10.75693, + "35": 10.78547, + "36": 10.7813, + "37": 10.76519, + "38": 10.74958, + "39": 10.72448, + "40": 10.72666, + "41": 10.6995, + "42": 10.69088, + "43": 10.67411, + "44": 10.68361, + "45": 10.67984, + "46": 10.66603, + "47": 10.65467, + "48": 10.63581, + "49": 10.61281, + "50": 10.62945 + } + }, + "num-zeros": { + "start_step": 1, + "end_step": 50, + "step_interval": 1, + "values": { + "1": 2889.0, + "2": 3128.0, + "3": 3038.0, + "4": 2975.0, + "5": 3008.0, + "6": 2944.0, + "7": 3192.0, + "8": 2804.0, + "9": 3108.0, + "10": 3113.0, + "11": 3013.0, + "12": 2778.0, + "13": 3176.0, + "14": 3048.0, + "15": 3095.0, + "16": 2873.0, + "17": 3075.0, + "18": 2890.0, + "19": 3059.0, + "20": 3130.0, + "21": 2987.0, + "22": 2948.0, + "23": 3003.0, + "24": 3074.0, + "25": 2917.0, + "26": 3079.0, + "27": 3090.0, + "28": 3109.0, + "29": 3010.0, + "30": 2961.0, + "31": 3008.0, + "32": 2927.0, + "33": 2897.0, + "34": 3081.0, + "35": 3063.0, + "36": 2906.0, + "37": 3172.0, + "38": 3042.0, + "39": 2863.0, + "40": 3092.0, + "41": 2729.0, + "42": 2655.0, + "43": 2848.0, + "44": 2830.0, + "45": 2888.0, + "46": 2871.0, + "47": 2949.0, + "48": 2811.0, + "49": 2900.0, + "50": 2886.0 + } + }, + "mem-allocated-bytes": { + "start_step": 1, + "end_step": 50, + "step_interval": 1, + "values": { + "1": 1384147968.0, + "2": 1384147968.0, + "3": 1384147968.0, + "4": 1384147968.0, + "5": 1384147968.0, + "6": 1384147968.0, + "7": 1384147968.0, + "8": 1384147968.0, + "9": 1384147968.0, + "10": 1384147968.0, + "11": 1384147968.0, + "12": 1384147968.0, + "13": 1384147968.0, + "14": 1384147968.0, + "15": 1384147968.0, + "16": 1384147968.0, + "17": 1384147968.0, + "18": 1384147968.0, + "19": 1384147968.0, + "20": 1384147968.0, + "21": 1384147968.0, + "22": 1384147968.0, + "23": 1384147968.0, + "24": 1384147968.0, + "25": 1384147968.0, + "26": 1384147968.0, + "27": 1384147968.0, + "28": 1384147968.0, + "29": 1384147968.0, + "30": 1384147968.0, + "31": 1384147968.0, + "32": 1384147968.0, + "33": 1384147968.0, + "34": 1384147968.0, + "35": 1384147968.0, + "36": 1384147968.0, + "37": 1384147968.0, + "38": 1384147968.0, + "39": 1384147968.0, + "40": 1384147968.0, + "41": 1384147968.0, + "42": 1384147968.0, + "43": 1384147968.0, + "44": 1384147968.0, + "45": 1384147968.0, + "46": 1384147968.0, + "47": 1384147968.0, + "48": 1384147968.0, + "49": 1384147968.0, + "50": 1384147968.0 + } + }, + "mem-max-allocated-bytes": { + "start_step": 1, + "end_step": 50, + "step_interval": 1, + "values": { + "1": 2412110336.0, + "2": 2956330496.0, + "3": 2956330496.0, + "4": 2956330496.0, + "5": 2956330496.0, + "6": 2956330496.0, + "7": 2956330496.0, + "8": 2956330496.0, + "9": 2956330496.0, + "10": 2956330496.0, + "11": 2956330496.0, + "12": 2956330496.0, + "13": 2956330496.0, + "14": 2956330496.0, + "15": 2956330496.0, + "16": 2956330496.0, + "17": 2956330496.0, + "18": 2956330496.0, + "19": 2956330496.0, + "20": 2956330496.0, + "21": 2956330496.0, + "22": 2956330496.0, + "23": 2956330496.0, + "24": 2956330496.0, + "25": 2956330496.0, + "26": 2956330496.0, + "27": 2956330496.0, + "28": 2956330496.0, + "29": 2956330496.0, + "30": 2956330496.0, + "31": 2956330496.0, + "32": 2956330496.0, + "33": 2956330496.0, + "34": 2956330496.0, + "35": 2956330496.0, + "36": 2956330496.0, + "37": 2956330496.0, + "38": 2956330496.0, + "39": 2956330496.0, + "40": 2956330496.0, + "41": 2956330496.0, + "42": 2956330496.0, + "43": 2956330496.0, + "44": 2956330496.0, + "45": 2956330496.0, + "46": 2956330496.0, + "47": 2956330496.0, + "48": 2956330496.0, + "49": 2956330496.0, + "50": 2956330496.0 + } + }, + "iteration-time": { + "start_step": 1, + "end_step": 50, + "step_interval": 1, + "values": { + "1": "nan", + "2": 2.51248, + "3": 0.08422, + "4": 0.07179, + "5": 0.07404, + "6": 0.08961, + "7": 0.07117, + "8": 0.06529, + "9": 0.07149, + "10": 0.06989, + "11": 0.06262, + "12": 0.07354, + "13": 0.0647, + "14": 0.06635, + "15": 0.05955, + "16": 0.06158, + "17": 0.06379, + "18": 0.06662, + "19": 0.05881, + "20": 0.0703, + "21": 0.06782, + "22": 0.06408, + "23": 0.06134, + "24": 0.06504, + "25": 0.06351, + "26": 0.06236, + "27": 0.0644, + "28": 0.06731, + "29": 0.06577, + "30": 0.07671, + "31": 0.07497, + "32": 0.05898, + "33": 0.06061, + "34": 0.06394, + "35": 0.06166, + "36": 0.05898, + "37": 0.05861, + "38": 0.06348, + "39": 0.06346, + "40": 0.06186, + "41": 0.06818, + "42": 0.06654, + "43": 0.06504, + "44": 0.05785, + "45": 0.05701, + "46": 0.05604, + "47": 0.05691, + "48": 0.06225, + "49": 0.05948, + "50": 0.0585 + } + } +} \ No newline at end of file diff --git a/tests/functional_tests/test_cases/moe/gpt3_mcore_te_tp1_pp1_te_4experts_groupedGEMM_op_fuser/model_config.yaml b/tests/functional_tests/test_cases/moe/gpt3_mcore_te_tp1_pp1_te_4experts_groupedGEMM_op_fuser/model_config.yaml new file mode 100644 index 00000000000..b092774efd3 --- /dev/null +++ b/tests/functional_tests/test_cases/moe/gpt3_mcore_te_tp1_pp1_te_4experts_groupedGEMM_op_fuser/model_config.yaml @@ -0,0 +1,62 @@ +ENV_VARS: + CUDA_DEVICE_MAX_CONNECTIONS: 1 + NVTE_ALLOW_NONDETERMINISTIC_ALGO: 0 + NCCL_ALGO: Ring + CUBLAS_WORKSPACE_CONFIG: :4096:8 + # Forward-compat with Blackwell: TE's cuDSL fused grouped MLP kernel + # additionally requires SM100, so on H100/A100 CI hardware this var is a no-op + # and TE falls back to its basic-op path. Setting it so the kernel is picked + # up automatically when Blackwell joins the CI matrix; the Megatron op-fuser + # wrapper (this PR's contribution) is exercised end-to-end on either path. + NVTE_CUTEDSL_FUSED_GROUPED_MLP: 1 +MODEL_ARGS: + --num-layers: 4 + --hidden-size: 512 + --num-attention-heads: 8 + --log-params-norm: true + --log-num-zeros-in-grad: true + --log-validation-ppl-to-tensorboard: true + --log-timers-to-tensorboard: true + --tensorboard-dir: ${TENSORBOARD_PATH} + --micro-batch-size: 4 + --global-batch-size: 32 + --seq-length: 1024 + --max-position-embeddings: 1024 + --train-iters: 50 + --timing-log-level: 0 + --lr-decay-iters: 320000 + --data-path: ${DATA_PATH}/text/common_pile/v01_filtered_data/my-gpt3_00_text_document + --vocab-file: ${DATA_PATH}/text/common_pile/v01_filtered_data/bpe/vocab.json + --merge-file: ${DATA_PATH}/text/common_pile/v01_filtered_data/bpe/merges.txt + --split: 949,50,1 + --distributed-backend: nccl + --lr: 0.00015 + --lr-decay-style: cosine + --min-lr: 1.0e-5 + --weight-decay: 1e-2 + --clip-grad: 1.0 + --lr-warmup-fraction: .01 + --log-interval: 1 + --eval-interval: 1000 + --eval-iters: 10 + --transformer-impl: transformer_engine + --tensor-model-parallel-size: 1 + --pipeline-model-parallel-size: 1 + --expert-model-parallel-size: 1 + --moe-grouped-gemm: true + --use-transformer-engine-op-fuser: true + --disable-bias-linear: true + --num-experts: 4 + --moe-router-load-balancing-type: sinkhorn + --moe-router-topk: 1 + --swiglu: true + --deterministic-mode: true + --no-gradient-accumulation-fusion: true + --attention-softmax-in-fp32: true + --use-mcore-models: true + --ckpt-format: torch_dist + --data-cache-path: ${DATA_CACHE_PATH} + --bf16: true + --no-bias-gelu-fusion: true + --log-memory-to-tensorboard: true +TEST_TYPE: regular diff --git a/tests/test_utils/recipes/h100/moe.yaml b/tests/test_utils/recipes/h100/moe.yaml index 6aeee162ab8..339964f1b49 100644 --- a/tests/test_utils/recipes/h100/moe.yaml +++ b/tests/test_utils/recipes/h100/moe.yaml @@ -151,6 +151,11 @@ products: - environment: [dev] scope: [mr, mr-github] platforms: [dgx_h100] + - test_case: [gpt3_mcore_te_tp1_pp1_te_4experts_groupedGEMM_op_fuser] + products: + - environment: [dev] + scope: [mr, mr-github] + platforms: [dgx_h100] - test_case: [gpt3_mcore_te_tp2_pp1_te_8experts_etp1_ep4] products: # - environment: [dev] diff --git a/tests/test_utils/recipes/h100/unit-tests.yaml b/tests/test_utils/recipes/h100/unit-tests.yaml index f90704b1264..80dbbcb6a93 100644 --- a/tests/test_utils/recipes/h100/unit-tests.yaml +++ b/tests/test_utils/recipes/h100/unit-tests.yaml @@ -43,6 +43,14 @@ spec: TEST_PATH="/opt/megatron-lm-legacy/" fi + # Enable TE's cuDSL fused grouped MLP path for MoE buckets. The kernel + # additionally requires SM100 (Blackwell), so on H100/A100 CI this is a + # no-op; setting it here means the kernel is picked up automatically when + # Blackwell hardware joins the unit-test matrix. + if [[ "$BUCKET" == *"transformer/moe"* || "$BUCKET" == *"test_moe_experts"* ]]; then + export NVTE_CUTEDSL_FUSED_GROUPED_MLP=1 + fi + bash $TEST_PATH/tests/unit_tests/run_ci_test.sh \ --tag $TAG \ diff --git a/tests/unit_tests/models/test_hybrid_moe_model.py b/tests/unit_tests/models/test_hybrid_moe_model.py index 3935964c975..243f43b4879 100644 --- a/tests/unit_tests/models/test_hybrid_moe_model.py +++ b/tests/unit_tests/models/test_hybrid_moe_model.py @@ -293,6 +293,10 @@ "inference_grouped_gemm_backend": "vllm", "inference_moe_disable_fused_quant_kernels": False, "inference_moe_token_dispatcher_type": "nvls", + "moe_mlp_glu_interleave_size": None, + "use_transformer_engine_op_fuser": False, + "moe_single_grouped_weight": False, + "moe_single_grouped_bias": False, } # Fields to ignore entirely (ephemeral, environment-specific, very large). SKIP_FIELDS = set() diff --git a/tests/unit_tests/optimizer/test_distrib_optimizer_grouped_quantized.py b/tests/unit_tests/optimizer/test_distrib_optimizer_grouped_quantized.py new file mode 100644 index 00000000000..59aebed27d9 --- /dev/null +++ b/tests/unit_tests/optimizer/test_distrib_optimizer_grouped_quantized.py @@ -0,0 +1,321 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +import torch + +from megatron.core.optimizer.distrib_optimizer import DistributedOptimizer + + +class _FakeGroupedQuantizedTensor: + def __init__(self, members, quantized_tensors=None): + self._members = members + self.quantized_tensors = quantized_tensors + self.quantizer = object() + + def split_into_quantized_tensors(self): + return self._members + + +def test_expand_quantized_param_shard_for_cast_splits_grouped_wrapper(): + optimizer = DistributedOptimizer.__new__(DistributedOptimizer) + members = [torch.empty(3), torch.empty(5), torch.empty(2)] + grouped_param = _FakeGroupedQuantizedTensor(members, quantized_tensors=members) + shard_main_param = torch.arange(6) + + expanded_params, expanded_main_params, expanded_offsets = ( + optimizer._expand_quantized_param_shard_for_cast( + grouped_param, shard_main_param, start_offset=2 + ) + ) + + assert len(expanded_params) == len(members) + assert all(expanded is member for expanded, member in zip(expanded_params, members)) + torch.testing.assert_close(expanded_main_params[0], torch.tensor([0])) + torch.testing.assert_close(expanded_main_params[1], torch.tensor([1, 2, 3, 4, 5])) + assert expanded_main_params[2] is None + assert expanded_offsets == [2, 0, None] + + +def test_expand_quantized_param_shard_for_cast_keeps_plain_param_unchanged(): + optimizer = DistributedOptimizer.__new__(DistributedOptimizer) + model_param = torch.empty(4) + shard_main_param = torch.arange(4) + + expanded_params, expanded_main_params, expanded_offsets = ( + optimizer._expand_quantized_param_shard_for_cast( + model_param, shard_main_param, start_offset=1 + ) + ) + + assert len(expanded_params) == 1 and expanded_params[0] is model_param + assert len(expanded_main_params) == 1 and expanded_main_params[0] is shard_main_param + assert expanded_offsets == [1] + + +def test_grouped_quantized_tensor_detection_allows_lazy_split_members(): + grouped_param = _FakeGroupedQuantizedTensor([torch.empty(1)], quantized_tensors=None) + + assert DistributedOptimizer._is_grouped_quantized_tensor(grouped_param) + assert DistributedOptimizer._is_distopt_quantized_param(grouped_param) + + +def test_grouped_quantized_tensor_detection_requires_quantizer(): + grouped_param = _FakeGroupedQuantizedTensor([torch.empty(1)], quantized_tensors=None) + grouped_param.quantizer = None + + assert not DistributedOptimizer._is_grouped_quantized_tensor(grouped_param) + assert not DistributedOptimizer._is_distopt_quantized_param(grouped_param) + + +def test_expand_quantized_param_shard_for_cast_lazy_splits_when_members_unset(): + optimizer = DistributedOptimizer.__new__(DistributedOptimizer) + members = [torch.empty(2), torch.empty(2)] + grouped_param = _FakeGroupedQuantizedTensor(members, quantized_tensors=None) + shard_main_param = torch.arange(4) + + expanded_params, expanded_main_params, expanded_offsets = ( + optimizer._expand_quantized_param_shard_for_cast( + grouped_param, shard_main_param, start_offset=0 + ) + ) + + assert expanded_params == members + torch.testing.assert_close(expanded_main_params[0], torch.tensor([0, 1])) + torch.testing.assert_close(expanded_main_params[1], torch.tensor([2, 3])) + assert expanded_offsets == [0, 0] + + +def test_expand_quantized_param_shard_for_cast_handles_none_shard_for_grouped_param(): + optimizer = DistributedOptimizer.__new__(DistributedOptimizer) + members = [torch.empty(3), torch.empty(2)] + grouped_param = _FakeGroupedQuantizedTensor(members, quantized_tensors=members) + + expanded_params, expanded_main_params, expanded_offsets = ( + optimizer._expand_quantized_param_shard_for_cast( + grouped_param, shard_main_param=None, start_offset=None + ) + ) + + assert expanded_params == members + assert expanded_main_params == [None, None] + assert expanded_offsets == [None, None] + + +# ----------------------------------------------------------------------------- +# Tests for `_normalize_state_dict_for_grouped_params`. +# +# The method's duck-type discriminator is `hasattr(module, 'num_gemms')` and +# `hasattr(module, '_split_grouped_checkpoint_tensor')`, so we don't need a real +# `TEGroupedLinear` instance — a `torch.nn.Module` subclass with those attrs +# stands in. Building real model trees lets us exercise `named_modules()` paths +# and the `module.`-prefix stripping logic. +# ----------------------------------------------------------------------------- + + +class _FakeTEGroupedLinear(torch.nn.Module): + """Stub matching the duck-type the optimizer keys on.""" + + def __init__( + self, num_gemms, *, single_grouped_weight, use_bias=False, single_grouped_bias=False + ): + super().__init__() + self.num_gemms = num_gemms + self.single_grouped_weight = single_grouped_weight + self.use_bias = use_bias + self.single_grouped_bias = single_grouped_bias + + def _split_grouped_checkpoint_tensor(self, tensor, checkpoint_key): + # Mirror the real method's first-dim split (sufficient for round-trip tests). + if tensor.shape[0] == self.num_gemms: + return list(tensor.unbind(dim=0)) + if tensor.shape[0] % self.num_gemms == 0: + return list(torch.chunk(tensor, self.num_gemms, dim=0)) + raise RuntimeError(f"cannot split {checkpoint_key}") + + +class _Wrapper(torch.nn.Module): + """Minimal model_chunk wrapping a child module under a configurable name path.""" + + def __init__(self, child, child_path: str = "experts.linear_fc1"): + super().__init__() + # Build the path one level at a time using nested Modules so that + # named_modules() emits an entry for each path segment (matching real + # mcore module trees that look like `decoder.layers.0.mlp.experts.linear_fc1`). + cursor = self + parts = child_path.split(".") + for part in parts[:-1]: + inner = torch.nn.Module() + setattr(cursor, part, inner) + cursor = inner + setattr(cursor, parts[-1], child) + + +def test_normalize_state_dict_indexed_to_grouped_weight(): + grouped = _FakeTEGroupedLinear(num_gemms=3, single_grouped_weight=True) + chunk = _Wrapper(grouped, child_path="experts.linear_fc1") + indexed = [torch.tensor([float(i), float(i) + 0.5]) for i in range(3)] + state_dict = {f"experts.linear_fc1.weight{i}": indexed[i] for i in range(3)} + + DistributedOptimizer._normalize_state_dict_for_grouped_params(state_dict, chunk) + + assert set(state_dict.keys()) == {"experts.linear_fc1.weight"} + torch.testing.assert_close(state_dict["experts.linear_fc1.weight"], torch.stack(indexed, dim=0)) + + +def test_normalize_state_dict_grouped_to_indexed_weight(): + grouped = _FakeTEGroupedLinear(num_gemms=3, single_grouped_weight=False) + chunk = _Wrapper(grouped, child_path="experts.linear_fc1") + grouped_tensor = torch.arange(9, dtype=torch.float32).view(3, 3) + state_dict = {"experts.linear_fc1.weight": grouped_tensor} + + DistributedOptimizer._normalize_state_dict_for_grouped_params(state_dict, chunk) + + assert set(state_dict.keys()) == { + "experts.linear_fc1.weight0", + "experts.linear_fc1.weight1", + "experts.linear_fc1.weight2", + } + for i in range(3): + torch.testing.assert_close(state_dict[f"experts.linear_fc1.weight{i}"], grouped_tensor[i]) + + +def test_normalize_state_dict_indexed_to_grouped_with_bias(): + grouped = _FakeTEGroupedLinear( + num_gemms=2, single_grouped_weight=True, use_bias=True, single_grouped_bias=True + ) + chunk = _Wrapper(grouped, child_path="layer.linear_fc1") + weights = [torch.tensor([1.0, 2.0]), torch.tensor([3.0, 4.0])] + biases = [torch.tensor([10.0]), torch.tensor([20.0])] + state_dict = { + "layer.linear_fc1.weight0": weights[0], + "layer.linear_fc1.weight1": weights[1], + "layer.linear_fc1.bias0": biases[0], + "layer.linear_fc1.bias1": biases[1], + } + + DistributedOptimizer._normalize_state_dict_for_grouped_params(state_dict, chunk) + + assert set(state_dict.keys()) == {"layer.linear_fc1.weight", "layer.linear_fc1.bias"} + torch.testing.assert_close(state_dict["layer.linear_fc1.weight"], torch.stack(weights, dim=0)) + torch.testing.assert_close(state_dict["layer.linear_fc1.bias"], torch.stack(biases, dim=0)) + + +def test_normalize_state_dict_skips_bias_when_use_bias_false(): + grouped = _FakeTEGroupedLinear( + num_gemms=2, single_grouped_weight=True, use_bias=False, single_grouped_bias=False + ) + chunk = _Wrapper(grouped, child_path="experts.linear_fc1") + state_dict = { + "experts.linear_fc1.weight0": torch.zeros(2), + "experts.linear_fc1.weight1": torch.ones(2), + # Stray bias keys that don't belong on this module — must remain untouched. + "experts.linear_fc1.bias0": torch.tensor([99.0]), + "experts.linear_fc1.bias1": torch.tensor([42.0]), + } + + DistributedOptimizer._normalize_state_dict_for_grouped_params(state_dict, chunk) + + assert "experts.linear_fc1.weight" in state_dict + assert state_dict["experts.linear_fc1.bias0"].item() == 99.0 + assert state_dict["experts.linear_fc1.bias1"].item() == 42.0 + + +def test_normalize_state_dict_strips_module_prefix_from_named_modules_path(): + """`named_modules()` may emit `module.foo.bar` (Float16Module wrapping); the + suffix lookup must use `foo.bar` after stripping the leading `module.`. + """ + + class _ModuleWrapper(torch.nn.Module): + def __init__(self, inner): + super().__init__() + self.module = inner # named_modules() now emits paths like `module.experts.linear_fc1`. + + grouped = _FakeTEGroupedLinear(num_gemms=2, single_grouped_weight=True) + inner = _Wrapper(grouped, child_path="experts.linear_fc1") + chunk = _ModuleWrapper(inner) + + indexed = [torch.tensor([1.0]), torch.tensor([2.0])] + state_dict = { + "experts.linear_fc1.weight0": indexed[0], + "experts.linear_fc1.weight1": indexed[1], + } + + DistributedOptimizer._normalize_state_dict_for_grouped_params(state_dict, chunk) + + assert set(state_dict.keys()) == {"experts.linear_fc1.weight"} + torch.testing.assert_close(state_dict["experts.linear_fc1.weight"], torch.stack(indexed, dim=0)) + + +def test_normalize_state_dict_skips_when_grouped_target_already_present(): + grouped = _FakeTEGroupedLinear(num_gemms=2, single_grouped_weight=True) + chunk = _Wrapper(grouped, child_path="experts.linear_fc1") + grouped_tensor = torch.tensor([[1.0, 2.0], [3.0, 4.0]]) + state_dict = { + "experts.linear_fc1.weight": grouped_tensor, + "experts.linear_fc1.weight0": torch.zeros(2), + "experts.linear_fc1.weight1": torch.ones(2), + } + + DistributedOptimizer._normalize_state_dict_for_grouped_params(state_dict, chunk) + + # No mutation: target layout already present. + torch.testing.assert_close(state_dict["experts.linear_fc1.weight"], grouped_tensor) + torch.testing.assert_close(state_dict["experts.linear_fc1.weight0"], torch.zeros(2)) + torch.testing.assert_close(state_dict["experts.linear_fc1.weight1"], torch.ones(2)) + + +def test_normalize_state_dict_skips_when_indexed_set_incomplete(): + """`single_grouped=True` requires ALL indexed keys to fold; partial set is left alone.""" + grouped = _FakeTEGroupedLinear(num_gemms=3, single_grouped_weight=True) + chunk = _Wrapper(grouped, child_path="experts.linear_fc1") + state_dict = { + "experts.linear_fc1.weight0": torch.zeros(2), + "experts.linear_fc1.weight2": torch.ones(2), # weight1 missing + } + + DistributedOptimizer._normalize_state_dict_for_grouped_params(state_dict, chunk) + + assert "experts.linear_fc1.weight" not in state_dict + assert set(state_dict.keys()) == {"experts.linear_fc1.weight0", "experts.linear_fc1.weight2"} + + +def test_normalize_state_dict_skips_ambiguous_indexed_match(): + """Two state_dict keys ending in the same indexed suffix (e.g. duplicate paths + across model chunks under one flat dict) must NOT be folded — the + `len(matches) == 1` guard in the optimizer leaves an incomplete `indexed_match_map`, + which fails the `len == num_gemms` check and the function bails out safely. + """ + grouped = _FakeTEGroupedLinear(num_gemms=2, single_grouped_weight=True) + chunk = _Wrapper(grouped, child_path="experts.linear_fc1") + state_dict = { + # Both keys end in `experts.linear_fc1.weight0` → endswith ambiguity. + "chunk_a.experts.linear_fc1.weight0": torch.tensor([1.0]), + "chunk_b.experts.linear_fc1.weight0": torch.tensor([2.0]), + # Single weight1 key — would-be match, but weight0 is ambiguous. + "chunk_a.experts.linear_fc1.weight1": torch.tensor([3.0]), + } + snapshot = dict(state_dict) + + DistributedOptimizer._normalize_state_dict_for_grouped_params(state_dict, chunk) + + # No remap: ambiguous weight0 means indexed_match_map is incomplete → skipped. + assert state_dict.keys() == snapshot.keys() + for k, v in snapshot.items(): + torch.testing.assert_close(state_dict[k], v) + + +def test_normalize_state_dict_ignores_non_te_grouped_modules(): + """Modules without `num_gemms`/`_split_grouped_checkpoint_tensor` (e.g. `torch.nn.Linear`) + must be skipped by the duck-type check — no traversal of their state_dict keys. + """ + plain = torch.nn.Linear(4, 4) + chunk = _Wrapper(plain, child_path="layer.linear") + # state_dict that, if naively matched, looks indexed. + state_dict = { + "layer.linear.weight0": torch.zeros(4, 4), + "layer.linear.weight1": torch.ones(4, 4), + } + snapshot = dict(state_dict) + + DistributedOptimizer._normalize_state_dict_for_grouped_params(state_dict, chunk) + + assert state_dict.keys() == snapshot.keys() diff --git a/tests/unit_tests/transformer/moe/test_grouped_mlp.py b/tests/unit_tests/transformer/moe/test_grouped_mlp.py index b1f15a9eef1..90f3cd1f6cf 100644 --- a/tests/unit_tests/transformer/moe/test_grouped_mlp.py +++ b/tests/unit_tests/transformer/moe/test_grouped_mlp.py @@ -1,9 +1,14 @@ # Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved. +import argparse +from types import SimpleNamespace + import pytest import torch import torch.nn.functional as F +import megatron.core.transformer.moe.experts as experts_module +from megatron.core.fusions.fused_bias_geglu import quick_gelu from megatron.core.models.gpt.gpt_layer_specs import ( get_gpt_layer_local_submodules, get_gpt_layer_with_transformer_engine_submodules, @@ -14,11 +19,551 @@ from megatron.core.transformer.spec_utils import get_submodules from megatron.core.transformer.transformer_config import TransformerConfig from megatron.core.utils import is_te_min_version -from megatron.training.arguments import parse_args +from megatron.training.arguments import _add_network_size_args, parse_args from megatron.training.initialize import _set_random_seed from tests.unit_tests.test_utilities import Utils +def test_op_fuser_transformer_config_args_are_exposed(): + parser = argparse.ArgumentParser() + _add_network_size_args(parser) + + args = parser.parse_args( + ["--use-transformer-engine-op-fuser", "--moe-mlp-glu-interleave-size", "16"] + ) + + assert args.use_transformer_engine_op_fuser is True + assert args.moe_mlp_glu_interleave_size == 16 + + +def test_remove_glu_interleaving_restores_contiguous_gate_and_linear_halves(): + interleaved = torch.tensor([[1, 2, 5, 6, 3, 4, 7, 8], [11, 12, 15, 16, 13, 14, 17, 18]]) + expected = torch.tensor([[1, 2, 3, 4, 5, 6, 7, 8], [11, 12, 13, 14, 15, 16, 17, 18]]) + + output = TEGroupedMLP._remove_glu_interleaving(interleaved, interleave_size=2) + + torch.testing.assert_close(output, expected) + + +def test_make_fused_ops_reuses_grouped_linear_weights_on_meta_device(monkeypatch): + class FakeGroupedLinear(torch.nn.Module): + def __init__( + self, + num_gemms, + in_features, + out_features, + *, + bias, + device, + dtype, + accumulate_into_main_grad, + single_grouped_weight, + single_grouped_bias=False, + delay_wgrad_compute=False, + ): + super().__init__() + self.num_gemms = num_gemms + self.in_features = in_features + self.out_features = out_features + self.use_bias = bias + self.device = device + self.dtype = dtype + self.fuse_wgrad_accumulation = accumulate_into_main_grad + self.single_grouped_weight = single_grouped_weight + self.single_grouped_bias = single_grouped_bias + self.delay_wgrad_compute = delay_wgrad_compute + + def need_backward_dw(self): + return False + + class FakeScaledSwiGLU(torch.nn.Module): + def __init__(self, glu_interleave_size): + super().__init__() + self.glu_interleave_size = glu_interleave_size + + class FakeSequential(list): + def register_forward_pre_hook(self, hook): + self.forward_pre_hook = hook + + fake_te = SimpleNamespace( + pytorch=SimpleNamespace( + GroupedLinear=FakeGroupedLinear, + ops=SimpleNamespace( + GroupedLinear=FakeGroupedLinear, + ScaledSwiGLU=FakeScaledSwiGLU, + Sequential=FakeSequential, + ), + ) + ) + monkeypatch.setattr(experts_module, "te", fake_te) + + module = TEGroupedMLP.__new__(TEGroupedMLP) + torch.nn.Module.__init__(module) + module.config = SimpleNamespace( + moe_mlp_glu_interleave_size=16, delay_wgrad_compute=False, activation_func_clamp_value=None + ) + module.activation_func = F.silu + module.config.gated_linear_unit = True + module.linear_fc1 = FakeGroupedLinear( + 2, + 4, + 8, + bias=True, + device="cuda", + dtype=torch.bfloat16, + accumulate_into_main_grad=True, + single_grouped_weight=False, + ) + module.linear_fc2 = FakeGroupedLinear( + 2, + 8, + 4, + bias=False, + device="cuda", + dtype=torch.bfloat16, + accumulate_into_main_grad=False, + single_grouped_weight=True, + ) + module.linear_fc1.weight0 = torch.nn.Parameter(torch.ones(8, 4)) + module.linear_fc1.weight1 = torch.nn.Parameter(torch.ones(8, 4) * 2) + module.linear_fc1.bias0 = torch.nn.Parameter(torch.zeros(8)) + module.linear_fc1.bias1 = torch.nn.Parameter(torch.ones(8)) + module.linear_fc2.weight = torch.nn.Parameter(torch.ones(4, 8)) + + ops = module._make_fused_ops() + + assert len(ops) == 3 + assert ops[0].device == "meta" + assert ops[0].weight0 is module.linear_fc1.weight0 + assert ops[0].weight1 is module.linear_fc1.weight1 + assert ops[0].bias0 is module.linear_fc1.bias0 + assert ops[0].bias1 is module.linear_fc1.bias1 + assert ops[1].glu_interleave_size == 16 + assert ops[2].device == "meta" + assert ops[2].weight is module.linear_fc2.weight + assert hasattr(ops, "forward_pre_hook") + + +def test_fused_forward_caches_ops_and_forwards_expected_arguments(): + class FakeFusedOps: + def __call__(self, hidden_states, fc1_tokens, probs, fc2_tokens): + self.args = (hidden_states, fc1_tokens, probs, fc2_tokens) + return hidden_states + 1 + + module = TEGroupedMLP.__new__(TEGroupedMLP) + # `_fused_forward` calls `skip_routed_expert_padding(config)` (added by PR 4071), which + # reads `moe_token_dispatcher_type` and `moe_flex_dispatcher_backend` after the + # `moe_router_padding_for_quantization` short-circuit fails. + module.config = SimpleNamespace( + fp8=False, + fp4=False, + moe_router_padding_for_quantization=False, + moe_token_dispatcher_type=None, + moe_flex_dispatcher_backend=None, + ) + module._fused_ops = None + fused_ops = FakeFusedOps() + module._make_fused_ops = lambda: fused_ops + hidden_states = torch.zeros(2, 4) + tokens_per_expert = torch.tensor([1, 1]) + probs = torch.ones(2) + + output = module._fused_forward(hidden_states, tokens_per_expert, probs) + + torch.testing.assert_close(output, torch.ones_like(hidden_states)) + assert module._fused_ops[0] is fused_ops + assert fused_ops.args[0] is hidden_states + assert fused_ops.args[1] is tokens_per_expert + assert fused_ops.args[2] is probs + assert fused_ops.args[3] is tokens_per_expert + + +def test_apply_bias_returns_input_unchanged_when_bias_is_none(): + intermediate = torch.arange(6, dtype=torch.float32).view(3, 2) + + output = TEGroupedMLP._apply_bias( + intermediate, bias_parallel=None, tokens_per_expert=[2, 1], permuted_probs=torch.ones(3) + ) + + assert output is intermediate + + +def test_apply_bias_combines_per_expert_bias_and_probs(): + intermediate = torch.tensor([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]], dtype=torch.float32) + bias_parallel = [torch.tensor([10.0, 20.0]), torch.tensor([100.0, 200.0])] + tokens_per_expert = [2, 1] + permuted_probs = torch.tensor([0.5, 0.5, 1.0]) + expected = torch.tensor([[6.0, 12.0], [8.0, 14.0], [105.0, 206.0]], dtype=torch.float32) + + output = TEGroupedMLP._apply_bias( + intermediate, bias_parallel, tokens_per_expert, permuted_probs + ) + + torch.testing.assert_close(output, expected) + assert output.dtype == intermediate.dtype + + +def test_make_fused_impl_pre_forward_hook_dispatches_submodule_hooks(): + module = TEGroupedMLP.__new__(TEGroupedMLP) + torch.nn.Module.__init__(module) + fc1_child = torch.nn.Linear(2, 2) + fc2_child = torch.nn.Linear(2, 2) + module.linear_fc1 = torch.nn.Sequential(fc1_child) + module.linear_fc2 = torch.nn.Sequential(fc2_child) + + calls = [] + + def fc1_hook(submodule, _inp): + calls.append(("fc1", submodule)) + return None + + def fc2_hook(submodule, _inp): + calls.append(("fc2", submodule)) + return None + + fc1_child.register_forward_pre_hook(fc1_hook) + fc2_child.register_forward_pre_hook(fc2_hook) + + hook = module._make_fused_impl_pre_forward_hook() + hook(object()) + + visited = {label for label, _ in calls} + assert visited == {"fc1", "fc2"} + + +def test_make_fused_impl_pre_forward_hook_rejects_input_modifying_hook(): + module = TEGroupedMLP.__new__(TEGroupedMLP) + torch.nn.Module.__init__(module) + fc1_child = torch.nn.Linear(2, 2) + module.linear_fc1 = torch.nn.Sequential(fc1_child) + module.linear_fc2 = torch.nn.Sequential(torch.nn.Linear(2, 2)) + + fc1_child.register_forward_pre_hook(lambda submodule, _inp: torch.zeros(1)) + + hook = module._make_fused_impl_pre_forward_hook() + + with pytest.raises(RuntimeError, match="modifies the input tensor"): + hook(object()) + + +def test_make_fused_ops_handles_single_grouped_weight_for_fc1(monkeypatch): + class FakeGroupedLinear(torch.nn.Module): + def __init__( + self, + num_gemms, + in_features, + out_features, + *, + bias, + device, + dtype, + accumulate_into_main_grad, + single_grouped_weight, + single_grouped_bias=False, + delay_wgrad_compute=False, + ): + super().__init__() + self.num_gemms = num_gemms + self.in_features = in_features + self.out_features = out_features + self.use_bias = bias + self.device = device + self.dtype = dtype + self.fuse_wgrad_accumulation = accumulate_into_main_grad + self.single_grouped_weight = single_grouped_weight + self.single_grouped_bias = single_grouped_bias + self.delay_wgrad_compute = delay_wgrad_compute + + def need_backward_dw(self): + return False + + class FakeScaledSwiGLU(torch.nn.Module): + def __init__(self, glu_interleave_size): + super().__init__() + self.glu_interleave_size = glu_interleave_size + + class FakeSequential(list): + def register_forward_pre_hook(self, hook): + self.forward_pre_hook = hook + + fake_te = SimpleNamespace( + pytorch=SimpleNamespace( + GroupedLinear=FakeGroupedLinear, + ops=SimpleNamespace( + GroupedLinear=FakeGroupedLinear, + ScaledSwiGLU=FakeScaledSwiGLU, + Sequential=FakeSequential, + ), + ) + ) + monkeypatch.setattr(experts_module, "te", fake_te) + + module = TEGroupedMLP.__new__(TEGroupedMLP) + torch.nn.Module.__init__(module) + module.config = SimpleNamespace( + moe_mlp_glu_interleave_size=8, + delay_wgrad_compute=False, + activation_func_clamp_value=None, + gated_linear_unit=True, + ) + module.activation_func = F.silu + module.linear_fc1 = FakeGroupedLinear( + 2, + 4, + 8, + bias=False, + device="cuda", + dtype=torch.bfloat16, + accumulate_into_main_grad=False, + single_grouped_weight=True, + ) + module.linear_fc2 = FakeGroupedLinear( + 2, + 8, + 4, + bias=True, + device="cuda", + dtype=torch.bfloat16, + accumulate_into_main_grad=True, + single_grouped_weight=False, + ) + module.linear_fc1.weight = torch.nn.Parameter(torch.ones(2, 8, 4)) + module.linear_fc2.weight0 = torch.nn.Parameter(torch.ones(4, 8)) + module.linear_fc2.weight1 = torch.nn.Parameter(torch.ones(4, 8) * 2) + module.linear_fc2.bias0 = torch.nn.Parameter(torch.zeros(4)) + module.linear_fc2.bias1 = torch.nn.Parameter(torch.ones(4)) + + ops = module._make_fused_ops() + + assert ops[0].weight is module.linear_fc1.weight + assert ops[1].glu_interleave_size == 8 + assert ops[2].weight0 is module.linear_fc2.weight0 + assert ops[2].weight1 is module.linear_fc2.weight1 + assert ops[2].bias0 is module.linear_fc2.bias0 + assert ops[2].bias1 is module.linear_fc2.bias1 + + +def _make_fake_te_namespace(): + """Build a fake TE namespace with the activation classes _make_fused_ops uses.""" + + class FakeGroupedLinear(torch.nn.Module): + def __init__( + self, + num_gemms, + in_features, + out_features, + *, + bias, + device, + dtype, + accumulate_into_main_grad, + single_grouped_weight, + single_grouped_bias=False, + delay_wgrad_compute=False, + ): + super().__init__() + self.num_gemms = num_gemms + self.in_features = in_features + self.out_features = out_features + self.use_bias = bias + self.device = device + self.dtype = dtype + self.fuse_wgrad_accumulation = accumulate_into_main_grad + self.single_grouped_weight = single_grouped_weight + self.single_grouped_bias = single_grouped_bias + self.delay_wgrad_compute = delay_wgrad_compute + + def need_backward_dw(self): + return False + + class FakeScaledSwiGLU(torch.nn.Module): + def __init__(self, glu_interleave_size): + super().__init__() + self.glu_interleave_size = glu_interleave_size + + class FakeScaledClampedQGeGLU(torch.nn.Module): + def __init__(self, glu_interleave_size, limit=None): + super().__init__() + self.glu_interleave_size = glu_interleave_size + self.limit = limit + + class FakeSequential(list): + def register_forward_pre_hook(self, hook): + self.forward_pre_hook = hook + + return ( + SimpleNamespace( + pytorch=SimpleNamespace( + GroupedLinear=FakeGroupedLinear, + ops=SimpleNamespace( + GroupedLinear=FakeGroupedLinear, + ScaledSwiGLU=FakeScaledSwiGLU, + ScaledClampedQGeGLU=FakeScaledClampedQGeGLU, + Sequential=FakeSequential, + ), + ) + ), + FakeGroupedLinear, + ) + + +def test_make_fused_ops_uses_clamped_qgeglu_for_quick_gelu(monkeypatch): + """quick_gelu + clamp value → ScaledClampedQGeGLU(limit=clamp).""" + fake_te, FakeGroupedLinear = _make_fake_te_namespace() + monkeypatch.setattr(experts_module, "te", fake_te) + + module = TEGroupedMLP.__new__(TEGroupedMLP) + torch.nn.Module.__init__(module) + module.config = SimpleNamespace( + moe_mlp_glu_interleave_size=4, + delay_wgrad_compute=False, + activation_func_clamp_value=7.0, + gated_linear_unit=True, + ) + module.activation_func = quick_gelu + common = dict( + device="cuda", + dtype=torch.bfloat16, + accumulate_into_main_grad=False, + single_grouped_weight=False, + ) + module.linear_fc1 = FakeGroupedLinear(2, 4, 8, bias=False, **common) + module.linear_fc2 = FakeGroupedLinear(2, 8, 4, bias=False, **common) + module.linear_fc1.weight0 = torch.nn.Parameter(torch.ones(8, 4)) + module.linear_fc1.weight1 = torch.nn.Parameter(torch.ones(8, 4)) + module.linear_fc2.weight0 = torch.nn.Parameter(torch.ones(4, 8)) + module.linear_fc2.weight1 = torch.nn.Parameter(torch.ones(4, 8)) + + ops = module._make_fused_ops() + + activation = ops[1] + assert type(activation).__name__ == "FakeScaledClampedQGeGLU" + assert activation.glu_interleave_size == 4 + assert activation.limit == 7.0 + + +def test_make_fused_ops_attaches_single_grouped_bias_for_fc1(monkeypatch): + """single_grouped_bias=True → bias attached as `bias` (not `bias{idx}`).""" + fake_te, FakeGroupedLinear = _make_fake_te_namespace() + monkeypatch.setattr(experts_module, "te", fake_te) + + module = TEGroupedMLP.__new__(TEGroupedMLP) + torch.nn.Module.__init__(module) + module.config = SimpleNamespace( + moe_mlp_glu_interleave_size=2, + delay_wgrad_compute=False, + activation_func_clamp_value=None, + gated_linear_unit=True, + ) + module.activation_func = F.silu + common = dict(device="cuda", dtype=torch.bfloat16, accumulate_into_main_grad=False) + # FC1 stores bias as a single grouped tensor; weights still per-GEMM. + module.linear_fc1 = FakeGroupedLinear( + 2, 4, 8, bias=True, single_grouped_weight=False, single_grouped_bias=True, **common + ) + module.linear_fc2 = FakeGroupedLinear( + 2, 8, 4, bias=False, single_grouped_weight=False, single_grouped_bias=False, **common + ) + module.linear_fc1.weight0 = torch.nn.Parameter(torch.ones(8, 4)) + module.linear_fc1.weight1 = torch.nn.Parameter(torch.ones(8, 4)) + module.linear_fc1.bias = torch.nn.Parameter(torch.zeros(2, 8)) + module.linear_fc2.weight0 = torch.nn.Parameter(torch.ones(4, 8)) + module.linear_fc2.weight1 = torch.nn.Parameter(torch.ones(4, 8)) + + ops = module._make_fused_ops() + + assert ops[0].weight0 is module.linear_fc1.weight0 + assert ops[0].weight1 is module.linear_fc1.weight1 + assert ops[0].bias is module.linear_fc1.bias # ← single grouped bias attached at "bias" + assert not hasattr( + ops[0], "bias0" + ), "bias should not be split into bias{idx} when single_grouped_bias=True" + + +def test_backward_dw_dispatches_fused_children_in_fc2_then_fc1_order(): + """delay_wgrad_compute=True (via wrapper) → backward_dw calls fused [2] then [0], + then triggers the wrapper-side wgrad/reduce hooks (PR 4311) so DDP's reduce-scatter + fires on the original linear_fc1/fc2 modules.""" + + class _Recorder: + def __init__(self, label): + self.label = label + self.calls = [] + + def backward_dw(self): + self.calls.append("backward_dw") + + class _FakeWrapper: + def __init__(self, label): + self.label = label + self.delay_wgrad_compute = True + self.backward_dw_calls = 0 + self.trigger_calls = 0 + + def backward_dw(self): + self.backward_dw_calls += 1 + + def _trigger_wgrad_accumulation_and_reduce_hooks(self): + self.trigger_calls += 1 + + class _FakeSequential: + def __init__(self, children): + self._children = children + + def children(self): + return iter(self._children) + + fc1_op = _Recorder("fc1") + activation_op = _Recorder("activation") + fc2_op = _Recorder("fc2") + fake_seq = _FakeSequential([fc1_op, activation_op, fc2_op]) + + module = TEGroupedMLP.__new__(TEGroupedMLP) + torch.nn.Module.__init__(module) + module._with_fused_impl = True + module._fused_ops = (fake_seq,) + module.linear_fc1 = _FakeWrapper("fc1_wrapper") + module.linear_fc2 = _FakeWrapper("fc2_wrapper") + module.config = SimpleNamespace(delay_wgrad_compute=False) # config flag is irrelevant + + module.backward_dw() + + assert fc2_op.calls == ["backward_dw"] + assert fc1_op.calls == ["backward_dw"] + assert activation_op.calls == [], "activation op must not be invoked" + assert module.linear_fc1.backward_dw_calls == 0, "wrapper backward_dw must be skipped" + assert module.linear_fc2.backward_dw_calls == 0 + # PR 4311: DDP reduce-scatter hooks live on the original wrappers and must be triggered + # explicitly because the fused path's backward_dw runs on the new GroupedLinear instances. + assert module.linear_fc2.trigger_calls == 1 + assert module.linear_fc1.trigger_calls == 1 + + +def test_backward_dw_falls_back_to_wrappers_when_delay_wgrad_off(): + """delay_wgrad_compute=False on wrappers → use the original wrapper backward_dw path.""" + + class _FakeWrapper: + def __init__(self): + self.delay_wgrad_compute = False + self.backward_dw_calls = 0 + + def backward_dw(self): + self.backward_dw_calls += 1 + + module = TEGroupedMLP.__new__(TEGroupedMLP) + torch.nn.Module.__init__(module) + module._with_fused_impl = True + module._fused_ops = None + module.linear_fc1 = _FakeWrapper() + module.linear_fc2 = _FakeWrapper() + module.config = SimpleNamespace(delay_wgrad_compute=False) + + module.backward_dw() + + assert module.linear_fc1.backward_dw_calls == 1 + assert module.linear_fc2.backward_dw_calls == 1 + + @pytest.mark.skipif( not is_te_min_version("1.9.0.dev0"), reason="TE Grouped MLP is only supported in TE 1.9.0.dev0 and later.", @@ -186,3 +731,183 @@ def test_gpu_forward_backward_with_no_tokens_allocated(self): for i in range(self.num_experts): assert getattr(self.grouped_mlp.experts.linear_fc1, f"weight{i}").grad is not None assert getattr(self.grouped_mlp.experts.linear_fc2, f"weight{i}").grad is not None + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + @pytest.mark.internal + def test_gpu_make_fused_ops_constructs_with_real_te(self): + """Verify `_make_fused_ops` builds a working TE op-fuser pipeline. + + Regression guard for the `device="meta"` shell + weight-reattachment + sequence — a TE-side change that allocates state in the constructor + would break this without any of the mock-based unit tests catching it. + """ + try: + from transformer_engine.pytorch.ops import GroupedLinear, ScaledSwiGLU + except ImportError: + pytest.skip("TE op fuser API not available") + import inspect + + if "single_grouped_weight" not in inspect.signature(GroupedLinear.__init__).parameters: + pytest.skip( + "Installed TE op fuser GroupedLinear lacks `single_grouped_weight` kwarg; " + "_make_fused_ops requires a TE build that exposes it." + ) + + Utils.destroy_model_parallel() + Utils.initialize_model_parallel(1, 1) + + tf_config = TransformerConfig( + num_layers=1, + hidden_size=self.hidden_size, + num_attention_heads=4, + num_moe_experts=self.num_experts, + use_cpu_initialization=False, + add_bias_linear=False, + gated_linear_unit=True, + activation_func=F.silu, + bias_activation_fusion=False, + bf16=True, + params_dtype=torch.bfloat16, + moe_router_load_balancing_type="sinkhorn", + moe_router_topk=1, + moe_grouped_gemm=True, + use_transformer_engine_op_fuser=True, + ) + _set_random_seed(seed_=123, data_parallel_random_init=False) + submodules = get_submodules( + get_gpt_layer_with_transformer_engine_submodules( + self.num_experts, moe_grouped_gemm=True + ).mlp + ) + assert isinstance(submodules, MoESubmodules) + layer = MoELayer(tf_config, submodules) + layer = Float16Module(layer.config, layer).module + layer.cuda() + experts = layer.experts + assert isinstance(experts, TEGroupedMLP) + assert experts._with_fused_impl + + ops = experts._make_fused_ops() + + assert len(ops) == 3 + assert isinstance(ops[0], GroupedLinear) + assert isinstance(ops[1], ScaledSwiGLU) + assert isinstance(ops[2], GroupedLinear) + # Weights of the wrapper ops must alias the underlying GroupedLinear + # parameters so optimizer updates are visible to the fused path. + for idx in range(experts.linear_fc1.num_gemms): + if not getattr(experts.linear_fc1, "single_grouped_weight", False): + assert getattr(ops[0], f"weight{idx}") is getattr( + experts.linear_fc1, f"weight{idx}" + ) + for idx in range(experts.linear_fc2.num_gemms): + if not getattr(experts.linear_fc2, "single_grouped_weight", False): + assert getattr(ops[2], f"weight{idx}") is getattr( + experts.linear_fc2, f"weight{idx}" + ) + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + @pytest.mark.internal + def test_gpu_fused_path_loss_decreases(self): + """End-to-end signal that the Megatron op-fuser wrapper actually trains. + + Builds a small grouped-MLP MoE layer with `use_transformer_engine_op_fuser=True`, + runs a few Adam steps driving the output toward a fixed target, and asserts the + loss decreases and stays finite. This catches regressions in the Megatron-side + wrapper plumbing — `_make_fused_ops` op-sequence construction, the fused-impl + pre-forward hook, weight aliasing into the TE op wrappers, and (delayed-)wgrad + backward — failure modes that pass our mock-based unit tests but break real + training. + + IMPORTANT: this test does NOT verify that TE's underlying cuDSL fused grouped + MLP kernel is actually invoked. That kernel additionally requires SM100 + (Blackwell) compute capability, plus `NVTE_CUTEDSL_FUSED_GROUPED_MLP` set, plus + a fusion-registration that happens at TE module import. On H100/A100 hardware + (which is the current CI matrix) the kernel is unavailable regardless and TE + falls back to its basic-op path — the Megatron wrapper still runs end-to-end + and the loss still decreases. So this test covers the wrapper code this PR + adds; it does not cover the cuDSL kernel selection. The functional test + `moe/gpt3_mcore_te_tp1_pp1_te_4experts_groupedGEMM_op_fuser` sets the env var + in its recipe as Blackwell forward-compat, but on current CI hardware it also + runs basic-op. + """ + try: + from transformer_engine.pytorch.ops import GroupedLinear # noqa: F401 + except ImportError: + pytest.skip("TE op fuser API not available") + import inspect + + if "single_grouped_weight" not in inspect.signature(GroupedLinear.__init__).parameters: + pytest.skip("Installed TE op fuser GroupedLinear lacks `single_grouped_weight` kwarg") + + Utils.destroy_model_parallel() + Utils.initialize_model_parallel(1, 1) + + tf_config = TransformerConfig( + num_layers=1, + hidden_size=self.hidden_size, + num_attention_heads=4, + num_moe_experts=self.num_experts, + use_cpu_initialization=False, + add_bias_linear=False, + gated_linear_unit=True, + activation_func=F.silu, + bias_activation_fusion=False, + bf16=True, + params_dtype=torch.bfloat16, + moe_router_load_balancing_type="sinkhorn", + moe_router_topk=1, + moe_grouped_gemm=True, + use_transformer_engine_op_fuser=True, + ) + _set_random_seed(seed_=123, data_parallel_random_init=False) + submodules = get_submodules( + get_gpt_layer_with_transformer_engine_submodules( + self.num_experts, moe_grouped_gemm=True + ).mlp + ) + layer = MoELayer(tf_config, submodules) + layer = Float16Module(layer.config, layer).module + layer.cuda() + assert isinstance(layer.experts, TEGroupedMLP) + assert layer.experts._with_fused_impl + + seq_len = 32 + batch_size = 2 + torch.manual_seed(7) + hidden_states = torch.rand( + (seq_len, batch_size, self.hidden_size), dtype=torch.bfloat16, device="cuda" + ) + # MoE blocks with small-init weights produce near-zero output in bf16, + # so a `||output||^2`-style loss starts at zero and never moves the + # weights. Drive the output toward a fixed non-zero target so the + # gradient is non-trivial even when the initial output is tiny. + with torch.no_grad(): + init_output, _ = layer(hidden_states) + target = torch.full_like(init_output, 0.5, dtype=torch.float32) + # Adam (not SGD) because bf16 weights quantize SGD updates of tiny + # gradients down to zero. Adam's per-parameter adaptive step keeps + # the update large enough to move bf16 weights regardless of grad + # magnitude — without it, even a working fused path produces only + # ~1e-6 loss movement per step, which is below trustable noise. + optimizer = torch.optim.Adam(layer.parameters(), lr=0.05) + + losses = [] + for _ in range(16): + optimizer.zero_grad(set_to_none=True) + output, _ = layer(hidden_states) + loss = (output.float() - target).pow(2).mean() + loss.backward() + optimizer.step() + losses.append(loss.detach().item()) + + assert all( + l == l and l != float("inf") for l in losses + ), f"Fused-path produced non-finite loss: {losses}" + # Require a meaningful relative drop. If the fused path silently no-ops + # weight updates (the historical wgrad-hook bug PR #4311 fixed) the + # loss is constant and this fires. + assert losses[-1] < losses[0] * 0.5, ( + f"Fused-path training did not reduce loss enough: " + f"initial={losses[0]:.6f}, final={losses[-1]:.6f}, full={losses}" + ) diff --git a/tests/unit_tests/transformer/test_transformer_engine_grouped_linear.py b/tests/unit_tests/transformer/test_transformer_engine_grouped_linear.py new file mode 100644 index 00000000000..467a6c31322 --- /dev/null +++ b/tests/unit_tests/transformer/test_transformer_engine_grouped_linear.py @@ -0,0 +1,226 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +import pytest +import torch + +from megatron.core.extensions import transformer_engine as te_ext + +pytestmark = pytest.mark.skipif( + not te_ext.HAVE_TE or not te_ext.is_te_min_version("1.9.0.dev0"), + reason="TE GroupedLinear is only supported in TE 1.9.0.dev0 and later.", +) + + +class _FakeGroupedCheckpointTensor: + def __init__(self, members, quantized_tensors=None): + self._members = members + self.quantized_tensors = quantized_tensors + + def split_into_quantized_tensors(self): + return self._members + + +def _grouped_linear_stub( + num_gemms, *, use_bias=False, single_grouped_weight=False, single_grouped_bias=False +): + module = te_ext.TEGroupedLinear.__new__(te_ext.TEGroupedLinear) + module.num_gemms = num_gemms + module.use_bias = use_bias + module.single_grouped_weight = single_grouped_weight + module.single_grouped_bias = single_grouped_bias + return module + + +def _empty_load_args(): + """Standard load_state_dict pre-hook trailing arguments.""" + return {}, True, [], [], [] + + +def test_split_grouped_checkpoint_tensor_uses_quantized_members(): + module = _grouped_linear_stub(num_gemms=2) + members = [torch.tensor([1, 2]), torch.tensor([3, 4])] + tensor = _FakeGroupedCheckpointTensor(members) + + splits = module._split_grouped_checkpoint_tensor(tensor, "weight") + + assert len(splits) == len(members) + assert all(split is member for split, member in zip(splits, members)) + + +def test_split_grouped_checkpoint_tensor_unbinds_grouped_first_dim(): + module = _grouped_linear_stub(num_gemms=3) + tensor = torch.arange(12).view(3, 4) + + splits = module._split_grouped_checkpoint_tensor(tensor, "weight") + + assert len(splits) == 3 + for gemm_idx, split in enumerate(splits): + torch.testing.assert_close(split, tensor[gemm_idx]) + + +def test_split_grouped_checkpoint_tensor_chunks_packed_first_dim(): + module = _grouped_linear_stub(num_gemms=3) + tensor = torch.arange(18).view(6, 3) + + splits = module._split_grouped_checkpoint_tensor(tensor, "weight") + + assert len(splits) == 3 + for split, expected in zip(splits, torch.chunk(tensor, 3, dim=0)): + torch.testing.assert_close(split, expected) + + +def test_split_grouped_checkpoint_tensor_rejects_bad_group_count(): + module = _grouped_linear_stub(num_gemms=3) + tensor = _FakeGroupedCheckpointTensor([torch.tensor([1]), torch.tensor([2])]) + + with pytest.raises(RuntimeError, match="has 2 groups, expected 3"): + module._split_grouped_checkpoint_tensor(tensor, "weight") + + +def test_split_grouped_checkpoint_tensor_rejects_unsplittable_first_dim(): + module = _grouped_linear_stub(num_gemms=3) + tensor = torch.arange(8).view(4, 2) + + with pytest.raises(RuntimeError, match="Cannot split checkpoint tensor"): + module._split_grouped_checkpoint_tensor(tensor, "weight") + + +def test_split_grouped_checkpoint_tensor_rejects_zero_dim(): + module = _grouped_linear_stub(num_gemms=2) + tensor = torch.tensor(7) + + with pytest.raises(RuntimeError, match="Cannot split checkpoint tensor"): + module._split_grouped_checkpoint_tensor(tensor, "weight") + + +def test_normalize_grouped_parameter_keys_indexed_to_grouped_weight_only(): + module = _grouped_linear_stub(num_gemms=3, use_bias=False, single_grouped_weight=True) + indexed = [torch.tensor([float(i), float(i) + 0.5]) for i in range(3)] + state_dict = {f"layer.weight{i}": indexed[i] for i in range(3)} + + module._normalize_grouped_parameter_keys(state_dict, "layer.", *_empty_load_args()) + + assert set(state_dict.keys()) == {"layer.weight"} + torch.testing.assert_close(state_dict["layer.weight"], torch.stack(indexed, dim=0)) + + +def test_normalize_grouped_parameter_keys_indexed_to_grouped_with_bias(): + module = _grouped_linear_stub( + num_gemms=2, use_bias=True, single_grouped_weight=True, single_grouped_bias=True + ) + weights = [torch.tensor([1.0, 2.0]), torch.tensor([3.0, 4.0])] + biases = [torch.tensor([10.0]), torch.tensor([20.0])] + state_dict = { + "layer.weight0": weights[0], + "layer.weight1": weights[1], + "layer.bias0": biases[0], + "layer.bias1": biases[1], + } + + module._normalize_grouped_parameter_keys(state_dict, "layer.", *_empty_load_args()) + + assert set(state_dict.keys()) == {"layer.weight", "layer.bias"} + torch.testing.assert_close(state_dict["layer.weight"], torch.stack(weights, dim=0)) + torch.testing.assert_close(state_dict["layer.bias"], torch.stack(biases, dim=0)) + + +def test_normalize_grouped_parameter_keys_grouped_to_indexed_weight_only(): + module = _grouped_linear_stub(num_gemms=3, use_bias=False, single_grouped_weight=False) + grouped = torch.arange(9, dtype=torch.float32).view(3, 3) + state_dict = {"layer.weight": grouped} + + module._normalize_grouped_parameter_keys(state_dict, "layer.", *_empty_load_args()) + + assert set(state_dict.keys()) == {"layer.weight0", "layer.weight1", "layer.weight2"} + for i in range(3): + torch.testing.assert_close(state_dict[f"layer.weight{i}"], grouped[i]) + + +def test_normalize_grouped_parameter_keys_grouped_to_indexed_with_bias(): + module = _grouped_linear_stub(num_gemms=2, use_bias=True, single_grouped_weight=False) + grouped_weight = torch.arange(8, dtype=torch.float32).view(2, 4) + grouped_bias = torch.tensor([[10.0, 20.0], [30.0, 40.0]]) + state_dict = {"layer.weight": grouped_weight, "layer.bias": grouped_bias} + + module._normalize_grouped_parameter_keys(state_dict, "layer.", *_empty_load_args()) + + assert set(state_dict.keys()) == { + "layer.weight0", + "layer.weight1", + "layer.bias0", + "layer.bias1", + } + torch.testing.assert_close(state_dict["layer.weight0"], grouped_weight[0]) + torch.testing.assert_close(state_dict["layer.weight1"], grouped_weight[1]) + torch.testing.assert_close(state_dict["layer.bias0"], grouped_bias[0]) + torch.testing.assert_close(state_dict["layer.bias1"], grouped_bias[1]) + + +def test_normalize_grouped_parameter_keys_skips_bias_when_use_bias_false(): + module = _grouped_linear_stub(num_gemms=2, use_bias=False, single_grouped_weight=True) + state_dict = { + "layer.weight0": torch.zeros(2), + "layer.weight1": torch.ones(2), + "layer.bias0": torch.tensor([99.0]), + "layer.bias1": torch.tensor([42.0]), + } + + module._normalize_grouped_parameter_keys(state_dict, "layer.", *_empty_load_args()) + + assert "layer.weight" in state_dict + assert state_dict["layer.bias0"].item() == 99.0 + assert state_dict["layer.bias1"].item() == 42.0 + + +def test_normalize_grouped_parameter_keys_returns_when_target_layout_already_present(): + # single_grouped_weight=True and grouped key already present → no-op + module = _grouped_linear_stub(num_gemms=2, use_bias=False, single_grouped_weight=True) + grouped = torch.tensor([[1.0, 2.0], [3.0, 4.0]]) + state_dict = {"layer.weight": grouped, "layer.weight0": torch.zeros(2)} + + module._normalize_grouped_parameter_keys(state_dict, "layer.", *_empty_load_args()) + + torch.testing.assert_close(state_dict["layer.weight"], grouped) + torch.testing.assert_close(state_dict["layer.weight0"], torch.zeros(2)) + + # single_grouped_weight=False and any indexed key present → no-op + module = _grouped_linear_stub(num_gemms=2, use_bias=False, single_grouped_weight=False) + grouped = torch.tensor([[5.0, 6.0], [7.0, 8.0]]) + state_dict = {"layer.weight": grouped, "layer.weight0": torch.tensor([99.0])} + + module._normalize_grouped_parameter_keys(state_dict, "layer.", *_empty_load_args()) + + torch.testing.assert_close(state_dict["layer.weight"], grouped) + torch.testing.assert_close(state_dict["layer.weight0"], torch.tensor([99.0])) + + +def test_normalize_grouped_parameter_keys_returns_when_indexed_set_incomplete(): + # single_grouped_weight=True needs ALL indexed keys to fold; partial → no-op + module = _grouped_linear_stub(num_gemms=3, use_bias=False, single_grouped_weight=True) + state_dict = {"layer.weight0": torch.zeros(2), "layer.weight2": torch.ones(2)} + + module._normalize_grouped_parameter_keys(state_dict, "layer.", *_empty_load_args()) + + assert "layer.weight" not in state_dict + assert set(state_dict.keys()) == {"layer.weight0", "layer.weight2"} + + +def test_normalize_grouped_parameter_keys_round_trips_via_chained_hooks(): + """Save in one layout, load in the other, then back: tensors survive intact.""" + members = [torch.tensor([1.0, 2.0, 3.0]), torch.tensor([4.0, 5.0, 6.0])] + + # Start from indexed checkpoint, target a single-grouped model. + state_dict = {f"layer.weight{i}": members[i] for i in range(2)} + folder = _grouped_linear_stub(num_gemms=2, use_bias=False, single_grouped_weight=True) + folder._normalize_grouped_parameter_keys(state_dict, "layer.", *_empty_load_args()) + + grouped = state_dict["layer.weight"] + torch.testing.assert_close(grouped, torch.stack(members, dim=0)) + + # Now use that grouped checkpoint with an indexed model. + splitter = _grouped_linear_stub(num_gemms=2, use_bias=False, single_grouped_weight=False) + splitter._normalize_grouped_parameter_keys(state_dict, "layer.", *_empty_load_args()) + + assert set(state_dict.keys()) == {"layer.weight0", "layer.weight1"} + for i in range(2): + torch.testing.assert_close(state_dict[f"layer.weight{i}"], members[i])