[None][fix] Load compressed-tensors NVFP4 weights in MiniMax-M3

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

@@ -13,14 +13,128 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 from collections.abc import Mapping
+from typing import Any
 
 import torch
 from torch import nn
 
 from tensorrt_llm._torch.models.modeling_utils import register_mapper
+from tensorrt_llm.logger import logger
 
+from ..base_weight_loader import ConsumableWeightsDict
 from ..base_weight_mapper import BaseWeightMapper
 
+_COMPRESSED_TENSORS_NVFP4_SUFFIXES = (".weight_packed", ".weight_global_scale",
+                                      ".input_global_scale")
+
+
+def _invert_compressed_tensors_global_scale(value) -> torch.Tensor:
+    """Convert a compressed-tensors global scale to TRT-LLM's divisor form.
+
+    A global scale of 0 marks an expert that saw no calibration data. It is
+    kept as 0 rather than inverted to Inf, which would otherwise win the MoE
+    per-layer max-reduction over experts and corrupt every expert's
+    input_scale; a 0 entry is harmless there and silences the dead expert.
+    """
+    t = value if isinstance(value, torch.Tensor) else value[...]
+    t = t.to(torch.float32)
+    out = torch.where(t > 0,
+                      1.0 / torch.clamp(t, min=torch.finfo(torch.float32).tiny),
+                      torch.zeros_like(t))
+    return torch.nan_to_num(out, nan=0.0, posinf=0.0, neginf=0.0).contiguous()
+
+
+def _read_compressed_tensors_global_scale(name: str, value, zero_scales: list,
+                                          invalid_scales: list) -> torch.Tensor:
+    """Read a global scale and record whether it has entries that cannot be inverted."""
+    t = value if isinstance(value, torch.Tensor) else value[...]
+    t32 = t.to(torch.float32)
+    if bool((~torch.isfinite(t32) | (t32 < 0)).any()):
+        invalid_scales.append(name)
+    elif bool((t32 == 0).any()):
+        zero_scales.append(name)
+    return t
+
+
+def _warn_uninvertible_global_scales(zero_scales: list,
+                                     invalid_scales: list) -> None:
+    """Log one summary line for the global scales that were set to 0."""
+    if not zero_scales and not invalid_scales:
+        return
+    parts = []
+    if invalid_scales:
+        parts.append(
+            f"{len(invalid_scales)} with negative or non-finite values, which "
+            f"means a broken checkpoint (e.g. {', '.join(invalid_scales[:3])})")
+    if zero_scales:
+        parts.append(
+            f"{len(zero_scales)} with 0, as stored for experts without "
+            f"calibration data (e.g. {', '.join(zero_scales[:3])})")
+    logger.warning(
+        "compressed-tensors NVFP4 global scales that cannot be inverted: " +
+        "; ".join(parts) + ". They are set to 0; a weight global scale of 0 "
+        "zeroes that expert's output.")
+
+
+def convert_compressed_tensors_nvfp4_weights(
+        weights: Mapping[str, Any]) -> Mapping[str, Any]:
+    """Rename llm-compressor "compressed-tensors" NVFP4 tensors to TRT-LLM's.
+
+    compressed-tensors stores
+      weight_packed        (uint8, packed FP4)
+      weight_scale         (FP8 per-group)
+      weight_global_scale  (FP32 scalar = (FP8_MAX*FP4_MAX)/amax_weight)
+      input_global_scale   (FP32 scalar = (FP8_MAX*FP4_MAX)/amax_input)
+    while TRT-LLM's NVFP4 Linear and MoE loaders expect
+      weight               (uint8 view of float4_e2m1x2)
+      weight_scale         (FP8, same)
+      weight_scale_2       (FP32 scalar = amax_weight/(FP8_MAX*FP4_MAX))
+      input_scale          (FP32 scalar = amax_input/(FP8_MAX*FP4_MAX))
+    so the global scales are inverted on the way through.
+
+    Args:
+        weights: checkpoint weights keyed by checkpoint name; a plain dict or
+            a ``ConsumableWeightsDict``.
+
+    Returns:
+        ``weights`` itself when it holds no compressed-tensors NVFP4 names.
+        Otherwise a converted mapping of the same kind; other keys are passed
+        through unchanged and the caller must not use ``weights`` afterwards
+        (see ``ConsumableWeightsDict.take_ownership``).
+
+    Raises:
+        ValueError: two checkpoint tensors map to the same TRT-LLM name (for
+            example both ``x.weight_packed`` and ``x.weight``), which would
+            otherwise keep whichever came last.
+    """
+    if not any(
+            name.endswith(_COMPRESSED_TENSORS_NVFP4_SUFFIXES)
+            for name in weights.keys()):
+        return weights
+
+    converted, sources = {}, {}
+    zero_scales, invalid_scales = [], []
+    for name, value in weights.items():
+        source = name
+        if name.endswith(".weight_packed"):
+            name = name[:-len(".weight_packed")] + ".weight"
+        elif name.endswith((".weight_global_scale", ".input_global_scale")):
+            value = _read_compressed_tensors_global_scale(
+                name, value, zero_scales, invalid_scales)
+            if name.endswith(".weight_global_scale"):
+                name = name[:-len(".weight_global_scale")] + ".weight_scale_2"
+            else:
+                name = name[:-len(".input_global_scale")] + ".input_scale"
+            value = _invert_compressed_tensors_global_scale(value)
+        if name in converted:
+            raise ValueError(
+                f"compressed-tensors NVFP4 checkpoint maps both {sources[name]!r} "
+                f"and {source!r} to {name!r}")
+        converted[name] = value
+        sources[name] = source
+    _warn_uninvertible_global_scales(zero_scales, invalid_scales)
+    return ConsumableWeightsDict.take_ownership(weights, converted)
+
 
 @register_mapper("MX")
 @register_mapper("HF")

