tensorrt_llm/_torch/models/checkpoints/hf/nemotron_h_weight_mapper.py (2)

37-39: TP-safety for “_scale” tensors

Blindly copying scales may mis-shard when scales are per-output-channel and need TP splitting (similar to A/D/dt_bias). Suggest splitting on dim 0 when divisible.

-            if "_scale" in key:
-                new_weights[key] = weights[name]
+            if "_scale" in key:
+                w = weights[name]
+                # Respect TP sharding for per-channel scales
+                if w.ndim > 0 and (w.shape[0] % tp_size == 0):
+                    w = split(w, tp_size, tp_rank)
+                new_weights[key] = w

---

1-1: Add NVIDIA header (2025)

This file is missing the required SPDX header with current year.

Add to the top of the file (non-diff snippet):

# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0

tensorrt_llm/_torch/models/modeling_nemotron_h.py (2)

16-16: Make exclude_modules normalization token-safe

Current regex may replace any substring “backbone” inside a longer token. Limit to whole-token to avoid accidental renames.

-                re.sub(r'(model\.layers\.)?backbone', 'model', K)
+                re.sub(r'(?<!\w)backbone(?!\w)', 'model', k)

If you specifically intend only top-level “backbone” or “model.layers.backbone”, consider asserting boundaries via path tokenization instead.

Also applies to: 259-261

---

1-15: Update header year to 2025

Header shows 2022–2024; update to 2022–2025 per guidelines.

-# SPDX-FileCopyrightText: Copyright (c) 2022-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2022-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.

tensorrt_llm/_torch/models/modeling_utils.py (2)

485-500: Broaden fused-module exclusion and collapse logic

Cover k/v synonyms for fused QKV and up_proj synonym for fused gate/up; reduce false negatives when exclude lists use different base names. Keep behavior identical for non-Linear modules.

-                    if isinstance(module, Linear):
-                        weight_mode = module.weights_loading_config.weight_mode
-                        if weight_mode == WeightMode.FUSED_GATE_UP_LINEAR:
-                            # gate_proj and up_proj share the same exclusion rule
-                            is_excluded = quant_config.is_module_excluded_from_quantization(
-                                name.replace('gate_up_proj', 'gate_proj'))
-                        elif weight_mode == WeightMode.FUSED_QKV_LINEAR:
-                            # q_proj, k_proj and v_proj share the same exclusion rule
-                            is_excluded = quant_config.is_module_excluded_from_quantization(
-                                name.replace('qkv', 'q'))
-                        else:
-                            is_excluded = quant_config.is_module_excluded_from_quantization(
-                                name)
+                    if isinstance(module, Linear):
+                        weight_mode = module.weights_loading_config.weight_mode
+                        candidates = [name]
+                        if weight_mode == WeightMode.FUSED_GATE_UP_LINEAR:
+                            # gate_proj and up_proj share the same exclusion rule
+                            candidates += [
+                                name.replace('gate_up_proj', 'gate_proj'),
+                                name.replace('gate_up_proj', 'up_proj'),
+                            ]
+                        elif weight_mode == WeightMode.FUSED_QKV_LINEAR:
+                            # q_proj, k_proj and v_proj share the same exclusion rule
+                            q_name = name.replace('qkv', 'q')
+                            candidates += [
+                                q_name,
+                                q_name.replace('.q_proj', '.k_proj'),
+                                q_name.replace('.q_proj', '.v_proj'),
+                            ]
+                        is_excluded = any(
+                            quant_config.is_module_excluded_from_quantization(n)
+                            for n in candidates
+                        )
                     else:
                         is_excluded = quant_config.is_module_excluded_from_quantization(
                             name)

---

1-1: Add NVIDIA header (2025)

Missing required SPDX header.

Add at top (non-diff snippet):

# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0

tensorrt_llm/_torch/modules/mamba/mamba2_mixer.py (2)

92-102: Defensive default for skip_create_weights_in_init

Accessing config.skip_create_weights_in_init directly can break older configs. Use getattr with False default.

         self.in_proj = Linear(
             d_model,
             d_in_proj,
             bias=bias,
             dtype=dtype,
             mapping=self.mapping,
             tensor_parallel_mode=TensorParallelMode.COLUMN,
             quant_config=config.get_quant_config(),
-            skip_create_weights_in_init=config.skip_create_weights_in_init,
+            skip_create_weights_in_init=getattr(config, "skip_create_weights_in_init", False),
             allreduce_strategy=config.allreduce_strategy)
