Spaces:
Running
Running
File size: 8,831 Bytes
ca2a6e7 | 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 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 | """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()
|