spec100m / code /train_medusa.py

squash to reclaim LFS quota

a8f07a3 about 1 month ago

17.9 kB

"""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()

	"""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()