| |
| """Bobic 1.5 Raye — 125M Transformer with GQA and NumberHead. |
| Architecture: |
| - 125.85M parameters |
| - 16 layers x 768 d_model |
| - GQA: 12 query heads, 4 KV heads (hd=64) |
| - SwiGLU FFN (hidden=2048) |
| - RoPE + RMSNorm |
| - NumberHead adapter (d=768) |
| - Untied lm_head (preserves full expressivity at 125M scale) |
| """ |
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
|
|
| CFG_RAYE = dict( |
| vocab=16384, |
| ctx=2048, |
| d_model=768, |
| n_layers=16, |
| n_heads=12, |
| n_kv_heads=4, |
| ffn=2048, |
| rope_base=10000.0, |
| num_pos=12, |
| ) |
|
|
|
|
| def rope_cache(ctx, hd, base, device): |
| inv = 1.0 / (base ** (torch.arange(0, hd, 2, device=device, dtype=torch.float32) / hd)) |
| t = torch.arange(ctx, device=device, dtype=torch.float32) |
| f = torch.outer(t, inv) |
| return torch.cos(f), torch.sin(f) |
|
|
|
|
| def apply_rope(q, k, cos, sin, offset=0): |
| T = q.shape[2] |
| c, s = cos[offset:offset+T].to(q.dtype), sin[offset:offset+T].to(q.dtype) |
| qf, kf = q.float(), k.float() |
| d = q.shape[-1] // 2 |
| q1, q2, k1, k2 = qf[..., :d], qf[..., d:], kf[..., :d], kf[..., d:] |
| out_q = torch.cat([q1 * c - q2 * s, q1 * s + q2 * c], -1).to(q.dtype) |
| out_k = torch.cat([k1 * c - k2 * s, k1 * s + k2 * c], -1).to(k.dtype) |
| return out_q, out_k |
|
|
|
|
| class RMSNorm(nn.Module): |
| def __init__(self, d): |
| super().__init__() |
| self.w = nn.Parameter(torch.ones(d)) |
| def forward(self, x): |
| return x * torch.rsqrt(x.float().pow(2).mean(-1, keepdim=True) + 1e-6).to(x.dtype) * self.w |
|
|
|
|
| class GQA(nn.Module): |
| """Grouped-Query Attention with 12 Q heads and 4 KV heads.""" |
| def __init__(self, cfg): |
| super().__init__() |
| d = cfg['d_model'] |
| self.H_q = cfg['n_heads'] |
| self.H_kv = cfg['n_kv_heads'] |
| self.hd = d // self.H_q |
| self.groups = self.H_q // self.H_kv |
|
|
| self.q = nn.Linear(d, self.H_q * self.hd, bias=False) |
| self.k = nn.Linear(d, self.H_kv * self.hd, bias=False) |
| self.v = nn.Linear(d, self.H_kv * self.hd, bias=False) |
| self.o = nn.Linear(self.H_q * self.hd, d, bias=False) |
|
|
| def forward(self, x, cos, sin): |
| B, T, _ = x.shape |
| q = self.q(x).view(B, T, self.H_q, self.hd).transpose(1, 2) |
| k = self.k(x).view(B, T, self.H_kv, self.hd).transpose(1, 2) |
| v = self.v(x).view(B, T, self.H_kv, self.hd).transpose(1, 2) |
|
|
| q, k = apply_rope(q, k, cos, sin) |
|
|
| |
| if self.groups > 1: |
| k = k.repeat_interleave(self.groups, dim=1) |
| v = v.repeat_interleave(self.groups, dim=1) |
|
|
| y = F.scaled_dot_product_attention(q, k, v, is_causal=True) |
| return self.o(y.transpose(1, 2).reshape(B, T, -1)) |
|
|
|
|
| class SwiGLU(nn.Module): |
| def __init__(self, cfg): |
| super().__init__() |
| d, f = cfg['d_model'], cfg['ffn'] |
| self.gate = nn.Linear(d, f, bias=False) |
| self.up = nn.Linear(d, f, bias=False) |
| self.down = nn.Linear(f, d, bias=False) |
| def forward(self, x): |
| return self.down(F.silu(self.gate(x)) * self.up(x)) |
|
|
|
|
| class Block(nn.Module): |
| def __init__(self, cfg): |
| super().__init__() |
| self.attn = GQA(cfg) |
| self.ff = SwiGLU(cfg) |
| self.ln1 = RMSNorm(cfg['d_model']) |
| self.ln2 = RMSNorm(cfg['d_model']) |
| def forward(self, x, cos, sin): |
| x = x + self.attn(self.ln1(x), cos, sin) |
| return x + self.ff(self.ln2(x)) |
|
|
|
|
| class NumberHead(nn.Module): |
| def __init__(self, cfg, digit_ids): |
| super().__init__() |
| self.register_buffer('digit_ids', torch.tensor(sorted(digit_ids), dtype=torch.long)) |
| self.pos = nn.Embedding(cfg['num_pos'], cfg['d_model']) |
| self.gate = nn.Parameter(torch.zeros(cfg['d_model'])) |
|
|
| def forward(self, h, idx, run_init=None): |
| B, T, d = h.shape |
| isd = (idx.unsqueeze(-1) == self.digit_ids).any(-1) |
| pos = torch.zeros(B, T, dtype=torch.long, device=h.device) |
| run = torch.zeros(B, dtype=torch.long, device=h.device) if run_init is None else run_init.to(h.device) |
|
|
| prev = torch.zeros(B, dtype=torch.bool, device=h.device) |
| for t in range(T): |
| cur = isd[:, t] |
| run = torch.where(cur & ~prev, torch.zeros_like(run), run) |
| pos[:, t] = torch.where(cur, run, torch.zeros_like(run)) |
| run = run + cur.to(torch.long) |
| prev = cur |
| pos = pos.clamp(max=self.pos.num_embeddings - 1) |
| add = self.gate * self.pos(pos) |
| out = h + torch.where(isd.unsqueeze(-1), add.to(h.dtype), torch.zeros_like(h)) |
| return out, run |
|
|
|
|
| class Bobic15Raye(nn.Module): |
| def __init__(self, cfg=CFG_RAYE, digit_ids=range(10)): |
| super().__init__() |
| self.cfg = cfg |
| self.tok = nn.Embedding(cfg['vocab'], cfg['d_model']) |
| self.blocks = nn.ModuleList([Block(cfg) for _ in range(cfg['n_layers'])]) |
| self.ln = RMSNorm(cfg['d_model']) |
| self.nhead = NumberHead(cfg, digit_ids) |
| self.head = nn.Linear(cfg['d_model'], cfg['vocab'], bias=False) |
|
|
| cos, sin = rope_cache(cfg['ctx'], cfg['d_model'] // cfg['n_heads'], cfg['rope_base'], 'cpu') |
| self.register_buffer('rope_cos', cos, persistent=False) |
| self.register_buffer('rope_sin', sin, persistent=False) |
|
|
| def forward(self, idx, targets=None): |
| x = self.tok(idx) |
| for b in self.blocks: |
| x = b(x, self.rope_cos, self.rope_sin) |
| h, _ = self.nhead(self.ln(x), idx) |
| logits = self.head(h) |
| loss = None |
| if targets is not None: |
| loss = F.cross_entropy(logits.reshape(-1, logits.size(-1)), targets.reshape(-1)) |
| return logits, loss |
|
|
| @torch.no_grad() |
| def generate(self, idx, n, temp=0.7): |
| self.eval() |
| for _ in range(n): |
| idx_cond = idx[:, -self.cfg['ctx']:] |
| logits, _ = self(idx_cond) |
| logits = logits[:, -1] / temp |
| p = F.softmax(logits, -1) |
| nxt = torch.multinomial(p, 1) |
| idx = torch.cat([idx, nxt], 1) |
| return idx |
|
|
|
|
| if __name__ == '__main__': |
| import time |
| dev = torch.device('cuda' if torch.cuda.is_available() else 'cpu') |
| print(f"Device: {dev} ({torch.cuda.get_device_name(0)})") |
| |
| m = Bobic15Raye().to(dev).to(torch.bfloat16) |
| n_params = sum(p.numel() for p in m.parameters()) |
| print(f"Bobic 1.5 Raye Total Parameters: {n_params:,} ({n_params/1e6:.2f}M)") |
|
|
| |
| BS, SEQ = 16, 512 |
| x = torch.randint(0, CFG_RAYE['vocab'], (BS, SEQ), device=dev) |
| torch.cuda.reset_peak_memory_stats() |
| torch.cuda.synchronize() |
| t0 = time.time() |
| |
| for _ in range(10): |
| _, loss = m(x, x) |
| loss.backward() |
| torch.cuda.synchronize() |
| dt = (time.time() - t0) / 10 |
| vram = torch.cuda.max_memory_allocated() / (1024**3) |
| tok_per_sec = (BS * SEQ) / dt |
|
|
| print(f"Train Step (BS={BS}, SEQ={SEQ}): {dt*1000:.1f} ms | {tok_per_sec:,.0f} tok/s | VRAM Peak: {vram:.2f} GiB") |
| |
| |
| g_ctx = torch.randint(0, CFG_RAYE['vocab'], (1, 64), device=dev) |
| t0 = time.time() |
| out = m.generate(g_ctx, 32, temp=0.5) |
| torch.cuda.synchronize() |
| print(f"Generation 32 tokens: {(time.time() - t0)*1000:.1f} ms") |
| print("Smoke test passed successfully!") |
|
|
