diff --git a/cpp/tensorrt_llm/thop/moeOp.cpp b/cpp/tensorrt_llm/thop/moeOp.cpp index 81c975241a2e..6e9503ce9ece 100644 --- a/cpp/tensorrt_llm/thop/moeOp.cpp +++ b/cpp/tensorrt_llm/thop/moeOp.cpp @@ -475,8 +475,19 @@ class FusedMoeRunner : public torch::CustomClassHolder { // Note: The weight shape for INT8 weight only quantization is different, e.g., fc2_expert_weights: // [num_experts, inter_size, hidden_size] - TORCH_CHECK(fc1_expert_weights.sizes()[2] == fc2_expert_weights.sizes()[1] * mInnerDimMultiplier * 2, - "fc1_expert_weights inter size must be 2 times fc2_expert_weights inter size."); + // Mirror the non-woq else-branch below: gated activations (Swiglu/Geglu) require fc1's + // intermediate dim to be 2x fc2's (one half each for gate and up), while non-gated + // activations (Relu2/Identity/ReLU/SiLU/Gelu, e.g. Nemotron-H) require them to be equal. + if (isGatedActivation(base_activation_type)) + { + TORCH_CHECK(fc1_expert_weights.sizes()[2] == fc2_expert_weights.sizes()[1] * mInnerDimMultiplier * 2, + "fc1_expert_weights inter size must be 2 times fc2_expert_weights inter size."); + } + else + { + TORCH_CHECK(fc1_expert_weights.sizes()[2] == fc2_expert_weights.sizes()[1] * mInnerDimMultiplier, + "fc1_expert_weights inter size must be equal to fc2_expert_weights inter size."); + } } else { @@ -749,8 +760,6 @@ class FusedMoeRunner : public torch::CustomClassHolder } TORCH_CHECK(fc1_expert_weights.sizes()[0] == fc2_expert_weights.sizes()[0], "fc1_expert_weights and fc2_expert_weights must have the same number of experts."); - TORCH_CHECK(fc1_expert_weights.sizes()[1] == fc2_expert_weights.sizes()[2] * mInnerDimMultiplier * 2, - "fc1_expert_weights inter size must be 2 times fc2_expert_weights inter size."); TORCH_CHECK(!input_sf.has_value() || isWMxfp4AMxfp8Quant() || isNvfp4Quant(), "Block-scaling factors provided for non block-scaling quantization"); @@ -761,6 +770,13 @@ class FusedMoeRunner : public torch::CustomClassHolder int64_t unpadded_hidden_size_val = unpadded_hidden_size.has_value() ? unpadded_hidden_size.value() : hidden_size; int64_t inter_size = fc2_expert_weights.sizes()[2] * mInnerDimMultiplier; + if (mUseINT8WoqPerChannel) + { + // Note: The weight shape for INT8 weight only quantization is different, e.g., fc2_expert_weights: + // [num_experts, inter_size, hidden_size] + hidden_size = fc2_expert_weights.sizes()[2] * mInnerDimMultiplier; + inter_size = fc2_expert_weights.sizes()[1]; + } int const num_experts_on_rank = fc2_expert_weights.sizes()[0]; auto const num_experts_total = static_cast(num_experts_on_rank * ep_size); auto parallelism_config @@ -794,6 +810,38 @@ class FusedMoeRunner : public torch::CustomClassHolder reinterpret_cast(swiglu_beta.has_value() ? swiglu_beta.value().const_data_ptr() : nullptr), reinterpret_cast(swiglu_limit.has_value() ? swiglu_limit.value().const_data_ptr() : nullptr)); + // Validate the fc1/fc2 inter-size relationship now that the activation type (gated vs + // non-gated) is finalized. INT8-woq uses a transposed weight layout, so its fc1/fc2 dim + // ordering differs from the non-woq path; both mirror the gated/non-gated split used in + // runMoe(). Gated activations (Swiglu/Geglu) require fc1's intermediate dim to be 2x fc2's; + // non-gated (Relu2/Identity/ReLU/SiLU/Gelu, e.g. Nemotron-H) require them to be equal. + if (mUseINT8WoqPerChannel) + { + if (isGatedActivation(base_activation_type)) + { + TORCH_CHECK(fc1_expert_weights.sizes()[2] == fc2_expert_weights.sizes()[1] * mInnerDimMultiplier * 2, + "fc1_expert_weights inter size must be 2 times fc2_expert_weights inter size."); + } + else + { + TORCH_CHECK(fc1_expert_weights.sizes()[2] == fc2_expert_weights.sizes()[1] * mInnerDimMultiplier, + "fc1_expert_weights inter size must be equal to fc2_expert_weights inter size."); + } + } + else + { + if (isGatedActivation(base_activation_type)) + { + TORCH_CHECK(fc1_expert_weights.sizes()[1] == fc2_expert_weights.sizes()[2] * mInnerDimMultiplier * 2, + "fc1_expert_weights inter size must be 2 times fc2_expert_weights inter size."); + } + else + { + TORCH_CHECK(fc1_expert_weights.sizes()[1] == fc2_expert_weights.sizes()[2] * mInnerDimMultiplier, + "fc1_expert_weights inter size must be equal to fc2_expert_weights inter size."); + } + } + setRunnerProfiles(profile_ids); auto stream = at::cuda::getCurrentCUDAStream(input.get_device()); diff --git a/tensorrt_llm/_torch/modules/fused_moe/quantization.py b/tensorrt_llm/_torch/modules/fused_moe/quantization.py index 2364aaf99eed..076454732a8d 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/quantization.py +++ b/tensorrt_llm/_torch/modules/fused_moe/quantization.py @@ -1333,14 +1333,14 @@ def create_weights(self, module: torch.nn.Module): # since the quantized weights have their own layout w3_w1_weight_shape = (module.expert_size_per_partition, module.hidden_size, - module.intermediate_size_per_partition * 2) + module.expand_intermediate_size_per_partition) w2_weight_shape = (module.expert_size_per_partition, module.intermediate_size_per_partition, module.hidden_size) fc31_weight_scale = nn.Parameter(torch.empty( module.expert_size_per_partition, - module.intermediate_size_per_partition * 2, + module.expand_intermediate_size_per_partition, dtype=module.dtype), requires_grad=False) module.register_parameter("fc31_weight_scale", fc31_weight_scale) @@ -1375,10 +1375,19 @@ def load_expert_w3_w1_weight(self, module: torch.nn.Module, w1_weight_shard = load_weight_shard(w1_weight, module.tp_size, module.tp_rank, TensorParallelMode.COLUMN) - w3_weight_shard = load_weight_shard(w3_weight, module.tp_size, - module.tp_rank, - TensorParallelMode.COLUMN) - w31_weight_shard = torch.cat([w3_weight_shard, w1_weight_shard], dim=0) + + # w3_weight (gate_proj) is empty for non-gated MoE (e.g. Nemotron-H squared-ReLU). + # Only concatenate the gate projection when present; otherwise the single + # up-projection fills the (non-doubled) intermediate buffer. The unquantized + # fused-MoE path handles non-gated experts the same way. + if w3_weight is not None and w3_weight.numel() > 0: + w3_weight_shard = load_weight_shard(w3_weight, module.tp_size, + module.tp_rank, + TensorParallelMode.COLUMN) + w31_weight_shard = torch.cat([w3_weight_shard, w1_weight_shard], + dim=0) + else: + w31_weight_shard = w1_weight_shard weight_dtype = torch.int8 @@ -1419,25 +1428,33 @@ def load_expert_w2_weight(self, module: torch.nn.Module, non_blocking=True) def load_quant_scales(self, module: torch.nn.Module, weights: Dict): - # fc31 scales - all_w3_scales = [ - load_weight_shard(weights[f"{expert_id}.w3.weight_scale"], - module.tp_size, module.tp_rank, - TensorParallelMode.COLUMN) - for expert_id in module.initial_local_expert_ids - ] + # fc31 scales. w1 (up_proj) is always present; w3 (gate_proj) is absent + # for non-gated MoE (e.g. Nemotron-H squared-ReLU). Only concatenate the + # gate-projection scales when the gate weights are present; otherwise the + # up-projection scales alone fill the (non-doubled) fc31 scale buffer. all_w1_scales = [ load_weight_shard(weights[f"{expert_id}.w1.weight_scale"], module.tp_size, module.tp_rank, TensorParallelMode.COLUMN) for expert_id in module.initial_local_expert_ids ] - w3_w1_scales = torch.cat( - [torch.stack(all_w3_scales), - torch.stack(all_w1_scales)], dim=-1) + has_w3_scales = all(f"{expert_id}.w3.weight_scale" in weights + for expert_id in module.initial_local_expert_ids) + if module.is_gated_activation and has_w3_scales: + all_w3_scales = [ + load_weight_shard(weights[f"{expert_id}.w3.weight_scale"], + module.tp_size, module.tp_rank, + TensorParallelMode.COLUMN) + for expert_id in module.initial_local_expert_ids + ] + w3_w1_scales = torch.cat( + [torch.stack(all_w3_scales), + torch.stack(all_w1_scales)], + dim=-1) + else: + w3_w1_scales = torch.stack(all_w1_scales) w3_w1_scales = w3_w1_scales.to(module.dtype) module.fc31_weight_scale.data.copy_(w3_w1_scales.contiguous()) - # fc2 scales all_w2_scales = [ load_weight_shard(weights[f"{expert_id}.w2.weight_scale"], diff --git a/tests/unittest/_torch/modules/moe/quantize_utils.py b/tests/unittest/_torch/modules/moe/quantize_utils.py index 2633e839284e..7ed5ace32426 100644 --- a/tests/unittest/_torch/modules/moe/quantize_utils.py +++ b/tests/unittest/_torch/modules/moe/quantize_utils.py @@ -2596,12 +2596,16 @@ def __init__( swiglu_beta: Optional[torch.Tensor] = None, swiglu_limit: Optional[torch.Tensor] = None, ): - assert activation_type == ActivationType.Swiglu, ( - "Only Swiglu activation is supported for W8A16RefGatedMLPFusedMoE" - ) + assert activation_type in ( + ActivationType.Swiglu, + ActivationType.Relu2, + ActivationType.Silu, + ), f"Unsupported activation for W8A16RefGatedMLPFusedMoE: {activation_type}" # Store the original quant_config for assertion in load_weights self._original_quant_config = model_config.quant_config if model_config else None - # Create experts without quantization config since we'll dequantize weights + # Create experts without quantization config since we'll dequantize weights. + # Forward activation_type so the base builds gated (Swiglu) or non-gated + # (squared-ReLU / SiLU) experts as appropriate. super().__init__( num_experts=num_experts, routing_method=routing_method, @@ -2610,6 +2614,7 @@ def __init__( dtype=dtype, model_config=ModelConfig(), # No quant_config bias=bias, + activation_type=activation_type, swiglu_alpha=swiglu_alpha, swiglu_beta=swiglu_beta, swiglu_limit=swiglu_limit, @@ -2628,24 +2633,33 @@ def load_weights(self, weights_list: List[Dict]): # Get quantized weights and scales w1 = weights[f"{expert}.w1.weight"] s1 = weights[f"{expert}.w1.weight_scale"] - w3 = weights[f"{expert}.w3.weight"] - s3 = weights[f"{expert}.w3.weight_scale"] w2 = weights[f"{expert}.w2.weight"] s2 = weights[f"{expert}.w2.weight_scale"] # Dequantize weights: w_dequant = (w.float() * scale).to(dtype) # Note: weights are (out_features, in_features), need transpose for matmul w1_dequant = (w1.T.contiguous().float() * s1).to(self.dtype).T.contiguous() - w3_dequant = (w3.T.contiguous().float() * s3).to(self.dtype).T.contiguous() w2_dequant = (w2.T.contiguous().float() * s2).to(self.dtype).T.contiguous() # Load as regular weights (no scales) - gate_up_proj_weights = [{}, {}] down_proj_weights = [{}] - gate_up_proj_weights[0]["weight"] = w1_dequant - gate_up_proj_weights[1]["weight"] = w3_dequant down_proj_weights[0]["weight"] = w2_dequant - self.experts[expert].gate_up_proj.load_weights(gate_up_proj_weights) + + # Gated experts pack gate (w3) + up (w1) into gate_up_proj; non-gated + # experts (squared-ReLU) have only the up (w1) projection. + if self._is_gated: + w3 = weights[f"{expert}.w3.weight"] + s3 = weights[f"{expert}.w3.weight_scale"] + w3_dequant = (w3.T.contiguous().float() * s3).to(self.dtype).T.contiguous() + gate_up_proj_weights = [{}, {}] + gate_up_proj_weights[0]["weight"] = w1_dequant + gate_up_proj_weights[1]["weight"] = w3_dequant + self.experts[expert].gate_up_proj.load_weights(gate_up_proj_weights) + else: + up_proj_weights = [{}] + up_proj_weights[0]["weight"] = w1_dequant + self.experts[expert].up_proj.load_weights(up_proj_weights) + self.experts[expert].down_proj.load_weights(down_proj_weights) def check_accuracy(self, output, ref_output, weight_dtype=torch.int8): @@ -2683,10 +2697,6 @@ def create_weights(self, **quant_kwargs) -> Dict[str, torch.Tensor]: w2_weight = torch.randint( -128, 127, (self.hidden_size, self.intermediate_size), dtype=torch.int8 ).cuda() - w3_weight = torch.randint( - -128, 127, (self.intermediate_size, self.hidden_size), dtype=torch.int8 - ).cuda() - # Per-channel scales w1_scale = ( torch.randn(self.intermediate_size, dtype=self.dtype, device="cuda") @@ -2696,10 +2706,20 @@ def create_weights(self, **quant_kwargs) -> Dict[str, torch.Tensor]: torch.randn(self.hidden_size, dtype=self.dtype, device="cuda") / self.intermediate_size ) - w3_scale = ( - torch.randn(self.intermediate_size, dtype=self.dtype, device="cuda") - / self.hidden_size - ) + + # Non-gated experts (e.g. Nemotron-H squared-ReLU) have no gate (w3) + # projection, so emit empty w3 tensors. Mirrors NVFP4QuantizeUtil. + if self._is_gated: + w3_weight = torch.randint( + -128, 127, (self.intermediate_size, self.hidden_size), dtype=torch.int8 + ).cuda() + w3_scale = ( + torch.randn(self.intermediate_size, dtype=self.dtype, device="cuda") + / self.hidden_size + ) + else: + w3_weight = torch.empty(0, dtype=torch.int8, device="cuda") + w3_scale = torch.empty(0, dtype=self.dtype, device="cuda") weights[f"{expert_id}.w1.weight"] = w1_weight weights[f"{expert_id}.w2.weight"] = w2_weight diff --git a/tests/unittest/_torch/modules/moe/test_moe_backend.py b/tests/unittest/_torch/modules/moe/test_moe_backend.py index e8c8d9511e2c..96af81d10f2c 100644 --- a/tests/unittest/_torch/modules/moe/test_moe_backend.py +++ b/tests/unittest/_torch/modules/moe/test_moe_backend.py @@ -791,7 +791,7 @@ def generate_element_wise_test_params() -> List: SEQ_LENS_TO_TEST, DTYPES_TO_TEST, [MoeBackendType.CUTLASS, MoeBackendType.TRTLLM], - [None, QuantAlgo.NVFP4], + [None, QuantAlgo.NVFP4, QuantAlgo.W8A16], ): if skip_reason: continue @@ -799,6 +799,9 @@ def generate_element_wise_test_params() -> List: continue if backend_type == MoeBackendType.TRTLLM and quant_algo is None: continue + # INT8 weight-only per-channel non-gated MoE is CUTLASS-path only. + if quant_algo == QuantAlgo.W8A16 and backend_type != MoeBackendType.CUTLASS: + continue test_id = f"act={activation_type.name}-{base_test_id}" param_values = ( dtype,