tensorrt_llm/_torch/models/modeling_laguna.py

@@ -50,7 +50,7 @@
 from ..moe.fused_moe.weight_owner import is_moe_weight_owner
 from ..speculative import SpecMetadata
 from ..utils import AuxStreamType
-from .checkpoints.hf.weight_mapper import HfWeightMapper
+from .checkpoints.hf.weight_mapper import HfWeightMapper, convert_compressed_tensors_nvfp4_weights
 from .modeling_speculative import SpecDecOneEngineForCausalLM
 from .modeling_utils import (
     DecoderModel,
@@ -637,59 +637,8 @@ def preprocess_weights(self, weights: dict) -> dict:
             weights,
         )
 
-        # Translate llm-compressor "compressed-tensors" NVFP4 naming into what
-        # TRT-LLM's NVFP4 linear / MoE loaders expect. llm-compressor stores
-        #   weight_packed        (uint8, packed FP4)
-        #   weight_scale         (FP8 per-group)
-        #   weight_global_scale  (FP32 scalar = (FP8_MAX*FP4_MAX)/amax_weight)
-        #   input_global_scale   (FP32 scalar = (FP8_MAX*FP4_MAX)/amax_input)
-        # TRT-LLM's loaders expect:
-        #   weight               (uint8 view of float4_e2m1x2)
-        #   weight_scale         (FP8, same)
-        #   weight_scale_2       (FP32 scalar = amax_weight/(FP8_MAX*FP4_MAX))
-        #   input_scale          (FP32 scalar = amax_input/(FP8_MAX*FP4_MAX))
-        # which is the reciprocal of the *_global_scale values.
         if self._config is not None and self._config.quant_config.quant_mode.has_nvfp4():
-
-            def _invert_global_scale(t: torch.Tensor) -> torch.Tensor:
-                """Convert compressed-tensors multiplier form to TRT-LLM
-                divisor form, protecting against uncalibrated experts.
-
-                llm-compressor stores global_scale=(FP8_MAX*FP4_MAX)/amax;
-                TRT-LLM expects {weight_scale_2, input_scale}=amax/(FP8_MAX*FP4_MAX),
-                i.e. the reciprocal. Some checkpoints store 0 for experts with
-                no calibration data (amax unavailable); a naive 1/x would
-                produce Inf which then contaminates the MoE per-layer
-                max-reduction over experts and zeroes out *every* expert's
-                input_scale. We preserve 0 -> 0 so those entries are harmless
-                in the max reduction and the corresponding expert's alpha
-                evaluates to 0 (silencing the dead expert).
-                """
-                t = t.to(torch.float32)
-                out = torch.where(
-                    t > 0,
-                    1.0 / torch.clamp(t, min=torch.finfo(torch.float32).tiny),
-                    torch.zeros_like(t),
-                )
-                # Belt-and-braces: any residual Inf/NaN becomes 0.
-                return torch.nan_to_num(out, nan=0.0, posinf=0.0, neginf=0.0).contiguous()
-
-            renamed: dict = {}
-            for name, value in weights.items():
-                if name.endswith(".weight_packed"):
-                    new_name = name[: -len(".weight_packed")] + ".weight"
-                    renamed[new_name] = value
-                elif name.endswith(".weight_global_scale"):
-                    new_name = name[: -len(".weight_global_scale")] + ".weight_scale_2"
-                    t = value[...] if not isinstance(value, torch.Tensor) else value
-                    renamed[new_name] = _invert_global_scale(t)
-                elif name.endswith(".input_global_scale"):
-                    new_name = name[: -len(".input_global_scale")] + ".input_scale"
-                    t = value[...] if not isinstance(value, torch.Tensor) else value
-                    renamed[new_name] = _invert_global_scale(t)
-                else:
-                    renamed[name] = value
-            weights = renamed
+            weights = convert_compressed_tensors_nvfp4_weights(weights)
 
         return weights
 

tensorrt_llm/_torch/models/modeling_minimaxm3.py

