[None][feat] Support Mistral Large3 LLM part (NVIDIA#9820)

Signed-off-by: bhsueh <[email protected]>

Author: byshiue

examples/llm-api/quickstart_advanced.py

@@ -23,6 +23,11 @@ def add_llm_args(parser):
                         type=str,
                         nargs="+",
                         help="A single or a list of text prompts.")
+    parser.add_argument('--checkpoint_format',
+                        type=str,
+                        default=None,
+                        choices=["HF", "mistral"],
+                        help="Model checkpoint format.")
     # Build config
     parser.add_argument("--max_seq_len",
                         type=int,
@@ -237,6 +242,7 @@ def setup_llm(args, **kwargs):
     llm = LLM(
         model=args.model_dir,
         backend='pytorch',
+        checkpoint_format=args.checkpoint_format,
         disable_overlap_scheduler=args.disable_overlap_scheduler,
         kv_cache_config=kv_cache_config,
         attn_backend=args.attention_backend,

examples/models/core/mistral_large_3/README.md

@@ -0,0 +1,53 @@
+# Mistral Large V3
+
+* Setup the model path
+
+```bash
+export mistral_large_3_model_path=<mistral_large_3_model_path>
+```
+
+## LLM-only run
+
+* Run the Mistral Large V3 by `quickstart_advanced.py`
+
+```bash
+mpirun -n 1 --allow-run-as-root --oversubscribe python3 examples/llm-api/quickstart_advanced.py \
+    --model_dir ${mistral_large_3_model_path} \
+    --tp_size 4 \
+    --moe_ep_size 4 \
+    --max_tokens 100 \
+    --checkpoint_format mistral \
+    --moe_backend TRTLLM
+```
+
+* Launch the trtllm-serve and send a request
+
+```bash
+echo "
+backend: pytorch
+tensor_parallel_size: 4
+moe_expert_parallel_size: 4
+enable_attention_dp: false
+kv_cache_config:
+  enable_block_reuse: true
+checkpoint_format: mistral
+" > serve.yml
+mpirun -n 1 --allow-run-as-root --oversubscribe python3 -m tensorrt_llm.commands.serve serve \
+    ${mistral_large_3_model_path} \
+    --host localhost --port 8001 --backend pytorch \
+    --extra_llm_api_options serve.yml \
+    --tokenizer ${mistral_large_3_model_path} \
+    2>&1 | tee serve_debug.log &
+
+curl http://localhost:8001/v1/completions \
+  -H "Content-Type: application/json" \
+  -d '{
+      "model": "${mistral_large_3_model_path}",
+      "prompt": "The capital of France is",
+      "max_tokens": 16,
+      "top_k": 16
+  }'
+
+# The result would be like
+{"id":"cmpl-7e342c1d722d4226a1bf3ed35d762c35","object":"text_completion","created":1764061351,"model":"${mistral_large_3_model_path}","choices":[{"index":0,"text":"The capital of France is **Paris**.\n\nParis is the largest city in France and","token_ids":null,"logprobs":null,"context_logits":null,"finish_reason":"length","stop_reason":null,"disaggregated_params":null,"avg_decoded_tokens_per_iter":1.0}],"usage":{"prompt_tokens":7,"total_tokens":23,"completion_tokens":16,"prompt_tokens_details":{"cached_tokens":1}},"prompt_token_ids":null}
+```

requirements.txt

@@ -75,3 +75,4 @@ numexpr<2.14.0 # WAR for attempted use of nonexistent numpy.typing
 partial_json_parser
 apache-tvm-ffi==0.1.4 # used for reduce nvidia-cutlass-dsl host overhead
 torch-c-dlpack-ext==0.1.3 # used for reduce nvidia-cutlass-dsl host overhead, optional package for improved torch tensor calling perf
+mistral-common==1.8.6

tensorrt_llm/_torch/models/checkpoints/__init__.py

@@ -12,11 +12,30 @@
 from .hf.qwen3_next_weight_mapper import Qwen3NextHfWeightMapper
 from .hf.weight_loader import HfWeightLoader
 from .hf.weight_mapper import HfWeightMapper
+from .mistral.checkpoint_loader import (MistralCheckpointLoader,
+                                        MistralLarge3CheckpointLoader)
+from .mistral.config_loader import MistralConfigLoader
+from .mistral.weight_mapper import (MistralLarge3WeightMapper,
+                                    MistralWeightMapper)
 
 __all__ = [
-    "HfConfigLoader", "HfWeightLoader", "HfWeightMapper",
-    "BaseCheckpointLoader", "HfCheckpointLoader", "NemotronHHfWeightMapper",
-    "Gemma3HfWeightMapper", "MixtralHfWeightMapper", "Llama4HfWeightMapper",
-    "Qwen2MoeHfWeightMapper", "Qwen3MoeHfWeightMapper", "Qwen2VLHfWeightMapper",
-    "Qwen3NextHfWeightMapper", "LlavaNextHfWeightMapper"
+    "HfConfigLoader",
+    "HfWeightLoader",
+    "HfWeightMapper",
+    "MistralConfigLoader",
+    "MistralWeightMapper",
+    "MistralCheckpointLoader",
+    "BaseCheckpointLoader",
+    "HfCheckpointLoader",
+    "NemotronHHfWeightMapper",
+    "Gemma3HfWeightMapper",
+    "MixtralHfWeightMapper",
+    "Llama4HfWeightMapper",
+    "Qwen2MoeHfWeightMapper",
+    "Qwen3MoeHfWeightMapper",
+    "Qwen2VLHfWeightMapper",
+    "Qwen3NextHfWeightMapper",
+    "LlavaNextHfWeightMapper",
+    "MistralLarge3CheckpointLoader",
+    "MistralLarge3WeightMapper",
 ]

tensorrt_llm/_torch/models/checkpoints/hf/weight_loader.py

@@ -19,6 +19,7 @@
 from tensorrt_llm.mapping import Mapping
 
 
+@register_checkpoint_weight_loader("mistral")
 @register_checkpoint_weight_loader("HF")
 class HfWeightLoader(BaseWeightLoader):
     """

tensorrt_llm/_torch/models/checkpoints/mistral/__init__.py

@@ -0,0 +1,75 @@
+from tensorrt_llm._torch.models.checkpoints.base_config_loader import BaseConfigLoader
+from tensorrt_llm._torch.models.checkpoints.base_weight_loader import BaseWeightLoader
+from tensorrt_llm._torch.models.checkpoints.base_weight_mapper import BaseWeightMapper
+from tensorrt_llm._torch.models.checkpoints.hf.checkpoint_loader import HfCheckpointLoader
+from tensorrt_llm._torch.models.checkpoints.mistral.config_loader import MistralConfigLoader
+from tensorrt_llm._torch.models.modeling_utils import register_checkpoint_loader
+
+
+@register_checkpoint_loader("mistral")
+class MistralCheckpointLoader(HfCheckpointLoader):
+    def __init__(
+        self,
+        *,
+        weight_loader: BaseWeightLoader | None = None,
+        weight_mapper: BaseWeightMapper | None = None,
+        config_loader: BaseConfigLoader | None = None,
+    ):
+        super().__init__(
+            weight_loader=weight_loader, weight_mapper=weight_mapper, config_loader=config_loader
+        )
+        self._checkpoint_format = "mistral"
+        self.mm_module_mapping = {
+            "vision_encoder": "vision_tower",
+            "pre_mm_projector_norm": "multi_modal_projector.norm",
+            "vision_language_adapter": "multi_modal_projector",
+            "patch_merger": "multi_modal_projector.patch_merger",
+        }
+
+    def preprocess_weights(self, weights: dict) -> dict:
+        """
+        Aggregate weights by module
+        """
+        hf_weights = {}
+
+        for key, value in weights.items():
+            modules = key.split(".")
+
+            if modules[0] not in self.mm_module_mapping.keys():
+                hf_weights["language_model." + key] = value
+
+            else:
+                modules[0] = self.mm_module_mapping[modules[0]]
+                hf_weights[".".join(modules)] = value
+
+        return hf_weights
+
+    def inverse_nvfp4_global_scales(self, weights):
+        for key in weights.keys():
+            if "global_scale" in key:
+                weights[key] = 1.0 / weights[key]
+
+    def load_weights(self, checkpoint_dir: str, **kwargs):
+        weights = super().weight_loader.load_weights(checkpoint_dir, **kwargs)
+        weights = self.preprocess_weights(weights)
+        # The definition of global_scale is different in Mistral, need to inverse the scale
+        self.inverse_nvfp4_global_scales(weights)
+        return weights
+
+    def get_default_config_loader(self) -> MistralConfigLoader:
+        return MistralConfigLoader()
+
+
+@register_checkpoint_loader("mistral_large_3")
+class MistralLarge3CheckpointLoader(MistralCheckpointLoader):
+    def __init__(
+        self,
+        *,
+        weight_loader: BaseWeightLoader | None = None,
+        weight_mapper: BaseWeightMapper | None = None,
+        config_loader: BaseConfigLoader | None = None,
+    ):
+        super().__init__(
+            weight_loader=weight_loader, weight_mapper=weight_mapper, config_loader=config_loader
+        )
+        self._checkpoint_format = "mistral_large_3"