607 lines
27 KiB
Python
607 lines
27 KiB
Python
"""E2 stage A (item 21): dense short-CoT supervision through the carry
|
|
whiteboard on GSM8K.
|
|
|
|
train_carry.py skeleton, one change that matters: CE targets are the
|
|
model's own VERIFIED terse scratchpad + answer (gsm_cot_data.json,
|
|
~30-60 tokens) instead of the ~3-token bare answer — the dense-output
|
|
ingredient the latent GSM arms structurally lacked.
|
|
Control: --feedforward = same supervision, adapter(e,e), no recurrence.
|
|
"""
|
|
|
|
import argparse
|
|
import json
|
|
import os
|
|
import math
|
|
import random
|
|
import sys
|
|
import time
|
|
from pathlib import Path
|
|
|
|
import torch
|
|
import torch.nn.functional as F
|
|
|
|
from carry_common import PAUSE_ID, carry_logits
|
|
from loop_common import BandLooper, MergeAdapter, chat_prompt, DIRECT_SUFFIX
|
|
|
|
sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
|
|
from jlens.core import load_model # noqa: E402
|
|
|
|
OUT = Path(os.environ.get("LOOP_OUT",
|
|
Path(__file__).resolve().parent.parent / "results-loop"))
|
|
STEPS = 600
|
|
BATCH = 4
|
|
LR = 1e-3
|
|
WARMUP = 20
|
|
K_PREFILL = 2
|
|
P_BY_LABEL = {"easy": 2, "hard": 6}
|
|
|
|
ap = argparse.ArgumentParser()
|
|
ap.add_argument("--feedforward", action="store_true")
|
|
ap.add_argument("--seed", type=int, default=0)
|
|
ap.add_argument("--lensnoise", default=None, metavar="RANK,SCALE",
|
|
help="E2-N1: inject noise into the carried state during "
|
|
"training, shaped by the top-RANK sensitivity "
|
|
"directions of jbar at the band entrance; noise norm "
|
|
"= SCALE * per-position state norm (e.g. 32,0.05)")
|
|
ap.add_argument("--drop-steps", type=int, default=0, metavar="D",
|
|
help="E2-L rung B: delete the first D scratchpad steps, "
|
|
"each replaced by --pause-per-step extra pauses")
|
|
ap.add_argument("--pause-per-step", type=int, default=10,
|
|
help="pauses per deleted step (median step = 10 tokens)")
|
|
ap.add_argument("--warm-start", default=None, metavar="ADAPTER_PT")
|
|
ap.add_argument("--steps", type=int, default=STEPS)
|
|
ap.add_argument("--lr", type=float, default=LR)
|
|
ap.add_argument("--tag-suffix", default="",
|
|
help="appended to TAG (distinguish control variants)")
|
|
ap.add_argument("--bandlora", type=int, default=0, metavar="RANK",
|
|
help="item 24: loop-only LoRA (lora_band.LoopLoRA) on every "
|
|
"band layer, uniform scale 1.0 — active only during "
|
|
"band re-runs, k=0 stays bit-exact")
|
|
ap.add_argument("--lora-lr", type=float, default=1e-3)
|
|
ap.add_argument("--lensteach", type=float, default=0.0, metavar="LAMBDA",
|
|
help="item 25: latent process supervision — lens-CE at the "
|
|
"replacement pauses against the DELETED step's tokens "
|
|
"(1:1 pause j <-> step token j), mixed at LAMBDA into "
|
|
"the output CE")
|
|
ap.add_argument("--lensteach-gen", type=float, default=0.0, metavar="LAMBDA",
|
|
help="item 26: result-staging supervision during generation "
|
|
"— at each visible scratchpad line's pre-'=' positions "
|
|
"(result not yet in causal context), lens-CE the "
|
|
"carried state against that line's result tokens")
|
|
ap.add_argument("--inner-iters", type=int, default=0, metavar="M",
|
|
help="item 27: M in-place band iterations at the anchor "
|
|
"(internal looping, no extra tokens); with --lensteach "
|
|
"the lens-CE aligns iteration i <-> deleted-step token "
|
|
"i instead of pause positions")
|
|
ap.add_argument("--base-pauses", type=int, default=-1, metavar="P",
|
|
help="override P_BY_LABEL with a fixed pause count; 0 = NO "
|
|
"pause tokens at all (inner iterations anchor on the "
|
|
"last prompt position)")
|
|
ap.add_argument("--traj-tf", type=float, default=0.0, metavar="LAMBDA",
|
|
help="item 29: teacher-forced TRANSITION loss — iteration i "
|
|
"gets teacher waypoint T[i-1] as merge input, its band "
|
|
"output is pulled onto T[i] (cosine); 10 independent "
|
|
"supervised one-step regressions")
|
|
ap.add_argument("--traj-fr", type=float, default=0.0, metavar="LAMBDA",
|
|
help="item 29: free-running trajectory loss — the burst's "
|
|
"own state s_i pulled onto waypoint T[i] (exposure-gap "
|
|
"term)")
|
|
ap.add_argument("--traj-span", choices=("step", "full"), default="step",
|
|
help="waypoints sampled evenly across the deleted step "
|
|
"(job 1, d=1) or the full CoT (job 2, answer-only)")
|
|
ap.add_argument("--teachstate", type=float, default=0.0, metavar="LAMBDA",
|
|
help="item 28 (Nils's variant): teacher-state distillation "
|
|
"— frozen warm-start adapter runs the FULL cot (step "
|
|
"visible); its band-exit state at the deleted step's "
|
|
"last token becomes the target; burst final iterate "
|
|
"trained to it by cosine, weighted LAMBDA")
|
|
ARGS = ap.parse_args()
|
|
STEPS, LR = ARGS.steps, ARGS.lr
|
|
TAG = ("carrycot_ff" if ARGS.feedforward else "carrycot") + (
|
|
f"_b{ARGS.drop_steps}" if ARGS.drop_steps else "") + (
|
|
f"_blr{ARGS.bandlora}" if ARGS.bandlora else "") + (
|
|
f"_lt{str(ARGS.lensteach).replace('.', '')}" if ARGS.lensteach else "") + (
|
|
f"_lg{str(ARGS.lensteach_gen).replace('.', '')}"
|
|
if ARGS.lensteach_gen else "") + (
|
|
f"_ii{ARGS.inner_iters}" if ARGS.inner_iters else "") + (
|
|
f"_ts{str(ARGS.teachstate).replace('.', '')}"
|
|
if ARGS.teachstate else "") + (
|
|
f"_tj{ARGS.traj_span[0]}{str(ARGS.traj_tf).replace('.', '')}"
|
|
f"_{str(ARGS.traj_fr).replace('.', '')}"
|
|
if (ARGS.traj_tf or ARGS.traj_fr) else "") + (
|
|
f"_p{ARGS.base_pauses}" if ARGS.base_pauses >= 0 else "") + (
|
|
f"_ln{ARGS.lensnoise.replace(',', '_')}" if ARGS.lensnoise else "") + (
|
|
f"_s{ARGS.seed}" if ARGS.seed else "") + ARGS.tag_suffix
|
|
|
|
|
|
def drop_cot_steps(cot, d):
|
|
"""Delete the first d scratchpad lines (front-first: the deleted
|
|
computation must ride the pause-chain before the visible remainder).
|
|
Returns (new_cot, n_deleted, deleted_text); unparseable cots pass
|
|
through intact."""
|
|
lines = [l for l in cot.split("\n") if l.strip()]
|
|
ans = [l for l in lines if l.startswith("Answer")]
|
|
steps = [l for l in lines if not l.startswith("Answer")]
|
|
if len(ans) != 1 or not steps:
|
|
return cot, 0, ""
|
|
n = min(d, len(steps))
|
|
return "\n".join(steps[n:] + ans), n, "\n".join(steps[:n])
|
|
|
|
|
|
class LensNoiseWrapper(torch.nn.Module):
|
|
"""Perturb the carried state s (not the anchor e) before the merge,
|
|
within the span of the lens's top-r readout-sensitive directions.
|
|
Train-time only (noise_on flag); eval and checkpoints use .base."""
|
|
|
|
def __init__(self, base, jbar_layer, rank, scale):
|
|
super().__init__()
|
|
self.base = base
|
|
self.scale = scale
|
|
self.noise_on = True
|
|
J = jbar_layer.float()
|
|
_, _, Vt = torch.linalg.svd(J, full_matrices=False)
|
|
self.register_buffer("V", Vt[:rank].T.contiguous()) # (d, r)
|
|
|
|
def forward(self, e, s):
|
|
if self.noise_on and self.scale > 0:
|
|
z = torch.randn(*s.shape[:-1], self.V.shape[1],
|
|
device=s.device, dtype=torch.float32)
|
|
n = z @ self.V.T
|
|
n = n * (s.float().norm(dim=-1, keepdim=True) * self.scale
|
|
/ (n.norm(dim=-1, keepdim=True) + 1e-6))
|
|
s = (s.float() + n).to(s.dtype)
|
|
return self.base(e, s)
|
|
|
|
|
|
def gen_staging_targets(tok, cot):
|
|
"""Item 26: per-line result-staging spans, computed in TOKEN space.
|
|
For each cot line with '=', the pre-'=' positions (result not yet in
|
|
causal context) target the line's result tokens; the Answer line's
|
|
'Answer:' positions target the answer tokens.
|
|
Returns [(rel_positions, result_token_ids), ...] relative to tok(cot)."""
|
|
ids = tok(cot, add_special_tokens=False)["input_ids"]
|
|
decoded = [tok.decode([t]) for t in ids]
|
|
out, line_start = [], 0
|
|
for i, d in enumerate(decoded + ["\n"]):
|
|
if "\n" not in d and i < len(ids):
|
|
continue
|
|
line = list(range(line_start, min(i, len(ids))))
|
|
line_start = i + 1
|
|
if not line:
|
|
continue
|
|
text = "".join(decoded[j] for j in line)
|
|
seps = [j for j in line if decoded[j].strip() == "="]
|
|
if seps:
|
|
sep = seps[-1]
|
|
elif text.strip().startswith("Answer"):
|
|
colons = [j for j in line if decoded[j].strip() == ":"]
|
|
if not colons:
|
|
continue
|
|
sep = colons[-1]
|
|
else:
|
|
continue
|
|
result = [ids[j] for j in line if j > sep
|
|
and decoded[j].strip()]
|
|
pre = [j for j in line if j <= sep]
|
|
if result and pre:
|
|
out.append((pre, result))
|
|
return out
|
|
|
|
|
|
def traj_burst_forward(looper, adapter, ids, msk, plens, m, T, tf_on, k):
|
|
"""Item 29 forward: optional teacher-forced transition predictions,
|
|
then the free-running burst (whose final state seeds the answer scan,
|
|
matching inference), then the visible-token carry. Returns
|
|
(logits, tf_preds, fr_states)."""
|
|
from torch.utils.checkpoint import checkpoint
|
|
from carry_common import prompt_prefill, build_step_updates, carry_steps
|
|
dev = ids.device
|
|
calls, _ = looper.capture(ids, msk, logits_to_keep=1)
|
|
e = looper._hin[looper.l0].detach()
|
|
B = ids.shape[0]
|
|
ar = torch.arange(ids.shape[1], device=dev)
|
|
pmask = ar[None, :] < plens[:, None].to(dev)
|
|
S, X = prompt_prefill(looper, adapter, e, calls, pmask, k)
|
|
rows = torch.arange(B, device=dev)
|
|
anchor = (plens - 1).to(dev)
|
|
s0 = S[rows, anchor]
|
|
|
|
def upd(S, X, seed):
|
|
x_new = adapter(e[rows, anchor], seed)
|
|
X = X.clone()
|
|
X[rows, anchor] = x_new.to(X.dtype)
|
|
S = checkpoint(lambda X_: looper.band(X_, calls), X,
|
|
use_reentrant=False)
|
|
return S, X
|
|
|
|
tf_preds, fr_states = [], []
|
|
if tf_on:
|
|
for i in range(m):
|
|
seed = s0 if i == 0 else T[:, i - 1].to(s0.dtype)
|
|
S, X = upd(S, X, seed)
|
|
tf_preds.append(S[rows, anchor])
|
|
for i in range(m):
|
|
seed = s0 if i == 0 else fr_states[-1]
|
|
S, X = upd(S, X, seed)
|
|
fr_states.append(S[rows, anchor])
|
|
total = msk.sum(-1)
|
|
updates = build_step_updates(plens.to(dev), total.to(dev), dev)
|
|
S, X = carry_steps(looper, adapter, e, calls, S, X, updates,
|
|
use_checkpoint=True)
|
|
return looper.suffix_logits(S, calls), tf_preds, fr_states
|
|
|
|
|
|
def lr_at(step):
|
|
if step < WARMUP:
|
|
return LR * (step + 1) / WARMUP
|
|
t = (step - WARMUP) / max(1, STEPS - WARMUP)
|
|
return 1e-4 + 0.5 * (LR - 1e-4) * (1 + math.cos(math.pi * t))
|
|
|
|
|
|
def build_batch(tok, items, p, device="cuda"):
|
|
seqs, labs, plens = [], [], []
|
|
for it in items:
|
|
pr = tok(chat_prompt(tok, it["question"], DIRECT_SUFFIX),
|
|
add_special_tokens=False)["input_ids"]
|
|
a = tok(it["cot"] + "<end_of_turn>",
|
|
add_special_tokens=False)["input_ids"]
|
|
pi = p + it.get("extra_pauses", 0)
|
|
seqs.append(pr + [PAUSE_ID] * pi + a)
|
|
labs.append([-100] * (len(pr) + pi) + a)
|
|
plens.append(len(pr))
|
|
T = max(len(s) for s in seqs)
|
|
pad = tok.pad_token_id or 0
|
|
ids = torch.full((len(seqs), T), pad, dtype=torch.long)
|
|
lab = torch.full((len(seqs), T), -100, dtype=torch.long)
|
|
msk = torch.zeros((len(seqs), T), dtype=torch.long)
|
|
for i, (s, l) in enumerate(zip(seqs, labs)):
|
|
ids[i, : len(s)] = torch.tensor(s)
|
|
lab[i, : len(s)] = torch.tensor(l)
|
|
msk[i, : len(s)] = 1
|
|
return (ids.to(device), msk.to(device), lab.to(device),
|
|
torch.tensor(plens, device=device))
|
|
|
|
|
|
@torch.no_grad()
|
|
def val_loss(looper, adapter, tok, items, p):
|
|
tot, n = 0.0, 0
|
|
for i in range(0, len(items), BATCH):
|
|
chunk = items[i : i + BATCH]
|
|
ids, msk, lab, plens = build_batch(tok, chunk, p)
|
|
ckw = {}
|
|
if ARGS.inner_iters:
|
|
extras = torch.tensor([c.get("extra_pauses", 0) for c in chunk],
|
|
device=plens.device)
|
|
ckw = dict(inner_iters=ARGS.inner_iters,
|
|
inner_at=plens + p + extras - 1)
|
|
logits = carry_logits(looper, adapter, ids, msk, plens, K_PREFILL,
|
|
feedforward=ARGS.feedforward, **ckw)
|
|
loss = F.cross_entropy(logits[:, :-1].flatten(0, 1).float(),
|
|
lab[:, 1:].flatten(), ignore_index=-100)
|
|
tot += loss.item() * len(ids)
|
|
n += len(ids)
|
|
return tot / n
|
|
|
|
|
|
def main():
|
|
rng = random.Random(ARGS.seed)
|
|
torch.manual_seed(ARGS.seed)
|
|
star = {it["idx"]: it for it in json.load(open(OUT / "star_data.json"))
|
|
if it["split"] == "train"}
|
|
cots = json.load(open(OUT / "gsm_cot_data.json"))
|
|
data = []
|
|
n_dropped = 0
|
|
for r in cots:
|
|
it = star.get(r["idx"])
|
|
if it is None:
|
|
continue
|
|
cot, ndel, deleted = (drop_cot_steps(r["cot"], ARGS.drop_steps)
|
|
if ARGS.drop_steps else (r["cot"], 0, ""))
|
|
n_dropped += ndel
|
|
data.append({"question": it["question"], "label": r["label"],
|
|
"cot": cot, "deleted": deleted,
|
|
"extra_pauses": ndel * ARGS.pause_per_step})
|
|
if ARGS.drop_steps:
|
|
print(f"rung B d={ARGS.drop_steps}: {n_dropped} steps deleted "
|
|
f"across {len(data)} items "
|
|
f"({ARGS.pause_per_step} pauses each)", flush=True)
|
|
|
|
model, tok = load_model(dtype=torch.bfloat16)
|
|
for pp in model.parameters():
|
|
pp.requires_grad_(False)
|
|
looper = BandLooper(model)
|
|
adapter = MergeAdapter(
|
|
d=model.config.get_text_config().hidden_size).cuda()
|
|
if ARGS.lensnoise:
|
|
r, sc = ARGS.lensnoise.split(",")
|
|
jbar = torch.load(Path(__file__).resolve().parent.parent
|
|
/ "results/jbar.pt", map_location="cpu")["Jbar"]
|
|
from loop_common import BAND
|
|
adapter = LensNoiseWrapper(adapter, jbar[BAND[0]], int(r),
|
|
float(sc)).cuda()
|
|
print(f"lens-noise: rank={r} scale={sc} on jbar L{BAND[0]}",
|
|
flush=True)
|
|
if ARGS.warm_start:
|
|
(adapter.base if ARGS.lensnoise else adapter).load_state_dict(
|
|
torch.load(ARGS.warm_start, map_location="cuda"))
|
|
print(f"warm-started from {ARGS.warm_start}", flush=True)
|
|
lora_params, lora_band_layers = [], []
|
|
if ARGS.bandlora:
|
|
from lora_band import inject_band_lora
|
|
from loop_common import BAND
|
|
lora_band_layers = list(range(BAND[0], BAND[1] + 1))
|
|
scales = {l: 1.0 for l in lora_band_layers}
|
|
lora_params = inject_band_lora(looper.tm, BAND[0], scales,
|
|
rank=ARGS.bandlora)
|
|
for p in lora_params:
|
|
p.data = p.data.cuda()
|
|
print(f"band-lora r={ARGS.bandlora}: "
|
|
f"{sum(p.numel() for p in lora_params)/1e6:.1f}M params, "
|
|
f"layers {lora_band_layers[0]}-{lora_band_layers[-1]}, "
|
|
f"lr={ARGS.lora_lr}", flush=True)
|
|
lens_teach = None
|
|
if ARGS.lensteach or ARGS.lensteach_gen:
|
|
from loop_common import BAND
|
|
J30 = torch.load(Path(__file__).resolve().parent.parent
|
|
/ "results/jbar.pt",
|
|
map_location="cuda")["Jbar"][BAND[1]].float()
|
|
tm = looper.tm
|
|
softcap = model.config.get_text_config().final_logit_softcapping
|
|
|
|
def lens_teach(h):
|
|
proj = h.float() @ J30.T
|
|
x = tm.norm(proj.to(tm.norm.weight.dtype))
|
|
lg = model.lm_head(x)
|
|
if softcap:
|
|
lg = softcap * torch.tanh(lg / softcap)
|
|
return lg
|
|
|
|
if ARGS.lensteach:
|
|
n_tgt = 0
|
|
for it in data:
|
|
if it["deleted"]:
|
|
cap = ARGS.inner_iters or it["extra_pauses"]
|
|
it["lens_targets"] = tok(
|
|
it["deleted"], add_special_tokens=False
|
|
)["input_ids"][:cap]
|
|
n_tgt += bool(it["lens_targets"])
|
|
print(f"lens-teach λ={ARGS.lensteach}: targets on {n_tgt} "
|
|
f"items (pause j <-> deleted-step token j, lens at "
|
|
f"L{BAND[1]})", flush=True)
|
|
if ARGS.lensteach_gen:
|
|
n_gt, n_spans = 0, 0
|
|
for it in data:
|
|
gt = gen_staging_targets(tok, it["cot"])
|
|
if gt:
|
|
it["gen_targets"] = gt
|
|
n_gt += 1
|
|
n_spans += len(gt)
|
|
print(f"lens-teach-gen λ={ARGS.lensteach_gen}: staging targets "
|
|
f"on {n_gt} items ({n_spans} line-spans; pre-'=' "
|
|
f"positions target the line result)", flush=True)
|
|
if ARGS.traj_tf or ARGS.traj_fr:
|
|
assert ARGS.warm_start and ARGS.inner_iters \
|
|
and ARGS.base_pauses == 0, \
|
|
"traj mode needs --warm-start, --inner-iters, --base-pauses 0"
|
|
data = [it for it in data if it["deleted"]]
|
|
m = ARGS.inner_iters
|
|
t0_ = time.time()
|
|
with torch.no_grad():
|
|
for i in range(0, len(data), BATCH):
|
|
chunk = data[i:i + BATCH]
|
|
full_items = [{"question": c["question"],
|
|
"cot": c["deleted"] + "\n" + c["cot"],
|
|
"extra_pauses": 0} for c in chunk]
|
|
ids_, msk_, _, plens_ = build_batch(tok, full_items, 0)
|
|
_, S_ = carry_logits(looper, adapter, ids_, msk_, plens_,
|
|
K_PREFILL, return_states=True)
|
|
for b, c in enumerate(chunk):
|
|
if ARGS.traj_span == "step":
|
|
span = len(tok(c["deleted"],
|
|
add_special_tokens=False)["input_ids"])
|
|
else:
|
|
span = len(tok(c["deleted"] + "\n" + c["cot"],
|
|
add_special_tokens=False)["input_ids"])
|
|
idxs = [round(j * (span - 1) / (m - 1)) for j in range(m)]
|
|
pos = [int(plens_[b]) + jj for jj in idxs]
|
|
c["traj_states"] = S_[b, pos].float().clone()
|
|
print(f"trajectory waypoints: {len(data)} items x {m} states "
|
|
f"({ARGS.traj_span} span, {time.time()-t0_:.0f}s; frozen "
|
|
f"warm-start teacher, full cot)", flush=True)
|
|
if ARGS.teachstate:
|
|
assert ARGS.warm_start and ARGS.inner_iters, \
|
|
"teachstate needs --warm-start (frozen teacher) + --inner-iters"
|
|
t0_ = time.time()
|
|
todo = [it for it in data if it["deleted"]]
|
|
pt = ARGS.base_pauses if ARGS.base_pauses >= 0 else 0
|
|
with torch.no_grad():
|
|
for i in range(0, len(todo), BATCH):
|
|
chunk = todo[i:i + BATCH]
|
|
full_items = [{"question": c["question"],
|
|
"cot": c["deleted"] + "\n" + c["cot"],
|
|
"extra_pauses": 0} for c in chunk]
|
|
ids_, msk_, _, plens_ = build_batch(tok, full_items, pt)
|
|
_, S_ = carry_logits(looper, adapter, ids_, msk_, plens_,
|
|
K_PREFILL, return_states=True)
|
|
for b, c in enumerate(chunk):
|
|
nstep = len(tok(c["deleted"],
|
|
add_special_tokens=False)["input_ids"])
|
|
pos = int(plens_[b]) + pt + nstep - 1
|
|
c["teacher_state"] = S_[b, pos].float().clone()
|
|
print(f"teacher states: {len(todo)} captured ({time.time()-t0_:.0f}s;"
|
|
f" frozen warm-start adapter, full cot, band-exit at the "
|
|
f"deleted step's last token)", flush=True)
|
|
groups = [{"params": list(adapter.parameters()), "lr": LR, "base": LR}]
|
|
if lora_params:
|
|
groups.append({"params": lora_params, "lr": ARGS.lora_lr,
|
|
"base": ARGS.lora_lr})
|
|
opt = torch.optim.AdamW(groups, weight_decay=0.01)
|
|
|
|
keep = [it for it in data
|
|
if len(tok(it["question"])["input_ids"])
|
|
+ len(tok(it["cot"])["input_ids"])
|
|
+ it.get("extra_pauses", 0) + 16 <= 460]
|
|
rng.shuffle(keep)
|
|
pool = {l: [it for it in keep if it["label"] == l]
|
|
for l in ("easy", "hard")}
|
|
val = {l: pool[l][:12] for l in pool}
|
|
pool = {l: pool[l][12:] for l in pool}
|
|
print(f"pool: easy={len(pool['easy'])} hard={len(pool['hard'])}",
|
|
flush=True)
|
|
if min(len(v) for v in pool.values()) < BATCH:
|
|
print("INSUFFICIENT POOL — aborting", flush=True)
|
|
return
|
|
|
|
log = []
|
|
t0 = time.time()
|
|
for step in range(STEPS):
|
|
lbl = ("easy", "hard")[step % 2]
|
|
batch = rng.sample(pool[lbl], BATCH)
|
|
p = (ARGS.base_pauses if ARGS.base_pauses >= 0
|
|
else P_BY_LABEL[lbl])
|
|
ids, msk, lab, plens = build_batch(tok, batch, p)
|
|
for g in opt.param_groups:
|
|
g["lr"] = g["base"] * lr_at(step) / LR
|
|
ckw, itstates, tfp, frs = {}, None, None, None
|
|
if ARGS.traj_tf or ARGS.traj_fr:
|
|
T = torch.stack([torch.as_tensor(b_["traj_states"])
|
|
for b_ in batch]).cuda()
|
|
logits, tfp, frs = traj_burst_forward(
|
|
looper, adapter, ids, msk, plens, ARGS.inner_iters, T,
|
|
bool(ARGS.traj_tf), K_PREFILL)
|
|
else:
|
|
if ARGS.inner_iters:
|
|
extras = torch.tensor([b_.get("extra_pauses", 0)
|
|
for b_ in batch],
|
|
device=plens.device)
|
|
assert (extras == extras[0]).all(), \
|
|
"inner-iters needs a batch-uniform pause block"
|
|
ckw = dict(inner_iters=ARGS.inner_iters,
|
|
inner_at=plens + p + extras - 1)
|
|
if ARGS.lensteach or ARGS.teachstate:
|
|
itstates = []
|
|
ckw["iter_states"] = itstates
|
|
if ARGS.lensteach or ARGS.lensteach_gen or ARGS.inner_iters:
|
|
logits, S = carry_logits(looper, adapter, ids, msk, plens,
|
|
K_PREFILL, use_checkpoint=True,
|
|
feedforward=ARGS.feedforward,
|
|
return_states=True, **ckw)
|
|
else:
|
|
logits = carry_logits(looper, adapter, ids, msk, plens,
|
|
K_PREFILL, use_checkpoint=True,
|
|
feedforward=ARGS.feedforward)
|
|
loss = F.cross_entropy(logits[:, :-1].flatten(0, 1).float(),
|
|
lab[:, 1:].flatten(), ignore_index=-100)
|
|
lce_val, lgen_val = 0.0, 0.0
|
|
if ARGS.lensteach and ARGS.inner_iters:
|
|
terms = []
|
|
for b, it in enumerate(batch):
|
|
tgt = it.get("lens_targets")
|
|
if not tgt or not itstates:
|
|
continue
|
|
m = min(len(tgt), len(itstates))
|
|
hs = torch.stack([itstates[i][b] for i in range(m)])
|
|
terms.append(F.cross_entropy(
|
|
lens_teach(hs).float(),
|
|
torch.tensor(tgt[:m], device=hs.device)))
|
|
if terms:
|
|
lce = torch.stack(terms).mean()
|
|
lce_val = lce.item()
|
|
loss = loss + ARGS.lensteach * lce
|
|
elif ARGS.lensteach:
|
|
terms = []
|
|
for b, it in enumerate(batch):
|
|
tgt = it.get("lens_targets")
|
|
if not tgt:
|
|
continue
|
|
s0 = int(plens[b]) + p
|
|
hs = S[b, s0 : s0 + len(tgt)]
|
|
terms.append(F.cross_entropy(
|
|
lens_teach(hs).float(),
|
|
torch.tensor(tgt, device=hs.device)))
|
|
if terms:
|
|
lce = torch.stack(terms).mean()
|
|
lce_val = lce.item()
|
|
loss = loss + ARGS.lensteach * lce
|
|
if ARGS.lensteach_gen:
|
|
gterms = []
|
|
for b, it in enumerate(batch):
|
|
base = int(plens[b]) + p + it.get("extra_pauses", 0)
|
|
for pre, res in it.get("gen_targets", []):
|
|
posl = torch.tensor([base + r for r in pre],
|
|
device=S.device)
|
|
lg = lens_teach(S[b, posl]).float()
|
|
gterms.append(torch.stack([
|
|
F.cross_entropy(
|
|
lg, torch.full((len(pre),), r, dtype=torch.long,
|
|
device=S.device))
|
|
for r in res]).mean())
|
|
if gterms:
|
|
gl = torch.stack(gterms).mean()
|
|
lgen_val = gl.item()
|
|
loss = loss + ARGS.lensteach_gen * gl
|
|
lts_val = 0.0
|
|
if ARGS.teachstate and itstates:
|
|
tterms = []
|
|
for b, it in enumerate(batch):
|
|
T = it.get("teacher_state")
|
|
if T is None:
|
|
continue
|
|
tterms.append(1 - F.cosine_similarity(
|
|
itstates[-1][b].float(), T, dim=0))
|
|
if tterms:
|
|
lt = torch.stack(tterms).mean()
|
|
lts_val = lt.item()
|
|
loss = loss + ARGS.teachstate * lt
|
|
ltf_val, lfr_val = 0.0, 0.0
|
|
if ARGS.traj_tf and tfp:
|
|
ltf = torch.stack([
|
|
(1 - F.cosine_similarity(p_.float(), T[:, i], dim=-1)).mean()
|
|
for i, p_ in enumerate(tfp)]).mean()
|
|
ltf_val = ltf.item()
|
|
loss = loss + ARGS.traj_tf * ltf
|
|
if ARGS.traj_fr and frs:
|
|
lfr = torch.stack([
|
|
(1 - F.cosine_similarity(s_.float(), T[:, i], dim=-1)).mean()
|
|
for i, s_ in enumerate(frs)]).mean()
|
|
lfr_val = lfr.item()
|
|
loss = loss + ARGS.traj_fr * lfr
|
|
opt.zero_grad(set_to_none=True)
|
|
loss.backward()
|
|
torch.nn.utils.clip_grad_norm_(
|
|
list(adapter.parameters()) + lora_params, 1.0)
|
|
opt.step()
|
|
log.append({"step": step, "loss": loss.item(), "lce": lce_val,
|
|
"lgen": lgen_val, "lts": lts_val, "ltf": ltf_val,
|
|
"lfr": lfr_val})
|
|
if step % 10 == 0:
|
|
print(f"step {step:4d} {lbl:4s} loss={loss.item():.4f} "
|
|
f"lce={lce_val:.3f} lgen={lgen_val:.3f} "
|
|
f"lts={lts_val:.3f} ltf={ltf_val:.3f} "
|
|
f"lfr={lfr_val:.3f} "
|
|
f"({(time.time()-t0)/(step+1):.1f}s/step)", flush=True)
|
|
if step % 200 == 199 or step == STEPS - 1:
|
|
if ARGS.lensnoise:
|
|
adapter.noise_on = False
|
|
for l in ("easy", "hard"):
|
|
v = val_loss(looper, adapter, tok, val[l],
|
|
ARGS.base_pauses
|
|
if ARGS.base_pauses >= 0
|
|
else P_BY_LABEL[l])
|
|
print(f" val@{step}: {l}={v:.3f}", flush=True)
|
|
if ARGS.lensnoise:
|
|
adapter.noise_on = True
|
|
sd = (adapter.base if ARGS.lensnoise else adapter).state_dict()
|
|
torch.save(sd, OUT / f"adapter_{TAG}_e{step+1}.pt")
|
|
if lora_params:
|
|
torch.save({"rank": ARGS.bandlora, "band": lora_band_layers,
|
|
"tensors": [p.detach().cpu()
|
|
for p in lora_params]},
|
|
OUT / f"lora_{TAG}_e{step+1}.pt")
|
|
json.dump(log, open(OUT / f"train_{TAG}_log.json", "w"))
|
|
print("done", flush=True)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|