esh-pfi-infrastructure/scripts/r49-corpus/train_voice_lora.py

"""R49 H02 pilot — author-voice LoRA on a dense Qwen3 carrier.

Pure continuation. No beat annotation, no Director, no orchestration loop --
that is H02's design, not a shortcut: if a carrier cannot hold the voice on plain
continuation, no amount of beat engineering rescues it, and the negative arrives
in hours rather than weeks.

Provenance is recorded from what the run RESOLVED to, never from what it
requested -- the attention implementation, the dtype, the device and the corpus
hash are all read back off the live objects after construction, because a config
value is a request and the playbook's §4 lesson is that two runs with the same
config and different backends produce different numbers and nobody notices.
"""
from __future__ import annotations
import argparse, hashlib, json, math, os, random, subprocess, sys, time
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"]


class Packed(Dataset):
    """Order-preserving packing into fixed-length blocks, one work-copy at a time.

    Documents are never packed across a work boundary. On a dense carrier an
    attention mask would handle it, but keeping the boundary costs nothing here
    and the constraint has to hold anyway if a hybrid carrier is ever revisited,
    where SSM state ignores the mask entirely.
    """
    def __init__(self, blocks): self.blocks = blocks
    def __len__(self): return len(self.blocks)
    def __getitem__(self, i):
        ids = torch.tensor(self.blocks[i], dtype=torch.long)
        return {"input_ids": ids, "labels": ids.clone(), "attention_mask": torch.ones_like(ids)}


def pack(tok, records, seq_len):
    by_stream = {}
    for r in records:
        by_stream.setdefault((r["work"], r["copy"]), []).append(r)
    blocks = []
    for key in sorted(by_stream):
        rows = sorted(by_stream[key], key=lambda r: r["chapter"])
        buf = []
        for r in rows:
            buf.extend(tok.encode(r["text"] + "\n\n", add_special_tokens=False))
            while len(buf) >= seq_len:
                blocks.append(buf[:seq_len]); buf = buf[seq_len:]
    return blocks


class LossLog(TrainerCallback):
    def __init__(self, path): self.path, self.series = path, []
    def on_log(self, args, state, control, logs=None, **kw):
        if logs and "loss" in logs:
            self.series.append({"step": state.global_step, **{k: v for k, v in logs.items()
                                                              if isinstance(v, (int, float))}})
            Path(self.path).write_text(json.dumps(self.series, indent=1))


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--corpus", required=True)
    ap.add_argument("--base", required=True)
    ap.add_argument("--out", required=True)
    ap.add_argument("--seq-len", type=int, default=4096)
    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)
    a = ap.parse_args()

    torch.manual_seed(a.seed); random.seed(a.seed)
    out = Path(a.out); out.mkdir(parents=True, exist_ok=True)

    tok = AutoTokenizer.from_pretrained(a.base)
    records, val_records, h = [], [], hashlib.sha256()
    for f in sorted(Path(a.corpus).glob("copies/*.jsonl")):
        h.update(f.read_bytes())
        for line in f.read_text(encoding="utf-8").splitlines():
            r = json.loads(line)
            (val_records if r["split"] == "val" else records).append(r)
    corpus_sha = h.hexdigest()[:16]
    print(f"[data] {len(records):,} train records, {len(val_records):,} val, corpus sha {corpus_sha}", flush=True)

    t0 = time.time()
    train_blocks = pack(tok, records, a.seq_len)
    val_blocks = pack(tok, val_records, a.seq_len)
    tr_tok = len(train_blocks) * a.seq_len
    print(f"[data] packed {len(train_blocks):,} train blocks ({tr_tok:,} tokens), "
          f"{len(val_blocks):,} val blocks, in {time.time()-t0:.0f}s", 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())

    # ⚠ Read back what the run RESOLVED to, not what it requested.
    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_blocks) / (a.batch * a.accum))
    prov = {"run": "r49-h02-pilot", "base": a.base, "corpus": a.corpus, "corpus_sha256_16": corpus_sha,
            "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,
            "train_blocks": len(train_blocks), "train_tokens": tr_tok, "val_blocks": len(val_blocks),
            "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")}
    (out / "provenance.json").write_text(json.dumps(prov, indent=2))
    print("[prov] " + json.dumps({k: prov[k] for k in
          ("corpus_sha256_16", "train_tokens", "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,
        # transformers 5.x dropped `warmup_ratio`; only `warmup_steps` survives, so
        # the 3% warmup is computed here rather than requested by a name that no
        # longer exists. Read the signature, do not assume the 4.x one.
        lr_scheduler_type="cosine", warmup_steps=max(1, int(0.03 * steps_per_epoch * int(a.epochs))),
        bf16=True, logging_steps=10,
        save_strategy="no", eval_strategy="epoch", report_to=[], seed=a.seed,
        gradient_checkpointing=True, dataloader_num_workers=2,
    )
    trainer = Trainer(model=model, args=args, train_dataset=Packed(train_blocks),
                      eval_dataset=Packed(val_blocks),
                      callbacks=[LossLog(out / "loss-series.json")])
    res = trainer.train()
    model.save_pretrained(out / "adapter")
    tok.save_pretrained(out / "adapter")

    prov["train_result"] = {k: v for k, v in res.metrics.items()}
    prov["finished_at"] = time.strftime("%Y-%m-%dT%H:%M:%S%z")
    (out / "provenance.json").write_text(json.dumps(prov, indent=2))
    saved = sorted(p.name for p in (out / "adapter").iterdir())
    print(f"[done] {res.metrics} -> {out/'adapter'} ({len(saved)} files)", flush=True)
    return 0


if __name__ == "__main__":
    sys.exit(main())