File size: 2,300 Bytes
24e7fa4 | 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 | 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 # padding_idx needs the real range; embeddings are unused (inputs_embeds)
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: # pragma: no cover - checkpoint not cached on this machine
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()
|