"""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=256, device="cuda", max_heads=None): """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. max_heads: only evaluate the first max_heads slots (curriculum mode — deep untrained heads would dominate the streak otherwise).""" K = cfg.medusa_heads if max_heads is None else min( max_heads, 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 - cfg.medusa_heads - 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] # draft whole groups; SCORE only the first K slots n_cond = min(((K + G - 1) // G) * G, cfg.medusa_cond_heads) n_par = min(cfg.medusa_par_heads, max(0, K - n_cond)) chain = model.spec_draft(h_a, hist, first_head=0, n_cond=n_cond, n_par=n_par) # 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[:, :K] == 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("--cond_max", type=int, default=0, help="only train conditioned heads 0..cond_max-1 " "(curriculum: widen progressively). 0 = all.") 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("--base-label", type=float, default=0.0, help="fraction of steps using the BASE model's own " "argmax as the CE target instead of the corpus " "token. Verification tests head==base, not " "head==corpus — the base drifts off-corpus after " "~50 tokens so corpus labels teach the wrong " "answer for deep heads. 1.0 = always on-policy.") ap.add_argument("--soft-label", type=float, default=0.0, help="fraction of steps using KL(base || head) " "distillation — the base's FULL distribution as " "soft targets (denser than argmax, teaches the " "heads the real distribution for sampling AND " "argmax). Soft-step overrides the tgt choice.") 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() best_streak = -1.0 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 if args.base_label > 0 or args.soft_label > 0: # base's own prediction at every position — the label that # verification actually tests against (argmax) and the # full distribution for distillation (soft) base_lg = model.lm_head(h).float() # [B,T,V] base_am = base_lg.argmax(-1) # [B,T] base_lp = F.log_softmax(base_lg, -1) # [B,T,V] else: base_am = None base_lp = None use_soft = base_lp is not None and \ torch.rand(()) < args.soft_label use_base = (base_am is not None and not use_soft and torch.rand(()) < args.base_label) 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). n_cond_pool = min(Kc, args.cond_max) if args.cond_max > 0 else Kc ks = torch.randperm(n_cond_pool, 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] H_hi = getattr(cfg, "medusa_hi_heads", 0) if H_hi > 0 and k < H_hi: # hi-rank near-field head: wide code + own vocab proj code = F.silu(inp @ model.spec_hi_w[k]) # [B,P,R] lg = code @ model.spec_hi_vocab + model.spec_hi_bias else: 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 token — corpus token, or the # base's own greedy predecessor when training on-policy if use_base or use_soft: prev_tok = base_am[:, pos + k - 1] if k > 0 else ids[:, pos] else: prev_tok = ids[:, pos + k] mb = model.spec_markov_emb[prev_tok] @ model.spec_markov_w2 lg = lg + mb # [B,P,V] tgt = (base_am[:, pos + k] if (use_base or use_soft) else 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[..., :cfg.medusa_rank]], dim=-1) # hi heads: use first r dims 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(()) # depth weighting: flat-ish (research: exp decay starves deep # heads — weight by bottleneck, not position). Mild decay only. w = max(0.995 ** k, 0.5) if use_soft: # KL(base || head): base's distribution at pos+k predicts # position pos+k+1 = the head's target slot soft = base_lp[:, pos + k] # [B,P,V] l = -(soft.exp() * F.log_softmax( lg.float(), -1)).sum(-1).mean() else: 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 = (base_am[:, pos + off - 1] if (use_base or use_soft) else ids[:, pos + off]) w = max(0.995 ** off, 0.5) # flat-ish: deep heads need signal too if use_soft: soft = base_lp[:, pos + off - 1] l = -(soft.exp() * F.log_softmax( lg.float(), -1)).sum(-1).mean() else: 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_loss.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, max_heads=args.cond_max or None) model.train() tag = "" if streak > best_streak: best_streak = streak torch.save({"model": model.state_dict(), "cfg": cfg.__dict__, "step": step + 1, "loss": total_loss, "streak": streak}, os.path.join(args.out, "best.pt")) tag = " (new best streak)" print(f" [eval] accept streak ~{streak:.1f} tokens " f"head-0 acc {acc1:.2f}{tag}", 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, max_heads=args.cond_max or None) 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()