diff --git a/tensorrt_llm/_torch/models/checkpoints/hf/weight_mapper.py b/tensorrt_llm/_torch/models/checkpoints/hf/weight_mapper.py index c180cb455fb5..9b4f486159f5 100644 --- a/tensorrt_llm/_torch/models/checkpoints/hf/weight_mapper.py +++ b/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") diff --git a/tensorrt_llm/_torch/models/modeling_laguna.py b/tensorrt_llm/_torch/models/modeling_laguna.py index 189f95c10573..1aa9420b59c9 100644 --- a/tensorrt_llm/_torch/models/modeling_laguna.py +++ b/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 diff --git a/tensorrt_llm/_torch/models/modeling_minimaxm3.py b/tensorrt_llm/_torch/models/modeling_minimaxm3.py index 311600ab67ca..1f27171184fa 100644 --- a/tensorrt_llm/_torch/models/modeling_minimaxm3.py +++ b/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) diff --git a/tests/unittest/_torch/models/checkpoints/hf/test_minimaxm3_weight_mapper.py b/tests/unittest/_torch/models/checkpoints/hf/test_minimaxm3_weight_mapper.py index 160bb485d5f8..0484c139fbee 100644 --- a/tests/unittest/_torch/models/checkpoints/hf/test_minimaxm3_weight_mapper.py +++ b/tests/unittest/_torch/models/checkpoints/hf/test_minimaxm3_weight_mapper.py @@ -178,10 +178,10 @@ 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. @@ -189,6 +189,13 @@ def test_load_weights_accepts_base_mapper_without_params_map() -> None: 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])) diff --git a/tests/unittest/_torch/models/checkpoints/hf/test_weight_mapper.py b/tests/unittest/_torch/models/checkpoints/hf/test_weight_mapper.py index c067e7722c68..08648f49317f 100644 --- a/tests/unittest/_torch/models/checkpoints/hf/test_weight_mapper.py +++ b/tests/unittest/_torch/models/checkpoints/hf/test_weight_mapper.py @@ -13,10 +13,16 @@ # See the License for the specific language governing permissions and # limitations under the License. +from unittest.mock import Mock + import pytest import torch -from tensorrt_llm._torch.models.checkpoints.hf.weight_mapper import HfWeightMapper +from tensorrt_llm._torch.models.checkpoints.base_weight_loader import ConsumableWeightsDict +from tensorrt_llm._torch.models.checkpoints.hf.weight_mapper import ( + HfWeightMapper, + convert_compressed_tensors_nvfp4_weights, +) from tensorrt_llm._torch.models.modeling_utils import duplicate_kv_weight @@ -32,3 +38,122 @@ def test_duplicate_kv_bias_preserves_head_boundaries(duplicate_kv) -> None: expected_bias = torch.tensor([10, 11, 10, 11, 20, 21, 20, 21]) torch.testing.assert_close(duplicated_bias, expected_bias) + + +class _LazySlice: + """Stands in for a lazily read safetensors slice, which ``[...]`` loads.""" + + def __init__(self, tensor: torch.Tensor) -> None: + self._tensor = tensor + + def __getitem__(self, index): + return self._tensor[index] + + +def test_compressed_tensors_nvfp4_names_and_global_scales_are_converted() -> None: + prefix = "model.layers.0.mlp.experts.0.w1" + packed = torch.tensor([[1, 2]], dtype=torch.uint8) + block_scale = torch.ones(1, 1).to(torch.float8_e4m3fn) + norm = torch.tensor([3.0]) + weights = { + f"{prefix}.weight_packed": packed, + f"{prefix}.weight_scale": block_scale, + # One entry is 0, as stored for an expert that saw no calibration data. + f"{prefix}.weight_global_scale": torch.tensor([4.0, 0.0]), + f"{prefix}.input_global_scale": _LazySlice(torch.tensor([8.0])), + "model.norm.weight": norm, + } + + converted = convert_compressed_tensors_nvfp4_weights(weights) + + assert set(converted.keys()) == { + f"{prefix}.weight", + f"{prefix}.weight_scale", + f"{prefix}.weight_scale_2", + f"{prefix}.input_scale", + "model.norm.weight", + } + assert converted[f"{prefix}.weight"] is packed + assert converted[f"{prefix}.weight_scale"] is block_scale + assert converted["model.norm.weight"] is norm + torch.testing.assert_close(converted[f"{prefix}.weight_scale_2"], torch.tensor([0.25, 0.0])) + torch.testing.assert_close(converted[f"{prefix}.input_scale"], torch.tensor([0.125])) + + +def test_compressed_tensors_nvfp4_uninvertible_global_scales_log_one_warning(monkeypatch) -> None: + import tensorrt_llm._torch.models.checkpoints.hf.weight_mapper as weight_mapper + + log = Mock() + monkeypatch.setattr(weight_mapper, "logger", log) + weights = { + "model.layers.0.mlp.experts.0.w1.weight_global_scale": torch.tensor([4.0]), + "model.layers.0.mlp.experts.1.w1.weight_global_scale": torch.tensor([0.0]), + "model.layers.0.mlp.experts.2.w1.weight_global_scale": torch.tensor([-1.0]), + "model.layers.0.mlp.experts.3.w1.input_global_scale": _LazySlice( + torch.tensor([float("nan")]) + ), + } + + converted = convert_compressed_tensors_nvfp4_weights(weights) + + for i in (1, 2): + torch.testing.assert_close( + converted[f"model.layers.0.mlp.experts.{i}.w1.weight_scale_2"], torch.tensor([0.0]) + ) + torch.testing.assert_close( + converted["model.layers.0.mlp.experts.3.w1.input_scale"], torch.tensor([0.0]) + ) + assert log.warning.call_count == 1 + message = log.warning.call_args[0][0] + assert "2 with negative or non-finite values" in message + assert "experts.2.w1.weight_global_scale" in message + assert "experts.3.w1.input_global_scale" in message + assert "1 with 0" in message + assert "experts.1.w1.weight_global_scale" in message + + log.reset_mock() + convert_compressed_tensors_nvfp4_weights( + {"model.layers.0.mlp.experts.0.w1.weight_global_scale": torch.tensor([4.0])} + ) + log.warning.assert_not_called() + + +def test_compressed_tensors_nvfp4_conversion_returns_other_checkpoints_unchanged() -> None: + prefix = "model.layers.0.mlp.experts.0.w1" + weights = { + f"{prefix}.weight": torch.tensor([[1, 2]], dtype=torch.uint8), + f"{prefix}.weight_scale_2": torch.tensor([0.25]), + f"{prefix}.input_scale": torch.tensor([0.125]), + } + + assert convert_compressed_tensors_nvfp4_weights(weights) is weights + + +@pytest.mark.parametrize( + "names", + [ + ("model.layers.0.self_attn.o_proj.weight_packed", "model.layers.0.self_attn.o_proj.weight"), + ( + "model.layers.0.mlp.experts.0.w1.weight_global_scale", + "model.layers.0.mlp.experts.0.w1.weight_scale_2", + ), + ], + ids=["weight", "weight_scale_2"], +) +def test_compressed_tensors_nvfp4_conversion_rejects_colliding_names(names) -> None: + weights = {name: torch.tensor([4.0]) for name in names} + + with pytest.raises(ValueError, match="maps both"): + convert_compressed_tensors_nvfp4_weights(weights) + + +def test_compressed_tensors_nvfp4_conversion_keeps_consumable_weights() -> None: + packed = torch.tensor([[1, 2]], dtype=torch.uint8) + source = ConsumableWeightsDict({"model.layers.0.self_attn.o_proj.weight_packed": packed}) + + converted = convert_compressed_tensors_nvfp4_weights(source) + + assert isinstance(converted, ConsumableWeightsDict) + assert converted["model.layers.0.self_attn.o_proj.weight"] is packed + # The converted dict owns the tensors, so consuming it can release them. + assert len(source) == 0