...
         self.out_proj = Linear(
             d_inner,
             d_model,
             bias=bias,
             dtype=dtype,
             mapping=self.mapping,
             tensor_parallel_mode=TensorParallelMode.ROW,
             quant_config=config.get_quant_config(),
-            skip_create_weights_in_init=config.skip_create_weights_in_init,
+            skip_create_weights_in_init=getattr(config, "skip_create_weights_in_init", False),
             allreduce_strategy=config.allreduce_strategy)

Also apply the same to conv1d init for consistency:

# conv1d ...
skip_create_weights_in_init=getattr(config, "skip_create_weights_in_init", False),

Also applies to: 143-153

---

1-15: Update header year to 2025

Adjust to 2022–2025.

-# SPDX-FileCopyrightText: Copyright (c) 2022-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-FileCopyrightText: Copyright (c) 2022-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

tensorrt_llm/_torch/models/modeling_nemotron_h.py (1)

259-261: Tighten the regex to avoid accidental substring replacements

Current pattern can replace "backbone" inside longer tokens (e.g., "backbone_extra"). Add word boundaries to target path segments only.

-                re.sub(r'(model\.layers\.)?backbone', 'model', k)
+                re.sub(r'(?<!\w)(?:model\.layers\.)?backbone(?!\w)', 'model', k)

tensorrt_llm/_torch/models/modeling_utils.py (1)

485-505: Avoid duplicate candidate names and guard replacements by substring

Minor cleanup: using a set prevents duplicates; gating replacements on substring presence avoids no-op strings.

-                    candidates = [name]
+                    candidates = {name}
                     if isinstance(module, Linear):
                         weight_mode = module.weights_loading_config.weight_mode
                         if weight_mode == WeightMode.FUSED_GATE_UP_LINEAR:
                             # sometimes gate and up proj are not packed in the checkpoint,
                             # but they still share the same exclusion rule
-                            candidates += [
-                                name.replace('gate_up_proj', 'gate_proj'),
-                                name.replace('gate_up_proj', 'up_proj')
-                            ]
+                            if 'gate_up_proj' in name:
+                                candidates.update({
+                                    name.replace('gate_up_proj', 'gate_proj'),
+                                    name.replace('gate_up_proj', 'up_proj'),
+                                })
                         elif weight_mode == WeightMode.FUSED_QKV_LINEAR:
                             # sometimes q_proj, k_proj and v_proj are not packed in the checkpoint,
                             # but they still share the same exclusion rule
-                            candidates += [
-                                name.replace('qkv_proj', 'q_proj'),
-                                name.replace('qkv_proj', 'k_proj'),
-                                name.replace('qkv_proj', 'v_proj')
-                            ]
-                    is_excluded = any(
+                            if 'qkv_proj' in name:
+                                candidates.update({
+                                    name.replace('qkv_proj', 'q_proj'),
+                                    name.replace('qkv_proj', 'k_proj'),
+                                    name.replace('qkv_proj', 'v_proj'),
+                                })
+                    is_excluded = any(
-                        quant_config.is_module_excluded_from_quantization(n)
-                        for n in candidates)
+                        quant_config.is_module_excluded_from_quantization(n)
+                        for n in candidates)

tensorrt_llm/_torch/modules/mamba/mamba2_mixer.py (2)

143-152: Mirror the getattr fallback here; consider aligning conv1d initializer too.

Make out_proj resilient the same way; and (optional) apply the same change to self.conv1d for consistency.

Apply this diff in-place:

-            skip_create_weights_in_init=config.skip_create_weights_in_init,
+            skip_create_weights_in_init=getattr(config, "skip_create_weights_in_init", False),

Optional (outside this hunk) for conv1d:

# in self.conv1d = Linear(...):
skip_create_weights_in_init=getattr(config, "skip_create_weights_in_init", False),

---

92-101: Add backward-compatible fallback for skip_create_weights_in_init

Linear.init already accepts skip_create_weights_in_init and ModelConfig sets skip_create_weights_in_init: bool = False by default; use getattr to guard against configs that lack this field:

-           skip_create_weights_in_init=config.skip_create_weights_in_init,
+           skip_create_weights_in_init=getattr(config, "skip_create_weights_in_init", False),

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro