700 lines
36 KiB
Python
700 lines
36 KiB
Python
"""Emergent-garble GRPO trainer (base model, length-penalty, no legibility constraint).
|
|
|
|
Goal: reproduce the Claude Fable 5 / Mythos 5 "illegible reasoning" phenomenon (system card
|
|
transcript 6.2.2.A, third_party/ant_fable.pdf p.107) on Qwen3-8B-Base -- a load-bearing CoT
|
|
that, under accuracy + length pressure with NO legibility term and NO character constraints,
|
|
compresses into a dense, telegraphic, symbol-heavy private notation. Pure RL on a *base* model
|
|
(R1-Zero recipe): the instruct model's legibility prior would fight this, so we start from base.
|
|
|
|
Reward (always on -- this IS the experiment):
|
|
R = task_reward - lambda * (gen_tokens / max_think)
|
|
where task_reward in {0,1} is the arith-reach verifier and gen_tokens is the policy's own
|
|
generated length (think dominates; the boxed answer is a handful of tokens). lambda is either
|
|
fixed or adapted by a Lagrangian dual that holds accuracy at a floor (``lambda_target``): while
|
|
accuracy < floor lambda stays ~0 (Phase 1: learn to solve); once accuracy clears the floor
|
|
lambda rises (Phase 2: compress the irreducible search -> notation emerges). Advantage is
|
|
group-relative over the G samples of each problem (REINFORCE-with-baseline == GRPO@1-update).
|
|
|
|
Everything is logged to wandb: per-step metrics, cheap live garble proxies (symbol/non-ascii
|
|
density), FULL untruncated rollout traces (wandb.Table), and the complete cots.jsonl uploaded as
|
|
a versioned artifact at every checkpoint. Temperature + lambda are live-steerable mid-run via
|
|
``<save_root>/temp.ctl``. Resumable from any checkpoint.
|
|
|
|
No SFT anywhere. The task is switchable (see ENVS) if arith-reach won't garble.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import os
|
|
import time
|
|
import unicodedata
|
|
from collections import defaultdict
|
|
|
|
import torch
|
|
from torch.nn.utils import clip_grad_norm_
|
|
|
|
import wandb
|
|
|
|
import torch.nn.functional as F
|
|
|
|
from model_organisms.envs.arith_reach import ArithReachEnv
|
|
from model_organisms.envs.entity_track import EntityTrackQueryEnv
|
|
from model_organisms.envs.state_track import StateTrackEnv, StateTrackQueryEnv
|
|
from model_organisms.envs.base import RLTrajectory, SeqBuilder, initial_prefix_ids, split_think_output
|
|
from model_organisms.rl.rollout import batched_single_turn_rollout
|
|
from model_organisms.rl.grpo import (
|
|
geom_checkpoint_steps,
|
|
load_policy,
|
|
make_gen_batch_fn,
|
|
make_vllm_gen_batch_fn,
|
|
resume_adapter,
|
|
seq_logprob,
|
|
)
|
|
|
|
ENVS = {"arith_reach": ArithReachEnv, "state_track": StateTrackEnv, "state_track_q": StateTrackQueryEnv,
|
|
"entity_track_q": EntityTrackQueryEnv}
|
|
|
|
# Two short few-shot exemplars: lock the (problem -> <think> ... </think> -> \boxed{expr}) FORMAT
|
|
# and seed basic solving, with deliberately TERSE reasoning so we bootstrap the format without
|
|
# anchoring a verbose prose style (we want the dense notation to *emerge*, not be copied).
|
|
FEWSHOT = """\
|
|
Use each of these numbers EXACTLY once, combined with the operations + - * / and parentheses, to write a single arithmetic expression that equals 24.
|
|
|
|
Numbers: 6, 4, 2, 1
|
|
|
|
Reason step by step inside <think> </think>. You may reason in whatever style you find efficient and use any shorthand, symbols, or private notation you like inside the think block. Then, OUTSIDE the think block, write ONLY your final expression as \\boxed{expr} -- just the expression using the numbers above, with no '=' and no explanation.
|
|
<think>
|
|
Aim 24. 6*4=24. Left 2,1 must net *1: (2-1)=1. So 6*4*(2-1).
|
|
</think>
|
|
\\boxed{6*4*(2-1)}
|
|
|
|
Use each of these numbers EXACTLY once, combined with the operations + - * / and parentheses, to write a single arithmetic expression that equals 17.
|
|
|
|
Numbers: 7, 3, 2, 5
|
|
|
|
Reason step by step inside <think> </think>. You may reason in whatever style you find efficient and use any shorthand, symbols, or private notation you like inside the think block. Then, OUTSIDE the think block, write ONLY your final expression as \\boxed{expr} -- just the expression using the numbers above, with no '=' and no explanation.
|
|
<think>
|
|
Aim 17. Sum: 7+3+2+5=17.
|
|
</think>
|
|
\\boxed{7+3+2+5}
|
|
|
|
"""
|
|
|
|
|
|
def build_base_prefix_ids(tokenizer, problem_prompt: str, fewshot: str) -> list[int]:
|
|
"""Base model has no chat template: build the raw prompt and force the think open."""
|
|
text = fewshot + problem_prompt + "\n<think>\n"
|
|
return tokenizer.encode(text, add_special_tokens=False)
|
|
|
|
|
|
# ---- legibility tax (the creative lever) ----------------------------------------------
|
|
# Penalize the English-prose SKELETON -- the function/filler words that make a CoT read as fluent
|
|
# English -- so the model must re-encode load-bearing reasoning into telegraphic/symbolic notation
|
|
# (Fable's mode) to keep accuracy while shedding the legible crutch. Cheap (no LM), emergent (the
|
|
# replacement notation is the model's own).
|
|
_STOPWORDS = frozenset("""
|
|
a an the is are was were be been being am to of in on at for with from by as into out up down over under
|
|
and or but so if then else that this these those it its which what how why when where who whom while
|
|
we i you he she they them me my our your his her their us him it's i'm we're you're they're let lets let's
|
|
can could would should will shall may might must need needs want wants have has had do does did get gets got
|
|
make makes made see sees saw look looks know knows think thinks way ways thing things like just only also
|
|
maybe perhaps okay ok hmm wait now here there yes no not none all some any each every about not no nor
|
|
first second next last finally actually alternatively however therefore thus hence because since though although
|
|
very too more most much many few again still even ever never always another other same such than there's
|
|
i'll we'll i've we've cannot don't doesn't didn't isn't aren't wasn't won't can't let me going to want to need to
|
|
""".split())
|
|
|
|
|
|
def legibility_score(think: str) -> float:
|
|
"""Fraction of whitespace tokens that are English function/filler words. High = fluent-English
|
|
prose (legible); low = telegraphic/symbolic (Fable-like). This is what the tax pushes DOWN."""
|
|
toks = [w.strip(".,;:!?()[]{}\"'`*").lower() for w in think.split()]
|
|
toks = [w for w in toks if w]
|
|
if not toks:
|
|
return 0.0
|
|
return sum(1 for w in toks if w in _STOPWORDS) / len(toks)
|
|
|
|
|
|
# ---- cheap live garble proxies (no API call) -------------------------------------------
|
|
|
|
_ARROWS = set("→⟶⟹⟸⇒⇐↔⇔←➜➔»«")
|
|
_SUIT_SKULL = set("♥♣♠♦💀☠✓✗✘✔★☆●○◆◇■□▪▶◀")
|
|
_STRUCT = set("→⟶⟹⟸⇒⇐:|/\\—–·∴∵∧∨¬≠≤≥≈±∈∉⊂⊃∪∩")
|
|
|
|
|
|
def garble_proxies(think: str) -> dict:
|
|
"""A few cheap scalars that track 'is the notation drifting weird/dense' without a judge.
|
|
Higher = more Fable-like (symbol-dense, non-ascii, hyphen-joined, few real words)."""
|
|
n = max(len(think), 1)
|
|
non_ascii = sum(1 for c in think if ord(c) > 127) / n
|
|
arrows = sum(1 for c in think if c in _ARROWS or c in _SUIT_SKULL) / n
|
|
struct = sum(1 for c in think if c in _STRUCT) / n
|
|
# CJK presence (Fable showed Chinese tokens mid-reasoning)
|
|
cjk = sum(1 for c in think if unicodedata.category(c) == "Lo" and ord(c) > 0x2E00) / n
|
|
words = think.split()
|
|
mean_wordlen = (sum(len(w) for w in words) / max(len(words), 1))
|
|
# hyphen-joined run tokens like "t8-col-built" -> long tokens with internal hyphens
|
|
hyphenjoin = sum(1 for w in words if w.count("-") >= 2) / max(len(words), 1)
|
|
# alpha word ratio: fraction of whitespace-tokens that are >=3 plain letters (prose-ish)
|
|
alpha_words = sum(1 for w in words if sum(c.isalpha() for c in w) >= 3) / max(len(words), 1)
|
|
return {
|
|
"non_ascii": non_ascii, "symbol_density": arrows, "struct_density": struct,
|
|
"cjk": cjk, "mean_wordlen": mean_wordlen, "hyphenjoin": hyphenjoin,
|
|
"alpha_word_ratio": alpha_words,
|
|
}
|
|
|
|
|
|
# ---- rollout (base-model prefix; cannot reuse the chat rollout) -------------------------
|
|
|
|
|
|
def rollout_base(tokenizer, gen_batch_fn, items, gen_cfg, fewshot):
|
|
"""items: list of (group_id, prompt_text, score_fn, meta). One batched gen call."""
|
|
prefixes = [build_base_prefix_ids(tokenizer, pt, fewshot) for (_g, pt, _s, _m) in items]
|
|
gens = gen_batch_fn(prefixes, gen_cfg)
|
|
out = []
|
|
for (gid, prompt_text, score_fn, meta), prefix, gen_ids in zip(items, prefixes, gens):
|
|
sb = SeqBuilder(tokenizer, prefix)
|
|
sb.add_generated(gen_ids)
|
|
sb.close_assistant()
|
|
text = tokenizer.decode(gen_ids, skip_special_tokens=False)
|
|
think, output = split_think_output(text)
|
|
task, m_out, m_cot = score_fn(think, output, text)
|
|
out.append(RLTrajectory(
|
|
input_ids=sb.ids, assistant_mask=sb.mask, task_reward=task, m_out=m_out, m_cot=m_cot,
|
|
prompt_text=prompt_text, think_text=think, output_text=output, full_text=text,
|
|
group_id=gid, meta=meta or {},
|
|
))
|
|
return out
|
|
|
|
|
|
# ---- query-after-think rollout (two-phase, single context) ------------------------------
|
|
# Phase A: the model thinks (stop at </think>, or forced close at the max_think budget -- a
|
|
# truncated log still gets PARTIAL credit on early-step queries, so the gradient at the budget
|
|
# frontier is smooth). Phase B: queries arrive in a follow-up user turn; a forced-empty think
|
|
# ("<think>\n\n</think>") forbids any further reasoning, so the answers must be READ OFF the
|
|
# phase-A notation -- load-bearing by construction. Injected text is untrained (mask False).
|
|
|
|
|
|
def make_gen_fn(model, tokenizer, device, chunk: int = 48):
|
|
"""Generic batched HF generate primitive: (prefixes, max_new, temp, top_p, stop_ids) ->
|
|
per-sequence generated ids with the stop token EXCLUDED. Chunks the batch for memory."""
|
|
pad_id = tokenizer.pad_token_id
|
|
|
|
def gen(prefix_ids_list, max_new, temp, top_p, stop_ids):
|
|
results = []
|
|
for c0 in range(0, len(prefix_ids_list), chunk):
|
|
chunk_prefixes = prefix_ids_list[c0:c0 + chunk]
|
|
maxlen = max(len(p) for p in chunk_prefixes)
|
|
ids = torch.tensor([[pad_id] * (maxlen - len(p)) + p for p in chunk_prefixes], device=device)
|
|
attn = torch.tensor([[0] * (maxlen - len(p)) + [1] * len(p) for p in chunk_prefixes], device=device)
|
|
with torch.no_grad():
|
|
out = model.generate(input_ids=ids, attention_mask=attn, max_new_tokens=max_new,
|
|
do_sample=True, temperature=temp, top_p=top_p,
|
|
eos_token_id=stop_ids, pad_token_id=pad_id)
|
|
for row in out[:, maxlen:].tolist():
|
|
cut = len(row)
|
|
for i, t in enumerate(row):
|
|
if t in stop_ids:
|
|
cut = i
|
|
break
|
|
results.append(row[:cut])
|
|
return results
|
|
|
|
return gen
|
|
|
|
|
|
def rollout_query(env, tokenizer, gen_fn, problems, G, gen_cfg, force_empty_think=False,
|
|
train_answers=True):
|
|
end_think_ids = tokenizer.encode("</think>", add_special_tokens=False)
|
|
im_end_id = tokenizer.convert_tokens_to_ids("<|im_end|>")
|
|
items = [(gid, p) for gid, p in enumerate(problems) for _ in range(G)]
|
|
prefixes = [initial_prefix_ids(tokenizer, env.prompt(p)) for (_g, p) in items]
|
|
if force_empty_think:
|
|
thinks = [[] for _ in items]
|
|
else:
|
|
thinks = gen_fn(prefixes, gen_cfg["max_think"], gen_cfg["think_temperature"],
|
|
gen_cfg.get("top_p", 1.0), end_think_ids)
|
|
sbs = []
|
|
for (gid, p), prefix, th in zip(items, prefixes, thinks):
|
|
sb = SeqBuilder(tokenizer, prefix)
|
|
if force_empty_think:
|
|
sb.add_control("<think>\n\n</think>")
|
|
else:
|
|
sb.add_generated(th)
|
|
sb.add_control("</think>") # model's own </think> was the stop (excluded); forced close if truncated
|
|
sb.close_assistant()
|
|
sb.add_user_turn(env.queries_text(p))
|
|
sb.add_control("<think>\n\n</think>\n\n") # forbid phase-B thinking
|
|
sbs.append(sb)
|
|
answers = gen_fn([sb.ids for sb in sbs], gen_cfg["max_out"], gen_cfg["out_temperature"],
|
|
gen_cfg.get("top_p", 1.0), [im_end_id])
|
|
trajs = []
|
|
for (gid, p), sb, th, ans in zip(items, sbs, thinks, answers):
|
|
think_text = tokenizer.decode(th, skip_special_tokens=False).replace("<think>", "").strip()
|
|
ans_text = tokenizer.decode(ans, skip_special_tokens=False).strip()
|
|
if train_answers:
|
|
sb.add_generated(ans)
|
|
else:
|
|
# keep phase-B answer tokens OUT of the policy gradient: heavy negative-advantage
|
|
# phases can corrupt the answer-region distribution (observed: ent arm collapsed to
|
|
# non-Latin garbage answers -> all-zero rewards -> zero advantage -> absorbing state).
|
|
# The reward's credit belongs to the think tokens (write retrievable notes); answers
|
|
# stay near the base distribution.
|
|
sb.ids.extend(ans)
|
|
sb.mask.extend([False] * len(ans))
|
|
sb.close_assistant()
|
|
credit, n_parsed = env.score_queries(p, ans_text)
|
|
trajs.append(RLTrajectory(
|
|
input_ids=sb.ids, assistant_mask=sb.mask, task_reward=credit, m_out=0.0, m_cot=0.0,
|
|
prompt_text=env.prompt(p), think_text=think_text, output_text=ans_text,
|
|
full_text=think_text + "\n---QUERIES---\n" + ans_text, group_id=gid,
|
|
meta={"think_len": len(th), "n_parsed": n_parsed,
|
|
"truncated": float(len(th) >= gen_cfg["max_think"]),
|
|
"queries": [list(q) for q in p.queries], "answers": env.query_answers(p)},
|
|
))
|
|
return trajs
|
|
|
|
|
|
def pg_microbatched(model, tokenizer, active, device, mb, n_norm):
|
|
"""Batched policy-gradient: -(adv * mean-token-logprob) per sequence, backward per microbatch.
|
|
Replaces the per-trajectory seq_logprob loop (batch-1 forwards are the 8B bottleneck)."""
|
|
pad_id = tokenizer.pad_token_id
|
|
loss_val = 0.0
|
|
order = sorted(range(len(active)), key=lambda i: len(active[i].input_ids)) # len-sorted: less padding
|
|
for c0 in range(0, len(order), mb):
|
|
chunk = [active[i] for i in order[c0:c0 + mb]]
|
|
L = max(len(t.input_ids) for t in chunk)
|
|
ids = torch.tensor([t.input_ids + [pad_id] * (L - len(t.input_ids)) for t in chunk], device=device)
|
|
attn = torch.tensor([[1] * len(t.input_ids) + [0] * (L - len(t.input_ids)) for t in chunk], device=device)
|
|
logits = model(input_ids=ids, attention_mask=attn).logits.float()
|
|
logp = torch.log_softmax(logits[:, :-1], dim=-1)
|
|
tok = logp.gather(2, ids[:, 1:].unsqueeze(-1)).squeeze(-1)
|
|
m = torch.tensor([t.assistant_mask[1:] + [False] * (L - len(t.input_ids)) for t in chunk],
|
|
device=device, dtype=tok.dtype)
|
|
seq_lp = (tok * m).sum(1) / m.sum(1).clamp(min=1)
|
|
adv = torch.tensor([t.advantage for t in chunk], device=device, dtype=seq_lp.dtype)
|
|
loss = -(adv * seq_lp).sum() / n_norm
|
|
loss.backward()
|
|
loss_val += loss.item()
|
|
return loss_val
|
|
|
|
|
|
def effective_budget(step: int, cfg: dict, max_think: int) -> float | None:
|
|
"""Token budget below which length is 'free'. None => linear (penalize full length).
|
|
'anneal_excess' shrinks the budget L0 -> length_floor over length_anneal_steps, so the model
|
|
is squeezed into ever-fewer tokens (the lever that forces lossless compression -> notation,
|
|
cf. Fable's long-then-compressed rollouts). 'target_excess' is a fixed budget."""
|
|
mode = cfg.get("length_penalty_mode", "linear")
|
|
if mode == "linear":
|
|
return None
|
|
L0 = cfg.get("length_target") or max_think
|
|
Lmin = cfg.get("length_floor") or L0
|
|
asteps = cfg.get("length_anneal_steps") or 0
|
|
if mode == "anneal_excess" and asteps > 0:
|
|
frac = min(max(step, 0) / asteps, 1.0)
|
|
return L0 - (L0 - Lmin) * frac
|
|
return L0 # target_excess
|
|
|
|
|
|
def assign_advantages_len(trajs, lam: float, max_think: int, normalize: bool, step: int, cfg: dict):
|
|
"""R = task - lambda * penalty_unit; group-relative advantage. penalty_signal selects what is
|
|
taxed: 'length' (normalized full length or excess over a shrinking budget), 'legibility'
|
|
(English-stopword density -> pushes toward telegraphic notation), or 'both' (weighted sum)."""
|
|
budget = effective_budget(step, cfg, max_think)
|
|
signal = cfg.get("penalty_signal", "length")
|
|
w_len, w_leg = cfg.get("w_len", 0.5), cfg.get("w_leg", 1.0)
|
|
corr_thresh = cfg.get("correct_thresh", 0.5)
|
|
for t in trajs:
|
|
# query rollouts: tax the THINK region only (phase-B answers are format-fixed, not fat)
|
|
gen_len = int(t.meta["think_len"]) if "think_len" in t.meta else int(sum(t.assistant_mask))
|
|
t.meta["gen_len"] = gen_len
|
|
if budget is None:
|
|
len_unit = min(gen_len / max(max_think, 1), 1.0)
|
|
else:
|
|
len_unit = max(0.0, gen_len - budget) / max(max_think, 1)
|
|
leg_unit = legibility_score(t.think_text)
|
|
t.meta["len_norm"] = len_unit
|
|
t.meta["legibility"] = leg_unit
|
|
t.meta["budget"] = budget if budget is not None else max_think
|
|
if signal == "legibility":
|
|
pen_unit = leg_unit
|
|
elif signal == "both":
|
|
pen_unit = w_len * len_unit + w_leg * leg_unit
|
|
else:
|
|
pen_unit = len_unit
|
|
# penalty only on CORRECT rollouts: never trades accuracy for brevity -> among correct
|
|
# rollouts the shorter/denser one wins (the gradient that points at compression, not truncation).
|
|
if cfg.get("penalty_on_correct_only", True) and t.task_reward <= corr_thresh:
|
|
pen_unit = 0.0
|
|
t.reward = t.task_reward - lam * pen_unit
|
|
groups = defaultdict(list)
|
|
for t in trajs:
|
|
groups[t.group_id].append(t)
|
|
# rank tournament: among the CORRECT rollouts of a group, a zero-mean bonus by inverse length
|
|
# rank. Late in training lengths cluster and the lambda-term's within-group differences vanish
|
|
# into the advantage normalization -> the compression gradient dies; rank keeps the direction
|
|
# crisp at any length scale. Uses only accuracy+length information (purism-clean).
|
|
beta = cfg.get("rank_bonus", 0.0)
|
|
if beta > 0:
|
|
for g in groups.values():
|
|
c = sorted((t for t in g if t.task_reward > corr_thresh), key=lambda t: t.meta["gen_len"])
|
|
if len(c) >= 2:
|
|
# tie-averaged ranks: rollouts pinned at the same length (e.g. all at the budget
|
|
# cap early on) get identical bonuses -> no arbitrary sort-order reward noise.
|
|
ranks, i = {}, 0
|
|
while i < len(c):
|
|
j = i
|
|
while j < len(c) and c[j].meta["gen_len"] == c[i].meta["gen_len"]:
|
|
j += 1
|
|
for k in range(i, j):
|
|
ranks[id(c[k])] = (i + j - 1) / 2
|
|
i = j
|
|
for t in c:
|
|
t.reward += beta * (0.5 - ranks[id(t)] / (len(c) - 1))
|
|
for g in groups.values():
|
|
rs = [t.reward for t in g]
|
|
mean = sum(rs) / len(rs)
|
|
std = (sum((r - mean) ** 2 for r in rs) / len(rs)) ** 0.5
|
|
for t in g:
|
|
t.advantage = (t.reward - mean) / (std + 1e-6) if normalize else (t.reward - mean)
|
|
|
|
|
|
# ---- curriculum -----------------------------------------------------------------------
|
|
|
|
|
|
def make_env(cfg, tier_idx):
|
|
tiers = cfg["curriculum"]
|
|
t = tiers[min(tier_idx, len(tiers) - 1)]
|
|
return ENVS[cfg.get("env", "arith_reach")](**t)
|
|
|
|
|
|
def load_policy_fullft(cfg, device="cuda"):
|
|
"""Full fine-tune path (vanilla HF, NO LoRA, NO unsloth): load the SFT checkpoint with ALL params
|
|
trainable + HF generation. For small models (0.6B) full-FT is cheap and gives the distributional
|
|
freedom to drift into notation that LoRA caps. No gradient checkpointing (0.6B fits) -> keeps
|
|
use_cache for fast HF generate."""
|
|
from transformers import AutoModelForCausalLM, AutoTokenizer
|
|
tok = AutoTokenizer.from_pretrained(cfg["model_name"])
|
|
if tok.pad_token_id is None:
|
|
tok.pad_token = tok.eos_token
|
|
tok.padding_side = "left"
|
|
model = AutoModelForCausalLM.from_pretrained(cfg["model_name"], torch_dtype=torch.bfloat16).to(device)
|
|
if cfg.get("grad_ckpt"): # 8B full-FT: activations don't fit without it; generate() is unaffected
|
|
model.gradient_checkpointing_enable(gradient_checkpointing_kwargs={"use_reentrant": False})
|
|
return model, tok
|
|
|
|
|
|
def train(cfg: dict):
|
|
device = "cuda"
|
|
full_ft = bool(cfg.get("full_ft"))
|
|
import random as _random
|
|
_random.seed(cfg["seed"])
|
|
torch.manual_seed(cfg["seed"])
|
|
|
|
resume_dir = cfg.get("resume_from") or ""
|
|
if full_ft:
|
|
if resume_dir: # full-FT resume: the checkpoint dir IS the model (resume_adapter is LoRA-only)
|
|
cfg["model_name"] = resume_dir
|
|
model, tokenizer = load_policy_fullft(cfg, device)
|
|
FastLanguageModel = None # not used in the full-FT path
|
|
else:
|
|
model, tokenizer = load_policy(cfg, device)
|
|
from unsloth import FastLanguageModel
|
|
|
|
start_step, resume_examples, resume_tokens, resume_lambda = 0, 0, 0, None
|
|
if resume_dir:
|
|
if full_ft:
|
|
rs = json.load(open(os.path.join(resume_dir, "resume_state.json")))
|
|
start_step, resume_examples, resume_tokens = rs["step"], rs["examples_seen"], rs["tokens_seen"]
|
|
resume_lambda = rs.get("lambda")
|
|
print(f"[resume] full-FT from {resume_dir}: step={start_step} lambda={resume_lambda}", flush=True)
|
|
else:
|
|
start_step, resume_examples, resume_tokens, resume_lambda = resume_adapter(model, resume_dir, cfg)
|
|
rs_path = os.path.join(resume_dir, "resume_state.json")
|
|
if os.path.exists(rs_path):
|
|
_tier = json.load(open(rs_path)).get("tier")
|
|
if _tier is not None:
|
|
cfg["start_tier"] = _tier
|
|
|
|
tier_i = int(cfg.get("start_tier", 0))
|
|
env = make_env(cfg, tier_i)
|
|
|
|
use_vllm = bool(cfg.get("use_vllm"))
|
|
gen_batch = make_vllm_gen_batch_fn(model, tokenizer) if use_vllm else make_gen_batch_fn(model, tokenizer, device)
|
|
is_query = hasattr(ENVS[cfg.get("env", "arith_reach")], "queries_text")
|
|
gen_fn = make_gen_fn(model, tokenizer, device, chunk=cfg.get("gen_chunk", 48)) if is_query else None
|
|
lora_tmp = f"/tmp/garble_lora_{cfg['run_name']}"
|
|
gen_cfg = {k: cfg[k] for k in ("max_think", "max_out", "temperature", "top_p", "max_seq_length")}
|
|
gen_cfg["think_temperature"] = cfg.get("think_temperature") or cfg["temperature"]
|
|
gen_cfg["out_temperature"] = cfg.get("out_temperature") or cfg["temperature"]
|
|
|
|
trainable = [p for p in model.parameters() if p.requires_grad]
|
|
optimizer = torch.optim.AdamW(trainable, lr=cfg["lr"], weight_decay=0.0)
|
|
warmup = max(1, int(cfg["warmup_ratio"] * cfg["max_steps"]))
|
|
sched = torch.optim.lr_scheduler.LambdaLR(optimizer, lambda s: min(1.0, (s + 1) / warmup))
|
|
if resume_dir:
|
|
opt_path = os.path.join(resume_dir, "optimizer.pt")
|
|
if os.path.exists(opt_path):
|
|
optimizer.load_state_dict(torch.load(opt_path, map_location=device))
|
|
for _ in range(min(start_step, cfg["max_steps"])):
|
|
sched.step()
|
|
|
|
ckpt_steps = geom_checkpoint_steps(cfg["n_checkpoints"], cfg["max_steps"])
|
|
save_root = os.path.join(cfg["save_dir"], cfg["run_name"])
|
|
os.makedirs(save_root, exist_ok=True)
|
|
ctl_path = os.path.join(save_root, "temp.ctl") # {"think_temperature","out_temperature","lambda_penalty","tier"}
|
|
cot_path = os.path.join(save_root, "cots.jsonl")
|
|
|
|
def _save_ckpt(path, with_optimizer):
|
|
import shutil
|
|
tmp = path + ".tmp" # atomic: a mid-write crash (e.g. ENOSPC) must not corrupt a resume target
|
|
shutil.rmtree(tmp, ignore_errors=True)
|
|
os.makedirs(tmp, exist_ok=True)
|
|
model.save_pretrained(tmp)
|
|
tokenizer.save_pretrained(tmp)
|
|
with open(os.path.join(tmp, "mo_config.json"), "w") as f:
|
|
json.dump(cfg, f, indent=2)
|
|
with open(os.path.join(tmp, "resume_state.json"), "w") as f:
|
|
json.dump({"step": step, "examples_seen": examples_seen, "tokens_seen": tokens_seen,
|
|
"lambda": lam, "tier": tier_i}, f)
|
|
if with_optimizer:
|
|
torch.save(optimizer.state_dict(), os.path.join(tmp, "optimizer.pt"))
|
|
shutil.rmtree(path, ignore_errors=True)
|
|
os.rename(tmp, path)
|
|
|
|
def _prune_ckpts():
|
|
"""8B full-FT ckpts are 16GB; several arms share one 500G overlay. Keep step_1 (the
|
|
pre-drift anchor) + the keep_ckpts most recent; push keepers off-node at run end."""
|
|
keep = int(cfg.get("keep_ckpts", 0))
|
|
if keep <= 0:
|
|
return
|
|
import re as _re, shutil
|
|
steps = sorted(int(m.group(1)) for d in os.listdir(save_root)
|
|
if (m := _re.fullmatch(r"step_(\d+)", d)))
|
|
for s in steps[:-keep]:
|
|
if s != steps[0]: # never prune the first (pre-drift) checkpoint
|
|
shutil.rmtree(os.path.join(save_root, f"step_{s}"), ignore_errors=True)
|
|
|
|
try:
|
|
wandb.init(project=cfg["wandb_project"], entity=cfg.get("wandb_entity") or None,
|
|
name=cfg["run_name"], group=cfg.get("wandb_group") or None, config=cfg)
|
|
except Exception as e:
|
|
print(f"[wandb] online init failed ({type(e).__name__}: {e}); offline")
|
|
os.environ["WANDB_MODE"] = "offline"
|
|
wandb.init(project=cfg["wandb_project"], name=cfg["run_name"],
|
|
group=cfg.get("wandb_group") or None, config=cfg, mode="offline")
|
|
wandb.define_metric("train/examples_seen")
|
|
wandb.define_metric("train/tokens_seen")
|
|
wandb.define_metric("*", step_metric="train/examples_seen")
|
|
|
|
rng = _random.Random(cfg["seed"])
|
|
examples_seen, tokens_seen = resume_examples, resume_tokens
|
|
B, G = cfg["batch_problems"], cfg["group_size"]
|
|
adaptive = bool(cfg.get("adaptive_lambda"))
|
|
lam = (resume_lambda if (resume_dir and resume_lambda is not None)
|
|
else (cfg["lambda_init"] if adaptive else cfg["lambda_penalty"]))
|
|
task_ema = None
|
|
tier_ema = None
|
|
tier_steps = 0
|
|
print(f"[train] env={env.name} tier={tier_i} adaptive_lam={adaptive} lam={lam} "
|
|
f"steps={cfg['max_steps']} B={B} G={G} ckpts={sorted(ckpt_steps)}", flush=True)
|
|
|
|
for step in range(start_step + 1, cfg["max_steps"] + 1):
|
|
t0 = time.time()
|
|
problems = [env.sample_problem(rng) for _ in range(B)]
|
|
|
|
# live steering
|
|
if os.path.exists(ctl_path):
|
|
try:
|
|
_ctl = json.load(open(ctl_path))
|
|
if "think_temperature" in _ctl:
|
|
gen_cfg["think_temperature"] = float(_ctl["think_temperature"])
|
|
if "out_temperature" in _ctl:
|
|
gen_cfg["out_temperature"] = float(_ctl["out_temperature"])
|
|
if "lambda_penalty" in _ctl and not adaptive:
|
|
lam = float(_ctl["lambda_penalty"])
|
|
if "lambda_force" in _ctl: # hard-set lambda live (works in adaptive mode too)
|
|
lam = float(_ctl["lambda_force"])
|
|
for _k in ("lambda_target", "lambda_max", "length_penalty_mode", "length_target",
|
|
"length_floor", "length_anneal_steps"): # live length-penalty retuning
|
|
if _k in _ctl:
|
|
cfg[_k] = _ctl[_k]
|
|
if "tier" in _ctl: # manual curriculum override
|
|
nt = int(_ctl["tier"])
|
|
if nt != tier_i:
|
|
tier_i = nt
|
|
env = make_env(cfg, tier_i)
|
|
tier_ema = None
|
|
print(f" [ctl] tier -> {tier_i}", flush=True)
|
|
except Exception:
|
|
pass
|
|
|
|
# ---- rollout ----
|
|
if use_vllm:
|
|
model.save_lora(lora_tmp)
|
|
gen_cfg["lora_request"] = model.load_lora(lora_tmp)
|
|
elif full_ft:
|
|
model.eval()
|
|
else:
|
|
FastLanguageModel.for_inference(model)
|
|
with torch.no_grad():
|
|
if is_query:
|
|
trajs = rollout_query(env, tokenizer, gen_fn, problems, G, gen_cfg,
|
|
train_answers=bool(cfg.get("train_answer_tokens", True)))
|
|
for t in trajs:
|
|
t.meta.setdefault("tier", tier_i)
|
|
else:
|
|
items = []
|
|
for gid, prob in enumerate(problems):
|
|
sf = lambda th, out, full, p=prob: env.score(p, th, out, full)
|
|
meta = {"tier": tier_i,
|
|
"target": getattr(prob, "target", getattr(prob, "answer", None)),
|
|
"n": getattr(prob, "n", getattr(prob, "R", None))}
|
|
items.extend([(gid, env.prompt(prob), sf, meta)] * G)
|
|
if cfg.get("base_model"):
|
|
trajs = rollout_base(tokenizer, gen_batch, items, gen_cfg, FEWSHOT)
|
|
else: # instruct/thinking model: standard chat rollout (apply_chat_template, enable_thinking)
|
|
trajs = batched_single_turn_rollout(tokenizer, gen_batch, items, gen_cfg, system=None)
|
|
|
|
# cadence eval (never end-only): no-think probe = the load-bearing gap, on a per-tier-fixed
|
|
# eval set. Forced-empty think -> queries must be answered from nothing -> credit should sit
|
|
# at the guess floor iff the think block is genuinely load-bearing.
|
|
eval_metrics = {}
|
|
if is_query and (step % cfg.get("eval_every", 10) == 0 or step == 1):
|
|
ev_rng = _random.Random(1234 + 1000 * tier_i)
|
|
ev_probs = [env.sample_problem(ev_rng) for _ in range(cfg.get("eval_problems", 16))]
|
|
with torch.no_grad():
|
|
nt = rollout_query(env, tokenizer, gen_fn, ev_probs, 1, gen_cfg, force_empty_think=True)
|
|
eval_metrics["eval/nothink_credit"] = sum(t.task_reward for t in nt) / len(nt)
|
|
t_roll = time.time()
|
|
|
|
step_task = sum(t.task_reward for t in trajs) / len(trajs)
|
|
# adaptive lambda: hold EMA(task) at the floor
|
|
if adaptive:
|
|
task_ema = step_task if task_ema is None else (
|
|
cfg["lambda_ema_alpha"] * task_ema + (1 - cfg["lambda_ema_alpha"]) * step_task)
|
|
lam = min(max(lam + cfg["lambda_lr"] * (task_ema - cfg["lambda_target"]), 0.0), cfg["lambda_max"])
|
|
assign_advantages_len(trajs, lam, cfg["max_think"], cfg["normalize_advantages"], step, cfg)
|
|
|
|
# ---- policy gradient ----
|
|
if full_ft:
|
|
model.train()
|
|
else:
|
|
FastLanguageModel.for_training(model)
|
|
optimizer.zero_grad()
|
|
active = [t for t in trajs if abs(t.advantage) > 1e-8]
|
|
n = max(len(active), 1)
|
|
loss_val = pg_microbatched(model, tokenizer, active, device, cfg.get("pg_microbatch", 1), n)
|
|
grad_norm = clip_grad_norm_(trainable, cfg["max_grad_norm"]) if active else torch.tensor(0.0)
|
|
optimizer.step()
|
|
sched.step()
|
|
|
|
# ---- metrics ----
|
|
step_tokens = sum(int(sum(t.assistant_mask)) for t in trajs)
|
|
examples_seen += len(trajs)
|
|
tokens_seen += step_tokens
|
|
task = step_task
|
|
gen_lens = [t.meta["gen_len"] for t in trajs]
|
|
correct = [t for t in trajs if t.task_reward > cfg.get("correct_thresh", 0.5)]
|
|
gen_len_correct = (sum(t.meta["gen_len"] for t in correct) / len(correct)) if correct else 0.0
|
|
reward = sum(t.reward for t in trajs) / len(trajs)
|
|
prox = defaultdict(float)
|
|
for t in trajs:
|
|
for k, v in garble_proxies(t.think_text).items():
|
|
prox[k] += v / len(trajs)
|
|
# garble proxies among CORRECT rollouts only (the load-bearing-and-garbled cell we care about)
|
|
prox_c = defaultdict(float)
|
|
if correct:
|
|
for t in correct:
|
|
for k, v in garble_proxies(t.think_text).items():
|
|
prox_c[k] += v / len(correct)
|
|
|
|
# curriculum advance with DWELL: EMA must clear the threshold AND the tier must have run
|
|
# >= tier_min_dwell steps (the 0.6B/A3B lesson: 2-step EMA jumps straight onto the next
|
|
# cliff and stall). Grow-T-at-fixed-budget = these tiers raise T while max_think stays put,
|
|
# so the squeeze tightens through the accuracy channel itself.
|
|
if cfg.get("auto_curriculum", True) and tier_i < len(cfg["curriculum"]) - 1:
|
|
tier_ema = task if tier_ema is None else 0.8 * tier_ema + 0.2 * task
|
|
tier_steps = tier_steps + 1
|
|
if tier_ema >= cfg.get("tier_advance_thresh", 0.7) and tier_steps >= cfg.get("tier_min_dwell", 8):
|
|
tier_i += 1
|
|
env = make_env(cfg, tier_i)
|
|
tier_ema = None
|
|
tier_steps = 0
|
|
print(f" [curriculum] step {step}: advanced to tier {tier_i}", flush=True)
|
|
|
|
log = {
|
|
"train/examples_seen": examples_seen, "train/tokens_seen": tokens_seen, "train/step": step,
|
|
"train/task_reward": task, "train/reward": reward, "train/lambda": lam, "train/tier": tier_i,
|
|
"train/gen_len_mean": sum(gen_lens) / len(gen_lens),
|
|
"train/gen_len_max": max(gen_lens), "train/gen_len_min": min(gen_lens),
|
|
"train/gen_len_correct": gen_len_correct,
|
|
"train/len_norm": sum(t.meta["len_norm"] for t in trajs) / len(trajs),
|
|
"train/budget": trajs[0].meta.get("budget", cfg["max_think"]),
|
|
"train/legibility": sum(t.meta.get("legibility", 0.0) for t in trajs) / len(trajs),
|
|
"train/legibility_correct": (sum(t.meta.get("legibility", 0.0) for t in correct) / len(correct)) if correct else 0.0,
|
|
"train/loss": loss_val, "train/grad_norm": float(grad_norm),
|
|
"train/lr": sched.get_last_lr()[0], "train/n_active": len(active),
|
|
"train/n_correct": len(correct),
|
|
"train/think_temp": gen_cfg["think_temperature"], "train/out_temp": gen_cfg["out_temperature"],
|
|
"train/task_ema": (task_ema if task_ema is not None else step_task),
|
|
"train/roll_time": t_roll - t0, "train/step_time": time.time() - t0,
|
|
}
|
|
if is_query:
|
|
log["train/truncated_frac"] = sum(t.meta.get("truncated", 0.0) for t in trajs) / len(trajs)
|
|
log["train/think_len_mean"] = sum(t.meta.get("think_len", 0) for t in trajs) / len(trajs)
|
|
nq = max(getattr(env, "n_queries", 1), 1)
|
|
log["train/parse_rate"] = sum(min(t.meta.get("n_parsed", 0), nq) for t in trajs) / (len(trajs) * nq)
|
|
log.update(eval_metrics)
|
|
if "eval/nothink_credit" in eval_metrics:
|
|
log["eval/loadbearing_gap"] = task - eval_metrics["eval/nothink_credit"]
|
|
for k, v in prox.items():
|
|
log[f"garble/{k}"] = v
|
|
for k, v in prox_c.items():
|
|
log[f"garble_correct/{k}"] = v
|
|
wandb.log(log)
|
|
|
|
# full traces -> cots.jsonl (untruncated)
|
|
with open(cot_path, "a") as f:
|
|
for t in trajs:
|
|
f.write(json.dumps({"step": step, "group": t.group_id, "task": t.task_reward,
|
|
"gen_len": t.meta["gen_len"], "reward": round(t.reward, 3), "tier": tier_i,
|
|
"target": t.meta.get("target"), "n": t.meta.get("n"),
|
|
"queries": t.meta.get("queries"), "answers": t.meta.get("answers"),
|
|
"think": t.think_text, "output": t.output_text}) + "\n")
|
|
|
|
if step % cfg.get("console_every", 2) == 0 or step == 1:
|
|
print(f" step {step}/{cfg['max_steps']} task={task:.3f} nc={len(correct)}/{len(trajs)} "
|
|
f"len={sum(gen_lens)/len(gen_lens):.0f}(c:{gen_len_correct:.0f}) lam={lam:.3f} "
|
|
f"sym={prox['symbol_density']:.3f} na={prox['non_ascii']:.3f} "
|
|
f"tier={tier_i} reward={reward:.3f} t={time.time()-t0:.1f}s", flush=True)
|
|
|
|
# FULL untruncated traces to a wandb Table (the user explicitly wants the traces in wandb)
|
|
if step % cfg.get("table_every", 10) == 0 or step == 1:
|
|
rows = [[step, t.group_id, round(t.task_reward, 2), t.meta["gen_len"], t.meta.get("target"),
|
|
t.think_text, t.output_text] for t in trajs[: 4 * G]]
|
|
wandb.log({"rollouts": wandb.Table(
|
|
columns=["step", "group", "task", "gen_len", "target", "think", "output"], data=rows)})
|
|
|
|
if step in ckpt_steps:
|
|
d = os.path.join(save_root, f"step_{step}")
|
|
_save_ckpt(d, with_optimizer=False)
|
|
_prune_ckpts()
|
|
# upload the FULL trace log as a versioned wandb artifact (everything in wandb)
|
|
try:
|
|
art = wandb.Artifact(f"cots_{cfg['run_name']}", type="rollouts")
|
|
art.add_file(cot_path)
|
|
wandb.log_artifact(art)
|
|
except Exception as e:
|
|
print(f" [artifact] upload failed: {e}", flush=True)
|
|
print(f" [ckpt] saved {d}", flush=True)
|
|
|
|
if cfg.get("save_every") and (step % cfg["save_every"] == 0 or step == cfg["max_steps"]):
|
|
_save_ckpt(os.path.join(save_root, "_latest"),
|
|
with_optimizer=bool(cfg.get("latest_optimizer", True)))
|
|
|
|
wandb.finish()
|
|
print(f"[train] done. examples_seen={examples_seen} tokens_seen={tokens_seen}", flush=True)
|
|
return save_root
|