File size: 4,444 Bytes
e1acc15 | 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 | """Small native-format fixtures; no downloaded or proprietary tensors."""
from __future__ import annotations
import torch
from torch import Tensor
from .model import Config
class SyntheticSource:
def __init__(self, c: Config, seed: int = 11):
self.fingerprint = f"synthetic-native-v1:{c}:{seed}"
self.tensors: dict[str, Tensor] = {}
self.dtypes: dict[str, str] = {}
g = torch.Generator().manual_seed(seed)
def add(name: str, shape: tuple[int, ...], kind: str = "BF16", ones: bool = False) -> None:
value = torch.ones(shape) if ones else torch.randn(shape, generator=g) * .025
if kind == "F8_E4M3":
value = (value * 128).to(torch.float8_e4m3fn).view(torch.uint8)
self.tensors[name.replace("weight", "scale")] = torch.full(
((shape[0]+31)//32, (shape[1]+31)//32), 120, dtype=torch.uint8)
self.dtypes[name.replace("weight", "scale")] = "F8_E8M0"
self.tensors[name] = value if kind == "F8_E4M3" else value.to(torch.float32 if kind == "F32" else torch.bfloat16)
self.dtypes[name] = kind
add("embed.weight", (c.vocab, c.dim))
add("head.weight", (c.vocab, c.dim))
add("mtp.0.main_proj.weight", (c.dim, c.dim * 3), "F8_E4M3")
add("mtp.0.main_norm.weight", (c.dim,), ones=True)
last = f"mtp.{c.stages-1}."
add(last + "norm.weight", (c.dim,), ones=True)
add(last + "markov_head.embed.weight", (c.vocab, c.markov))
add(last + "markov_head.head.weight", (c.vocab, c.markov))
add(last + "confidence_head.proj.weight", (1, c.dim+c.markov))
for i in range(c.stages):
p = f"mtp.{i}."
for name, shape in {
"attn.wq_a": (c.qr, c.dim), "attn.wq_b": (c.heads*c.hd, c.qr),
"attn.wkv": (c.hd, c.dim), "attn.wo_a": (c.groups*c.outr, c.heads*c.hd//c.groups),
"attn.wo_b": (c.dim, c.groups*c.outr),
"ffn.shared_experts.w1": (c.inter, c.dim), "ffn.shared_experts.w2": (c.dim, c.inter),
"ffn.shared_experts.w3": (c.inter, c.dim),
}.items():
add(p + name + ".weight", shape, "F8_E4M3")
for name, size in {"attn_norm": c.dim, "ffn_norm": c.dim, "attn.q_norm": c.qr, "attn.kv_norm": c.hd}.items():
add(p+name+".weight", (size,), ones=True)
for kind in ("attn", "ffn"):
add(p+f"hc_{kind}_fn", ((2+c.hc)*c.hc, c.hc*c.dim), "F32")
add(p+f"hc_{kind}_base", ((2+c.hc)*c.hc,), "F32")
add(p+f"hc_{kind}_scale", (3,), "F32")
add(p+"attn.attn_sink", (c.heads,), "F32")
add(p+"ffn.gate.weight", (c.experts, c.dim))
add(p+"ffn.gate.bias", (c.experts,), "F32")
for e in range(c.experts):
for name, (n, k) in {"w1": (c.inter, c.dim), "w2": (c.dim, c.inter), "w3": (c.inter, c.dim)}.items():
ep = f"{p}ffn.experts.{e}.{name}"
self.tensors[ep+".weight"] = torch.randint(0, 256, (n, k//2), generator=g, dtype=torch.uint8)
self.tensors[ep+".scale"] = torch.full((n, k//32), 120, dtype=torch.uint8)
self.dtypes[ep+".weight"], self.dtypes[ep+".scale"] = "I8", "F8_E8M0"
def get(self, name: str) -> Tensor:
return self.tensors[name].clone()
def info(self, name: str) -> dict[str, object]:
return {"dtype": self.dtypes[name], "shape": tuple(self.tensors[name].shape)}
def verify(self, name: str) -> None:
if name not in self.tensors:
raise KeyError(name)
def tiny_config() -> Config:
return Config(dim=64, inter=64, stages=3, vocab=96, noise=95, heads=2, hd=32, rd=16, qr=32,
groups=2, outr=32, experts=4, topk=3, markov=32, window=8)
def inputs(c: Config, batch: int = 2, P: int = 12) -> dict[str, Tensor]:
g = torch.Generator().manual_seed(37)
pos = torch.arange(max(0, P-c.window+1), P+1)[None].expand(batch, -1).clone()
shape = pos.shape
return {"taps": torch.randn(*shape, c.dim*3, generator=g).to(torch.bfloat16),
"engine_mx": torch.zeros(*shape, c.dim), "has_taps": torch.ones(shape, dtype=torch.bool),
"positions": pos, "mask": torch.ones(shape, dtype=torch.bool),
"anchor": torch.arange(batch) + 3, "P": torch.full((batch,), P, dtype=torch.long)}
|