[https://nvbugs/5441729][test] Fix test_modeling_llama_min_latency.py failures

The test_modeling_llama_min_latency.py::test_llama_allclose_to_hf tests
are failing with latest HF transformers due to a bug in their code.

A PR has been submitted to fix it in upstream repo:
huggingface/transformers#40609

Signed-off-by: Po-Han Huang <[email protected]>

Author: nvpohanh

tensorrt_llm/_torch/models/modeling_llama.py

@@ -1002,16 +1002,28 @@ def __init__(self, model_config: ModelConfig[Llama4Config], *args,
 
         self.dtype = self.pretrained_config.text_config.torch_dtype
 
-    def load_weights(self):
+    def load_weights(self, weights: Dict):
         module_dict = nn.ModuleDict({
             "vision_model":
             Llama4VisionModel(self.pretrained_config.vision_config),
             "multi_modal_projector":
             Llama4MultiModalProjector(self.pretrained_config),
         })
-        load_sharded_checkpoint(module_dict,
-                                self.pretrained_config._name_or_path,
-                                strict=False)
+
+        # If the named params are present in the weights, load them directly.
+        param_names = [name for name, _ in module_dict.named_parameters()]
+        if all(name in weights for name in param_names):
+            vision_encoder_weights = {
+                name: weights[name]
+                for name in param_names
+            }
+            module_dict.load_state_dict(vision_encoder_weights)
+
+        # Otherwise, load the weights from the checkpoint.
+        else:
+            load_sharded_checkpoint(module_dict,
+                                    self.pretrained_config._name_or_path,
+                                    strict=False)
 
         self.vision_model = module_dict["vision_model"].to(self.device)
         self.mm_projector = module_dict["multi_modal_projector"].to(self.device)
@@ -1294,7 +1306,7 @@ def infer_max_seq_len(self):
 
     def load_weights(self, weights: Dict, weight_mapper: BaseWeightMapper):
         if not DISAGG:
-            self.mm_encoder.load_weights()
+            self.mm_encoder.load_weights(weights)
 
         # Temporarily detach mm_encoder so the TRT-LLM loader doesn't try to load it
         had_mm_encoder = hasattr(self, "mm_encoder")

tests/unittest/_torch/modeling/test_modeling_llama_min_latency.py

@@ -266,10 +266,12 @@ def test_llama_allclose_to_hf(self, scenario: AllCloseScenario) -> None:
         attention_backend = "TRTLLM"
         metadata_cls = get_attention_backend(attention_backend).Metadata
 
-        if transformers.__version__ >= "4.55.0":
+        if transformers.__version__ >= "4.55.0" \
+            and transformers.__version__ < "4.56.1":
             self.skipTest(
-                "The transformers 4.55.0 has accuracy issues while 4.33.1 works fine. "
-                "https://nvbugspro.nvidia.com/bug/5441729")
+                "The transformers between 4.55.0 and 4.56.1 have accuracy "
+                "issues for Llama4. See: "
+                "https://github.com/huggingface/transformers/pull/40609")
 
         torch.random.manual_seed(0)
         config_dict = deepcopy(LLAMA_4_MAVERICK_TWO_LAYER_CONFIG)