From ab7890f87ca6e1824ca7b999d41f9e2e04dd7c10 Mon Sep 17 00:00:00 2001
From: William Zhang <133824995+2ez4bz@users.noreply.github.com>
Date: Mon, 6 Apr 2026 11:01:50 -0700
Subject: [PATCH] [TRTLLM-11768][fix] Config updates to enable NVFP4

* Why?

The Nemotron Nano VL model checkpoints for NVFP4 could not be loaded
into TRT-LLM.

* What?

Makes the necessary config parsing changes to fix this.

Signed-off-by: William Zhang <133824995+2ez4bz@users.noreply.github.com>
---
 .../_torch/models/modeling_nemotron_nano.py   | 27 +++++++++++++++++++
 1 file changed, 27 insertions(+)

diff --git a/tensorrt_llm/_torch/models/modeling_nemotron_nano.py b/tensorrt_llm/_torch/models/modeling_nemotron_nano.py
index f2f1943e000c..2737c6d2c0a2 100644
--- a/tensorrt_llm/_torch/models/modeling_nemotron_nano.py
+++ b/tensorrt_llm/_torch/models/modeling_nemotron_nano.py
@@ -1251,6 +1251,8 @@ def __init__(self, model_config: ModelConfig):
 
         llm_model_config = copy.deepcopy(model_config)
         llm_model_config.pretrained_config = llm_model_config.pretrained_config.llm_config
+        self._update_config_for_quantization(llm_model_config)
+
         self.llm = AutoModelForCausalLM.from_config(llm_model_config)
 
         self.vocab_size = llm_model_config.pretrained_config.vocab_size
@@ -1467,6 +1469,31 @@ def forward(
         logger.debug(f"output shape: {output_prob.shape}")
         return output_prob
 
+    @staticmethod
+    def _update_config_for_quantization(llm_model_config: ModelConfig) -> None:
+        # Strip the VL wrapper prefix from exclude_modules and
+        # quant_config_dict so patterns match the inner LLM's module names
+        # (e.g. "language_model.backbone.layers.0.mixer.conv1d" becomes
+        # "backbone.layers.0.mixer.conv1d").
+        _LM_PREFIX = "language_model."
+        if llm_model_config.quant_config.exclude_modules is not None:
+            llm_model_config.quant_config.exclude_modules = [
+                m[len(_LM_PREFIX) :] if m.startswith(_LM_PREFIX) else m
+                for m in llm_model_config.quant_config.exclude_modules
+            ]
+        if llm_model_config.quant_config_dict is not None:
+            # NOTE: without `_frozen` toggling, `ModelConfig` cannot have its attributes
+            # modified.
+            old_frozen = llm_model_config._frozen
+            llm_model_config._frozen = False
+            try:
+                llm_model_config.quant_config_dict = {
+                    k[len(_LM_PREFIX) :] if k.startswith(_LM_PREFIX) else k: v
+                    for k, v in llm_model_config.quant_config_dict.items()
+                }
+            finally:
+                llm_model_config._frozen = old_frozen
+
 
 def _rearrange_img(x: torch.Tensor, patch_size: int) -> torch.Tensor:
     py = x.shape[-2] // patch_size