"""Export the pinned MacroStories weights into a one-token, cached ONNX decoder. The independent inference adapter is verified against the original HF Python implementation, including empty caches, all four recurrent passes and long context. """ from pathlib import Path import argparse import hashlib import json import shutil import numpy as np import onnx import onnxruntime as ort import torch from huggingface_hub import snapshot_download from safetensors.torch import load_file from transformers import AutoModelForCausalLM, AutoTokenizer MODEL_ID = "raincandy-u/MacroStories" REVISION = "dc1ef998bf5331f7f70a77c815ed3e4f643a0ba7" ROOT = Path(__file__).resolve().parents[1] class CachedDecoder(torch.nn.Module): def __init__(self, state, config): super().__init__() self.weights = torch.nn.ParameterDict({ name.replace(".", "__"): torch.nn.Parameter(value, requires_grad=False) for name, value in state.items() if "early_exit_gate" not in name }) self.eps = config["rms_norm_eps"] self.register_buffer("inv_freq", 1.0 / (config["rope_theta"] ** (torch.arange(0, 16, 2).float() / 16))) def weight(self, name): return self.weights[name.replace(".", "__")] def layer(self, name): return self.weight("model.layers.0." + name + ".weight") def norm(self, hidden, weight): return hidden * torch.rsqrt(hidden.square().mean(-1, keepdim=True) + self.eps) * weight def rotate(self, vector, cosine, sine): return vector * cosine + torch.cat((-vector[..., 8:], vector[..., :8]), dim=-1) * sine def forward(self, input_ids, position, past_key_values): embed = self.weight("model.embed_tokens.weight") hidden = torch.nn.functional.embedding(input_ids, embed)[0] angles = position.float() * self.inv_freq angles = torch.cat((angles, angles), dim=-1) cosine, sine = angles.cos(), angles.sin() present = [] first_value = None for step in range(4): normed = self.norm(hidden, self.layer("input_layernorm")) query = torch.nn.functional.linear(normed, self.layer("self_attn.q_proj")).reshape(2, 16) key = torch.nn.functional.linear(normed, self.layer("self_attn.k_proj")) value = torch.nn.functional.linear(normed, self.layer("self_attn.v_proj")) if step == 0: first_value = value else: value = 0.5 * value + 0.5 * first_value query, key = self.rotate(query, cosine, sine), self.rotate(key, cosine, sine) keys = torch.cat((past_key_values[step, 0], key.unsqueeze(0)), dim=0) values = torch.cat((past_key_values[step, 1], value.unsqueeze(0)), dim=0) present.append(torch.stack((keys, values))) scores = torch.matmul(query, keys.transpose(0, 1)) * 0.25 attended = torch.matmul(torch.softmax(scores, dim=-1), values).reshape(32) projected = torch.nn.functional.linear(attended, self.layer("self_attn.o_proj")) hidden = hidden + self.norm(projected, self.layer("input_layernorm_2")) normed = self.norm(hidden, self.layer("post_attention_layernorm")) gate = torch.nn.functional.silu(torch.nn.functional.linear(normed, self.layer("mlp.gate_proj"))) up = torch.nn.functional.linear(normed, self.layer("mlp.up_proj")) down = torch.nn.functional.linear(gate * up, self.layer("mlp.down_proj")) hidden = hidden + self.norm(down, self.layer("post_attention_layernorm_2")) hidden = self.norm(hidden, self.weight("model.norm.weight")) return torch.nn.functional.linear(hidden, embed), torch.stack(present) def main(): parser = argparse.ArgumentParser() parser.add_argument("--source", type=Path, help="Existing snapshot; otherwise download pinned upstream revision.") args = parser.parse_args() source = args.source or Path(snapshot_download(MODEL_ID, revision=REVISION)) config = json.loads((source / "config.json").read_text()) state = load_file(source / "model.safetensors") assert sum(t.numel() for t in state.values()) == 19969 destination = ROOT / "models/macrostories" (destination / "onnx").mkdir(parents=True, exist_ok=True) for name in ["config.json", "tokenizer.json", "tokenizer_config.json", "special_tokens_map.json", "generation_config.json"]: shutil.copyfile(source / name, destination / name) adapter = CachedDecoder(state, config).eval() path = destination / "onnx/model.onnx" with torch.inference_mode(): torch.onnx.export( adapter, (torch.tensor([1]), torch.tensor([3]), torch.zeros(4, 2, 3, 16)), path, input_names=["input_ids", "position", "past_key_values"], output_names=["logits", "present_key_values"], opset_version=17, dynamic_axes={"past_key_values": {2: "past_length"}, "present_key_values": {2: "present_length"}}, dynamo=False, ) onnx.checker.check_model(onnx.load(path)) reference = AutoModelForCausalLM.from_pretrained(source, trust_remote_code=True, attn_implementation="eager").eval() tokenizer = AutoTokenizer.from_pretrained(source) torch.set_num_threads(1) session = ort.InferenceSession(str(path), providers=["CPUExecutionProvider"]) sequences = [ [1], [1] + tokenizer.encode('[PERSON_NAME_1] wanted to water the flower.'), [1] + tokenizer.encode('"I want to help this flower," [PERSON_NAME_1] said.'), [1] + tokenizer.encode('[PERSON_NAME_2] carried the bucket. ' * 36), ] max_error = 0.0 fixtures = [] with torch.inference_mode(): for ids in sequences: expected = reference(torch.tensor([ids]), use_cache=False).logits[0].numpy() past = np.zeros((4, 2, 0, 16), dtype=np.float32) observed = [] for pos, token in enumerate(ids): logits, past = session.run(None, { "input_ids": np.array([token], dtype=np.int64), "position": np.array([pos], dtype=np.int64), "past_key_values": past, }) observed.append(logits) observed = np.array(observed) error = float(np.max(np.abs(expected - observed))) max_error = max(max_error, error) np.testing.assert_allclose(observed, expected, atol=1e-4, rtol=1e-4) assert np.array_equal(observed.argmax(-1), expected.argmax(-1)) fixtures.append({"ids": ids, "last_logits": expected[-1].tolist()}) # Exercise RoPE and caches at the actual context boundary. long_ids = ([1] + tokenizer.encode('[PERSON_NAME_1] carried the bucket. ' * 400))[:2048] expected_last = reference(torch.tensor([long_ids]), use_cache=False).logits[0, -1].numpy() past = np.zeros((4, 2, 0, 16), dtype=np.float32) for pos, token in enumerate(long_ids): logits, past = session.run(None, {"input_ids": np.array([token], dtype=np.int64), "position": np.array([pos], dtype=np.int64), "past_key_values": past}) np.testing.assert_allclose(logits, expected_last, atol=2e-4, rtol=1e-4) boundary_error = float(np.max(np.abs(logits - expected_last))) fixture_dir = ROOT / "tests/fixtures" fixture_dir.mkdir(parents=True, exist_ok=True) (fixture_dir / "reference.json").write_text(json.dumps(fixtures)) tokenize_cases = ["", '[PERSON_NAME_1] wanted to water the flower.', 'The kite was stuck in the tree.', '"I want to help this flower," [PERSON_NAME_1] said.', " [PERSON_NAME_1]'s book", 'Explain quantum computing in German!'] (fixture_dir / "tokenizer.json").write_text(json.dumps([ {"text": text, "ids": tokenizer.encode(text), "decoded": tokenizer.decode(tokenizer.encode(text), skip_special_tokens=False)} for text in tokenize_cases ])) provenance = { "model_id": MODEL_ID, "revision": REVISION, "parameters": 19969, "source_weights_bytes": (source / "model.safetensors").stat().st_size, "source_weights_sha256": hashlib.sha256((source / "model.safetensors").read_bytes()).hexdigest(), "onnx_bytes": path.stat().st_size, "onnx_sha256": hashlib.sha256(path.read_bytes()).hexdigest(), "opset": 17, "precision": "fp32", "validation_max_abs_error": max_error, "validation_context_boundary_abs_error": boundary_error, "validated_positions": sum(len(ids) for ids in sequences) + len(long_ids), "torch_version": torch.__version__, "transformers_version": "4.55.0", } (destination / "provenance.json").write_text(json.dumps(provenance, indent=2) + "\n") print(json.dumps(provenance, indent=2)) if __name__ == "__main__": main()