File size: 3,484 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 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 | import unittest
import numpy as np
import onnxruntime as ort
import torch
from transformers import AutoConfig
from transformers.models.qwen3.modeling_qwen3 import Qwen3RotaryEmbedding
from tools.maya_export import export_backbone as EB
from tools.maya_export.common import ckpt_paths, work_dir
# The full fp32 fused graph is 5.6 GB and cannot be loaded under the 4 GB limit, so the merge is checked on the
# first two layers (same merge code, layer count 2). The full model is then the same graph over 28 layers.
N = 2
BUILD = work_dir() / "backbone_build" / "fp32"
TOL = 1e-4
def _sess(path):
so = ort.SessionOptions()
so.intra_op_num_threads = 4
return ort.InferenceSession(str(path), so, providers=["CPUExecutionProvider"])
class TestBackboneMerge(unittest.TestCase):
@classmethod
def setUpClass(cls):
paths = [BUILD / f"layer_{i:02d}.onnx" for i in range(N)] + [BUILD / "head.onnx"]
if not all(p.exists() for p in paths):
raise unittest.SkipTest("per-layer fp32 graphs not built (run export_backbone --precision fp32)")
cls.layers = [_sess(BUILD / f"layer_{i:02d}.onnx") for i in range(N)]
cls.head = _sess(BUILD / "head.onnx")
cls.fused_path = BUILD / f"fused{N}_fp32.onnx"
if not cls.fused_path.exists():
EB.merge([BUILD / f"layer_{i:02d}.onnx" for i in range(N)], BUILD / "head.onnx", cls.fused_path, "fp32")
cls.fused = _sess(cls.fused_path)
cfg = AutoConfig.from_pretrained(ckpt_paths()["backbone_dir"])
cfg._attn_implementation = "sdpa"
cls.rot = Qwen3RotaryEmbedding(config=cfg).eval()
def _rope(self, T, past):
cos, sin = EB.rope(self.rot, T, past)
return cos.numpy(), sin.numpy()
def test_fused_matches_layers(self):
rng = np.random.default_rng(0)
steps = [(8, 0), (1, 8), (1, 9), (1, 10), (1, 11)] # 8 prefill tokens, then 4 decode steps
kv_fused = {f"past_{kv}_{i}": np.zeros((1, EB.N_KV_HEADS, 0, EB.HEAD_DIM), np.float32)
for i in range(N) for kv in ("k", "v")}
kv_chain = [(np.zeros((1, EB.N_KV_HEADS, 0, EB.HEAD_DIM), np.float32),
np.zeros((1, EB.N_KV_HEADS, 0, EB.HEAD_DIM), np.float32)) for _ in range(N)]
for T, past in steps:
x = (rng.standard_normal((1, T, EB.D_MODEL)) * 0.02).astype(np.float32)
cos, sin = self._rope(T, past)
mask = EB.masks(T, past).numpy()
feed = {"embeds": x, "cos": cos, "sin": sin, "mask": mask, **kv_fused}
hid_f, cb_f, *present = self.fused.run(None, feed)
for i in range(N):
kv_fused[f"past_k_{i}"] = present[2 * i]
kv_fused[f"past_v_{i}"] = present[2 * i + 1]
h = x
for i in range(N):
out, nk, nv = self.layers[i].run(None, {"h": h, "cos": cos, "sin": sin, "mask": mask,
"past_k": kv_chain[i][0], "past_v": kv_chain[i][1]})
kv_chain[i] = (nk, nv)
h = out
hid_c, cb_c = self.head.run(None, {"hid": h})
rel_h = float(np.abs(hid_f - hid_c).max() / np.abs(hid_c).max())
rel_c = float(np.abs(cb_f - cb_c).max() / np.abs(cb_c).max())
self.assertLessEqual(rel_h, TOL, f"hidden T={T} past={past}")
self.assertLessEqual(rel_c, TOL, f"cb0 T={T} past={past}")
if __name__ == "__main__":
unittest.main()
|