diff --git a/megatron/core/distributed/param_and_grad_buffer.py b/megatron/core/distributed/param_and_grad_buffer.py index 6c36f119e19..de47ccc0d29 100644 --- a/megatron/core/distributed/param_and_grad_buffer.py +++ b/megatron/core/distributed/param_and_grad_buffer.py @@ -243,6 +243,9 @@ def __init__( # or bucket.grad_data. self.cached_param_buffer_shard_list = [None] * len(self.buckets) self.cached_grad_buffer_shard_list = [None] * len(self.buckets) + # Track grad mode used to create cached param views. Rebuild if mode changes to avoid + # mixing no_grad-created views with in-place updates in grad-enabled mode. + self._cached_param_buffer_shards_grad_enabled = None def reset(self): """ @@ -403,23 +406,28 @@ def start_param_sync(self, force_sync: bool = False): # Standard distributed optimizer path: use _coalescing_manager. # all_gather_into_tensor writes directly into a contiguous output buffer and # does not need a copy-back step, so coalescing works correctly. - with _coalescing_manager( - self.intra_distributed_optimizer_instance_group, async_ops=async_op - ) as cm: - for idx, bucket in enumerate(self.buckets): - if self.cached_param_buffer_shard_list[idx] is None: - self.cached_param_buffer_shard_list[idx] = shard_buffer( - bucket.param_data, self.intra_distributed_optimizer_instance_size + current_grad_enabled = torch.is_grad_enabled() + if self._cached_param_buffer_shards_grad_enabled != current_grad_enabled: + self.cached_param_buffer_shard_list = [None] * len(self.buckets) + self._cached_param_buffer_shards_grad_enabled = current_grad_enabled + with torch.no_grad(): + with _coalescing_manager( + self.intra_distributed_optimizer_instance_group, async_ops=async_op + ) as cm: + for idx, bucket in enumerate(self.buckets): + if self.cached_param_buffer_shard_list[idx] is None: + self.cached_param_buffer_shard_list[idx] = shard_buffer( + bucket.param_data, self.intra_distributed_optimizer_instance_size + ) + local_data_view = self.cached_param_buffer_shard_list[idx][ + self.intra_distributed_optimizer_instance_rank + ] + dist_all_gather_func( + bucket.param_data, + local_data_view, + group=self.intra_distributed_optimizer_instance_group, + async_op=async_op, ) - local_data_view = self.cached_param_buffer_shard_list[idx][ - self.intra_distributed_optimizer_instance_rank - ] - dist_all_gather_func( - bucket.param_data, - local_data_view, - group=self.intra_distributed_optimizer_instance_group, - async_op=async_op, - ) if async_op: self.param_gather_handle = cm else: diff --git a/megatron/core/extensions/transformer_engine.py b/megatron/core/extensions/transformer_engine.py index 17358f8a921..2487392b5f8 100644 --- a/megatron/core/extensions/transformer_engine.py +++ b/megatron/core/extensions/transformer_engine.py @@ -1787,6 +1787,50 @@ def __init__( setattr(weight, "partition_dim", part_dim) setattr(weight, "partition_stride", 1) + 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.""" + + def maybe_remap_param(param_name: str) -> 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 getattr(self, "single_grouped_parameter", False): + 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") + if self.use_bias: + maybe_remap_param("bias") + + self._register_load_state_dict_pre_hook( + normalize_grouped_parameter_keys, with_module=True + ) + def merge_extra_states( self, state_dict, @@ -1877,6 +1921,31 @@ def merge_extra_states( self._register_load_state_dict_pre_hook(merge_extra_states, with_module=True) + 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: @@ -1981,6 +2050,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) @@ -1988,11 +2072,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 eeda383a75d..0bd1c0db670 100644 --- a/megatron/core/optimizer/distrib_optimizer.py +++ b/megatron/core/optimizer/distrib_optimizer.py @@ -361,7 +361,10 @@ def _build_model_and_main_param_groups( if model_param.type() in ['torch.cuda.HalfTensor', 'torch.cuda.BFloat16Tensor']: # Generate sharded model param. - if is_float8tensor(model_param) and config.fp8_recipe != "delayed": + if ( + cls._is_distopt_quantized_param(model_param) + and config.fp8_recipe != "delayed" + ): # MXFP8Tensor and BlockwiseQTensor don't support view(-1) shard_model_param = None else: @@ -381,7 +384,7 @@ 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_float8tensor(model_param): + if 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() @@ -913,6 +916,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. @@ -2145,7 +2212,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 @@ -2167,7 +2234,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: @@ -2367,7 +2434,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) @@ -2382,7 +2449,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() @@ -2459,8 +2526,29 @@ def _copy_main_params_to_model_params(self): if self.config.use_precision_aware_optimizer_no_fp8_or_ds_fp8: return + 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, ) # Utility method for copying group params. @@ -2480,7 +2568,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 else: @@ -2592,8 +2680,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 34e9fb17a02..0f272d6adf7 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 @@ -242,10 +243,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 +290,191 @@ 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. + 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 + + # 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 + if self.linear_fc1.need_backward_dw() or self.linear_fc2.need_backward_dw(): + return False # Delayed weight gradient computation is not supported + + # Check activation + if self.activation_func != F.silu or not self.config.gated_linear_unit: + return False # Expected SwiGLU activation + + return True + + def _make_fused_ops(self) -> torch.nn.Module: + """Construct fused module for FC1, activation, and FC2.""" + + # 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_parameter = self.linear_fc1.single_grouped_parameter + fc1_weight_dtype = ( + self.linear_fc1.weight.dtype + if fc1_single_grouped_parameter + else self.linear_fc1.weight0.dtype + ) + fc2_single_grouped_parameter = self.linear_fc2.single_grouped_parameter + fc2_weight_dtype = ( + self.linear_fc2.weight.dtype + if fc2_single_grouped_parameter + else self.linear_fc2.weight0.dtype + ) + + # TODO:ksivamani: Why meta device? + 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=torch.cuda.current_device(), + dtype=fc1_weight_dtype, + accumulate_into_main_grad=self.linear_fc1.fuse_wgrad_accumulation, + single_grouped_parameter=fc1_single_grouped_parameter, + ) + + # Copy the weights from GroupedLinear module to GroupedLinear op. + if fc1_single_grouped_parameter: + setattr(op, "weight", getattr(self.linear_fc1, "weight")) + + for idx in range(self.linear_fc1.num_gemms): + if not fc1_single_grouped_parameter: + setattr(op, f"weight{idx}", getattr(self.linear_fc1, f"weight{idx}")) + if self.linear_fc1.use_bias: + setattr(op, f"bias{idx}", getattr(self.linear_fc1, f"bias{idx}")) + ops.append(op) + + # Activation and post-multiply probs + op = te.pytorch.ops.ScaledSwiGLU( + glu_interleave_size=self.config.moe_mlp_glu_interleave_size + ) + ops.append(op) + + # FC2 + has_bias = self.linear_fc2.use_bias + 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=torch.cuda.current_device(), + dtype=fc2_weight_dtype, + accumulate_into_main_grad=self.linear_fc2.fuse_wgrad_accumulation, + single_grouped_parameter=fc2_single_grouped_parameter, + ) + + # Copy the weights from GroupedLinear module to GroupedLinear op. + if fc2_single_grouped_parameter: + setattr(op, "weight", getattr(self.linear_fc2, "weight")) + + for idx in range(self.linear_fc2.num_gemms): + if not fc2_single_grouped_parameter: + setattr(op, f"weight{idx}", getattr(self.linear_fc2, f"weight{idx}")) + if self.linear_fc2.use_bias: + setattr(op, f"bias{idx}", getattr(self.linear_fc2, f"bias{idx}")) + 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. + """ - 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, - ) - 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 + + 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 self.config.moe_router_padding_for_quantization: + # Padding has already been applied in router + 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 ) - 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 - ) - 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 + # 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 def forward( self, @@ -341,17 +493,30 @@ 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 self.config.moe_router_padding_for_quantization: + # Padding has already been applied in router + 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 +541,100 @@ 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 + ) + + def remove_glu_interleaving(x: torch.Tensor) -> torch.Tensor: + """Reorder tensor so gate and linear units are contiguous. + + Should only be applied if the activation function is + an interleaved GLU. + + """ + shape = x.size() + interleave_size = self.config.moe_mlp_glu_interleave_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 + + 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 = remove_glu_interleaving(intermediate_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 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 = remove_glu_interleaving(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 + ) + + 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 +648,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 diff --git a/megatron/core/transformer/moe/moe_utils.py b/megatron/core/transformer/moe/moe_utils.py index 4c424c74b0b..6fdc6458004 100644 --- a/megatron/core/transformer/moe/moe_utils.py +++ b/megatron/core/transformer/moe/moe_utils.py @@ -1402,9 +1402,12 @@ 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) return 16 diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py index ebeb63ad51d..68e9a9dcfe3 100644 --- a/megatron/core/transformer/transformer_config.py +++ b/megatron/core/transformer/transformer_config.py @@ -449,6 +449,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 #################### @@ -807,6 +811,15 @@ class TransformerConfig(ModelParallelConfig): """Number of SMs to use for HybridEP. In pure NVL scenarios, 16 SMs can generally achieve good bandwidth.""" + 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 ################## diff --git a/tests/unit_tests/models/test_mamba_moe_model.py b/tests/unit_tests/models/test_mamba_moe_model.py index 534ed103efa..12b734612bc 100644 --- a/tests/unit_tests/models/test_mamba_moe_model.py +++ b/tests/unit_tests/models/test_mamba_moe_model.py @@ -287,6 +287,8 @@ "moe_router_force_biased": None, "inference_grouped_gemm_backend": "auto", "inference_moe_disable_fused_quant_kernels": False, + "moe_mlp_glu_interleave_size": None, + "use_transformer_engine_op_fuser": False, } # Fields to ignore entirely (ephemeral, environment-specific, very large). SKIP_FIELDS = set()