@@ -79,6 +79,7 @@
 from ..utils import AuxStreamType, EventType, get_model_extra_attrs, is_torch_compiling
 from .checkpoints.base_weight_mapper import BaseWeightMapper
 from .checkpoints.hf.minimaxm3_weight_mapper import MINIMAX_M3_PARAMS_MAP, MiniMaxM3HfWeightMapper
+from .checkpoints.hf.weight_mapper import convert_compressed_tensors_nvfp4_weights
 from .modeling_speculative import SpecDecOneEngineForCausalLM
 from .modeling_utils import (
     DecoderModel,
@@ -2993,6 +2994,10 @@ def load_weights(
         params_map: Optional[Dict[str, str]] = None,
         allow_partial_loading: bool = False,
     ) -> None:
+        # Accept compressed-tensors NVFP4 names before any loader below reads
+        # the dict. Keyed on the tensor names, not on the global quant config:
+        # that is MIXED_PRECISION when NVFP4 covers only some modules.
+        weights = convert_compressed_tensors_nvfp4_weights(weights)
         # Pack the opt-in five-way projection with MiniMax-specific TP
         # sharding before the generic mapper handles the compatibility path.
         packed_projection_modules = _load_qkv_index_proj_weights(self, weights)

tests/unittest/_torch/models/checkpoints/hf/test_minimaxm3_weight_mapper.py

@@ -178,17 +178,24 @@ def test_gate_bias_params_map() -> None:
     assert renamed["model.layers.3.block_sparse_moe.gate.weight"] is gate_weight
 
 
-def test_load_weights_accepts_base_mapper_without_params_map() -> None:
+def _make_layerless_model(model_class: type = MiniMaxM3ForCausalLM) -> MiniMaxM3ForCausalLM:
     config = SimpleNamespace(num_key_value_heads=_NUM_KV_HEADS, num_attention_heads=8)
     model_config = ModelConfig(pretrained_config=config, mapping=Mapping())
-    model = object.__new__(MiniMaxM3ForCausalLM)
+    model = object.__new__(model_class)
     torch.nn.Module.__init__(model)
     model.model_config = model_config
     # This mapper-only fixture has no decoder layers to refresh after loading.
     model.model = torch.nn.Module()
     model.model.layers = torch.nn.ModuleList()
     # Set by SpecDecOneEngineForCausalLM.__init__, which this test bypasses.
     model.spec_config = None
+    # Set by MiniMaxM3VLForConditionalGeneration.__init__ for text-only configs.
+    model.vision_tower = None
+    return model
+
+
+def test_load_weights_accepts_base_mapper_without_params_map() -> None:
+    model = _make_layerless_model()
     mapper = HfWeightMapper()
     source_name = "model.layers.3.block_sparse_moe.e_score_correction_bias"
     target_name = "model.layers.3.block_sparse_moe.gate.e_score_correction_bias"
@@ -201,3 +208,41 @@ def test_load_weights_accepts_base_mapper_without_params_map() -> None:
     assert call_kwargs["weight_mapper"] is mapper
     renamed = mapper.rename_by_params_map(call_kwargs["params_map"], {source_name: bias})
     assert renamed[target_name] is bias
+
+
+@pytest.mark.parametrize(
+    ("model_class", "checkpoint_prefix"),
+    [
+        (MiniMaxM3ForCausalLM, ""),
+        (MiniMaxM3VLForConditionalGeneration, "language_model."),
+    ],
+    ids=["text", "vl"],
+)
+def test_load_weights_converts_compressed_tensors_nvfp4_experts(
+    model_class: type, checkpoint_prefix: str
+) -> None:
+    model = _make_layerless_model(model_class)
+    expert = "model.layers.3.block_sparse_moe.experts.0.w1"
+    packed = torch.tensor([[1, 2]], dtype=torch.uint8)
+    block_scale = torch.ones(1, 1).to(torch.float8_e4m3fn)
+    weights = {
+        f"{checkpoint_prefix}{expert}.weight_packed": packed,
+        f"{checkpoint_prefix}{expert}.weight_scale": block_scale,
+        f"{checkpoint_prefix}{expert}.weight_global_scale": torch.tensor([4.0]),
+        f"{checkpoint_prefix}{expert}.input_global_scale": torch.tensor([0.0]),
+    }
+
+    with patch.object(DecoderModelForCausalLM, "load_weights") as base_load_weights:
+        model.load_weights(weights, weight_mapper=HfWeightMapper())
+
+    loaded = base_load_weights.call_args.kwargs["weights"]
+    assert set(loaded.keys()) == {
+        f"{expert}.weight",
+        f"{expert}.weight_scale",
+        f"{expert}.weight_scale_2",
+        f"{expert}.input_scale",
+    }
+    assert loaded[f"{expert}.weight"] is packed
+    assert loaded[f"{expert}.weight_scale"] is block_scale
+    torch.testing.assert_close(loaded[f"{expert}.weight_scale_2"], torch.tensor([0.25]))
+    torch.testing.assert_close(loaded[f"{expert}.input_scale"], torch.tensor([0.0]))