"""Phase 2: train the hybrid speculative heads on the frozen base model. Architecture (v6): - Kc conditioned heads in groups of G: head k predicts token at offset k+1 given h_t + compressed embeddings of the previous group's G tokens. Teacher-forced during training (real tokens as conditioning). - Kp parallel heads predict offsets Kc+1..Kc+Kp from h_t alone. - All heads share one vocab projection (spec_vocab [r,V]) + bias. Key ideas vs old rank-1 medusa: - Every head emits a real context-dependent distribution (not 2 fixed tokens) - Conditioning carries the chain state forward -> deep heads stay accurate - Verification-ready: heads trained to match base's greedy predictions Memory strategy: h computed under no_grad (base frozen), then each sampled head's [B,T,V] logits are backward()'d separately so only one head's graph lives at a time. Head + position subsampling bounds the step cost. Usage: python train_medusa.py --resume checkpoints/base_500m/final.pt \ --data mixture500m --steps 4000 --batch 8 --seq 1024 """ from __future__ import annotations import argparse import math import os import time import torch import torch.nn.functional as F from config import Config from model import build_model from data_pipeline import load_tokens from train import batched, TrainingController def build_cond_index(Kc: int, G: int, T: int, device) -> torch.Tensor: """cond_idx[k, t, j] = position of conditioning token j for head k at anchor t: t + (g-1)*G + j + 1 where g = k//G. Group 0 pulls the last G real tokens (offsets -G+1..0).""" ks = torch.arange(Kc, device=device) ts = torch.arange(T, device=device) g = ks // G # [Kc] j = torch.arange(G, device=device) idx = (ts[None, :, None] + ((g - 1) * G)[:, None, None] + j[None, None, :] + 1) # [Kc, T, G] return idx.clamp(0, T - 1) @torch.no_grad() def eval_streak(model, cfg, data, seq=1024, n_anchors=64, device="cuda"): """Free-running draft vs base-greedy streak (standard acceptance proxy). For each anchor t: draft a chain with self-fed conditioning, count how long head preds match the base model's argmax at each position.""" K = cfg.medusa_heads G = cfg.medusa_cond_group start = torch.randint(0, data.numel() - seq - 1, (1,), device=device) ids = data[start:start + seq].unsqueeze(0) # [1, T] h = model(ids) # [1, T, d] base_pred = model.lm_head(h).argmax(-1)[0] # [T] pred for pos t+1 lo, hi = G - 1, seq - K - 2 anchors = torch.randint(lo, hi, (n_anchors,), device=device).sort().values h_a = h[0, anchors] # [A, d] hist = torch.stack([ids[0, t - G + 1:t + 1] for t in anchors]) # [A, G] chain = model.spec_draft(h_a, hist, first_head=0) # [A, K] # base_pred[t+k] is base's greedy for position t+k+1 = draft slot k tgt = torch.stack([base_pred[t:t + K] for t in anchors]) # [A, K] match = (chain == tgt) streaks = match.cumprod(1).sum(1).float() # first-mismatch len return streaks.mean().item(), match[:, 0].float().mean().item() def main(): ap = argparse.ArgumentParser() ap.add_argument("--data", type=str, default="mixture500m") ap.add_argument("--steps", type=int, default=4000) ap.add_argument("--seq", type=int, default=1024) ap.add_argument("--batch", type=int, default=8) ap.add_argument("--lr", type=float, default=1e-3) ap.add_argument("--warmup", type=int, default=100) ap.add_argument("--cond_per_step", type=int, default=48) ap.add_argument("--par_per_step", type=int, default=16) ap.add_argument("--pos_per_head", type=int, default=512, help="positions sampled per head per step") ap.add_argument("--ss-prob", type=float, default=0.0, help="scheduled sampling: prob a head's conditioning uses " "the previous group's own predictions instead of " "ground truth (group 0 always uses real tokens)") ap.add_argument("--ss-start", type=int, default=0, help="step at which scheduled sampling kicks in") ap.add_argument("--conf-w", type=float, default=0.1, help="weight of the confidence-head BCE loss " "(dspark acceptance predictor; 0 disables)") ap.add_argument("--out", type=str, default="checkpoints/spec_v6") ap.add_argument("--checkpoint_every", type=int, default=500) ap.add_argument("--eval_every", type=int, default=200) ap.add_argument("--resume", type=str, required=True) ap.add_argument("--grad-clip", type=float, default=1.0) args = ap.parse_args() device = "cuda" torch.backends.cuda.matmul.allow_tf32 = True torch.backends.cudnn.allow_tf32 = True cfg = Config.v5_500m() assert cfg.medusa_rank > 1, "v6 head arch requires medusa_rank > 1" model = build_model(cfg, device) ckpt = torch.load(args.resume, map_location=device) missing, unexpected = model.load_state_dict(ckpt["model"], strict=False) print(f"resumed base: {args.resume} (step {ckpt.get('step')}, " f"loss {ckpt.get('loss'):.4f})") print(f" missing: {len(missing)} (new spec params), " f"unexpected: {len(unexpected)} (old r=1 heads)") # Freeze everything except spec_* params for name, p in model.named_parameters(): p.requires_grad = name.startswith("spec_") trainable = [p for p in model.parameters() if p.requires_grad] print(f"trainable: {sum(p.numel() for p in trainable)/1e6:.1f}M spec params") data = load_tokens(args.data).to(device) print(f"data: {args.data} ({data.numel():,} tokens)") opt = torch.optim.AdamW(trainable, lr=args.lr, betas=(0.9, 0.95), weight_decay=0.01, fused=True) ctrl = TrainingController(args.lr, args.warmup, args.steps) os.makedirs(args.out, exist_ok=True) B, T = args.batch, args.seq Kc, Kp = cfg.medusa_cond_heads, cfg.medusa_par_heads G, e = cfg.medusa_cond_group, cfg.medusa_emb_rank cond_idx = build_cond_index(Kc, G, T, device) # [Kc, T, G] print(f"\n=== Phase 2: spec-head training ===") print(f" heads: {Kc} conditioned (G={G}) + {Kp} parallel = {cfg.medusa_heads}") print(f" steps: {args.steps} batch {B}x{T} " f"heads/step: {args.cond_per_step}c+{args.par_per_step}p " f"x{args.pos_per_head}pos\n") model.train() loader = batched(data, T, B, device) t_start = time.perf_counter() for step in range(args.steps): ids = next(loader) opt.zero_grad(set_to_none=True) with torch.no_grad(): h = model(ids) # [B,T,d] frozen base total_loss = 0.0 # --- conditioned heads (teacher forcing on previous group) --- # NOTE: E = tok_emb @ spec_tok_proj is recomputed per head so each # backward() frees its own graph (shared E would need retain_graph # and keep every head's [B,P,V] logits alive -> OOM). ks = torch.randperm(Kc, device=device)[:args.cond_per_step] for k in ks.tolist(): g = k // G lo = G - 1 if g == 0 else 0 hi = T - k - 1 if hi <= lo: continue pos = torch.randint(lo, hi, (args.pos_per_head,), device=device) E = model.tok_emb(ids) @ model.spec_tok_proj # [B,T,e] cond_pos = cond_idx[k][pos] # [P,G] if (args.ss_prob > 0 and g > 0 and step >= args.ss_start and torch.rand(()) < args.ss_prob): # Scheduled sampling: condition on the PREVIOUS group's own # predictions (computed no-grad with teacher-forced inputs), # not ground truth — fixes the train/test conditioning gap # that caps free-running accept streaks. with torch.no_grad(): cond_tok = torch.empty(B, pos.numel(), G, dtype=torch.long, device=device) for j in range(G): pk = (g - 1) * G + j # prev-group head pc = cond_idx[pk][pos] # [P,G] pcond = E[:, pc].reshape(B, -1, G * e) pinp = torch.cat([h[:, pos], pcond], dim=-1) pcode = F.silu(pinp @ model.spec_cond_w[pk]) plg = pcode @ model.spec_vocab + model.spec_bias cond_tok[:, :, j] = plg.argmax(-1) cond = (model.tok_emb(cond_tok) @ model.spec_tok_proj ).reshape(B, -1, G * e) else: cond = E[:, cond_pos].reshape(B, -1, G * e) inp = torch.cat([h[:, pos], cond], dim=-1) # [B,P,d+G*e] code = F.silu(inp @ model.spec_cond_w[k]) # [B,P,r] lg = code @ model.spec_vocab + model.spec_bias # [B,P,V] # dspark markov: bias by previous (teacher-forced) token mb = (model.spec_markov_emb[ids[:, pos + k]] @ model.spec_markov_w2) # [B,P,V] lg = lg + mb tgt = ids[:, pos + k + 1] # [B,P] # confidence head: predict P(this slot accepted) vs real token if args.conf_w > 0: cin = torch.cat([inp, code], dim=-1) # [B,P,d+Ge+r] conf = torch.sigmoid(cin @ model.spec_conf_w).squeeze(-1) conf_t = (lg.argmax(-1) == tgt).float().detach() conf_loss = F.binary_cross_entropy(conf.float(), conf_t) else: conf_loss = lg.new_zeros(()) w = max(0.98 ** k, 0.05) l = F.cross_entropy(lg.reshape(-1, cfg.vocab_size).float(), tgt.reshape(-1)) (w * (l + args.conf_w * conf_loss) / args.cond_per_step).backward() total_loss += l.item() * w / args.cond_per_step # --- parallel tail heads (unconditioned, deep offsets) --- js = torch.randperm(Kp, device=device)[:args.par_per_step] for j in js.tolist(): off = Kc + j + 1 hi = T - off if hi <= 0: continue pos = torch.randint(0, hi, (args.pos_per_head,), device=device) code = F.silu(h[:, pos] @ model.spec_par_w[j]) # [B,P,r] lg = code @ model.spec_vocab + model.spec_bias # [B,P,V] tgt = ids[:, pos + off] w = max(0.98 ** off, 0.05) l = F.cross_entropy(lg.reshape(-1, cfg.vocab_size).float(), tgt.reshape(-1)) (w * l / args.par_per_step).backward() total_loss += l.item() * w / args.par_per_step grad_norm = torch.nn.utils.clip_grad_norm_(trainable, args.grad_clip) opt.step() if step % 20 == 0 or step == args.steps - 1: lr, should_ckpt, msg = ctrl.update(total_loss, grad_norm.item(), step) for pg in opt.param_groups: pg["lr"] = lr elapsed = time.perf_counter() - t_start print(f"step {step:5d} lr {lr:.2e} loss {total_loss:.4f} " f"grad {grad_norm.item():.2f} {elapsed:.0f}s{msg}", flush=True) if should_ckpt: torch.save({"model": model.state_dict(), "cfg": cfg.__dict__, "step": step + 1, "loss": total_loss}, os.path.join(args.out, "best.pt")) else: lr = ctrl.lr_at(step + 1) for pg in opt.param_groups: pg["lr"] = lr if (step + 1) % args.eval_every == 0: model.eval() streak, acc1 = eval_streak(model, cfg, data, seq=T, device=device) model.train() print(f" [eval] accept streak ~{streak:.1f} tokens " f"head-0 acc {acc1:.2f}", flush=True) if (step + 1) % args.checkpoint_every == 0: torch.save({"model": model.state_dict(), "cfg": cfg.__dict__, "step": step + 1, "loss": total_loss}, os.path.join(args.out, f"step_{step+1}.pt")) print(f" checkpoint: step_{step+1}.pt", flush=True) torch.save({"model": model.state_dict(), "cfg": cfg.__dict__, "step": args.steps, "loss": total_loss}, os.path.join(args.out, "final.pt")) streak, acc1 = eval_streak(model, cfg, data, seq=T, device=device) print(f"\n=== Done ===") print(f" best loss: {ctrl.best_loss:.4f}") print(f" final accept streak ~{streak:.1f} tokens head-0 acc {acc1:.2f}") if __name__ == "__main__": main()