(self, n_layers: int = 2, include_lm_head: bool = True)
| 85 | """Smoke tests for the HF → tinygrad.apps.llm state-dict keymap.""" |
| 86 | |
| 87 | def _fake_hf_weights(self, n_layers: int = 2, include_lm_head: bool = True): |
| 88 | keys = [ |
| 89 | "model.embed_tokens.weight", |
| 90 | "model.norm.weight", |
| 91 | ] |
| 92 | if include_lm_head: |
| 93 | keys.append("lm_head.weight") |
| 94 | for l in range(n_layers): |
| 95 | keys += [ |
| 96 | f"model.layers.{l}.input_layernorm.weight", |
| 97 | f"model.layers.{l}.post_attention_layernorm.weight", |
| 98 | f"model.layers.{l}.self_attn.q_proj.weight", |
| 99 | f"model.layers.{l}.self_attn.k_proj.weight", |
| 100 | f"model.layers.{l}.self_attn.v_proj.weight", |
| 101 | f"model.layers.{l}.self_attn.o_proj.weight", |
| 102 | f"model.layers.{l}.self_attn.q_norm.weight", |
| 103 | f"model.layers.{l}.self_attn.k_norm.weight", |
| 104 | f"model.layers.{l}.mlp.gate_proj.weight", |
| 105 | f"model.layers.{l}.mlp.up_proj.weight", |
| 106 | f"model.layers.{l}.mlp.down_proj.weight", |
| 107 | ] |
| 108 | # sentinel objects so we can verify identity-based aliasing |
| 109 | return {k: object() for k in keys} |
| 110 | |
| 111 | def test_keymap_renames_every_hf_key(self): |
| 112 | hf = self._fake_hf_weights(n_layers=2) |
no test coverage detected