diff --git a/tensorrt_llm/_torch/model_config.py b/tensorrt_llm/_torch/model_config.py index 78fbd82da645..9afc065094fc 100644 --- a/tensorrt_llm/_torch/model_config.py +++ b/tensorrt_llm/_torch/model_config.py @@ -45,6 +45,11 @@ _DEEPSEEK_V4_ARCHITECTURES = {"DeepseekV4ForCausalLM"} _DEEPSEEK_V4_ROUTED_EXPERT_WEIGHT = "layers.0.ffn.experts.0.w1.weight" +_MINIMAX_M3_ARCHITECTURES = { + "MiniMaxM3SparseForCausalLM", + "MiniMaxM3SparseForConditionalGeneration", +} + def _unified_kv_pool_includes_mamba( is_disagg: bool, spec_config: Optional['SpeculativeConfig']) -> bool: @@ -674,6 +679,81 @@ def _set_deepseek_v4_routed_moe_quant_config(pretrained_config, experts_quant_config.quant_algo) return layer_quant_config + @staticmethod + def _set_minimax_m3_layer_quant_config(pretrained_config, + layer_quant_config): + """Normalize the Minimax M3 MIXED_PRECISION per-layer quant config. + + Two fix-ups are applied: + + 1. Strip the ``language_model.`` prefix from every per-layer key. + The M3 VL checkpoint stores keys like + ``language_model.model.layers.0.self_attn.o_proj -> MXFP8`` in + ``hf_quant_config.json``, but the TRT-LLM module tree names the text + decoder ``model.layers.0.self_attn.o_proj`` (no ``language_model.`` + prefix -- the loader strips it). ``apply_layerwise_quant_config`` + matches *standalone* Linears (e.g. ``o_proj``, ``down_proj``) with an + **exact** ``name == key`` comparison, so the prefixed keys never match + and those layers silently fall back to the global ``MIXED_PRECISION`` + config -> loaded unquantized (MXFP8 ``weight_scale`` dropped) -> the + attention/MLP output magnitude explodes. Stripping the prefix makes + the exact match succeed. (Fused qkv/gate_up Linears and the Attention + wrapper use substring matches and happened to work regardless.) + + 2. Inject a single coarse ``model.layers.N.block_sparse_moe.experts`` + entry per MoE layer so ``MiniMaxM3MoE._get_experts_quant_config`` can + select the NVFP4 backend for the routed experts (the fine-grained + per-linear NVFP4 expert keys can't be used directly). + + Does nothing when there is no per-layer config (e.g. the uniform MXFP8 + or a BF16 checkpoint). + """ + from tensorrt_llm.models.modeling_utils import QuantAlgo + if layer_quant_config is None: + return layer_quant_config + + # (1) Strip the ``language_model.`` prefix so exact-match per-layer + # quant assignment works for standalone base Linears. + _LM_PREFIX = "language_model." + layer_quant_config = { + (k[len(_LM_PREFIX):] if k.startswith(_LM_PREFIX) else k): v + for k, v in layer_quant_config.items() + } + + # (2) Inject coarse NVFP4 expert entries (only when routed experts are + # NVFP4). + has_nvfp4_experts = any( + "block_sparse_moe.experts" in k and isinstance(v, QuantConfig) + and v.quant_algo == QuantAlgo.NVFP4 + for k, v in layer_quant_config.items()) + if not has_nvfp4_experts: + return layer_quant_config + + experts_quant_config = QuantConfig() + experts_quant_config.quant_algo = QuantAlgo.NVFP4 + # TODO: remove the hardcoded group_size and read it from the per-linear + # NVFP4 expert entries in hf_quant_config.json instead. 16 is correct + # for standard NVFP4 today, but this is a latent bug if a checkpoint + # ever ships a different group size. + experts_quant_config.group_size = 16 + + text_config = getattr(pretrained_config, "text_config", + pretrained_config) + if isinstance(text_config, dict): + moe_layer_freq = text_config.get("moe_layer_freq", []) + else: + moe_layer_freq = getattr(text_config, "moe_layer_freq", []) + + for layer_idx, freq in enumerate(moe_layer_freq): + if int(freq) != 0: + layer_quant_config[ + f"model.layers.{layer_idx}.block_sparse_moe.experts"] = experts_quant_config + + logger.info( + "Detected Minimax M3 NVFP4 routed MoE checkpoint; using NVFP4 " + "for routed experts.") + return layer_quant_config + @staticmethod def load_quant_config_from_dtypes_json(dtypes_json_file, moe_backend: str): quant_config = QuantConfig() @@ -1060,6 +1140,10 @@ def _recursive_update_config(config: transformers.PretrainedConfig, kwargs.get('spec_config', None), require_layout=require_deepseek_v4_routed_moe_layout) + if architecture in _MINIMAX_M3_ARCHITECTURES: + layer_quant_config = cls._set_minimax_m3_layer_quant_config( + pretrained_config, layer_quant_config) + model_config = cls(pretrained_config=pretrained_config, quant_config=quant_config, quant_config_dict=layer_quant_config, diff --git a/tensorrt_llm/_torch/models/modeling_minimaxm3.py b/tensorrt_llm/_torch/models/modeling_minimaxm3.py index 0b8b9f3e00ad..86da05822e97 100644 --- a/tensorrt_llm/_torch/models/modeling_minimaxm3.py +++ b/tensorrt_llm/_torch/models/modeling_minimaxm3.py @@ -31,6 +31,7 @@ from tensorrt_llm.functional import AllReduceStrategy, PositionEmbeddingType from tensorrt_llm.mapping import Mapping +from tensorrt_llm.models.modeling_utils import QuantConfig from ..attention_backend import AttentionMetadata from ..attention_backend.interface import PositionalEmbeddingParams, RopeParams @@ -352,6 +353,24 @@ class MiniMaxM3MoE(nn.Module): that the previous independent-reduction wiring incurred. """ + @staticmethod + def _get_experts_quant_config(model_config: "ModelConfig", layer_idx: int) -> QuantConfig: + """Return the per-layer quant config for the routed experts. + + For MIXED_PRECISION checkpoints (MXFP8 base + NVFP4 experts), + ``ModelConfig._set_minimax_m3_moe_quant_config`` pre-populates + ``quant_config_dict`` with coarse entries keyed by + ``model.layers.N.block_sparse_moe.experts``. Falls back to the + global ``quant_config`` when no per-layer entry exists (e.g. BF16 + or user-supplied global NVFP4 config). + """ + if getattr(model_config, "quant_config_dict", None) is None: + return model_config.quant_config + return model_config.quant_config_dict.get( + f"model.layers.{layer_idx}.block_sparse_moe.experts", + model_config.quant_config, + ) + def __init__( self, model_config: "ModelConfig[PretrainedConfig]", @@ -415,6 +434,7 @@ def __init__( # (matches DeepSeekV3). Under Attention DP the fused MoE already # skips its in-op all-reduce, so ``reduce_results=False`` is # also the correct flag there. + experts_quant_config = MiniMaxM3MoE._get_experts_quant_config(model_config, layer_idx) self.experts = create_moe( routing_method=self.gate.routing_method, num_experts=self.num_experts, @@ -422,6 +442,7 @@ def __init__( reduce_results=False, model_config=model_config, layer_idx=layer_idx, + override_quant_config=experts_quant_config, swiglu_alpha=self.swiglu_alpha, swiglu_beta=self.swiglu_beta, swiglu_limit=self.swiglu_limit, diff --git a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_trtllm_gen.py b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_trtllm_gen.py index 820d92a95e4a..d224dc4e73c8 100644 --- a/tensorrt_llm/_torch/modules/fused_moe/fused_moe_trtllm_gen.py +++ b/tensorrt_llm/_torch/modules/fused_moe/fused_moe_trtllm_gen.py @@ -325,6 +325,10 @@ def _to_trtllm_gen_activation_type(self, activation_type: ActivationType) -> int: if activation_type == ActivationType.Swiglu: return 0 + elif activation_type == ActivationType.SwigluBias: + # SwigluBias uses the same SwiGlu kernel path (ActType::SwiGlu == 0); + # the per-expert alpha/beta/clamp_limit are passed as separate tensors. + return 0 elif activation_type == ActivationType.Relu2: return 1 elif activation_type == ActivationType.Silu: diff --git a/tests/integration/defs/accuracy/references/gsm8k.yaml b/tests/integration/defs/accuracy/references/gsm8k.yaml index 2c32b05995d3..fd7b4096b7d4 100644 --- a/tests/integration/defs/accuracy/references/gsm8k.yaml +++ b/tests/integration/defs/accuracy/references/gsm8k.yaml @@ -447,6 +447,9 @@ MiniMaxAI/MiniMax-M3: MiniMaxAI/MiniMax-M3-MXFP8: - quant_algo: MXFP8 accuracy: 89 +nvidia/MiniMax-M3-NVFP4: + - quant_algo: MIXED_PRECISION + accuracy: 88 nvidia/NVIDIA-Nemotron-Nano-9B-v2: - accuracy: 85.027 - quant_algo: FP8 diff --git a/tests/integration/defs/accuracy/references/mmlu.yaml b/tests/integration/defs/accuracy/references/mmlu.yaml index a33ceedb4a8d..d4012b2c6c32 100644 --- a/tests/integration/defs/accuracy/references/mmlu.yaml +++ b/tests/integration/defs/accuracy/references/mmlu.yaml @@ -278,6 +278,9 @@ MiniMaxAI/MiniMax-M3: MiniMaxAI/MiniMax-M3-MXFP8: - quant_algo: MXFP8 accuracy: 85 +nvidia/MiniMax-M3-NVFP4: + - quant_algo: MIXED_PRECISION + accuracy: 83 moonshotai/Kimi-K2-Instruct: - quant_algo: FP8_BLOCK_SCALES accuracy: 87.65 diff --git a/tests/integration/defs/accuracy/test_llm_api_pytorch.py b/tests/integration/defs/accuracy/test_llm_api_pytorch.py index 641ed8c11ca4..a13e745957f0 100644 --- a/tests/integration/defs/accuracy/test_llm_api_pytorch.py +++ b/tests/integration/defs/accuracy/test_llm_api_pytorch.py @@ -7233,10 +7233,11 @@ class TestMiniMaxM3(LlmapiAccuracyTestHarness): # config requires ``trust_remote_code`` and the runtime requires the # MiniMax-M3 sparse attention backend + matching KV-cache manager v2 # (selected by ``sparse_attention_config``). Blackwell-only (SM100+). - # Two checkpoints are covered: the upstream BF16 checkpoint and the + # Three checkpoints are covered: the upstream BF16 checkpoint, the # NVIDIA MXFP8 checkpoint (weights in MXFP8, activations / KV cache - # remain BF16; the runtime path is identical aside from the quant - # config). + # remain BF16), and the NVIDIA NVFP4 checkpoint (MXFP8 base layers + + # NVFP4 routed experts); the runtime path is identical aside from the + # quant config. MODEL_NAME = "MiniMaxAI/MiniMax-M3" MODEL_PATH = f"{llm_models_root()}/MiniMax-M3" @@ -7261,16 +7262,20 @@ def test_auto_dtype(self, tp_size, ep_size): task = GSM8K(self.MODEL_NAME) task.evaluate(llm) - @pytest.mark.skip_less_device(8) + @pytest.mark.skip_less_device(4) @pytest.mark.skip_less_device_memory(140000) - @parametrize_with_ids("tp_size,ep_size", [(8, 8)]) + @parametrize_with_ids("tp_size,ep_size", [(4, 4)]) def test_mxfp8(self, tp_size, ep_size): # MXFP8 checkpoint: weights are MXFP8 (e4m3 + UE8M0 1x32 block # scales) with MXFP8 dynamic activations; the KV cache stays in # BF16 and the sparse attention path is unchanged from BF16. model_name = "MiniMaxAI/MiniMax-M3-MXFP8" model_path = f"{llm_models_root()}/MiniMax-M3-MXFP8" - kv_cache_config = KvCacheConfig(free_gpu_memory_fraction=0.6, + # Halving TP from the BF16 reference (TP=8) doubles per-rank + # model + KV footprint; cap KV cache at 0.4 of free memory and + # constrain batch / token budget so the runtime allocator stays + # under the PyTorch cap. + kv_cache_config = KvCacheConfig(free_gpu_memory_fraction=0.4, enable_block_reuse=False) sparse_attention_config = MiniMaxM3SparseAttentionConfig() with LLM(model_path, @@ -7279,6 +7284,8 @@ def test_mxfp8(self, tp_size, ep_size): kv_cache_config=kv_cache_config, sparse_attention_config=sparse_attention_config, max_seq_len=4096, + max_batch_size=32, + max_num_tokens=4096, trust_remote_code=True) as llm: assert llm.args.quant_config.quant_algo == QuantAlgo.MXFP8 task = MMLU(model_name) @@ -7286,6 +7293,33 @@ def test_mxfp8(self, tp_size, ep_size): task = GSM8K(model_name) task.evaluate(llm) + @pytest.mark.skip_less_device(4) + @pytest.mark.skip_less_device_memory(140000) + @parametrize_with_ids("tp_size,ep_size", [(4, 4)]) + def test_nvfp4(self, tp_size, ep_size): + # NVFP4 checkpoint: MXFP8 base layers with NVFP4 routed experts + # (MIXED_PRECISION checkpoint); the KV cache stays in BF16 and the + # sparse attention path is unchanged from BF16. + model_name = "nvidia/MiniMax-M3-NVFP4" + model_path = f"{llm_models_root()}/MiniMax-M3-NVFP4" + kv_cache_config = KvCacheConfig(free_gpu_memory_fraction=0.6, + enable_block_reuse=False) + sparse_attention_config = MiniMaxM3SparseAttentionConfig() + moe_config = MoeConfig(backend="CUTLASS") + with LLM(model_path, + tensor_parallel_size=tp_size, + moe_expert_parallel_size=ep_size, + kv_cache_config=kv_cache_config, + sparse_attention_config=sparse_attention_config, + moe_config=moe_config, + max_seq_len=4096, + trust_remote_code=True) as llm: + assert llm.args.quant_config.quant_algo == QuantAlgo.MIXED_PRECISION + task = MMLU(model_name) + task.evaluate(llm) + task = GSM8K(model_name) + task.evaluate(llm) + @skip_pre_blackwell class TestGLM5FP8(LlmapiAccuracyTestHarness): diff --git a/tests/integration/test_lists/qa/llm_function_core.txt b/tests/integration/test_lists/qa/llm_function_core.txt index d2a91b70a423..108afd235ab6 100644 --- a/tests/integration/test_lists/qa/llm_function_core.txt +++ b/tests/integration/test_lists/qa/llm_function_core.txt @@ -664,7 +664,8 @@ accuracy/test_llm_api_pytorch.py::TestLlama3_3_70BInstruct::test_nvfp4_tp4[torch accuracy/test_llm_api_pytorch.py::TestMiniMaxM2::test_4gpus[attention_dp=False-cuda_graph=True-overlap_scheduler=True-tp_size=4-ep_size=4] accuracy/test_llm_api_pytorch.py::TestMiniMaxM2_5::test_4gpus[attention_dp=False-cuda_graph=True-overlap_scheduler=True-tp_size=4-ep_size=4] accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_auto_dtype[tp_size=8-ep_size=8] TIMEOUT (180) -accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_mxfp8[tp_size=8-ep_size=8] TIMEOUT (180) +accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_mxfp8[tp_size=4-ep_size=4] TIMEOUT (180) +accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_nvfp4[tp_size=4-ep_size=4] TIMEOUT (180) accuracy/test_llm_api_pytorch.py::TestMinistral8BInstruct::test_auto_dtype accuracy/test_llm_api_pytorch.py::TestMinistral8BInstruct::test_fp8 accuracy/test_llm_api_pytorch.py::TestMistralLarge3_675B::test_fp8[latency_moe_deepgemm]