#!/usr/bin/env python3 """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, # 12 query heads -> hd=64 n_kv_heads=4, # 4 KV heads (GQA with ratio 3:1) ffn=2048, # SwiGLU hidden dim 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 # 12 // 4 = 3 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) # (B, H_q, T, hd) k = self.k(x).view(B, T, self.H_kv, self.hd).transpose(1, 2) # (B, H_kv, T, hd) v = self.v(x).view(B, T, self.H_kv, self.hd).transpose(1, 2) # (B, H_kv, T, hd) q, k = apply_rope(q, k, cos, sin) # Expand KV heads to match Q groups for SDPA if self.groups > 1: k = k.repeat_interleave(self.groups, dim=1) # (B, H_q, T, hd) v = v.repeat_interleave(self.groups, dim=1) # (B, H_q, T, hd) 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)") # 1. Forward-backward smoke test 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") # 2. Generation test 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!")