item 29 pre-registered (Nils's design): trajectory teacher-forcing — 10 waypoint transitions supervised independently (T[i-1]→T[i]), TF and TF+FR arms; job 1 d=1 step-span, job 2 answer-only full-CoT span
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
Nils
Claude Fable 5
+131
-21
@@ -77,6 +77,18 @@ 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 "
|
||||
@@ -94,6 +106,9 @@ TAG = ("carrycot_ff" if ARGS.feedforward else "carrycot") + (
|
||||
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
|
||||
@@ -173,6 +188,49 @@ def gen_staging_targets(tok, cot):
|
||||
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
|
||||
@@ -321,6 +379,35 @@ def main():
|
||||
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"
|
||||
@@ -375,26 +462,34 @@ def main():
|
||||
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 = {}, None
|
||||
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)
|
||||
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:
|
||||
logits = carry_logits(looper, adapter, ids, msk, plens,
|
||||
K_PREFILL, use_checkpoint=True,
|
||||
feedforward=ARGS.feedforward)
|
||||
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
|
||||
@@ -458,17 +553,32 @@ def main():
|
||||
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})
|
||||
"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} "
|
||||
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:
|
||||
|
||||
Reference in New Issue
Block a user