From ffb70be757c58bf6dda894b08af267d5b0d0c151 Mon Sep 17 00:00:00 2001 From: Mihai Chiorean Date: Fri, 14 Aug 2026 14:42:11 -0700 Subject: [PATCH 1/2] [#17723][feat] Parse multiple compressed-tensors config groups Signed-off-by: Mihai Chiorean --- tensorrt_llm/_torch/model_config.py | 56 ++- tensorrt_llm/models/quant_config_utils.py | 332 ++++++++++++++---- .../models/test_quant_config_utils.py | 324 ++++++++++++++++- 3 files changed, 630 insertions(+), 82 deletions(-) diff --git a/tensorrt_llm/_torch/model_config.py b/tensorrt_llm/_torch/model_config.py index c1f10ee89158..fc351194202c 100644 --- a/tensorrt_llm/_torch/model_config.py +++ b/tensorrt_llm/_torch/model_config.py @@ -590,8 +590,24 @@ def load_hf_quant_config(hf_quant_config, moe_backend, checkpoint_dir=None): # NOTE: This is for llm-compressor's quantized checkpoints. elif hf_quant_config.get("quant_method") == "compressed-tensors": - update_quant_config_from_compressed_tensors(quant_config, - hf_quant_config) + # Multi-group ("mixed-precision") checkpoints resolve to + # MIXED_PRECISION plus one QuantConfig per quantized module, which + # needs the checkpoint's module names. Single-group checkpoints do + # not, so the tensor index is only read when there is more than one + # group. + module_names = None + if len(hf_quant_config.get("config_groups") or {}) > 1: + module_names = ModelConfig._read_checkpoint_module_names( + checkpoint_dir) + layer_quant_config = update_quant_config_from_compressed_tensors( + quant_config, hf_quant_config, module_names) + if (quant_config.quant_algo == QuantAlgo.MIXED_PRECISION + and layer_quant_config is None): + raise ValueError( + "compressed-tensors checkpoint quantizes modules " + "differently per config group, but its per-layer quant " + "config could not be resolved because no safetensors " + f"tensor index was found under {checkpoint_dir}.") elif hf_quant_config.get("quant_method") == "nvfp4": quant_config.quant_algo = QuantAlgo.NVFP4 group_size = hf_quant_config.get("group_size", 16) @@ -613,6 +629,42 @@ def _read_safetensors_header(path: Path) -> Dict[str, Any]: header_size = struct.unpack(" Optional[List[str]]: + """Read the module names, in checkpoint (HF) namespace, of every stored tensor. + + compressed-tensors ``config_groups`` select modules with regexes, so + resolving a multi-group (mixed-precision) checkpoint into per-layer + quant configs needs the checkpoint's module list. + + Args: + checkpoint_dir: Local checkpoint directory, may be None. + + Returns: + Deduplicated module names, or None when no safetensors tensor + index could be read. + """ + if checkpoint_dir is None: + return None + + checkpoint_path = Path(checkpoint_dir) + index_path = checkpoint_path / "model.safetensors.index.json" + tensor_names = [] + if index_path.exists(): + with open(index_path) as f: + tensor_names = list(json.load(f).get("weight_map", {})) + else: + for shard in sorted(checkpoint_path.glob("*.safetensors")): + tensor_names.extend(ModelConfig._read_safetensors_header(shard)) + + # Drop the parameter name ("weight", "weight_packed", ...) and the + # "__metadata__" header entry, neither of which is a module. + module_names = [ + name.rsplit('.', 1)[0] for name in tensor_names if '.' in name + ] + return list(dict.fromkeys(module_names)) or None + @staticmethod def _get_safetensors_header_for_tensor(checkpoint_dir: str, tensor_name: str) -> Optional[Dict]: diff --git a/tensorrt_llm/models/quant_config_utils.py b/tensorrt_llm/models/quant_config_utils.py index 95d7d9955f37..b56de9c95c68 100644 --- a/tensorrt_llm/models/quant_config_utils.py +++ b/tensorrt_llm/models/quant_config_utils.py @@ -13,54 +13,110 @@ # See the License for the specific language governing permissions and # limitations under the License. -from typing import Any, Mapping +import re +from collections.abc import Iterable, Mapping +from typing import Any +from tensorrt_llm.logger import logger from tensorrt_llm.models.modeling_utils import QuantConfig from tensorrt_llm.quantization.mode import QuantAlgo +# compressed-tensors marks a ``targets``/``ignore`` entry as a regex with this +# prefix; every other entry is an exact module-name (or class-name) match. +_TARGET_REGEX_PREFIX = "re:" -def update_quant_config_from_compressed_tensors( - quant_config: QuantConfig, hf_quant_config: Mapping[str, Any] -) -> None: - """Mutate QuantConfig from an llm-compressor compressed-tensors config.""" - config_groups = hf_quant_config.get("config_groups") - if config_groups is None: - raise ValueError(f"config_groups is not set in {hf_quant_config}.") - # compressed-tensors keys config_groups by group name: custom recipes use - # "group_0"; named preset schemes (e.g. FP8_DYNAMIC) use the scheme name. - # This path applies a single algo globally, so resolve to one group. - group_config = config_groups.get("group_0") - if group_config is None: - if len(config_groups) != 1: - raise ValueError( - f"Expected 'group_0' or exactly one config group, got {sorted(config_groups)}." - ) - group_config = next(iter(config_groups.values())) +def _matches_target(module_name: str, target: str) -> bool: + """Whether ``module_name`` matches a compressed-tensors target/ignore entry. + + Mirrors ``compressed_tensors.utils.match.match_name``: ``re:``-prefixed + entries are matched with ``re.match`` (anchored at the start only, *not* + ``re.fullmatch``), everything else must equal the module name. Class-name + targets (e.g. ``"Linear"``) are not resolvable here because they need the + instantiated module tree, which does not exist at config-parse time. + + Args: + module_name: Module name in the checkpoint (HF) namespace. + target: One ``targets`` or ``ignore`` entry. + + Returns: + True if the entry matches the module name. + """ + if target.startswith(_TARGET_REGEX_PREFIX): + return re.match(target.removeprefix(_TARGET_REGEX_PREFIX), module_name) is not None + return target == module_name + + +def _ordered_targets(targets: Iterable[str]) -> list[str]: + r"""Order ``config_groups`` targets by compressed-tensors match precedence. + + Mirrors ``compressed_tensors.utils.match.match_targets``: exact-name + targets sort before ``re:`` targets, each block in ascending lexicographic + order. ``_scheme_from_targets`` then takes the *first* matching target, so + this ordering -- not pattern specificity -- decides which ``config_group`` + owns a module that several groups target. + + Qwen3.8-27B-NVFP4 depends on exactly this: its FP8 group targets + ``re:.*layers\.(56|...|63)\.mlp\.(gate|up|down)_proj$`` while its NVFP4 + group targets ``re:.*mlp\.(gate|up|down)_proj$``. Both match the last + eight MLP blocks, and ``layers`` sorts before ``mlp``, so those blocks are + FP8. + + Args: + targets: Target entries collected across all ``config_groups``. + + Returns: + The targets, most-specific-first per compressed-tensors' ordering. + """ + return sorted(targets, key=lambda target: (_TARGET_REGEX_PREFIX in target, target)) + + +def _exclude_modules_from_hf_config(hf_quant_config: Mapping[str, Any]) -> list[str]: + """Collect the modules the producer left unquantized. + + Args: + hf_quant_config: ``quantization_config`` from the checkpoint's config. + + Returns: + ``modules_to_not_convert`` merged with ``ignore``, order-preserving and + deduplicated. + """ + hf_exclude_modules = hf_quant_config.get("modules_to_not_convert", None) + ignore = list(hf_quant_config.get("ignore", [])) + if hf_exclude_modules is None: + return ignore + return list(dict.fromkeys(list(hf_exclude_modules) + ignore)) + + +def _quant_config_from_config_group( + group_config: Mapping[str, Any], group_format: str | None +) -> QuantConfig: + """Resolve one compressed-tensors ``config_group`` to a TRT-LLM QuantConfig. + + Only ``quant_algo`` and ``group_size`` are set; the KV-cache algorithm and + the exclude-module list are global properties of the checkpoint and stay + with the caller. + + Args: + group_config: One entry of the checkpoint's ``config_groups``. + group_format: The group's ``format``, falling back to the checkpoint's + top-level ``format`` when the group does not declare one. + + Returns: + A QuantConfig carrying the group's ``quant_algo`` and ``group_size``. + + Raises: + ValueError: The group uses a bit width, strategy or group size that + TRT-LLM does not support. + """ weights_quant_config = group_config["weights"] weights_quant_strategy = weights_quant_config["strategy"] - # kv_cache_scheme (llm-compressor): FP8 per-tensor KV cache. Handled - # before the weight-algo branches so recipes that early-return (MXFP4 - # pack-quantized) still pick up the KV-cache quantization. - kv_cache_scheme = hf_quant_config.get("kv_cache_scheme") - if kv_cache_scheme is not None: - if kv_cache_scheme.get("num_bits") == 8 and kv_cache_scheme.get("type") == "float": - if quant_config.kv_cache_quant_algo in (None, QuantAlgo.FP8): - quant_config.kv_cache_quant_algo = QuantAlgo.FP8 - else: - raise ValueError( - f"Specified kv_cache_quant_algo={quant_config.kv_cache_quant_algo}, " - "conflicting with FP8 KV cache from HF quant config." - ) - else: - raise ValueError(f"Unsupported kv_cache_scheme: {kv_cache_scheme}.") - # MXFP4 pack-quantized (weight-only): FP4 E2M1 weights packed two per # uint8 with per-32-group uint8 E8M0 scales and no activation # quantization (e.g. Kimi K3 routed experts). Handled before reading the # input-activation strategy, which is null for weight-only recipes. - if hf_quant_config.get("format") == "mxfp4-pack-quantized" or ( + if group_format == "mxfp4-pack-quantized" or ( weights_quant_config["num_bits"] == 4 and weights_quant_config.get("type") == "float" and weights_quant_strategy == "group" @@ -69,13 +125,7 @@ def update_quant_config_from_compressed_tensors( group_size = weights_quant_config["group_size"] if group_size != 32: raise ValueError(f"Unsupported group_size: {group_size}. Supported: 32 for MXFP4.") - quant_config.quant_algo = QuantAlgo.W4A16_MXFP4 - quant_config.group_size = group_size - hf_exclude_modules = hf_quant_config.get("modules_to_not_convert", None) - quant_config.exclude_modules = list( - set((hf_exclude_modules or []) + hf_quant_config.get("ignore", [])) - ) - return + return QuantConfig(quant_algo=QuantAlgo.W4A16_MXFP4, group_size=group_size) inputs_quant_config = group_config["input_activations"] inputs_quant_strategy = inputs_quant_config["strategy"] @@ -84,30 +134,30 @@ def update_quant_config_from_compressed_tensors( if weights_quant_strategy == "channel": if inputs_quant_strategy != "token": raise ValueError(f"Unsupported inputs_quant_strategy: {inputs_quant_strategy}.") - quant_config.quant_algo = QuantAlgo.FP8_PER_CHANNEL_PER_TOKEN - elif weights_quant_strategy == "block": + return QuantConfig(quant_algo=QuantAlgo.FP8_PER_CHANNEL_PER_TOKEN) + if weights_quant_strategy == "block": if inputs_quant_strategy != "group": raise ValueError(f"Unsupported inputs_quant_strategy: {inputs_quant_strategy}.") - quant_config.quant_algo = QuantAlgo.FP8_BLOCK_SCALES group_size = inputs_quant_config["group_size"] # TRT-LLM only supports group_size=128 for FP8_BLOCK_SCALES. if group_size != 128: raise ValueError(f"Unsupported group_size: {group_size}. Supported: 128.") - quant_config.group_size = group_size + return QuantConfig(quant_algo=QuantAlgo.FP8_BLOCK_SCALES, group_size=group_size) - else: - raise ValueError( - f"Unsupported weights_quant_strategy: {weights_quant_strategy}. " - "Supported strategies: 'channel', 'block'." - ) - elif ( + raise ValueError( + f"Unsupported weights_quant_strategy: {weights_quant_strategy}. " + "Supported strategies: 'channel', 'block'." + ) + + if ( weights_quant_config["num_bits"] == 4 and weights_quant_config.get("type") == "float" and weights_quant_strategy == "tensor_group" ): - # llm-compressor NVFP4: weights FP4 with FP8 per-group scales - # (group_size=16), scaled by an FP32 global scale. + # llm-compressor NVFP4 quantizes both weights and activations. Model + # adapters must normalize module names independently of the algorithm; + # SM120/121 execute this W4A4 contract through the FP4 CUTLASS path. if inputs_quant_strategy != "tensor_group": raise ValueError( f"Unsupported inputs_quant_strategy for NVFP4: {inputs_quant_strategy}." @@ -115,18 +165,172 @@ def update_quant_config_from_compressed_tensors( group_size = weights_quant_config["group_size"] if group_size != 16: raise ValueError(f"Unsupported group_size: {group_size}. Supported: 16 for NVFP4.") - quant_config.quant_algo = QuantAlgo.NVFP4 - quant_config.group_size = group_size - else: + return QuantConfig(quant_algo=QuantAlgo.NVFP4, group_size=group_size) + + raise ValueError( + f"Unsupported quant_bits: {weights_quant_config['num_bits']}. " + "Supported: 8 (FP8) or 4 (NVFP4)." + ) + + +def _build_layer_quant_configs( + hf_quant_config: Mapping[str, Any], + module_names: Iterable[str], + kv_cache_quant_algo: QuantAlgo | None, +) -> dict[str, QuantConfig]: + """Resolve a multi-group checkpoint to one QuantConfig per quantized module. + + Args: + hf_quant_config: ``quantization_config`` from the checkpoint's config. + module_names: Module names (HF namespace) present in the checkpoint. + kv_cache_quant_algo: Global KV-cache algorithm, copied onto every + per-module config the way the modelopt MIXED_PRECISION path does. + + Returns: + QuantConfigs keyed by module name, covering exactly the modules the + producer quantized. Modules absent from the mapping inherit the global + ``MIXED_PRECISION`` config, i.e. they stay unquantized. + + Raises: + ValueError: A ``config_group`` has no ``targets``, uses a scheme + TRT-LLM does not support, or a target matched no checkpoint module. + """ + top_level_format = hf_quant_config.get("format") + target_quant_configs: dict[str, QuantConfig] = {} + for group_name, group_config in hf_quant_config["config_groups"].items(): + targets = group_config.get("targets") + if not targets: + raise ValueError( + f"config_group '{group_name}' has no 'targets'. Targets are required to " + "resolve a compressed-tensors checkpoint with multiple config groups." + ) + group_quant_config = _quant_config_from_config_group( + group_config, group_config.get("format", top_level_format) + ) + group_quant_config.kv_cache_quant_algo = kv_cache_quant_algo + for target in targets: + if target in target_quant_configs: + logger.warning( + f"compressed-tensors target '{target}' appears in more than one " + "config group; the last group takes precedence." + ) + target_quant_configs[target] = group_quant_config + + ignore = hf_quant_config.get("ignore", []) + ordered_targets = _ordered_targets(target_quant_configs) + + layer_quant_configs: dict[str, QuantConfig] = {} + matched_targets: set[str] = set() + for module_name in module_names: + matching_targets = [ + target for target in ordered_targets if _matches_target(module_name, target) + ] + matched_targets.update(matching_targets) + if any(_matches_target(module_name, entry) for entry in ignore): + continue + if matching_targets: + layer_quant_configs[module_name] = target_quant_configs[ + matching_targets[0] + ].model_copy() + + unmatched_targets = sorted(set(ordered_targets) - matched_targets) + if unmatched_targets: raise ValueError( - f"Unsupported quant_bits: {weights_quant_config['num_bits']}. " - "Supported: 8 (FP8) or 4 (NVFP4)." + f"config_groups targets matched no checkpoint module: {unmatched_targets}. " + "compressed-tensors targets that select modules by class name are not supported." ) + return layer_quant_configs - hf_exclude_modules = hf_quant_config.get("modules_to_not_convert", None) - if hf_exclude_modules is not None: + +def update_quant_config_from_compressed_tensors( + quant_config: QuantConfig, + hf_quant_config: Mapping[str, Any], + module_names: Iterable[str] | None = None, +) -> dict[str, QuantConfig] | None: + """Mutate QuantConfig from an llm-compressor compressed-tensors config. + + A checkpoint with a single ``config_group`` maps onto one global + ``quant_algo``. A checkpoint with several groups (``"format": + "mixed-precision"``) quantizes different modules differently and maps onto + ``QuantAlgo.MIXED_PRECISION`` plus one QuantConfig per quantized module. + Qwen3.8-27B-NVFP4 is the motivating example: ``group_0`` is FP8 + (attention, GDN projections, ``lm_head`` and the last eight MLP blocks) and + ``group_1`` is NVFP4 (every other MLP block). + + Args: + quant_config: Mutated in place with the checkpoint's global settings. + hf_quant_config: ``quantization_config`` from the checkpoint's config. + module_names: Module names (HF namespace) present in the checkpoint. + Multi-group checkpoints need them to resolve ``config_groups`` + targets; single-group checkpoints ignore them. + + Returns: + Per-module QuantConfigs keyed by module name for a multi-group + checkpoint, or ``None`` when the checkpoint has a single group or + ``module_names`` was not supplied. + + Raises: + ValueError: ``config_groups`` is missing or empty, the KV-cache scheme + is unsupported or conflicts with an explicit request, or a group + uses a scheme TRT-LLM does not support. + """ + config_groups = hf_quant_config.get("config_groups") + if config_groups is None: + raise ValueError(f"config_groups is not set in {hf_quant_config}.") + if not config_groups: + raise ValueError(f"config_groups is empty in {hf_quant_config}.") + + # kv_cache_scheme (llm-compressor): FP8 per-tensor KV cache. Handled + # before the weight-algo branches so recipes that bail out early (MXFP4 + # pack-quantized) still pick up the KV-cache quantization. + kv_cache_scheme = hf_quant_config.get("kv_cache_scheme") + if kv_cache_scheme is not None: + if kv_cache_scheme.get("num_bits") == 8 and kv_cache_scheme.get("type") == "float": + if quant_config.kv_cache_quant_algo in (None, QuantAlgo.FP8): + quant_config.kv_cache_quant_algo = QuantAlgo.FP8 + else: + raise ValueError( + f"Specified kv_cache_quant_algo={quant_config.kv_cache_quant_algo}, " + "conflicting with FP8 KV cache from HF quant config." + ) + else: + raise ValueError(f"Unsupported kv_cache_scheme: {kv_cache_scheme}.") + + exclude_modules = _exclude_modules_from_hf_config(hf_quant_config) + + if len(config_groups) > 1: + quant_config.quant_algo = QuantAlgo.MIXED_PRECISION + # The per-layer map is authoritative for a mixed checkpoint. Producer + # ``ignore`` entries are non-recursive and have already been applied + # while building that map; copying them into TRT-LLM's recursive + # exclude list could incorrectly shadow quantized children. quant_config.exclude_modules = list( - set(hf_exclude_modules + hf_quant_config.get("ignore", [])) + hf_quant_config.get("modules_to_not_convert", None) or [] + ) + if module_names is None: + # The global algo is fully determined without the module list, but + # the per-module configs are not. Callers that only need + # ``quant_config`` (e.g. reporting the checkpoint's quantization + # back through LlmArgs) are fine; a caller that builds the model + # from this config must pass ``module_names``. + logger.debug( + f"compressed-tensors checkpoint has multiple config groups " + f"({sorted(config_groups)}) but no checkpoint module names were supplied; " + "per-layer quantization configs were not resolved." + ) + return None + layer_quant_configs = _build_layer_quant_configs( + hf_quant_config, module_names, quant_config.kv_cache_quant_algo ) - else: - quant_config.exclude_modules = hf_quant_config.get("ignore", []) + return layer_quant_configs + + # compressed-tensors keys config_groups by group name: custom recipes use + # "group_0"; named preset schemes (e.g. FP8_DYNAMIC) use the scheme name. + group_config = next(iter(config_groups.values())) + group_quant_config = _quant_config_from_config_group( + group_config, group_config.get("format", hf_quant_config.get("format")) + ) + quant_config.quant_algo = group_quant_config.quant_algo + quant_config.group_size = group_quant_config.group_size + quant_config.exclude_modules = exclude_modules + return None diff --git a/tests/unittest/models/test_quant_config_utils.py b/tests/unittest/models/test_quant_config_utils.py index d10c6da77709..13acb2830044 100644 --- a/tests/unittest/models/test_quant_config_utils.py +++ b/tests/unittest/models/test_quant_config_utils.py @@ -135,24 +135,18 @@ def test_update_quant_config_from_compressed_tensors_parses_scheme_named_group() assert quant_config.exclude_modules == ["lm_head"] -def test_update_quant_config_from_compressed_tensors_rejects_ambiguous_groups(): - # Multiple groups, none named "group_0" -> ambiguous which to apply globally. - config = _compressed_tensors_config( - weights={ - "num_bits": 8, - "strategy": "channel", - }, - input_activations={ - "num_bits": 8, - "strategy": "token", - }, +def test_update_quant_config_from_compressed_tensors_single_group_has_no_layer_configs(): + # Regression: a single-group checkpoint stays a global quant_algo, and + # supplying module names must not turn it into a per-layer config. + quant_config = QuantConfig() + layer_quant_configs = update_quant_config_from_compressed_tensors( + quant_config, + _compressed_tensors_config(ignore=["lm_head"]), + module_names=["model.layers.0.mlp.down_proj", "lm_head"], ) - group = config["config_groups"].pop("group_0") - config["config_groups"]["FP8_DYNAMIC"] = group - config["config_groups"]["FP8_STATIC"] = dict(group) - with pytest.raises(ValueError, match="exactly one config group"): - update_quant_config_from_compressed_tensors(QuantConfig(), config) + assert layer_quant_configs is None + assert quant_config.quant_algo == QuantAlgo.NVFP4 @pytest.mark.parametrize( @@ -311,3 +305,301 @@ def test_update_quant_config_from_compressed_tensors_mxfp4_with_fp8_kv_cache(): # The MXFP4 branch returns early; kv_cache_scheme must still be honored. assert quant_config.kv_cache_quant_algo == QuantAlgo.FP8 assert set(quant_config.exclude_modules) == {"lm_head"} + + +def test_update_quant_config_from_compressed_tensors_group_format_overrides_top_level(): + # Multi-group checkpoints declare "format" per group; a group-level format + # must win over the checkpoint's top-level one. + config = _compressed_tensors_config( + weights={ + "num_bits": 4, + "type": "float", + "strategy": "group", + "group_size": 32, + }, + format="mixed-precision", + ) + config["config_groups"]["group_0"]["format"] = "mxfp4-pack-quantized" + + quant_config = QuantConfig() + update_quant_config_from_compressed_tensors(quant_config, config) + + assert quant_config.quant_algo == QuantAlgo.W4A16_MXFP4 + assert quant_config.group_size == 32 + + +# Module names and config groups shaped like unsloth/Qwen3.8-27B-NVFP4, whose +# text decoder interleaves GDN ("linear_attn") and full-attention blocks. The +# real checkpoint has 64 blocks with the FP8 MLP tail at 56-63; here it has 4 +# with the tail at 2-3. +_QWEN38_FP8_MLP_LAYERS = (2, 3) + +_QWEN38_MODULE_NAMES = [ + # GDN block: in_proj_a/in_proj_b stay bf16, the rest is FP8. + "model.language_model.layers.0.linear_attn", + "model.language_model.layers.0.linear_attn.in_proj_a", + "model.language_model.layers.0.linear_attn.in_proj_b", + "model.language_model.layers.0.linear_attn.in_proj_qkv", + "model.language_model.layers.0.linear_attn.in_proj_z", + "model.language_model.layers.0.linear_attn.out_proj", + "model.language_model.layers.0.input_layernorm", + "model.language_model.layers.0.mlp.gate_proj", + "model.language_model.layers.0.mlp.up_proj", + "model.language_model.layers.0.mlp.down_proj", + # Full-attention block. + "model.language_model.layers.1.self_attn.q_proj", + "model.language_model.layers.1.self_attn.k_proj", + "model.language_model.layers.1.self_attn.v_proj", + "model.language_model.layers.1.self_attn.o_proj", + "model.language_model.layers.1.mlp.gate_proj", + "model.language_model.layers.1.mlp.down_proj", + # FP8 MLP tail. + "model.language_model.layers.2.mlp.gate_proj", + "model.language_model.layers.2.mlp.down_proj", + "model.language_model.layers.3.mlp.gate_proj", + "model.language_model.layers.3.mlp.down_proj", + # Never quantized by this recipe. + "model.visual.blocks.0.attn.qkv", + "mtp.layers.0.mlp.down_proj", + "lm_head", +] + +_QWEN38_IGNORE = [ + # Non-recursive in compressed-tensors: the GDN module itself is unquantized + # but its in_proj_qkv/in_proj_z/out_proj children are FP8. + "model.language_model.layers.0.linear_attn", + "model.language_model.layers.0.linear_attn.in_proj_a", + "model.language_model.layers.0.linear_attn.in_proj_b", + "model.visual.blocks.0.attn.qkv", + "re:^mtp.*", +] + + +def _qwen38_dense_config(**overrides): + """A compressed-tensors config shaped like unsloth/Qwen3.8-27B-NVFP4. + + ``group_0`` is FP8 per-channel/per-token (attention, GDN projections, + ``lm_head`` and the MLP tail); ``group_1`` is NVFP4 (every other MLP). + Both groups target the tail MLPs, so the checkpoint only loads correctly + if target precedence resolves those to ``group_0``. + """ + fp8_mlp_layers = "|".join(str(layer) for layer in _QWEN38_FP8_MLP_LAYERS) + config = { + "quant_method": "compressed-tensors", + "format": "mixed-precision", + "kv_cache_scheme": { + "num_bits": 8, + "type": "float", + "strategy": "tensor", + }, + "config_groups": { + "group_0": { + "format": "float-quantized", + "targets": [ + r"re:.*self_attn\.(q|k|v|o)_proj$", + r"re:.*linear_attn\.(in_proj_qkv|in_proj_z|out_proj)$", + r"re:.*lm_head", + rf"re:.*layers\.({fp8_mlp_layers})\.mlp\.(gate|up|down)_proj$", + ], + "weights": { + "num_bits": 8, + "type": "float", + "strategy": "channel", + }, + "input_activations": { + "num_bits": 8, + "type": "float", + "strategy": "token", + }, + }, + "group_1": { + "format": "nvfp4-pack-quantized", + "targets": [r"re:.*mlp\.(gate|up|down)_proj$"], + "weights": { + "num_bits": 4, + "type": "float", + "strategy": "tensor_group", + "group_size": 16, + }, + "input_activations": { + "num_bits": 4, + "type": "float", + "strategy": "tensor_group", + "group_size": 16, + }, + }, + }, + "ignore": list(_QWEN38_IGNORE), + } + config.update(overrides) + return config + + +def test_update_quant_config_from_compressed_tensors_qwen38_dense_mixed_precision(): + quant_config = QuantConfig() + layer_quant_configs = update_quant_config_from_compressed_tensors( + quant_config, _qwen38_dense_config(), _QWEN38_MODULE_NAMES + ) + + assert quant_config.quant_algo == QuantAlgo.MIXED_PRECISION + assert quant_config.kv_cache_quant_algo == QuantAlgo.FP8 + + algos = {name: cfg.quant_algo for name, cfg in layer_quant_configs.items()} + prefix = "model.language_model.layers" + expected = { + # Early MLP blocks: W4A4 NVFP4 (only group_1 matches). + f"{prefix}.0.mlp.gate_proj": QuantAlgo.NVFP4, + f"{prefix}.0.mlp.up_proj": QuantAlgo.NVFP4, + f"{prefix}.0.mlp.down_proj": QuantAlgo.NVFP4, + f"{prefix}.1.mlp.gate_proj": QuantAlgo.NVFP4, + f"{prefix}.1.mlp.down_proj": QuantAlgo.NVFP4, + # Tail MLP blocks: matched by both groups, group_0 wins on precedence. + f"{prefix}.2.mlp.gate_proj": QuantAlgo.FP8_PER_CHANNEL_PER_TOKEN, + f"{prefix}.2.mlp.down_proj": QuantAlgo.FP8_PER_CHANNEL_PER_TOKEN, + f"{prefix}.3.mlp.gate_proj": QuantAlgo.FP8_PER_CHANNEL_PER_TOKEN, + f"{prefix}.3.mlp.down_proj": QuantAlgo.FP8_PER_CHANNEL_PER_TOKEN, + # Attention, GDN projections and lm_head: FP8. + f"{prefix}.1.self_attn.q_proj": QuantAlgo.FP8_PER_CHANNEL_PER_TOKEN, + f"{prefix}.1.self_attn.k_proj": QuantAlgo.FP8_PER_CHANNEL_PER_TOKEN, + f"{prefix}.1.self_attn.v_proj": QuantAlgo.FP8_PER_CHANNEL_PER_TOKEN, + f"{prefix}.1.self_attn.o_proj": QuantAlgo.FP8_PER_CHANNEL_PER_TOKEN, + f"{prefix}.0.linear_attn.in_proj_qkv": QuantAlgo.FP8_PER_CHANNEL_PER_TOKEN, + f"{prefix}.0.linear_attn.in_proj_z": QuantAlgo.FP8_PER_CHANNEL_PER_TOKEN, + f"{prefix}.0.linear_attn.out_proj": QuantAlgo.FP8_PER_CHANNEL_PER_TOKEN, + "lm_head": QuantAlgo.FP8_PER_CHANNEL_PER_TOKEN, + } + assert algos == expected + + # NVFP4 entries keep their group size; every entry carries the global KV + # cache algo, matching the modelopt MIXED_PRECISION path. + assert layer_quant_configs[f"{prefix}.0.mlp.down_proj"].group_size == 16 + assert all(cfg.kv_cache_quant_algo == QuantAlgo.FP8 for cfg in layer_quant_configs.values()) + + +def test_update_quant_config_from_compressed_tensors_mixed_precision_skips_ignored_modules(): + quant_config = QuantConfig() + layer_quant_configs = update_quant_config_from_compressed_tensors( + quant_config, _qwen38_dense_config(), _QWEN38_MODULE_NAMES + ) + + prefix = "model.language_model.layers" + # Producer-ignored modules, and modules no config group targets, stay out + # of the mapping: they inherit the global MIXED_PRECISION config. + for name in ( + f"{prefix}.0.linear_attn", + f"{prefix}.0.linear_attn.in_proj_a", + f"{prefix}.0.linear_attn.in_proj_b", + f"{prefix}.0.input_layernorm", + "model.visual.blocks.0.attn.qkv", + "mtp.layers.0.mlp.down_proj", + ): + assert name not in layer_quant_configs + + +def test_update_quant_config_from_compressed_tensors_mixed_ignore_stays_in_layer_map(): + quant_config = QuantConfig() + layer_quant_configs = update_quant_config_from_compressed_tensors( + quant_config, _qwen38_dense_config(), _QWEN38_MODULE_NAMES + ) + + prefix = "model.language_model.layers" + # compressed-tensors ignore entries are non-recursive and are applied + # while constructing the authoritative layer map. Copying them into + # TRT-LLM's recursive exclude list would shadow quantized children. + assert quant_config.exclude_modules == [] + assert f"{prefix}.0.linear_attn" not in layer_quant_configs + assert f"{prefix}.0.linear_attn.in_proj_a" not in layer_quant_configs + assert f"{prefix}.0.linear_attn.out_proj" in layer_quant_configs + + # No module with a per-layer config may be excluded by the final config. + assert not [ + name + for name in layer_quant_configs + if quant_config.is_module_excluded_from_quantization(name) + ] + + +def test_update_quant_config_from_compressed_tensors_keeps_modules_to_not_convert(): + # modules_to_not_convert is written in TRT-LLM's pattern language, so it + # keeps TRT-LLM's recursive semantics and must survive the ignore-entry + # filtering even when it shadows a module a config group targets. + prefix = "model.language_model.layers" + quant_config = QuantConfig() + update_quant_config_from_compressed_tensors( + quant_config, + _qwen38_dense_config(modules_to_not_convert=[f"{prefix}.1.self_attn"]), + _QWEN38_MODULE_NAMES, + ) + + assert f"{prefix}.1.self_attn" in quant_config.exclude_modules + assert quant_config.is_module_excluded_from_quantization(f"{prefix}.1.self_attn.q_proj") + + +def test_update_quant_config_from_compressed_tensors_mixed_precision_without_module_names(): + # The global algo is resolvable without the checkpoint's module list; the + # per-layer configs are not. + quant_config = QuantConfig() + layer_quant_configs = update_quant_config_from_compressed_tensors( + quant_config, _qwen38_dense_config() + ) + + assert layer_quant_configs is None + assert quant_config.quant_algo == QuantAlgo.MIXED_PRECISION + assert quant_config.kv_cache_quant_algo == QuantAlgo.FP8 + assert quant_config.exclude_modules == [] + + +def test_update_quant_config_from_compressed_tensors_exact_target_beats_regex_target(): + # compressed-tensors orders exact-name targets before "re:" targets, so an + # exact target wins even when a regex target also matches. + config = _qwen38_dense_config() + config["config_groups"]["group_0"]["targets"] = ["model.layers.0.mlp.down_proj"] + config["config_groups"]["group_1"]["targets"] = [r"re:.*mlp\.(gate|up|down)_proj$"] + config["ignore"] = [] + + layer_quant_configs = update_quant_config_from_compressed_tensors( + QuantConfig(), + config, + ["model.layers.0.mlp.down_proj", "model.layers.1.mlp.down_proj"], + ) + + assert layer_quant_configs["model.layers.0.mlp.down_proj"].quant_algo == ( + QuantAlgo.FP8_PER_CHANNEL_PER_TOKEN + ) + assert layer_quant_configs["model.layers.1.mlp.down_proj"].quant_algo == QuantAlgo.NVFP4 + + +def test_update_quant_config_from_compressed_tensors_mixed_precision_requires_targets(): + config = _qwen38_dense_config() + del config["config_groups"]["group_1"]["targets"] + + with pytest.raises(ValueError, match="has no 'targets'"): + update_quant_config_from_compressed_tensors(QuantConfig(), config, _QWEN38_MODULE_NAMES) + + +def test_update_quant_config_from_compressed_tensors_rejects_any_unmatched_target(): + # A class-name target cannot be resolved from tensor names. It must fail + # even when another target produces a non-empty layer map. + config = _qwen38_dense_config() + config["config_groups"]["group_0"]["targets"].append("Linear") + with pytest.raises(ValueError, match=r"matched no checkpoint module: \['Linear'\]"): + update_quant_config_from_compressed_tensors(QuantConfig(), config, _QWEN38_MODULE_NAMES) + + +def test_update_quant_config_from_compressed_tensors_copies_per_module_configs(): + layer_quant_configs = update_quant_config_from_compressed_tensors( + QuantConfig(), _qwen38_dense_config(), _QWEN38_MODULE_NAMES + ) + + assert len({id(config) for config in layer_quant_configs.values()}) == len(layer_quant_configs) + + +def test_update_quant_config_from_compressed_tensors_rejects_empty_config_groups(): + with pytest.raises(ValueError, match="config_groups is empty"): + update_quant_config_from_compressed_tensors( + QuantConfig(), + { + "quant_method": "compressed-tensors", + "config_groups": {}, + }, + ) From 96d0ed7709479cdc1e118663b46f39c5d2233b29 Mon Sep 17 00:00:00 2001 From: Mihai Chiorean Date: Fri, 14 Aug 2026 15:07:17 -0700 Subject: [PATCH 2/2] [#17723][fix] Resolve Hub tensor indexes for mixed quantization Signed-off-by: Mihai Chiorean --- tensorrt_llm/_torch/model_config.py | 30 ++++++++-- .../models/test_quant_config_utils.py | 60 +++++++++++++++++++ 2 files changed, 85 insertions(+), 5 deletions(-) diff --git a/tensorrt_llm/_torch/model_config.py b/tensorrt_llm/_torch/model_config.py index fc351194202c..b8a71d3caad9 100644 --- a/tensorrt_llm/_torch/model_config.py +++ b/tensorrt_llm/_torch/model_config.py @@ -632,7 +632,7 @@ def _read_safetensors_header(path: Path) -> Dict[str, Any]: @staticmethod def _read_checkpoint_module_names( checkpoint_dir: Optional[str]) -> Optional[List[str]]: - """Read the module names, in checkpoint (HF) namespace, of every stored tensor. + """Read checkpoint-namespace module names for every stored tensor. compressed-tensors ``config_groups`` select modules with regexes, so resolving a multi-group (mixed-precision) checkpoint into per-layer @@ -648,14 +648,34 @@ def _read_checkpoint_module_names( if checkpoint_dir is None: return None - checkpoint_path = Path(checkpoint_dir) - index_path = checkpoint_path / "model.safetensors.index.json" tensor_names = [] - if index_path.exists(): + checkpoint_path = Path(checkpoint_dir) + if checkpoint_path.is_dir(): + index_path = checkpoint_path / "model.safetensors.index.json" + shard_paths = sorted(checkpoint_path.glob("*.safetensors")) + else: + try: + cached_index = transformers.utils.hub.cached_file( + checkpoint_dir, "model.safetensors.index.json") + except OSError: + cached_index = None + index_path = Path( + cached_index) if cached_index is not None else None + shard_paths = [] + if index_path is None: + try: + cached_shard = transformers.utils.hub.cached_file( + checkpoint_dir, "model.safetensors") + except OSError: + cached_shard = None + if cached_shard is not None: + shard_paths.append(Path(cached_shard)) + + if index_path is not None and index_path.exists(): with open(index_path) as f: tensor_names = list(json.load(f).get("weight_map", {})) else: - for shard in sorted(checkpoint_path.glob("*.safetensors")): + for shard in shard_paths: tensor_names.extend(ModelConfig._read_safetensors_header(shard)) # Drop the parameter name ("weight", "weight_packed", ...) and the diff --git a/tests/unittest/models/test_quant_config_utils.py b/tests/unittest/models/test_quant_config_utils.py index 13acb2830044..71a0ab725ad1 100644 --- a/tests/unittest/models/test_quant_config_utils.py +++ b/tests/unittest/models/test_quant_config_utils.py @@ -13,8 +13,12 @@ # See the License for the specific language governing permissions and # limitations under the License. +import json +from unittest.mock import patch + import pytest +from tensorrt_llm._torch.model_config import ModelConfig from tensorrt_llm.models.modeling_utils import QuantConfig from tensorrt_llm.models.quant_config_utils import update_quant_config_from_compressed_tensors from tensorrt_llm.quantization.mode import QuantAlgo @@ -603,3 +607,59 @@ def test_update_quant_config_from_compressed_tensors_rejects_empty_config_groups "config_groups": {}, }, ) + + +def test_read_checkpoint_module_names_from_local_index(tmp_path): + index_path = tmp_path / "model.safetensors.index.json" + index_path.write_text( + json.dumps( + { + "weight_map": { + "model.layers.0.mlp.gate_proj.weight_packed": "model-1.safetensors", + "model.layers.0.mlp.gate_proj.weight_scale": "model-1.safetensors", + "lm_head.weight": "model-2.safetensors", + } + } + ) + ) + + assert ModelConfig._read_checkpoint_module_names(str(tmp_path)) == [ + "model.layers.0.mlp.gate_proj", + "lm_head", + ] + + +def test_read_checkpoint_module_names_resolves_hub_index(tmp_path): + index_path = tmp_path / "model.safetensors.index.json" + index_path.write_text( + json.dumps({"weight_map": {"model.layers.0.self_attn.q_proj.weight": "model.safetensors"}}) + ) + + with patch( + "tensorrt_llm._torch.model_config.transformers.utils.hub.cached_file", + return_value=str(index_path), + ) as cached_file: + module_names = ModelConfig._read_checkpoint_module_names("org/model") + + cached_file.assert_called_once_with("org/model", "model.safetensors.index.json") + assert module_names == ["model.layers.0.self_attn.q_proj"] + + +def test_model_config_builds_layer_map_for_compressed_tensors(tmp_path): + (tmp_path / "model.safetensors.index.json").write_text( + json.dumps( + { + "weight_map": { + f"{module_name}.weight": "model.safetensors" + for module_name in _QWEN38_MODULE_NAMES + } + } + ) + ) + + quant_config, layer_quant_configs = ModelConfig.load_hf_quant_config( + _qwen38_dense_config(), "AUTO", checkpoint_dir=str(tmp_path) + ) + + assert quant_config.quant_algo == QuantAlgo.MIXED_PRECISION + assert layer_quant_configs["lm_head"].quant_algo == QuantAlgo.FP8_PER_CHANNEL_PER_TOKEN