From 3440b3af882cd197fa7f11d20cca18e5463ffa9c Mon Sep 17 00:00:00 2001 From: Yifan Jiang <19356972+yifjiang@users.noreply.github.com> Date: Fri, 16 Oct 2026 13:59:07 -0700 Subject: [PATCH 1/3] [None][fix] Load compressed-tensors NVFP4 weights in MiniMax-M3 MiniMaxM3ForCausalLM.load_weights hands checkpoint tensors to the generic loader under their original names. Checkpoints quantized to NVFP4 with llm-compressor's compressed-tensors format name the packed weights `weight_packed` and store `weight_global_scale` and `input_global_scale` in multiplier form, so the NVFP4 loaders find no `weight`. The fused-MoE loader, for example, finds no `.w1.weight` and fails `assert w1_weight is not None and w3_weight is not None`. LagunaHfWeightMapper already translates these names. Move that translation into a shared helper, convert_compressed_tensors_nvfp4_weights in checkpoints/hf/weight_mapper.py: weight_packed -> weight, and weight_global_scale / input_global_scale -> weight_scale_2 / input_scale as reciprocals, with 0 kept as 0 for uncalibrated experts. The helper returns its input unchanged when no such names are present and keeps a ConsumableWeightsDict a ConsumableWeightsDict. Laguna calls the helper behind its existing NVFP4 gate. MiniMax-M3 calls it at the top of load_weights, before the five-way projection packer and the generic loader read the dict. It is keyed on tensor names rather than on the global quant config, which is MIXED_PRECISION when NVFP4 covers only some modules. The VL class reaches it through the text loader after the language_model. prefix is stripped. Signed-off-by: Yifan Jiang <19356972+yifjiang@users.noreply.github.com> Co-Authored-By: Claude Opus 5.5 --- .../models/checkpoints/hf/weight_mapper.py | 66 ++++++++++++++++++ tensorrt_llm/_torch/models/modeling_laguna.py | 55 +-------------- .../_torch/models/modeling_minimaxm3.py | 5 ++ .../hf/test_minimaxm3_weight_mapper.py | 49 ++++++++++++- .../checkpoints/hf/test_weight_mapper.py | 69 ++++++++++++++++++- 5 files changed, 188 insertions(+), 56 deletions(-) diff --git a/tensorrt_llm/_torch/models/checkpoints/hf/weight_mapper.py b/tensorrt_llm/_torch/models/checkpoints/hf/weight_mapper.py index c180cb455fb5..0897474e2422 100644 --- a/tensorrt_llm/_torch/models/checkpoints/hf/weight_mapper.py +++ b/tensorrt_llm/_torch/models/checkpoints/hf/weight_mapper.py @@ -13,14 +13,80 @@ # 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 ..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 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``). + """ + if not any( + name.endswith(_COMPRESSED_TENSORS_NVFP4_SUFFIXES) + for name in weights.keys()): + return weights + + converted = {} + for name, value in weights.items(): + if name.endswith(".weight_packed"): + name = name[:-len(".weight_packed")] + ".weight" + elif name.endswith(".weight_global_scale"): + name = name[:-len(".weight_global_scale")] + ".weight_scale_2" + value = _invert_compressed_tensors_global_scale(value) + elif name.endswith(".input_global_scale"): + name = name[:-len(".input_global_scale")] + ".input_scale" + value = _invert_compressed_tensors_global_scale(value) + converted[name] = value + 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..6b33bc12d028 100644 --- a/tests/unittest/_torch/models/checkpoints/hf/test_weight_mapper.py +++ b/tests/unittest/_torch/models/checkpoints/hf/test_weight_mapper.py @@ -16,7 +16,11 @@ 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 +36,66 @@ 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_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 + + +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 From f0849e2e6028654b6f9f391d1ee45a43ed87fdbd Mon Sep 17 00:00:00 2001 From: Yifan Jiang <19356972+yifjiang@users.noreply.github.com> Date: Sat, 24 Oct 2026 18:20:49 -0700 Subject: [PATCH 2/3] [None][fix] Warn once about compressed-tensors NVFP4 global scales set to 0 Review follow-up. A global scale that is 0, negative or non-finite becomes 0 when converted. 0 is intended for experts without calibration data, but a negative or NaN scale means a broken checkpoint, and a weight global scale of 0 zeroes that expert's output. Nothing showed either case. The helper now counts them and logs one summary warning with a few example tensor names. Signed-off-by: Yifan Jiang <19356972+yifjiang@users.noreply.github.com> Co-Authored-By: Claude Opus 5.5 --- .../models/checkpoints/hf/weight_mapper.py | 47 +++++++++++++++++-- .../checkpoints/hf/test_weight_mapper.py | 40 ++++++++++++++++ 2 files changed, 82 insertions(+), 5 deletions(-) diff --git a/tensorrt_llm/_torch/models/checkpoints/hf/weight_mapper.py b/tensorrt_llm/_torch/models/checkpoints/hf/weight_mapper.py index 0897474e2422..df760bda04f8 100644 --- a/tensorrt_llm/_torch/models/checkpoints/hf/weight_mapper.py +++ b/tensorrt_llm/_torch/models/checkpoints/hf/weight_mapper.py @@ -19,6 +19,7 @@ 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 @@ -43,6 +44,38 @@ def _invert_compressed_tensors_global_scale(value) -> torch.Tensor: 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. @@ -75,16 +108,20 @@ def convert_compressed_tensors_nvfp4_weights( return weights converted = {} + zero_scales, invalid_scales = [], [] for name, value in weights.items(): if name.endswith(".weight_packed"): name = name[:-len(".weight_packed")] + ".weight" - elif name.endswith(".weight_global_scale"): - name = name[:-len(".weight_global_scale")] + ".weight_scale_2" - value = _invert_compressed_tensors_global_scale(value) - elif name.endswith(".input_global_scale"): - name = name[:-len(".input_global_scale")] + ".input_scale" + 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) converted[name] = value + _warn_uninvertible_global_scales(zero_scales, invalid_scales) return ConsumableWeightsDict.take_ownership(weights, converted) 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 6b33bc12d028..7ba4481412e1 100644 --- a/tests/unittest/_torch/models/checkpoints/hf/test_weight_mapper.py +++ b/tests/unittest/_torch/models/checkpoints/hf/test_weight_mapper.py @@ -13,6 +13,8 @@ # See the License for the specific language governing permissions and # limitations under the License. +from unittest.mock import Mock + import pytest import torch @@ -78,6 +80,44 @@ def test_compressed_tensors_nvfp4_names_and_global_scales_are_converted() -> Non 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 = { From 68a7207e095ee4644022caaf23fbf0ca72c2f436 Mon Sep 17 00:00:00 2001 From: Yifan Jiang <19356972+yifjiang@users.noreply.github.com> Date: Tue, 27 Oct 2026 12:15:54 -0700 Subject: [PATCH 3/3] [None][fix] Reject compressed-tensors NVFP4 names that convert to the same key convert_compressed_tensors_nvfp4_weights renames weight_packed to weight and the global scales to weight_scale_2 / input_scale. A checkpoint holding both x.weight_packed and x.weight (or both a global scale and its converted name) silently kept whichever came last, which could then disagree with its scales. Raise a ValueError naming both source tensors instead, and test both cases. Co-Authored-By: Claude Opus 5.5 Signed-off-by: Yifan Jiang <19356972+yifjiang@users.noreply.github.com> --- .../models/checkpoints/hf/weight_mapper.py | 13 ++++++++++++- .../checkpoints/hf/test_weight_mapper.py | 18 ++++++++++++++++++ 2 files changed, 30 insertions(+), 1 deletion(-) diff --git a/tensorrt_llm/_torch/models/checkpoints/hf/weight_mapper.py b/tensorrt_llm/_torch/models/checkpoints/hf/weight_mapper.py index df760bda04f8..9b4f486159f5 100644 --- a/tensorrt_llm/_torch/models/checkpoints/hf/weight_mapper.py +++ b/tensorrt_llm/_torch/models/checkpoints/hf/weight_mapper.py @@ -101,15 +101,21 @@ def convert_compressed_tensors_nvfp4_weights( 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 = {} + 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")): @@ -120,7 +126,12 @@ def convert_compressed_tensors_nvfp4_weights( 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) 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 7ba4481412e1..08648f49317f 100644 --- a/tests/unittest/_torch/models/checkpoints/hf/test_weight_mapper.py +++ b/tests/unittest/_torch/models/checkpoints/hf/test_weight_mapper.py @@ -129,6 +129,24 @@ def test_compressed_tensors_nvfp4_conversion_returns_other_checkpoints_unchanged 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})