tests

Signed-off-by: Olya Kozlova <okozlova@nvidia.com>

Author: evezhier

tests/unittest/_torch/models/checkpoints/hf/test_weight_loader.py

@@ -11,7 +11,7 @@ class MyError(Exception):
 
 
 @pytest.mark.parametrize(
-    "dir_name, safetensor_filenames, expected_safetensor_filenames",
+    "dir_name, safetensor_filenames, expected_safetensor_filenames, use_consolidated",
     [
         (
             "foo",
@@ -21,6 +21,18 @@ class MyError(Exception):
                 "consolidated.safetensors",
             ],
             ["model-00001-of-00002.safetensors", "model-000002-of-00002.safetensors"],
+            False,
+        ),
+        # If use_consolidated specified explicitly.
+        (
+            "foo",
+            [
+                "model-00001-of-00002.safetensors",
+                "model-000002-of-00002.safetensors",
+                "consolidated.safetensors",
+            ],
+            ["consolidated.safetensors"],
+            True,
         ),
         (
             "foo",
@@ -29,12 +41,14 @@ class MyError(Exception):
                 "foo-consolidated.safetensors",
             ],
             [f"model-0000{i}-of-00010.safetensors" for i in range(1, 11)],
+            False,
         ),
         # If there is only a consolidated safetensor, that one should still be used.
         (
             "foo",
             ["consolidated.safetensors"],
             ["consolidated.safetensors"],
+            False,
         ),
         # If the directory contains "consolidated" in its name, but its contents are sharded tensors.
         (
@@ -45,6 +59,7 @@ class MyError(Exception):
                 "consolidated.safetensors",
             ],
             ["model-00001-of-00002.safetensors", "model-000002-of-00002.safetensors"],
+            False,
         ),
     ],
 )
@@ -53,6 +68,7 @@ def test_load_weights_ignores_consolidated_ckpt_when_sharded_ckpt_exists(
     dir_name: str,
     safetensor_filenames: list[str],
     expected_safetensor_filenames: list[str],
+    use_consolidated: bool,
 ):
     checkpoint_dir = tmp_path / dir_name
     checkpoint_dir.mkdir()
@@ -70,7 +86,9 @@ def test_load_weights_ignores_consolidated_ckpt_when_sharded_ckpt_exists(
         mock.patch.object(loader, "prefetch_files") as prefetch_files,
         pytest.raises(MyError),
     ):
-        loader.load_weights(checkpoint_dir=str(checkpoint_dir), mapping=Mapping())
+        loader.load_weights(
+            checkpoint_dir=str(checkpoint_dir), mapping=Mapping(), use_consolidated=use_consolidated
+        )
 
     prefetch_files.assert_called_once()
     prefetched_files = prefetch_files.call_args[0][0]

tests/unittest/_torch/models/checkpoints/mistral/test_weight_mapper.py

@@ -0,0 +1,47 @@
+import pytest
+import torch
+
+from tensorrt_llm._torch.models.checkpoints.mistral.weight_mapper import MistralWeightMapper
+
+
+@pytest.fixture
+def expected_renames():
+    return {
+        # Top-level embeddings and output projections
+        "tok_embeddings.weight": "model.embed_tokens.weight",
+        "output.weight": "lm_head.weight",
+        "norm.weight": "model.norm.weight",
+        # Per-layer attention projection weights (pixtral_mapping + mistral_llm_mapping)
+        "layers.0.attention.wq.weight": "model.layers.0.self_attn.q_proj.weight",
+        "layers.0.attention.wk.weight": "model.layers.0.self_attn.k_proj.weight",
+        "layers.0.attention.wv.weight": "model.layers.0.self_attn.v_proj.weight",
+        "layers.0.attention.wo.weight": "model.layers.0.self_attn.o_proj.weight",
+        # Per-layer MLP weights
+        "layers.0.feed_forward.w1.weight": "model.layers.0.mlp.gate_proj.weight",
+        "layers.0.feed_forward.w2.weight": "model.layers.0.mlp.down_proj.weight",
+        "layers.0.feed_forward.w3.weight": "model.layers.0.mlp.up_proj.weight",
+        # Layernorms
+        "layers.0.attention_norm.weight": "model.layers.0.input_layernorm.weight",
+        "layers.0.ffn_norm.weight": "model.layers.0.post_attention_layernorm.weight",
+        # Quantization scales: compound key must win over individual token
+        "layers.0.attention.kv_fake_quantizer.qscale_act": "model.layers.0.self_attn.kv_scale",
+        "layers.0.attention.qscale_act": "model.layers.0.self_attn.input_scale",
+        # Unknown keys must pass through unchanged
+        "some.unknown.tensor": "some.unknown.tensor",
+    }
+
+
+def test_rename_by_params_map(expected_renames):
+    mapper = MistralWeightMapper()
+    dummy = torch.tensor(0.0)
+    input_weights = {k: dummy for k in expected_renames}
+
+    result = mapper.rename_by_params_map(mapper.mistral_llm_mapping, input_weights)
+
+    mismatches = {k: v for k, v in expected_renames.items() if v not in result}
+    assert not mismatches, (
+        f"Keys not renamed as expected (input -> expected):\n"
+        + "\n".join(f"  {k!r} -> {v!r}" for k, v in mismatches.items())
+        + f"\nActual keys: {sorted(result.keys())}"
+    )
+    assert type(result) is dict