| |
| |
| |
|
|
| """ |
| 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, |
| |
| |
| "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 |
|
|
