File size: 3,449 Bytes
deaf348 | 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 | """vela3-fast UltraBERT-380M encoder as a fixed-length, eager-attention graph (ids + mask -> final-norm hidden states).
Same contract as the CPU path of modeling_ultrabert.UltrabertModel.forward: local layers attend to keys with |i - j| <= 128, global
layers to all keys, padding masks keys only; RoPE theta per layer (10k local, 2.56M global). Masks use -1e4 instead of finfo.min
(fp16-safe); cos / sin and the two masks are precomputed for the bucket length."""
import sys
import torch
from torch import nn
MASK_VALUE = -1e4
sys.path.insert(0, "src")
def load_decision(path="src"):
from vela3_encoder import DecisionEncoder
return DecisionEncoder.load(path, device="cpu")
class Vela3EncoderExport(nn.Module):
def __init__(self, enc, length):
super().__init__()
cfg = enc.config
cfg._attn_implementation = "eager"
self.length, self.layers, self.embed, self.norm = length, enc.layers, enc.embed_tokens, enc.norm
self.windows, self.thetas = list(cfg.attention_windows), list(cfg.rope_theta)
pos = torch.arange(length).unsqueeze(0)
probe = torch.zeros(1, length, cfg.hidden_size)
self.theta_key = {}
for i, base in enumerate(sorted(set(self.thetas))):
cos, sin = enc.rotary_emb[str(base).replace(".", "_")](probe, pos)
self.register_buffer(f"cos{i}", cos.detach().clone()); self.register_buffer(f"sin{i}", sin.detach().clone())
self.theta_key[base] = i
idx = torch.arange(length)
dist = (idx[:, None] - idx[None, :]).abs()
self.window_key = {}
for i, w in enumerate(sorted({w for w in self.windows if w is not None})):
self.register_buffer(f"band{i}", torch.where(dist <= w, 0.0, MASK_VALUE).view(1, 1, length, length))
self.window_key[w] = i
def forward(self, input_ids, attention_mask):
pad = (1.0 - attention_mask.float()).view(1, 1, 1, self.length) * MASK_VALUE
bands = {w: getattr(self, f"band{i}") + pad for w, i in self.window_key.items()}
x = self.embed(input_ids.long())
for layer, w, t in zip(self.layers, self.windows, self.thetas):
k = self.theta_key[t]
x = layer(x, attention_mask=pad if w is None else bands[w],
position_embeddings=(getattr(self, f"cos{k}"), getattr(self, f"sin{k}")))
return self.norm(x)
if __name__ == "__main__": # parity of the export graph vs the release encoder (sdpa, fp32 CPU) on a real decision input
import json
import pyarrow.parquet as pq
m = load_decision()
m.encoder.config._attn_implementation = "eager"
r = pq.read_table("data/typed/all/test-00000-of-00001.parquet").to_pylist()[0]
state, qs = json.loads(r["state"]), json.loads(r["questions"])
q = next(q for q in qs.values() if q["type"] == "choice")
ids = m.encode(state, q)["ids"]; n = len(ids)
with torch.no_grad():
ref = m.encoder(input_ids=torch.tensor([ids]), attention_mask=torch.ones(1, n, dtype=torch.long)).last_hidden_state[0]
for L in (n, 512):
g = Vela3EncoderExport(m.encoder, L).eval()
x = torch.zeros(1, L, dtype=torch.int32); x[0, :n] = torch.tensor(ids); mk = torch.zeros(1, L, dtype=torch.int32); mk[0, :n] = 1
h = g(x, mk)[0, :n]
print(f"n={n} L={L} max|dh|={(h - ref).abs().max():.2e} min cos={nn.functional.cosine_similarity(h, ref, dim=-1).min():.6f}")
|