| """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__": |
| 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}") |
|
|
