File size: 17,917 Bytes
a8f07a3 | 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 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321 322 323 324 325 326 327 328 329 330 331 332 333 334 335 336 337 338 339 340 341 342 343 344 345 346 347 348 349 350 351 352 353 354 355 356 | """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()
|