From 67ed6a46c22222fcf671c0e1c187147d10afef08 Mon Sep 17 00:00:00 2001 From: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com> Date: Tue, 18 Nov 2025 01:28:37 -0800 Subject: [PATCH] [FMDL-1328][feat] Add support for nano-v3 and super-v3 with pytorch backend Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com> --- cpp/tensorrt_llm/thop/moeOp.cpp | 22 +- .../hf/nemotron_h_weight_mapper.py | 38 ++- .../hf/qwen3_next_weight_mapper.py | 2 +- .../_torch/models/modeling_nemotron_h.py | 207 ++++++++++++++-- .../_torch/modules/fused_moe/create_moe.py | 10 +- .../modules/fused_moe/fused_moe_cutlass.py | 5 +- .../modules/fused_moe/fused_moe_vanilla.py | 46 +++- .../_torch/modules/fused_moe/interface.py | 8 +- .../_torch/modules/fused_moe/quantization.py | 47 ++-- tensorrt_llm/_torch/utils.py | 24 ++ .../test_lists/test-db/l0_h100.yml | 2 + .../modeling/test_modeling_nemotron_h.py | 224 +++++++++++------- 12 files changed, 491 insertions(+), 144 deletions(-) diff --git a/cpp/tensorrt_llm/thop/moeOp.cpp b/cpp/tensorrt_llm/thop/moeOp.cpp index 2cc1fb279779..953de1c58f19 100644 --- a/cpp/tensorrt_llm/thop/moeOp.cpp +++ b/cpp/tensorrt_llm/thop/moeOp.cpp @@ -435,7 +435,8 @@ class FusedMoeRunner : public torch::CustomClassHolder WorkspaceInfo const& workspace_info = getWorkspaceInfo(num_rows, hidden_size, inter_size, num_experts_total, static_cast(experts_per_token), base_activation_type, parallelism_config, min_latency_mode, stream); - auto const quant_params = getQuantParams(num_experts_on_rank, hidden_size, inter_size, quant_scales); + auto const quant_params + = getQuantParams(num_experts_on_rank, hidden_size, inter_size, quant_scales, base_activation_type); kernels::MoeMinLatencyParams min_latency_params{}; // TODO: support lora in the future @@ -613,7 +614,8 @@ class FusedMoeRunner : public torch::CustomClassHolder WorkspaceInfo const& workspace_info = getWorkspaceInfo(num_rows, hidden_size, inter_size, num_experts_total, static_cast(experts_per_token), base_activation_type, parallelism_config, min_latency_mode, stream); - auto const quant_params = getQuantParams(num_experts_on_rank, hidden_size, inter_size, quant_scales); + auto const quant_params + = getQuantParams(num_experts_on_rank, hidden_size, inter_size, quant_scales, base_activation_type); // TODO: support lora in the future ::tensorrt_llm::kernels::LoraParams lora_params{}; @@ -859,8 +861,10 @@ class FusedMoeRunner : public torch::CustomClassHolder } kernels::QuantParams getQuantParams(int64_t const num_experts_on_rank, int64_t const hidden_size, - int64_t const inter_size, torch::optional> const& quant_scales) const + int64_t const inter_size, torch::optional> const& quant_scales, + ActivationType base_activation_type) const { + int expand_ratio = isGatedActivation(base_activation_type) ? 2 : 1; if (isFp8Quant()) { TORCH_CHECK(quant_scales.has_value(), "Expecting quant scales for fp8 quantization"); @@ -925,12 +929,12 @@ class FusedMoeRunner : public torch::CustomClassHolder && fc1_weight_block.sizes()[1] == TmaWarpSpecializedGroupedGemmInput::alignToSfDim( inter_size, TmaWarpSpecializedGroupedGemmInput::MinNDimAlignmentMXFPX) - * 2 + * expand_ratio && fc1_weight_block.sizes()[2] * FP8_PER_INT32 * TmaWarpSpecializedGroupedGemmInput::MXFPXBlockScaleVectorSize == TmaWarpSpecializedGroupedGemmInput::alignToSfDim( hidden_size, TmaWarpSpecializedGroupedGemmInput::MinKDimAlignmentMXFPX), - "fc1 weight block size must be (num_experts_on_rank, inter_size * 2, hidden_size // 4 // " + "fc1 weight block size must be (num_experts_on_rank, inter_size * expand_ratio, hidden_size // 4 // " "block_scale_vector_size)"); TORCH_CHECK(fc1_global.sizes()[0] == num_experts_on_rank, "fc1 global size must be (num_experts_on_rank,)"); TORCH_CHECK(fc2_act_global.dim() == 0 || fc2_act_global.sizes()[0] == num_experts_on_rank, @@ -978,12 +982,12 @@ class FusedMoeRunner : public torch::CustomClassHolder && fc1_weight_block.sizes()[1] == TmaWarpSpecializedGroupedGemmInput::alignToSfDim( inter_size, TmaWarpSpecializedGroupedGemmInput::MinNDimAlignmentMXFPX) - * 2 + * expand_ratio && fc1_weight_block.sizes()[2] * FP8_PER_INT32 * TmaWarpSpecializedGroupedGemmInput::MXFPXBlockScaleVectorSize == TmaWarpSpecializedGroupedGemmInput::alignToSfDim( hidden_size, TmaWarpSpecializedGroupedGemmInput::MinKDimAlignmentMXFPX), - "fc1 weight block size must be (num_experts_on_rank, inter_size * 2, hidden_size // 4 // " + "fc1 weight block size must be (num_experts_on_rank, inter_size * expand_ratio, hidden_size // 4 // " "block_scale_vector_size)"); TORCH_CHECK(fc1_global.sizes()[0] == num_experts_on_rank, "fc1 global size must be (num_experts_on_rank,)"); TORCH_CHECK(fc2_weight_block.sizes()[0] == num_experts_on_rank @@ -1044,12 +1048,12 @@ class FusedMoeRunner : public torch::CustomClassHolder && fc1_weight_block.sizes()[1] == TmaWarpSpecializedGroupedGemmInput::alignToSfDim( inter_size, TmaWarpSpecializedGroupedGemmInput::MinKDimAlignmentNVFP4) - * 2 + * expand_ratio && fc1_weight_block.sizes()[2] * FP8_PER_INT32 * TmaWarpSpecializedGroupedGemmInput::NVFP4BlockScaleVectorSize == TmaWarpSpecializedGroupedGemmInput::alignToSfDim( hidden_size, TmaWarpSpecializedGroupedGemmInput::MinKDimAlignmentNVFP4), - "fc1 weight block size must be (num_experts_on_rank, inter_size * 2, hidden_size // 4 // " + "fc1 weight block size must be (num_experts_on_rank, inter_size * expand_ratio, hidden_size // 4 // " "block_scale_vector_size)"); TORCH_CHECK(fc1_global.sizes()[0] == num_experts_on_rank, "fc1 global size must be (num_experts_on_rank,)"); TORCH_CHECK(fc2_act_global.dim() == 0 || fc2_act_global.sizes()[0] == num_experts_on_rank, diff --git a/tensorrt_llm/_torch/models/checkpoints/hf/nemotron_h_weight_mapper.py b/tensorrt_llm/_torch/models/checkpoints/hf/nemotron_h_weight_mapper.py index 170f57d42cd7..869e11d7ea82 100644 --- a/tensorrt_llm/_torch/models/checkpoints/hf/nemotron_h_weight_mapper.py +++ b/tensorrt_llm/_torch/models/checkpoints/hf/nemotron_h_weight_mapper.py @@ -2,8 +2,8 @@ from tensorrt_llm._torch.models.checkpoints.hf.weight_mapper import \ HfWeightMapper -from tensorrt_llm._torch.models.modeling_nemotron_h import split from tensorrt_llm._torch.models.modeling_utils import register_mapper +from tensorrt_llm._torch.utils import split @register_mapper("HF", "NemotronHForCausalLM") @@ -34,7 +34,8 @@ def preprocess_weights(self, weights: dict) -> dict: if "A_log" in key: key = key.replace("A_log", "A") - if "_scale" in key: + if ("mixer.in_proj" in key + or "mixer.out_proj" in key) and "_scale" in key: new_weights[key] = weights[name] elif "A" in key: w = split(weights[name], tp_size, tp_rank) @@ -94,6 +95,39 @@ def preprocess_weights(self, weights: dict) -> dict: elif "mixer.norm.weight" in key: w = split(weights[name], tp_size, tp_rank) new_weights[key] = w + # Remap MoE expert weights. + elif "mixer.experts." in key: + if self.config.moe_backend == 'VANILLA': + new_weights[key] = weights[name] + else: + if "up_proj" in key: + w1_key = key.replace("up_proj", "w1") + w3_key = key.replace("up_proj", "w3") + # Don't need to handle with input_scale and weight_scale_2 since they are scalar for fp8 and nvfp4 models. + if "input_scale" in key or "weight_scale_2" in key: + new_weights[w3_key] = weights[name] + new_weights[w1_key] = weights[name] + elif "weight_scale" in key: + # NVFP4 case. + if weights[name].shape: + new_weights[w3_key] = weights[ + name][:weights[name].shape[0] // 2] + new_weights[w1_key] = weights[name][ + weights[name].shape[0] // 2:] + # FP8 case. + else: + new_weights[w3_key] = weights[name] + new_weights[w1_key] = weights[name] + else: + new_weights[w3_key] = weights[name][:weights[name]. + shape[0] // 2] + new_weights[w1_key] = weights[name][weights[name]. + shape[0] // 2:] + elif "down_proj" in key: + key = key.replace("down_proj", "w2") + new_weights[key] = weights[name] + else: + raise ValueError(f"Unknown MoE weight: {key}") else: new_weights[key] = weights[name] diff --git a/tensorrt_llm/_torch/models/checkpoints/hf/qwen3_next_weight_mapper.py b/tensorrt_llm/_torch/models/checkpoints/hf/qwen3_next_weight_mapper.py index e8df2da4156b..e2d3203a13d4 100644 --- a/tensorrt_llm/_torch/models/checkpoints/hf/qwen3_next_weight_mapper.py +++ b/tensorrt_llm/_torch/models/checkpoints/hf/qwen3_next_weight_mapper.py @@ -6,8 +6,8 @@ from tensorrt_llm._torch.model_config import ModelConfig from tensorrt_llm._torch.models.checkpoints.hf.qwen2_moe_weight_mapper import \ Qwen2MoeHfWeightMapper -from tensorrt_llm._torch.models.modeling_nemotron_h import split from tensorrt_llm._torch.models.modeling_utils import register_mapper +from tensorrt_llm._torch.utils import split from tensorrt_llm.models.modeling_utils import DecoderModelForCausalLM diff --git a/tensorrt_llm/_torch/models/modeling_nemotron_h.py b/tensorrt_llm/_torch/models/modeling_nemotron_h.py index 55fe5dc0d5fd..36e6d1b49bb7 100644 --- a/tensorrt_llm/_torch/models/modeling_nemotron_h.py +++ b/tensorrt_llm/_torch/models/modeling_nemotron_h.py @@ -14,44 +14,33 @@ # limitations under the License. import re -from typing import Optional +from typing import Dict, Optional import torch from torch import nn -from torch.nn import functional as F from transformers import AutoConfig, PretrainedConfig from tensorrt_llm._torch.models.checkpoints.base_weight_mapper import \ BaseWeightMapper from tensorrt_llm._torch.modules.mamba.mamba2_metadata import Mamba2Metadata +from tensorrt_llm._torch.utils import ActivationType, relu2 from ..attention_backend import AttentionMetadata from ..model_config import ModelConfig from ..modules.attention import Attention from ..modules.decoder_layer import DecoderLayer from ..modules.embedding import Embedding +from ..modules.fused_moe import MoEWeightLoadingMode, create_moe +from ..modules.linear import Linear from ..modules.mamba.mamba2_mixer import Mamba2Mixer from ..modules.mlp import MLP +from ..modules.multi_stream_utils import maybe_execute_in_parallel from ..modules.rms_norm import RMSNorm +from ..utils import AuxStreamType, EventType from .modeling_utils import (DecoderModel, DecoderModelForCausalLM, register_auto_model) -def split(x: torch.Tensor, - tp_size: int, - idx: int, - dim: int = 0) -> torch.Tensor: - assert x.shape[dim] % tp_size == 0 - split_size = x.shape[dim] // tp_size - if tp_size == 1: - return x - return torch.split(x, split_size, dim=dim)[idx] - - -def relu2(x: torch.Tensor) -> torch.Tensor: - return torch.square(F.relu(x)) - - class NemotronHConfig(PretrainedConfig): model_type = "nemotron_h" @@ -120,6 +109,158 @@ def forward( attn_metadata=attn_metadata) +# Ref code: https://huggingface.co/nvidia/Nemotron-Nano-3-30B-A3.5B-dev-1024/blob/main/modeling_nemotron_h.py#L818 +class NemotronHMOE(nn.Module): + + def __init__( + self, + model_config: ModelConfig[PretrainedConfig], + layer_idx: int, + aux_stream_dict: Dict[AuxStreamType, torch.cuda.Stream], + ): + super().__init__() + + # Import here to avoid circular dependency. + from .modeling_deepseekv3 import DeepseekV3Gate + + self.activation_type = ActivationType.Relu2 + self.reduce_results = True + + config = model_config.pretrained_config + self.hidden_dim = config.hidden_size + self.ffn_dim = config.intermediate_size + self.layer_idx = layer_idx + self.moe_intermediate_size = config.moe_intermediate_size[0] \ + if isinstance(config.moe_intermediate_size, list) else config.moe_intermediate_size + self.use_latent_moe: bool = getattr(config, "moe_latent_size", + None) is not None + self.moe_hidden_size: int = config.moe_latent_size if self.use_latent_moe else config.hidden_size + self.mlp_bias = config.mlp_bias if hasattr(config, + 'mlp_bias') else False + self.moe_n_group = config.n_group + self.num_experts = config.n_routed_experts + self.hidden_size = config.hidden_size + self.num_shared_experts = config.n_shared_experts + self.top_k = config.num_experts_per_tok + self.enable_attention_dp = model_config.mapping.enable_attention_dp + self.routed_scaling_factor = config.routed_scaling_factor + + # Setup shared expert MLP. + if config.n_shared_experts is None or config.n_shared_experts == 0: + self.shared_experts = None + else: + shared_expert_intermediate_size = ( + config.moe_shared_expert_intermediate_size * + config.n_shared_experts) + self.shared_experts = MLP( + hidden_size=config.hidden_size, + intermediate_size=shared_expert_intermediate_size, + bias=self.mlp_bias, + activation=relu2, + dtype=config.torch_dtype, + config=model_config, + layer_idx=self.layer_idx, + ) + # Setup MoE gate. + self.gate = DeepseekV3Gate( + self.hidden_size, + self.num_experts, + top_k=self.top_k, + n_group=self.moe_n_group, + topk_group=config.topk_group, + routed_scaling_factor=self.routed_scaling_factor, + dtype=config.torch_dtype, + fuse_routing_kernel=True, + apply_routing=False, + moe_backend=model_config.moe_backend) + + # Setup MoE experts. + self.experts = create_moe( + routing_method=self.gate.routing_method, + num_experts=self.num_experts, + hidden_size=self.moe_hidden_size, + intermediate_size=self.moe_intermediate_size, + aux_stream_dict=aux_stream_dict, + dtype=config.torch_dtype, + reduce_results=self.reduce_results, + model_config=model_config, + layer_idx=self.layer_idx, + weight_loading_mode=MoEWeightLoadingMode.VANILLA, + bias=self.mlp_bias, + activation_type=self.activation_type, + ) + + # Setup latent projection layers. + if self.use_latent_moe: + self.fc1_latent_proj = Linear( + in_features=self.hidden_size, + out_features=self.moe_hidden_size, + bias=self.mlp_bias, + dtype=config.torch_dtype, + ) + self.fc2_latent_proj = Linear( + in_features=self.moe_hidden_size, + out_features=self.hidden_size, + bias=self.mlp_bias, + dtype=config.torch_dtype, + ) + else: + self.fc1_latent_proj = None + self.fc2_latent_proj = None + + self.aux_stream_shared = aux_stream_dict[AuxStreamType.MoeShared] + self.event_dict = { + key: torch.cuda.Event() + for key in [EventType.Main, EventType.MoeShared] + } + + def forward( + self, + hidden_states: torch.Tensor, + attn_metadata: AttentionMetadata, + **kwargs, + ) -> torch.Tensor: + assert hidden_states.shape[-1] == self.hidden_dim + orig_shape = hidden_states.shape + hidden_states = hidden_states.view(-1, self.hidden_dim) + + def _compute_shared_output(): + if self.shared_experts is not None: + shared_expert_output = self.shared_experts(hidden_states) + else: + shared_expert_output = 0 + return shared_expert_output + + def _compute_routed_output(): + router_logits = self.gate(hidden_states) + + routed_hidden_states = hidden_states + if self.use_latent_moe: + routed_hidden_states = self.fc1_latent_proj( + routed_hidden_states) + + all_rank_num_tokens = attn_metadata.all_rank_num_tokens + final_hidden_states = self.experts( + routed_hidden_states, + router_logits, + all_rank_num_tokens=all_rank_num_tokens, + use_dp_padding=False) + + if self.use_latent_moe: + final_hidden_states = self.fc2_latent_proj(final_hidden_states) + + return final_hidden_states + + routed_output, shared_output = maybe_execute_in_parallel( + _compute_routed_output, _compute_shared_output, + self.event_dict[EventType.Main], + self.event_dict[EventType.MoeShared], self.aux_stream_shared) + + final_hidden_states = shared_output + routed_output + + return final_hidden_states.view(orig_shape) + + class NemotronHLayer(DecoderLayer): def __init__( @@ -130,6 +271,7 @@ def __init__( # - -> MLPLayer # * -> TransformerLayer layer_type: str, + aux_stream_dict: Dict[AuxStreamType, torch.cuda.Stream], ): super().__init__() @@ -160,6 +302,10 @@ def __init__( self.mixer = MLPLayer(model_config, layer_idx) elif layer_type == "*": self.mixer = TransformerLayer(model_config, layer_idx) + elif layer_type == "E": + self.mixer = NemotronHMOE(model_config, + layer_idx=layer_idx, + aux_stream_dict=aux_stream_dict) else: raise ValueError(f"{layer_type} is not supported") @@ -186,6 +332,18 @@ def __init__(self, model_config: ModelConfig[NemotronHConfig]): super().__init__(model_config) config = self.model_config.pretrained_config + aux_stream_list = [torch.cuda.Stream() for _ in range(3)] + self.aux_stream_dict = { + # TODO: add attention stream. + # AuxStreamType.Attention: aux_stream_list[0], + AuxStreamType.MoeShared: + aux_stream_list[0], + AuxStreamType.MoeChunkingOverlap: + aux_stream_list[1], + AuxStreamType.MoeBalancer: + aux_stream_list[2], + } + # calculate embeddings self.embed_tokens = Embedding( config.vocab_size, @@ -196,7 +354,11 @@ def __init__(self, model_config: ModelConfig[NemotronHConfig]): # create layers layers = [] for layer_idx, layer_type in enumerate(config.hybrid_override_pattern): - layers.append(NemotronHLayer(model_config, layer_idx, layer_type)) + layers.append( + NemotronHLayer(model_config, + layer_idx, + layer_type, + aux_stream_dict=self.aux_stream_dict)) self.layers = nn.ModuleList(layers) # final norm @@ -251,6 +413,15 @@ def __init__( self, model_config: ModelConfig[NemotronHConfig], ): + # rms_norm_eps might be named differently in the config. + if hasattr(model_config.pretrained_config, "rms_norm_eps"): + rms_epsilon = model_config.pretrained_config.rms_norm_eps + elif hasattr(model_config.pretrained_config, "layer_norm_epsilon"): + rms_epsilon = model_config.pretrained_config.layer_norm_epsilon + else: + raise ValueError("layer_norm_epsilon or rms_norm_eps is not set") + model_config.pretrained_config.rms_norm_eps = rms_epsilon + if not model_config.mapping.tp_size in [1, 2, 4, 8]: raise ValueError("TP has to be either 1, 2, 4 or 8") diff --git a/tensorrt_llm/_torch/modules/fused_moe/create_moe.py b/tensorrt_llm/_torch/modules/fused_moe/create_moe.py index f8ce1a0c3108..136f85549e8e 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/create_moe.py +++ b/tensorrt_llm/_torch/modules/fused_moe/create_moe.py @@ -7,7 +7,7 @@ from tensorrt_llm.models.modeling_utils import QuantConfig from ...model_config import ModelConfig -from ...utils import AuxStreamType +from ...utils import ActivationType, AuxStreamType from .configurable_moe import ConfigurableMoE from .fused_moe_cute_dsl import CuteDslFusedMoE from .fused_moe_cutlass import CutlassFusedMoE @@ -80,6 +80,7 @@ def create_moe_backend( swiglu_limit: Optional[torch.Tensor] = None, init_load_balancer: bool = True, without_comm: bool = False, + activation_type: ActivationType = ActivationType.Swiglu, ) -> MoE: """ Create MoE backend instance with validation. @@ -101,6 +102,7 @@ def create_moe_backend( swiglu_alpha: SwiGLU alpha parameter swiglu_beta: SwiGLU beta parameter swiglu_limit: SwiGLU limit parameter + activation_type: Activation type Returns: MoE: MoE backend instance @@ -183,6 +185,7 @@ def create_moe_backend( swiglu_beta=swiglu_beta, swiglu_limit=swiglu_limit, init_load_balancer=init_load_balancer, + activation_type=activation_type, ) elif moe_cls == WideEPMoE: return moe_cls( @@ -212,6 +215,8 @@ def create_moe_backend( model_config=model_config, weight_loading_mode=weight_loading_mode, apply_router_weight_on_input=apply_router_weight_on_input, + layer_idx=layer_idx, + activation_type=activation_type, ) elif moe_cls == CuteDslFusedMoE: return moe_cls( @@ -280,6 +285,7 @@ def create_moe( swiglu_alpha: Optional[torch.Tensor] = None, swiglu_beta: Optional[torch.Tensor] = None, swiglu_limit: Optional[torch.Tensor] = None, + activation_type: ActivationType = ActivationType.Swiglu, ) -> MoE: """ Create MoE instance with automatic parameter inference from model_config. @@ -301,6 +307,7 @@ def create_moe( swiglu_alpha: SwiGLU alpha parameter swiglu_beta: SwiGLU beta parameter swiglu_limit: SwiGLU limit parameter + activation_type: Activation type Returns: MoE: MoE instance @@ -384,4 +391,5 @@ def create_moe( swiglu_alpha=swiglu_alpha, swiglu_beta=swiglu_beta, swiglu_limit=swiglu_limit, + activation_type=activation_type, ) diff --git a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_cutlass.py b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_cutlass.py index 8b82a9ea6c9b..c06fd515ff27 100755 --- a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_cutlass.py +++ b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_cutlass.py @@ -11,7 +11,8 @@ from ...distributed import allgather from ...expert_statistic import ExpertStatistic from ...model_config import ModelConfig -from ...utils import AuxStreamType, EventType, Fp4QuantizedTensor, ceil_div +from ...utils import (ActivationType, AuxStreamType, EventType, + Fp4QuantizedTensor, ceil_div) from .interface import AlltoallMethodType, MoE from .quantization import UnquantizedFusedMoEMethod @@ -75,6 +76,7 @@ def __init__( swiglu_beta: Optional[torch.Tensor] = None, swiglu_limit: Optional[torch.Tensor] = None, init_load_balancer: bool = True, + activation_type: ActivationType = ActivationType.Swiglu, ): super().__init__( @@ -92,6 +94,7 @@ def __init__( swiglu_limit=swiglu_limit, layer_idx=layer_idx, init_load_balancer=init_load_balancer, + activation_type=activation_type, ) # Store original hidden size before any potential padding diff --git a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_vanilla.py b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_vanilla.py index 3e8cc6ec8879..c0b85e299be6 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_vanilla.py +++ b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_vanilla.py @@ -10,7 +10,9 @@ from ...distributed import allgather, reducescatter from ...model_config import ModelConfig +from ...utils import ActivationType, is_gated_activation, relu2 from ..gated_mlp import GatedMLP +from ..mlp import MLP from .interface import MoEWeightLoadingMode from .routing import BaseMoeRoutingMethod @@ -31,6 +33,8 @@ def __init__( VANILLA, apply_router_weight_on_input: bool = False, pack_weights: bool = False, + layer_idx: Optional[int] = None, + activation_type: ActivationType = ActivationType.Swiglu, ): from ...distributed import AllReduce @@ -41,6 +45,20 @@ def __init__( self.intermediate_size = intermediate_size self.weight_loading_mode = weight_loading_mode self.pack_weights = pack_weights + self.layer_idx = layer_idx + + self.activation_type = activation_type + self.is_gated_activation = is_gated_activation(activation_type) + # Limit support for VanillaMoE to non-gated activations for now. + if not self.is_gated_activation: + if pack_weights: + raise ValueError( + "pack_weights must be False for non-gated activations. Otherwise please update `create_weights`." + ) + if self.activation_type != ActivationType.Relu2: + raise ValueError( + f"Unsupported activation type: {self.activation_type} for non-gated activations. Only Relu2 is supported." + ) self.dtype = dtype self.reduce_results = reduce_results @@ -107,14 +125,26 @@ def create_experts(self, module_list: nn.ModuleList = None): ) for expert_idx in range(self.num_experts): if self.expert_start <= expert_idx < self.expert_end: - module_list[expert_idx] = GatedMLP( - hidden_size=self.hidden_size, - intermediate_size=self.intermediate_size, - bias=False, - dtype=self.dtype, - config=model_config, - reduce_output=False, - ) + if self.activation_type == ActivationType.Relu2: + module_list[expert_idx] = MLP( + hidden_size=self.hidden_size, + intermediate_size=self.intermediate_size, + bias=False, + activation=relu2, + dtype=self.dtype, + config=model_config, + layer_idx=self.layer_idx, + ) + else: + module_list[expert_idx] = GatedMLP( + hidden_size=self.hidden_size, + intermediate_size=self.intermediate_size, + bias=False, + dtype=self.dtype, + config=model_config, + reduce_output=False, + layer_idx=self.layer_idx, + ) else: # use identity as placeholder for unused experts module_list[expert_idx] = nn.Identity() diff --git a/tensorrt_llm/_torch/modules/fused_moe/interface.py b/tensorrt_llm/_torch/modules/fused_moe/interface.py index cd064d812185..ae837672f3a6 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/interface.py +++ b/tensorrt_llm/_torch/modules/fused_moe/interface.py @@ -9,7 +9,8 @@ from ...distributed.ops import reducescatter from ...model_config import ModelConfig from ...utils import (ActivationType, AuxStreamType, Fp4QuantizedTensor, - get_model_extra_attrs, is_torch_compiling) + get_model_extra_attrs, is_gated_activation, + is_torch_compiling) from .routing import BaseMoeRoutingMethod @@ -166,6 +167,11 @@ def __init__( self.layer_idx = layer_idx self.layer_idx_str = str(layer_idx) if layer_idx is not None else None self.activation_type = int(activation_type) + # Note: + # - for gated activations, there should be with gate and up projections, so the intermediate size should be expanded by 2. + # - for non-gated activations, there is only one up projection (no gate projection), so the intermediate size should not be expanded. + self.is_gated_activation = is_gated_activation(activation_type) + self.intermediate_size_expand_ratio = 2 if self.is_gated_activation else 1 self._register_layer(model_config) diff --git a/tensorrt_llm/_torch/modules/fused_moe/quantization.py b/tensorrt_llm/_torch/modules/fused_moe/quantization.py index 388b1a51d6f7..c3ff7d037d5e 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/quantization.py +++ b/tensorrt_llm/_torch/modules/fused_moe/quantization.py @@ -220,7 +220,8 @@ def create_weights( if module.bias: if w3_w1_bias_shape is None: w3_w1_bias_shape = (module.expert_size_per_partition, - module.intermediate_size_per_partition * 2) + module.intermediate_size_per_partition * + module.intermediate_size_expand_ratio) if w2_bias_shape is None: w2_bias_shape = (module.expert_size_per_partition, module.hidden_size) @@ -515,7 +516,8 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase): def create_weights(self, module: torch.nn.Module): weight_dtype = module.dtype w3_w1_weight_shape = (module.expert_size_per_partition, - module.intermediate_size_per_partition * 2, + module.intermediate_size_per_partition * + module.intermediate_size_expand_ratio, module.hidden_size) w2_weight_shape = ( module.expert_size_per_partition, @@ -580,13 +582,11 @@ def requantize_expert_w3_w1_weight_fp8_qdq(module: torch.nn.Module, w3_weight_scale = w3_weight_scale[...].reshape([]) max_w3_w1_weight_scale = max(w1_weight_scale, w3_weight_scale) + split_length = module.intermediate_size_per_partition * module.intermediate_size_expand_ratio // 2 w3_weight = dst_w3_w1_weight.narrow( - dim=0, start=0, - length=module.intermediate_size_per_partition).to(dtype=module.dtype) + dim=0, start=0, length=split_length).to(dtype=module.dtype) w1_weight = dst_w3_w1_weight.narrow( - dim=0, - start=module.intermediate_size_per_partition, - length=module.intermediate_size_per_partition).to(dtype=module.dtype) + dim=0, start=split_length, length=split_length).to(dtype=module.dtype) dequant_w3_weight = w3_weight * w3_weight_scale dequant_w1_weight = w1_weight * w1_weight_scale requant_w3_weight = (dequant_w3_weight / max_w3_w1_weight_scale).to( @@ -594,13 +594,10 @@ def requantize_expert_w3_w1_weight_fp8_qdq(module: torch.nn.Module, requant_w1_weight = (dequant_w1_weight / max_w3_w1_weight_scale).to( torch.float8_e4m3fn) - dst_w3_w1_weight.narrow( - dim=0, start=0, - length=module.intermediate_size_per_partition).copy_(requant_w3_weight) - dst_w3_w1_weight.narrow( - dim=0, - start=module.intermediate_size_per_partition, - length=module.intermediate_size_per_partition).copy_(requant_w1_weight) + dst_w3_w1_weight.narrow(dim=0, start=0, + length=split_length).copy_(requant_w3_weight) + dst_w3_w1_weight.narrow(dim=0, start=split_length, + length=split_length).copy_(requant_w1_weight) class FP8QDQFusedMoEMethod(FusedMoEMethodBase): @@ -609,7 +606,8 @@ def create_weights(self, module: torch.nn.Module): weight_dtype = torch.float8_e4m3fn w3_w1_weight_shape = (module.expert_size_per_partition, - module.intermediate_size_per_partition * 2, + module.intermediate_size_per_partition * + module.intermediate_size_expand_ratio, module.hidden_size) w2_weight_shape = ( module.expert_size_per_partition, @@ -1669,7 +1667,8 @@ def create_weights(self, module.scaling_vector_size = scaling_vector_size # Divide by 16 because we use int64 to pack 16 fp4 values w3_w1_weight_shape = (module.expert_size_per_partition, - module.intermediate_size_per_partition * 2, + module.intermediate_size_per_partition * + module.intermediate_size_expand_ratio, module.hidden_size // weight_vec_size) w2_weight_shape = (module.expert_size_per_partition, module.hidden_size, module.intermediate_size_per_partition // @@ -1679,7 +1678,8 @@ def create_weights(self, # column parallel w3_w1_weight_scale = nn.Parameter( torch.ones(module.expert_size_per_partition, - module.intermediate_size_per_partition * 2, + module.intermediate_size_per_partition * + module.intermediate_size_expand_ratio, module.hidden_size // module.scaling_vector_size // block_scales_vec_size, dtype=block_scales_dtype), @@ -1950,16 +1950,17 @@ def load_expert_w3_w1_weight_scale_nvfp4( device=device) # Keep weights in device buffer # w3 - dst_w3_weight_scale = dst_w3_w1_weight_scale.narrow( - dim=0, start=0, length=module.intermediate_size_per_partition) + split_length = module.intermediate_size_per_partition * module.intermediate_size_expand_ratio // 2 + dst_w3_weight_scale = dst_w3_w1_weight_scale.narrow(dim=0, + start=0, + length=split_length) dst_w3_weight_scale.copy_( w3_weight_scale.view(dst_w3_weight_scale.dtype)) # w1 - dst_w1_weight_scale = dst_w3_w1_weight_scale.narrow( - dim=0, - start=module.intermediate_size_per_partition, - length=module.intermediate_size_per_partition) + dst_w1_weight_scale = dst_w3_w1_weight_scale.narrow(dim=0, + start=split_length, + length=split_length) dst_w1_weight_scale.copy_( w1_weight_scale.view(dst_w1_weight_scale.dtype)) diff --git a/tensorrt_llm/_torch/utils.py b/tensorrt_llm/_torch/utils.py index 9301578b4a2d..5f77a4c7a1b2 100644 --- a/tensorrt_llm/_torch/utils.py +++ b/tensorrt_llm/_torch/utils.py @@ -6,6 +6,7 @@ from typing import Dict, List import torch +from torch.nn import functional as F from tensorrt_llm._utils import TensorWrapper, convert_to_torch_tensor from tensorrt_llm.mapping import Mapping @@ -46,6 +47,14 @@ class ActivationType(IntEnum): Relu2 = 8 +# IMPORTANT: when adding a new activation type, please update this function. +# And make sure it aligned with cpp/tensorrt_llm/kernels/cutlass_kernels/include/moe_gemm_kernels.h::isGatedActivation function. +def is_gated_activation(activation_type: ActivationType) -> bool: + return activation_type in [ + ActivationType.Swiglu, ActivationType.SwigluBias, ActivationType.Geglu + ] + + def set_torch_compiling(enable: bool): global is_torch_compiling_flag is_torch_compiling_flag = enable @@ -370,3 +379,18 @@ def wrapper(*args, **kwargs): return wrapper return decorator(func) if func else decorator + + +def split(x: torch.Tensor, + tp_size: int, + idx: int, + dim: int = 0) -> torch.Tensor: + assert x.shape[dim] % tp_size == 0 + split_size = x.shape[dim] // tp_size + if tp_size == 1: + return x + return torch.split(x, split_size, dim=dim)[idx] + + +def relu2(x: torch.Tensor) -> torch.Tensor: + return torch.square(F.relu(x)) diff --git a/tests/integration/test_lists/test-db/l0_h100.yml b/tests/integration/test_lists/test-db/l0_h100.yml index 070e7e77d302..7ef07bc51c05 100644 --- a/tests/integration/test_lists/test-db/l0_h100.yml +++ b/tests/integration/test_lists/test-db/l0_h100.yml @@ -33,6 +33,8 @@ l0_h100: - unittest/_torch/modeling -k "modeling_nemotron" - unittest/_torch/modeling -k "modeling_gemma3" - unittest/_torch/modeling -k "modeling_gpt_oss" + - unittest/_torch/modeling/test_modeling_nemotron_h.py::test_nemotron_h_correctness[Nemotron-Nano-3-30B-A3.5B-dev-1024-mamba_ssm_cache_dtype:float32] + - unittest/_torch/modeling/test_modeling_nemotron_h.py::test_nemotron_h_correctness[Nemotron-Nano-3-30B-A3.5B-dev-1024-mamba_ssm_cache_dtype:None] - unittest/disaggregated/test_disagg_utils.py - unittest/disaggregated/test_router.py - unittest/disaggregated/test_remoteDictionary.py diff --git a/tests/unittest/_torch/modeling/test_modeling_nemotron_h.py b/tests/unittest/_torch/modeling/test_modeling_nemotron_h.py index 21ab127a3669..c8d88cccee76 100644 --- a/tests/unittest/_torch/modeling/test_modeling_nemotron_h.py +++ b/tests/unittest/_torch/modeling/test_modeling_nemotron_h.py @@ -1,7 +1,7 @@ import pytest import torch from utils.llm_data import llm_models_root -from utils.util import skip_gpu_memory_less_than +from utils.util import similar, skip_gpu_memory_less_than from tensorrt_llm import LLM from tensorrt_llm.llmapi import KvCacheConfig @@ -31,14 +31,15 @@ def extract_decode_logprobs(result: RequestOutput, return get_logprobs(token_ids, logits) -def create_nemotron_h_llm(use_cuda_graph, +def create_nemotron_h_llm(model_folder, + use_cuda_graph, disable_overlap_scheduler, max_batch_size, mamba_ssm_cache_dtype=None, enable_chunked_prefill=False, - max_num_tokens=None): + max_num_tokens=8192): """Create LLM with specific overlap scheduler setting""" - model_dir = f"{llm_models_root(check=True)}/Nemotron-H-8B-Base-8K" + model_dir = f"{llm_models_root(check=True)}/{model_folder}" return LLM( model=model_dir, tensor_parallel_size=1, @@ -57,12 +58,15 @@ def create_nemotron_h_llm(use_cuda_graph, ) -@skip_gpu_memory_less_than( - (2 * 8 + 1) * 2**30) # 8B, bf16, plus 1 GB for good measure @pytest.mark.parametrize("mamba_ssm_cache_dtype", [None, "float32"], ids=lambda n: f"mamba_ssm_cache_dtype:{n}") -def test_nemotron_h_correctness(mamba_ssm_cache_dtype): - # This test is close to memory limit on A30 (with 24GB), so empty cache first +@pytest.mark.parametrize("model_folder", [ + pytest.param("Nemotron-H-8B-Base-8K", + marks=skip_gpu_memory_less_than((2 * 8 + 1) * 2**30)), + pytest.param("Nemotron-Nano-3-30B-A3.5B-dev-1024", + marks=skip_gpu_memory_less_than((2 * 30 + 1) * 2**30)), +]) +def test_nemotron_h_correctness(mamba_ssm_cache_dtype, model_folder): torch.cuda.empty_cache() text_prompts = [ @@ -72,78 +76,128 @@ def test_nemotron_h_correctness(mamba_ssm_cache_dtype): num_prompts = len(text_prompts) nemotron_h = create_nemotron_h_llm( + model_folder=model_folder, use_cuda_graph=False, disable_overlap_scheduler=False, max_batch_size=num_prompts, mamba_ssm_cache_dtype=mamba_ssm_cache_dtype) - expected_completions = [ - " bright, with endless possibilities for innovation and growth", - " the head of state and head of government of", - ] + if model_folder == "Nemotron-H-8B-Base-8K": + expected_completions = [ + " bright, with endless possibilities for innovation and growth", + " the head of state and head of government of", + ] - # reference logprobs for first prompt from mcore for prompt minus first token - # TODO(oargov): generate a reference on-the-fly once we have confidence in the HF impl - prefill_logprobs_ref_mcore = torch.tensor([ - -7.415980815887451, -0.36192911863327026, -2.8658294677734375, - -2.316344738006592 - ]) - - # reference logprobs from initial implementation (commit 5ce1102a02bd2938c0c8334138371f081f55fcc1 on single RTX 6000) - initial_impl_atol = 0.2 - batching_atol = 0.2 - - prefill_logprobs_ref_initial_no_batching = [ - torch.tensor([ - -7.4359540939331055, - -0.37661877274513245, - -2.8925108909606934, - -2.268364906311035, - ]), - torch.tensor([ - -8.759482383728027, - -1.656238079071045, - -0.5448741912841797, - -1.7702054977416992, - -0.05832016468048096, - -1.460732102394104, - ]) - ] - prefill_logprobs_ref_initial_with_batching = [ - torch.tensor([ - -7.401950836181641, -0.38696032762527466, -2.8725428581237793, - -2.2654521465301514 - ]), - torch.tensor([ - -8.73007583618164, -1.6853574514389038, -0.5468529462814331, - -1.7846013307571411, -0.053610533475875854, -1.4385275840759277 + # reference logprobs for first prompt from mcore for prompt minus first token + # TODO(oargov): generate a reference on-the-fly once we have confidence in the HF impl + prefill_logprobs_ref_mcore = torch.tensor([ + -7.415980815887451, -0.36192911863327026, -2.8658294677734375, + -2.316344738006592 ]) - ] - decode_logprobs_ref_initial_no_batching = [ - torch.tensor([ - -2.2722280025482178, -0.5124826431274414, -0.7916123270988464, - -2.1908130645751953, -0.059298671782016754, -0.5125972032546997, - -0.3856367766857147, -0.055953752249479294, -1.1059765815734863 - ]), - torch.tensor([ - -1.329713225364685, -1.5038213729858398, -0.021283088251948357, - -0.38457369804382324, -0.3582419157028198, -0.16527847945690155, - -0.0044861179776489735, -0.059462934732437134, -0.041099339723587036 - ]) - ] - decode_logprobs_ref_initial_with_batching = [ - torch.tensor([ - -2.2877156734466553, -0.46699056029319763, -0.7909849286079407, - -2.1276988983154297, -0.062114741653203964, -0.5291495323181152, - -0.38685765862464905, -0.05595658719539642, -1.1020748615264893 - ]), - torch.tensor([ - -1.3567769527435303, -1.5647790431976318, -0.022344056516885757, - -0.38503751158714294, -0.3581986725330353, -0.18398350477218628, - -0.004726295825093985, -0.05941498652100563, -0.04291720315814018 - ]) - ] + # reference logprobs from initial implementation (commit 5ce1102a02bd2938c0c8334138371f081f55fcc1 on single RTX 6000) + initial_impl_atol = 0.2 + batching_atol = 0.2 + + prefill_logprobs_ref_initial_no_batching = [ + torch.tensor([ + -7.4359540939331055, + -0.37661877274513245, + -2.8925108909606934, + -2.268364906311035, + ]), + torch.tensor([ + -8.759482383728027, + -1.656238079071045, + -0.5448741912841797, + -1.7702054977416992, + -0.05832016468048096, + -1.460732102394104, + ]) + ] + prefill_logprobs_ref_initial_with_batching = [ + torch.tensor([ + -7.401950836181641, -0.38696032762527466, -2.8725428581237793, + -2.2654521465301514 + ]), + torch.tensor([ + -8.73007583618164, -1.6853574514389038, -0.5468529462814331, + -1.7846013307571411, -0.053610533475875854, -1.4385275840759277 + ]) + ] + + decode_logprobs_ref_initial_no_batching = [ + torch.tensor([ + -2.2722280025482178, -0.5124826431274414, -0.7916123270988464, + -2.1908130645751953, -0.059298671782016754, -0.5125972032546997, + -0.3856367766857147, -0.055953752249479294, -1.1059765815734863 + ]), + torch.tensor([ + -1.329713225364685, -1.5038213729858398, -0.021283088251948357, + -0.38457369804382324, -0.3582419157028198, -0.16527847945690155, + -0.0044861179776489735, -0.059462934732437134, + -0.041099339723587036 + ]) + ] + decode_logprobs_ref_initial_with_batching = [ + torch.tensor([ + -2.2877156734466553, -0.46699056029319763, -0.7909849286079407, + -2.1276988983154297, -0.062114741653203964, -0.5291495323181152, + -0.38685765862464905, -0.05595658719539642, -1.1020748615264893 + ]), + torch.tensor([ + -1.3567769527435303, -1.5647790431976318, -0.022344056516885757, + -0.38503751158714294, -0.3581986725330353, -0.18398350477218628, + -0.004726295825093985, -0.05941498652100563, + -0.04291720315814018 + ]) + ] + elif model_folder == "Nemotron-Nano-3-30B-A3.5B-dev-1024": + + expected_completions = [ + " bright, with endless possibilities for innovation and growth", + " the head of state and head of government of", + ] + + # Copied from prefill_logprobs_no_batching[0] directly. + prefill_logprobs_ref_mcore = torch.tensor( + [-8.5145, -0.8952, -2.3531, -1.6690]) + + # reference logprobs from initial implementation (commit e4e42e0ec30227866ce30fc9c93d5e49352bb79c on single H200). + initial_impl_atol = 2.0 + batching_atol = 2.0 + + prefill_logprobs_ref_initial_no_batching = [ + torch.tensor([-8.5145, -0.8952, -2.3531, -1.6690]), + torch.tensor([-9.9306, -1.4935, -0.4787, -1.4945, -0.0195, -1.5253]) + ] + prefill_logprobs_ref_initial_with_batching = [ + torch.tensor([-8.5221, -0.8114, -2.4334, -1.6909]), + torch.tensor([-9.9466, -1.5095, -0.5282, -1.4701, -0.0185, -1.4108]) + ] + + decode_logprobs_ref_initial_no_batching = [ + torch.tensor([ + -9.2718e-01, -9.7786e-01, -7.5823e-01, -3.3243e-01, -8.7978e-01, + -3.2046e-02, -9.5047e-01, -9.2678e-01, -2.5973e-04 + ]), + torch.tensor([ + -1.6836, -0.8289, -0.0063, -0.5166, -0.1798, -0.6075, -1.0987, + -0.9075, -0.0025 + ]) + ] + decode_logprobs_ref_initial_with_batching = [ + torch.tensor([ + -9.0849e-01, -9.3238e-01, -8.2788e-01, -3.5542e-01, -9.0881e-01, + -3.4794e-02, -9.4975e-01, -9.2631e-01, -2.4041e-04 + ]), + torch.tensor([ + -1.6331, -0.7666, -0.0063, -0.5110, -0.1617, -0.6578, -1.1073, + -1.1447, -0.0024 + ]) + ] + else: + raise ValueError(f"Invalid model folder: {model_folder}") try: sampling_params = SamplingParams(max_tokens=9, @@ -201,9 +255,14 @@ def test_nemotron_h_correctness(mamba_ssm_cache_dtype): atol=initial_impl_atol, rtol=0.0) - # compare expected completion - assert completions_batching[i] == expected_completions[i] - assert completions_no_batching[i] == expected_completions[i] + if model_folder == "Nemotron-H-8B-Base-8K": + # compare expected completion + assert completions_batching[i] == expected_completions[i] + assert completions_no_batching[i] == expected_completions[i] + else: + assert similar(completions_batching[i], + completions_no_batching[i], + threshold=0.5) # compare decode logprobs with initial implementation torch.testing.assert_close( @@ -267,7 +326,8 @@ def test_nemotron_h_cuda_graph_overlap_scheduler(): return_generation_logits=True) # Test without cg and overlap scheduler disabled - with create_nemotron_h_llm(use_cuda_graph=False, + with create_nemotron_h_llm(model_folder="Nemotron-H-8B-Base-8K", + use_cuda_graph=False, disable_overlap_scheduler=True, max_batch_size=16) as llm: outputs_no_cg_no_overlap = llm.generate(prompts, @@ -275,14 +335,16 @@ def test_nemotron_h_cuda_graph_overlap_scheduler(): use_tqdm=True) # Test with cg and overlap scheduler disabled - with create_nemotron_h_llm(use_cuda_graph=True, + with create_nemotron_h_llm(model_folder="Nemotron-H-8B-Base-8K", + use_cuda_graph=True, disable_overlap_scheduler=True, max_batch_size=16) as llm: outputs_with_cg_no_overlap = llm.generate( prompts, sampling_params=sampling_config, use_tqdm=True) # Test with cg and overlap scheduler enabled - with create_nemotron_h_llm(use_cuda_graph=True, + with create_nemotron_h_llm(model_folder="Nemotron-H-8B-Base-8K", + use_cuda_graph=True, disable_overlap_scheduler=False, max_batch_size=16) as llm: outputs_with_cg_with_overlap = llm.generate( @@ -362,14 +424,16 @@ def test_nemotron_h_chunked_prefill(): return_context_logits=True, return_generation_logits=True) - with create_nemotron_h_llm(use_cuda_graph=False, + with create_nemotron_h_llm(model_folder="Nemotron-H-8B-Base-8K", + use_cuda_graph=False, disable_overlap_scheduler=True, max_batch_size=16) as llm: outputs = llm.generate(prompts, sampling_params=sampling_config, use_tqdm=True) - with create_nemotron_h_llm(use_cuda_graph=False, + with create_nemotron_h_llm(model_folder="Nemotron-H-8B-Base-8K", + use_cuda_graph=False, disable_overlap_scheduler=True, max_batch_size=16, enable_chunked_prefill=True,