| import unittest |
|
|
| import torch |
| from transformers import AutoConfig |
| from transformers.models.qwen3.modeling_qwen3 import Qwen3Model |
|
|
| from tools.maya_export import streamed_backbone as SB |
| from tools.maya_export.common import ckpt_paths |
|
|
| N_LAYERS = 2 |
|
|
|
|
| def _hf_two_layer(cfg): |
| """Upstream transformers Qwen3Model with the first N_LAYERS layers loaded from the checkpoint (fp32).""" |
| c = AutoConfig.from_pretrained(cfg.name_or_path) |
| c.num_hidden_layers = N_LAYERS |
| c.layer_types = c.layer_types[:N_LAYERS] |
| c.vocab_size = 200000 |
| c._attn_implementation = "sdpa" |
| model = Qwen3Model(c).eval() |
| weights = SB._Weights(SB.shard_files(ckpt_paths()["ckpt"])) |
| state = {} |
| for i in range(N_LAYERS): |
| prefix = SB.LAYER_PREFIX.format(i=i) |
| for name, _ in model.layers[i].named_parameters(): |
| state[f"layers.{i}.{name}"] = weights.get(prefix + name).float() |
| state["norm.weight"] = weights.get(SB.NORM_KEY).float() |
| model.load_state_dict(state, strict=False) |
| return model |
|
|
|
|
| class TestStreamedBackbone(unittest.TestCase): |
| @classmethod |
| def setUpClass(cls): |
| try: |
| cls.backbone_dir = ckpt_paths()["backbone_dir"] |
| cls.shards = SB.shard_files(ckpt_paths()["ckpt"]) |
| except Exception as e: |
| raise unittest.SkipTest(f"checkpoint not available: {e}") |
|
|
| def test_prefill_then_steps_match_full_forward(self): |
| torch.manual_seed(0) |
| stream = SB.StreamedBackbone(self.backbone_dir, self.shards, n_layers=N_LAYERS) |
| ref = _hf_two_layer(stream.cfg) |
| x = torch.randn(1, 9, stream.cfg.hidden_size) * 0.1 |
|
|
| out = stream.forward(x[:, :6]) |
| with torch.inference_mode(): |
| full = ref(inputs_embeds=x[:, :6]).last_hidden_state |
| self.assertLess(float((out - full).abs().max()), 1e-4) |
|
|
| for t in range(6, 9): |
| out = stream.forward(x[:, t : t + 1]) |
| with torch.inference_mode(): |
| full = ref(inputs_embeds=x[:, : t + 1]).last_hidden_state[:, -1:] |
| self.assertLess(float((out - full).abs().max()), 1e-4, f"step {t}") |
|
|
|
|
| if __name__ == "__main__": |
| unittest.main() |
|
|
