File size: 3,730 Bytes
3e0ec22 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 | # SPDX-FileCopyrightText: © 2026 Tenstorrent USA, Inc.
#
# SPDX-License-Identifier: Apache-2.0
"""
Lightweight (CPU-only, no model download) regression tests for the
multimodal HF-key-remapping pipeline in load_checkpoints.py.
"""
from types import SimpleNamespace
import torch
from models.tt_transformers.tt.load_checkpoints import (
convert_hf_to_meta_mllama,
map_hf_to_meta_keys_mllama,
split_hf_keys,
standardize_hf_keys_multimodal,
)
def _make_mllama_config(num_hidden_layers=4, cross_attention_layers=None):
"""Build a minimal config object accepted by map_hf_to_meta_keys_mllama."""
if cross_attention_layers is None:
cross_attention_layers = [1]
return SimpleNamespace(
text_config=SimpleNamespace(
num_hidden_layers=num_hidden_layers,
cross_attention_layers=cross_attention_layers,
)
)
def _make_sample_mllama_state_dict():
"""Return a minimal HF-format state_dict that covers the projector keys
plus the embed_tokens and lm_head keys required by map_hf_to_meta_keys_mllama."""
t = torch.zeros(1)
return {
"model.multi_modal_projector.weight": t,
"model.multi_modal_projector.bias": t,
"model.vision_model.layernorm_pre.weight": t,
"model.vision_model.layernorm_pre.bias": t,
# Both must be present so standardize_hf_keys (called inside
# standardize_hf_keys_multimodal) doesn't delete embed_tokens.
"lm_head.weight": torch.zeros(16, 4),
"model.embed_tokens.weight": torch.zeros(16, 4),
}
class TestMllamaProjectorKeyRemap:
"""Ensure model.multi_modal_projector.* keys survive the two-stage
multimodal pipeline and land as vision_model.vision_projection.*."""
def test_projector_keys_after_full_pipeline(self):
state_dict = _make_sample_mllama_state_dict()
config = _make_mllama_config()
state_dict = standardize_hf_keys_multimodal(state_dict)
state_dict = split_hf_keys(state_dict)
state_dict = map_hf_to_meta_keys_mllama(state_dict, config)
assert "vision_model.vision_projection.weight" in state_dict
assert "vision_model.vision_projection.bias" in state_dict
assert not any("multi_modal_projector" in k for k in state_dict)
def test_projector_keys_via_convert_hf_to_meta_mllama(self):
"""Standardize_hf_keys_multimodal() -> convert_hf_to_meta_mllama().
Asserts model.multi_modal_projector.weight ends up as
vision_model.vision_projection.weight."""
state_dict = _make_sample_mllama_state_dict()
config = _make_mllama_config()
head_dim = 64
state_dict = standardize_hf_keys_multimodal(state_dict)
state_dict = convert_hf_to_meta_mllama(state_dict, head_dim, config)
assert "vision_model.vision_projection.weight" in state_dict
assert "vision_model.vision_projection.bias" in state_dict
assert not any("multi_modal_projector" in k for k in state_dict)
def test_projector_keys_without_standardize(self):
"""map_hf_to_meta_keys_mllama should also work when called directly
with the original model.-prefixed keys (backward compat)."""
t = torch.zeros(1)
state_dict = {
"model.multi_modal_projector.weight": t,
"model.multi_modal_projector.bias": t,
"model.embed_tokens.weight": torch.zeros(16, 4),
}
config = _make_mllama_config()
state_dict = split_hf_keys(state_dict)
state_dict = map_hf_to_meta_keys_mllama(state_dict, config)
assert "vision_model.vision_projection.weight" in state_dict
assert "vision_model.vision_projection.bias" in state_dict
|