268 lines
14 KiB
Python
268 lines
14 KiB
Python
"""BabyYarros Option C — LoRA SFT on beat→paragraph pairs, loss on the RESPONSE only.
|
||
|
||
Sibling of `r49-corpus/train_voice_lora.py`, not a mode of it: that script is the
|
||
reproducibility anchor for every raw-text arm (Brontë 0.6B/1.7B/4B, Yarros base and
|
||
instruct) and adding a branch to it would put those runs one refactor away from moving.
|
||
Everything that is not the data path is held IDENTICAL to the raw-text instruct arm --
|
||
rank 32, alpha 64, lr 1e-4, batch 1 x accum 8, seed 4919, cosine with 3% warmup, eval and
|
||
save every 25 -- so the only variable between the two arms is the shape of the data.
|
||
|
||
⚠ LOSS IS MASKED TO THE ASSISTANT TURN. The raw-text arms train on every token because
|
||
every token is Yarros. Here the system prompt and the beat are NOT Yarros -- the beat is
|
||
machine-written -- and training on them teaches the carrier to generate beats, which is
|
||
capacity spent on the half of the exchange the product supplies. Prompt tokens get label
|
||
-100 and the collator checks that something survives: an example with zero unmasked labels
|
||
contributes no gradient and would otherwise train silently as a no-op.
|
||
|
||
⚠ NO PACKING. The raw-text path packs into fixed 4096-token blocks, which is right for
|
||
continuous prose and wrong here: packing would run loss across an example boundary and
|
||
teach the model that a response is followed by another system prompt. Pairs are padded
|
||
per example instead, and the cost -- padding waste -- is trivial at ~400 tokens each.
|
||
|
||
⚠ THE HELD-OUT LOSS FROM THIS RUN IS NOT COMPARABLE TO THE RAW-TEXT ARMS' 2.5263 / 2.6114.
|
||
Different objective over different tokens. It is good for exactly one thing, picking the
|
||
best checkpoint within this run, which the 4B rung showed is not the end-of-run adapter.
|
||
The cross-arm comparison has to be behavioural (on-beat / in-band / ran-on) plus delta_cb.
|
||
"""
|
||
from __future__ import annotations
|
||
import argparse, hashlib, json, math, random, subprocess, time
|
||
import pathlib
|
||
from pathlib import Path
|
||
|
||
import torch
|
||
from torch.utils.data import Dataset
|
||
from transformers import (AutoModelForCausalLM, AutoTokenizer, Trainer, TrainingArguments,
|
||
TrainerCallback)
|
||
from peft import LoraConfig, get_peft_model
|
||
|
||
TARGETS = ["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"]
|
||
|
||
# ⚠ THE SYSTEM PROMPT MUST MATCH THE ONE THE PAIRS WERE BUILT UNDER, and the pair file now
|
||
# records it. Reading it from the pairs' provenance rather than hardcoding it here is the fix
|
||
# for the defect that shaped the whole BabyYarros pilot: the trainer said "ONE paragraph" while
|
||
# the data was a median of four, and the carrier believed the data. A literal in this file is a
|
||
# second place for that to drift.
|
||
SYS_YARROS_FROZEN = (
|
||
"You expand a single story beat into ONE paragraph of prose in the manner of Rebecca "
|
||
"Yarros — contemporary first-person PRESENT-tense narration, emotionally charged, sensory "
|
||
"and physical, the voice of new-adult romantasy. Render the beat itself; do not move past "
|
||
"it, do not add a new scene, do not comment. Output the paragraph only, 90–140 words.")
|
||
|
||
|
||
def resolve_sys(pairs_path: str) -> tuple[str, str]:
|
||
"""Return (system_prompt, source). Prefer the pair build's own provenance."""
|
||
prov = pathlib.Path(str(pairs_path) + ".provenance.json")
|
||
if prov.exists():
|
||
d = json.loads(prov.read_text(encoding="utf-8"))
|
||
s = d.get("system_prompt")
|
||
if s:
|
||
return s, f"pairs provenance ({d.get('register', '?')})"
|
||
return SYS_YARROS_FROZEN, "frozen Yarros literal (no provenance found)"
|
||
|
||
|
||
def user_msg(rec: dict) -> str:
|
||
if rec.get("context"):
|
||
return (f"Continuing from:\n\n{rec['context']}\n\nBeat: {rec['beat']}")
|
||
return f"Beat: {rec['beat']}"
|
||
|
||
|
||
class Pairs(Dataset):
|
||
"""Renders each pair through the chat template and masks the prompt.
|
||
|
||
The prompt length is measured by tokenising the SAME prefix the full render begins
|
||
with -- add_generation_prompt=True on the messages minus the assistant turn -- rather
|
||
than by searching for a delimiter string. A delimiter search is what breaks silently
|
||
when a template changes its spacing, and it would mask the wrong span without erroring.
|
||
"""
|
||
|
||
def __init__(self, tok, records, seq_len, sys_prompt):
|
||
self.rows = []
|
||
dropped = 0
|
||
for r in records:
|
||
msgs = [{"role": "system", "content": sys_prompt}, {"role": "user", "content": user_msg(r)}]
|
||
# ⚠ transformers 5.16 returns a BatchEncoding from apply_chat_template(tokenize=True),
|
||
# not a list of ids. Taking len() of it yields 2 (the number of keys), so every
|
||
# example failed the `len(full) <= len(prefix)` test and the whole dataset was
|
||
# dropped. The REFUSING guard below is what surfaced it -- a silent version of this
|
||
# bug trains on nothing and reports a loss curve anyway.
|
||
prefix = tok.apply_chat_template(msgs, tokenize=True,
|
||
add_generation_prompt=True)["input_ids"]
|
||
full = tok.apply_chat_template(
|
||
msgs + [{"role": "assistant", "content": r["response"]}],
|
||
tokenize=True, add_generation_prompt=False)["input_ids"]
|
||
# The mask is only correct if the generation prefix is a TRUE prefix of the full
|
||
# render. Asserted rather than assumed: a template revision that reorders or
|
||
# re-spaces the header would shift the boundary and mask the wrong span, which
|
||
# trains without erroring and looks exactly like a normal run.
|
||
if list(full[:len(prefix)]) != list(prefix):
|
||
raise SystemExit("REFUSING: generation prefix is not a prefix of the full render "
|
||
"-- the loss mask would cover the wrong tokens")
|
||
if len(full) > seq_len or len(full) <= len(prefix):
|
||
dropped += 1
|
||
continue
|
||
labels = [-100] * len(prefix) + list(full[len(prefix):])
|
||
self.rows.append((full, labels))
|
||
self.dropped = dropped
|
||
if not self.rows:
|
||
raise SystemExit("REFUSING: every pair was dropped -- template or seq_len is wrong")
|
||
|
||
def __len__(self):
|
||
return len(self.rows)
|
||
|
||
def __getitem__(self, i):
|
||
ids, labels = self.rows[i]
|
||
return {"input_ids": ids, "labels": labels}
|
||
|
||
|
||
def collate(batch, pad_id):
|
||
n = max(len(b["input_ids"]) for b in batch)
|
||
ids, labels, mask = [], [], []
|
||
for b in batch:
|
||
k = n - len(b["input_ids"])
|
||
ids.append(b["input_ids"] + [pad_id] * k)
|
||
labels.append(b["labels"] + [-100] * k)
|
||
mask.append([1] * len(b["input_ids"]) + [0] * k)
|
||
lab = torch.tensor(labels)
|
||
if int((lab != -100).sum()) == 0:
|
||
raise SystemExit("REFUSING: a batch has no unmasked labels -- it would train as a no-op")
|
||
return {"input_ids": torch.tensor(ids), "labels": lab, "attention_mask": torch.tensor(mask)}
|
||
|
||
|
||
class LossLog(TrainerCallback):
|
||
def __init__(self, path):
|
||
self.path, self.series = path, []
|
||
|
||
def on_log(self, args, state, control, logs=None, **kw):
|
||
if logs:
|
||
self.series.append({"step": state.global_step, **logs})
|
||
self.path.write_text(json.dumps(self.series, indent=2))
|
||
|
||
|
||
def main() -> int:
|
||
ap = argparse.ArgumentParser()
|
||
ap.add_argument("--pairs", required=True)
|
||
ap.add_argument("--val-pairs", required=True)
|
||
ap.add_argument("--base", required=True)
|
||
ap.add_argument("--out", required=True)
|
||
ap.add_argument("--seq-len", type=int, default=1536)
|
||
ap.add_argument("--rank", type=int, default=32)
|
||
ap.add_argument("--lr", type=float, default=1e-4)
|
||
ap.add_argument("--epochs", type=float, default=3.0)
|
||
ap.add_argument("--batch", type=int, default=1)
|
||
ap.add_argument("--accum", type=int, default=8)
|
||
ap.add_argument("--seed", type=int, default=4919)
|
||
ap.add_argument("--eval-steps", type=int, default=25)
|
||
ap.add_argument("--save-steps", type=int, default=25)
|
||
# ⚠ Must exceed the number of saves the run will make. The pilot took 9 evals and the
|
||
# loss minimum was at the 6th; a full run at 1672 steps saving every 50 makes 33, and
|
||
# the hardcoded limit of 12 would have PRUNED an early minimum before it could be read.
|
||
# The 4B rung already proved the best checkpoint is not the last one -- a retention
|
||
# policy that silently deletes it turns that lesson into a trap rather than a guard.
|
||
ap.add_argument("--save-total-limit", type=int, default=12)
|
||
a = ap.parse_args()
|
||
|
||
torch.manual_seed(a.seed); random.seed(a.seed)
|
||
out = Path(a.out); out.mkdir(parents=True, exist_ok=True)
|
||
|
||
SYS, sys_src = resolve_sys(a.pairs)
|
||
print(f"[data] system prompt from {sys_src}")
|
||
tok = AutoTokenizer.from_pretrained(a.base)
|
||
if tok.chat_template is None:
|
||
raise SystemExit("REFUSING: carrier has no chat template -- it is not an instruct build")
|
||
if tok.pad_token_id is None:
|
||
tok.pad_token = tok.eos_token
|
||
|
||
h = hashlib.sha256()
|
||
def load(p):
|
||
b = Path(p).read_bytes(); h.update(b)
|
||
return [json.loads(l) for l in b.decode("utf-8").splitlines() if l.strip()]
|
||
train_recs, val_recs = load(a.pairs), load(a.val_pairs)
|
||
pairs_sha = h.hexdigest()[:16]
|
||
|
||
# The val pairs must come from corpus passages the train pairs never saw. A pair set
|
||
# split off the same passages would report a held-out loss over text the adapter had
|
||
# already fit, which is the failure that looks most like success.
|
||
tr_keys = {(r["work"], r["chapter"], r["response"][:80]) for r in train_recs}
|
||
overlap = [r for r in val_recs if (r["work"], r["chapter"], r["response"][:80]) in tr_keys]
|
||
if overlap:
|
||
raise SystemExit(f"REFUSING: {len(overlap)} val pairs share a passage with train")
|
||
if any(r.get("split") != "val" for r in val_recs):
|
||
raise SystemExit("REFUSING: val pairs are not all from the corpus val split")
|
||
|
||
train_ds = Pairs(tok, train_recs, a.seq_len, SYS)
|
||
val_ds = Pairs(tok, val_recs, a.seq_len, SYS)
|
||
with_ctx = sum(1 for r in train_recs if r.get("context"))
|
||
print(f"[data] {len(train_ds)} train pairs ({train_ds.dropped} dropped), "
|
||
f"{len(val_ds)} val ({val_ds.dropped} dropped), "
|
||
f"{with_ctx}/{len(train_recs)} carry context, pairs sha {pairs_sha}", flush=True)
|
||
|
||
model = AutoModelForCausalLM.from_pretrained(a.base, dtype=torch.bfloat16,
|
||
attn_implementation="sdpa").to("cuda")
|
||
model = get_peft_model(model, LoraConfig(r=a.rank, lora_alpha=2 * a.rank, lora_dropout=0.0,
|
||
bias="none", task_type="CAUSAL_LM",
|
||
target_modules=TARGETS))
|
||
model.gradient_checkpointing_enable(); model.enable_input_require_grads()
|
||
trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
|
||
total = sum(p.numel() for p in model.parameters())
|
||
|
||
resolved = {
|
||
"attn_implementation": getattr(model.config, "_attn_implementation", "?"),
|
||
"dtype": str(next(model.parameters()).dtype),
|
||
"device": torch.cuda.get_device_name(0),
|
||
"torch": torch.__version__,
|
||
"adapted_modules": sum(1 for n, _ in model.named_modules() if n.endswith("lora_A.default")),
|
||
}
|
||
try:
|
||
repo = Path(__file__).resolve().parents[2]
|
||
git = subprocess.run(["git", "-C", str(repo), "rev-parse", "--short", "HEAD"],
|
||
capture_output=True, text=True).stdout.strip()
|
||
dirty = bool(subprocess.run(["git", "-C", str(repo), "status", "--porcelain"],
|
||
capture_output=True, text=True).stdout.strip())
|
||
except Exception:
|
||
git, dirty = "?", True
|
||
|
||
steps_per_epoch = math.ceil(len(train_ds) / (a.batch * a.accum))
|
||
prov = {"run": "r49-babyyarros-pairs-pilot", "base": a.base,
|
||
"pairs": a.pairs, "val_pairs": a.val_pairs, "pairs_sha256_16": pairs_sha,
|
||
"objective": "chat-template SFT, loss on assistant turn only",
|
||
"loss_comparable_to_raw_text_arms": False,
|
||
"train_pairs": len(train_ds), "val_pairs_n": len(val_ds),
|
||
"pairs_with_context": with_ctx,
|
||
"seq_len": a.seq_len, "lora_rank": a.rank, "lora_alpha": 2 * a.rank, "targets": TARGETS,
|
||
"lr": a.lr, "epochs": a.epochs, "batch": a.batch, "grad_accum": a.accum, "seed": a.seed,
|
||
"trainable_params": trainable, "total_params": total,
|
||
"trainable_pct": round(100 * trainable / total, 3),
|
||
"steps_per_epoch": steps_per_epoch, "planned_steps": steps_per_epoch * int(a.epochs),
|
||
"resolved": resolved, "harness_commit": git, "harness_dirty_at_launch": dirty,
|
||
"launched_at": time.strftime("%Y-%m-%dT%H:%M:%S%z"), "system_prompt": SYS,
|
||
"system_prompt_source": sys_src}
|
||
(out / "provenance.json").write_text(json.dumps(prov, indent=2))
|
||
print("[prov] " + json.dumps({k: prov[k] for k in
|
||
("pairs_sha256_16", "train_pairs", "planned_steps", "trainable_pct",
|
||
"harness_dirty_at_launch")}), flush=True)
|
||
print("[prov] resolved: " + json.dumps(resolved), flush=True)
|
||
|
||
args = TrainingArguments(
|
||
output_dir=str(out / "checkpoints"), per_device_train_batch_size=a.batch,
|
||
gradient_accumulation_steps=a.accum, num_train_epochs=a.epochs, learning_rate=a.lr,
|
||
lr_scheduler_type="cosine",
|
||
warmup_steps=max(1, int(0.03 * steps_per_epoch * int(a.epochs))),
|
||
bf16=True, logging_steps=10,
|
||
save_strategy="steps", save_steps=a.save_steps, save_total_limit=a.save_total_limit,
|
||
eval_strategy="steps", eval_steps=a.eval_steps,
|
||
report_to=[], seed=a.seed,
|
||
gradient_checkpointing=True, dataloader_num_workers=2,
|
||
remove_unused_columns=False,
|
||
)
|
||
trainer = Trainer(model=model, args=args, train_dataset=train_ds, eval_dataset=val_ds,
|
||
data_collator=lambda b: collate(b, tok.pad_token_id),
|
||
callbacks=[LossLog(out / "loss-series.json")])
|
||
trainer.train()
|
||
model.save_pretrained(out / "adapter")
|
||
print("[done] adapter saved; the BEST checkpoint is in checkpoints/, not necessarily this one")
|
||
return 0
|
||
|
||
|
||
if __name__ == "__main__":
|
||
raise SystemExit(main())
|