code/fastfill/train.py

"""LoRA SFT for FastFill v1 (adapted from scripts/sft_train.py).

Differences from the OptiScene script: pre-split files instead of a random split, loss on the
assistant answer only, over-length samples dropped (never truncated mid-JSON), EOS = <|im_end|>
supervised, Qwen3 thinking disabled in the chat template, final adapter saved; run_manifest.json (the build's
MANIFEST + these arguments) is written into --out so every checkpoint says which data it saw.
Hyper-parameters: LoRA r16/alpha32 as in the OptiScene paper, but lr 1e-4 x 2 epochs (~6k steps at batch 32) instead
of the paper's 5e-6 x 10 epochs, which is a full-fine-tuning rate, ~20x too low for a rank-16 adapter.

Effective batch = bs x grad_accum x n_gpus; keep it at 32 when changing the GPU count:
  1 GPU:   CUDA_VISIBLE_DEVICES=0 python -m fastfill.train --model Qwen/Qwen3-8B --data DATA --out outputs/ff-v1 --bs 4 --grad_accum 8
  8 GPUs:  torchrun --nproc_per_node 8 -m fastfill.train --model Qwen/Qwen3-8B --data DATA --out outputs/ff-v1 --bs 4 --grad_accum 1
Resume after an interruption: add --resume (continues from the last checkpoint in --out).
Sequence length: --max_len is the model's limit (Qwen3-8B: 40960), not a memory setting; a batch is padded to its
longest sample, so memory follows the data. Measure it before choosing --bs / --grad_accum:
  --dry_run                       token length and object count percentiles with the real tokenizer, the samples
                                  over --max_len listed (they are dropped from training, never truncated)
  --mem_test 8192,16384,32768,max one forward/backward/optimizer step per length on one GPU with a batch of --bs rows
                                  padded like a training batch (row 0 at that length, the others at half, see
                                  mem_batch); peak memory printed next to the static part (weights + LoRA params/grads
                                  + AdamW states, read after a warm-up step); 'max' = the longest train or dev row,
                                  because Trainer evaluates dev every --save_steps
  --eval_max_len N                leave dev rows longer than N tokens out of that evaluation (eval loss then covers
                                  the remaining subset); pick N from the mem_test table
"""
import argparse
import collections
import hashlib
import json
import math
import os

import torch
from peft import LoraConfig, get_peft_model
from transformers import (AutoModelForCausalLM, AutoTokenizer, DataCollatorForSeq2Seq, Trainer,
                          TrainingArguments, set_seed)

from fastfill.scene import prompt_text


def end_token(tok):
    return "<|im_end|>" if "<|im_end|>" in tok.get_vocab() else tok.eos_token


def encode(tok, rows, max_len, chunk=2000):
    """Tokenize (prompt, answer) pairs in batches (per-row tokenizer calls are ~300x slower); labels mask the
    prompt; rows longer than max_len are dropped, never truncated (a cut JSON answer teaches broken output).
    -> (samples, dropped row indices)."""
    out, dropped, end = [], [], end_token(tok)
    total = len(rows)
    for i in range(0, total, chunk):
        if total > 5000 and (i // chunk) % 10 == 0:
            print(f"  [Tokenize] {i}/{total} samples ({i*100//total}%)...", flush=True)
        part = rows[i:i + chunk]
        ps = tok([prompt_text(tok, m[:-1]) for m in part], add_special_tokens=False)["input_ids"]
        ans = tok([m[-1]["content"] + end for m in part], add_special_tokens=False)["input_ids"]
        for k, (p, a) in enumerate(zip(ps, ans)):
            if len(p) + len(a) > max_len:
                dropped.append(i + k)
            else:
                out.append({"input_ids": p + a, "attention_mask": [1] * (len(p) + len(a)), "labels": [-100] * len(p) + a})
    return out, dropped


def select(recs, exclude_flags=(), source_weight=None):
    """The sampling layer: rows with any flag in exclude_flags are left out; a row of a source with weight w is
    used int(w) times plus once more with probability frac(w) (stable per uid), so w=0.5 keeps half a source and
    w=2 doubles it. Default: every row once (the data's own proportions)."""
    out, gone = [], collections.Counter()
    for r in recs:
        hit = {k for k, v in (r.get("flags") or {}).items() if v} & set(exclude_flags)
        if hit:
            gone["flag:" + sorted(hit)[0]] += 1
            continue
        w = (source_weight or {}).get(r.get("source"), 1.0)
        u = int(hashlib.sha1(r["uid"].encode()).hexdigest(), 16) % 10 ** 6 / 10 ** 6
        n = int(w) + (u < w - int(w))
        gone["weight"] += n == 0
        out += [r] * n
    return out, gone


def load(tok, path, max_len, exclude_flags=(), source_weight=None):
    print(f"Loading and tokenizing {path}...", flush=True)
    recs, gone = select([json.loads(line) for line in open(path)], exclude_flags, source_weight)
    rows, dropped = encode(tok, [r["messages"] for r in recs], max_len)
    comp = collections.Counter(r.get("source") for r in recs)
    print(f"{path}: {len(rows)} samples, dropped {len(dropped)} longer than {max_len} tokens"
          + (f", left out {dict(gone)}" if +gone else "") + f"; by source {dict(comp.most_common())}", flush=True)
    return rows


def length_report(tok, path, max_len):
    """Real-tokenizer profile of one file: tokens per sample (total / answer) and objects per room, the samples over
    max_len listed by uid so they can be looked at instead of silently dropped."""
    recs = [json.loads(line) for line in open(path)]
    rows, over = encode(tok, [r["messages"] for r in recs], float("inf"))
    lens = [len(r["input_ids"]) for r in rows]
    ans = [sum(l != -100 for l in r["labels"]) for r in rows]
    nobj = [len(json.loads(r["messages"][1]["content"])["objects"]) for r in recs]
    nfix = [len(json.loads(r["messages"][1]["content"]).get("fixed", [])) for r in recs]
    pct = lambda v: {k: sorted(v)[min(int(f * len(v)), len(v) - 1)] for k, f in
                     (("p50", .5), ("p90", .9), ("p95", .95), ("p99", .99), ("max", 1.0))}
    rep = {"file": path, "samples": len(rows), "tokens_total": sum(lens), "tokens": pct(lens), "answer_tokens": pct(ans),
           "objects": pct(nobj), "fixed": pct(nfix),
           "over_max_len": [{"uid": recs[i]["uid"], "source": recs[i]["source"], "tokens": lens[i], "objects": nobj[i]}
                            for i in range(len(rows)) if lens[i] > max_len],
           "buckets": {f"<= {b}": sum(l <= b for l in lens) for b in (4096, 8192, 16384, 32768, max_len)}}
    return rep


def mem_batch(tok, longest, spec, bs):
    """The batch DataCollatorForSeq2Seq(padding=True) builds for a --bs step whose longest row has the requested length:
    'max' = the longest train/dev row itself, a number = that many tokens (the longest row's tokens repeated, every
    token supervised: an upper bound). The other bs-1 rows are the same tokens cut to half and right-padded (pad id,
    attention_mask 0, labels -100), so SDPA materializes its (bs,1,L,L) mask as in almost every real batch.
    -> {input_ids, attention_mask, labels}: bs lists of L ints each."""
    if spec == "max":
        ids, lab = longest["input_ids"], longest["labels"]
    else:
        ids = (longest["input_ids"] * (int(spec) // len(longest["input_ids"]) + 1))[:int(spec)]
        lab = ids
    L, h = len(ids), len(ids) // 2
    rows = [(ids, lab)] + [(ids[:h], lab[:h])] * (bs - 1)
    return {"input_ids": [i + [tok.pad_token_id] * (L - len(i)) for i, _ in rows],
            "attention_mask": [[1] * len(i) + [0] * (L - len(i)) for i, _ in rows],
            "labels": [l + [-100] * (L - len(l)) for _, l in rows]}


def mem_test(model, tok, train, dev, specs, bs):
    """One training step per requested length on the current GPU in the training configuration: a batch of bs rows
    padded like the collator's (mem_batch), gradient checkpointing on. 'max' = the longest train or dev row (Trainer
    evaluates dev every save_steps). A warm-up step (512 tokens) first allocates the LoRA grads and AdamW states, so
    static = weights + LoRA params/grads + AdamW states and the per-length column is what the step adds
    (activations, logits, attention mask). OOM at one length moves on to the next."""
    gpu = torch.device("cuda")
    model.to(gpu)
    model.train()
    model.gradient_checkpointing_enable(gradient_checkpointing_kwargs={"use_reentrant": False})
    opt = torch.optim.AdamW([p for p in model.parameters() if p.requires_grad], lr=1e-4)
    longest = max(train + dev, key=lambda r: len(r["input_ids"]))
    print(f"longest train row {max(len(r['input_ids']) for r in train)} tokens, "
          f"longest dev row {max(len(r['input_ids']) for r in dev)} tokens (eval set)", flush=True)

    def step(spec):                       # -> peak bytes of one forward/backward/optimizer step, None on OOM
        b = None
        torch.cuda.reset_peak_memory_stats()
        try:                              # the batch's own allocation can OOM too
            b = {k: torch.tensor(v, device=gpu) for k, v in mem_batch(tok, longest, spec, bs).items()}
            model(**b).loss.backward()
            opt.step()
            return torch.cuda.max_memory_allocated()
        except torch.cuda.OutOfMemoryError:
            return None
        finally:
            opt.zero_grad(set_to_none=False)  # grads stay resident, as in training
            del b
            torch.cuda.empty_cache()

    if step("512") is None:               # warm-up: creates the LoRA grads and the AdamW moments
        print(f"OOM already at 512 tokens x bs {bs} on {torch.cuda.get_device_name()}: the model does not fit this card "
              f"in this configuration; no table", flush=True)
        return
    static = torch.cuda.memory_allocated()
    print(f"static (weights + LoRA params/grads + AdamW states) {static / 2 ** 30:.1f} GB", flush=True)
    for spec in specs:
        n = len(longest["input_ids"]) if spec == "max" else int(spec)
        peak = step(spec)
        if peak is None:
            print(f"len {n:6d}  bs {bs}  OOM on {torch.cuda.get_device_name()} "
                  f"({torch.cuda.get_device_properties(gpu).total_memory / 2 ** 30:.0f} GB)", flush=True)
        else:
            print(f"len {n:6d}  bs {bs}  peak {peak / 2 ** 30:5.1f} GB  = static {static / 2 ** 30:.1f} + "
                  f"activations/logits/mask {(peak - static) / 2 ** 30:.1f} GB", flush=True)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--model", required=True, help="HF id or local path, e.g. Qwen/Qwen3-8B")
    ap.add_argument("--data", required=True, help="dir with train.jsonl / dev.jsonl from fastfill.build")
    ap.add_argument("--out", required=True)
    ap.add_argument("--max_len", type=int, default=40960, help="model sequence limit; longer samples are dropped")
    ap.add_argument("--lr", type=float, default=1e-4)
    ap.add_argument("--epochs", type=float, default=2)
    ap.add_argument("--bs", type=int, default=4, help="per-device batch size")
    ap.add_argument("--grad_accum", type=int, default=8, help="per GPU; divide by the GPU count to keep batch 32")
    ap.add_argument("--lora_r", type=int, default=16)
    ap.add_argument("--lora_alpha", type=int, default=32)
    ap.add_argument("--full_ft", action="store_true", help="full fine-tuning instead of LoRA (use with --deepspeed)")
    ap.add_argument("--deepspeed", default=None)
    ap.add_argument("--attn", default="sdpa", help="sdpa | flash_attention_2")
    ap.add_argument("--save_steps", type=int, default=500)
    ap.add_argument("--report_to", default="none")
    ap.add_argument("--dry_run", action="store_true", help="tokenize and print the length profile only (no GPU)")
    ap.add_argument("--mem_test", default=None, help="e.g. 8192,16384,32768,max: measure peak memory per length, then exit")
    ap.add_argument("--exclude_flags", nargs="*", default=["oob_objects", "overlapping_furniture", "fixed_collision"],
                    help="train rows carrying any of these quality flags are left out (see the row's 'flags'). Default: "
                         "the three whose reference answer serve.py would refuse (objects > 10 cm through the walls; "
                         "copies of one piece through each other; furniture through a door / column / stair box). "
                         "Pass --exclude_flags with no value to train on all")
    ap.add_argument("--source_weight", nargs="*", default=[],
                    help="SOURCE=w: use each train row of SOURCE w times on average (0.5 halves it, 2 doubles it)")
    ap.add_argument("--eval_max_len", type=int, default=0,
                    help="dev rows longer than this are left out of the eval at every save_steps (0 = keep all)")
    ap.add_argument("--resume", action="store_true", help="continue from the last checkpoint in --out")
    a = ap.parse_args()
    local_rank = int(os.environ.get("LOCAL_RANK", 0))
    if torch.cuda.is_available():
        torch.cuda.set_device(local_rank)
    a.source_weight = {k: float(v) for k, v in (x.split("=") for x in a.source_weight)}
    known = {json.loads(line)["source"] for line in open(f"{a.data}/train.jsonl")} if a.source_weight else set()
    if set(a.source_weight) - known:
        raise SystemExit(f"--source_weight: unknown source(s) {sorted(set(a.source_weight) - known)}; known: {sorted(known)}")

    print(f"Loading tokenizer from {a.model}...", flush=True)
    tok = AutoTokenizer.from_pretrained(a.model)
    if tok.pad_token is None:
        tok.pad_token = tok.eos_token
    if a.dry_run:
        for split in ("train", "dev", "test"):
            rep = length_report(tok, f"{a.data}/{split}.jsonl", a.max_len)
            print(json.dumps(rep, indent=1))
        train = load(tok, f"{a.data}/train.jsonl", a.max_len, a.exclude_flags, a.source_weight)
        print(tok.decode(train[0]["input_ids"]))
        print("supervised part:", repr(tok.decode([t for t in train[0]["labels"] if t != -100])))
        return
    train = load(tok, f"{a.data}/train.jsonl", a.max_len, a.exclude_flags, a.source_weight)
    dev = load(tok, f"{a.data}/dev.jsonl", a.max_len)
    if a.eval_max_len:                      # protect the eval at every save_steps; eval loss then covers a subset
        kept = [r for r in dev if len(r["input_ids"]) <= a.eval_max_len]
        print(f"eval set: {len(dev) - len(kept)} dev rows longer than {a.eval_max_len} tokens left out, {len(kept)} kept", flush=True)
        dev = kept
    if int(os.environ.get("RANK", 0)) == 0 and not a.mem_test:   # provenance of this run (build MANIFEST + args)
        os.makedirs(a.out, exist_ok=True)
        man = f"{a.data}/MANIFEST.json"
        rec = {"args": vars(a), "data_manifest": json.load(open(man)) if os.path.exists(man) else None}
        path = f"{a.out}/run_manifest.json"
        if a.resume and os.path.exists(path):   # a resume is recorded beside the run it continues; a new run replaces
            k = 1
            while os.path.exists(f"{a.out}/run_manifest.resume{k}.json"):
                k += 1
            path = f"{a.out}/run_manifest.resume{k}.json"
        json.dump(rec, open(path, "w"), indent=1)

    model = AutoModelForCausalLM.from_pretrained(
        a.model, dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32, attn_implementation=a.attn)
    model.config.use_cache = False
    set_seed(42)                            # before get_peft_model: same LoRA init every run (v1.0 vs v1.1 pairing)
    if not a.full_ft:
        model.enable_input_require_grads()  # needed for gradient checkpointing with frozen embeddings
        model = get_peft_model(model, LoraConfig(
            task_type="CAUSAL_LM", r=a.lora_r, lora_alpha=a.lora_alpha, lora_dropout=0.05,
            target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"]))
        model.print_trainable_parameters()
    if a.mem_test:
        mem_test(model, tok, train, dev, a.mem_test.split(","), a.bs)
        return

    # 3% warmup as an integer step count: transformers 4.x rejects fractional warmup_steps, >=5.15 has no warmup_ratio
    # one process that sees several GPUs runs DataParallel: its batch is bs x n_gpu, like DDP's
    world = int(os.environ.get("WORLD_SIZE", 0)) or max(1, torch.cuda.device_count())
    total_steps = math.ceil(math.ceil(len(train) / (a.bs * world)) / a.grad_accum) * a.epochs
    # batch similar lengths together (~1/3 less padding); the argument was renamed in transformers 5.14
    by_length = ({"group_by_length": True} if "group_by_length" in TrainingArguments.__dataclass_fields__
                 else {"train_sampling_strategy": "group_by_length"})
    args = TrainingArguments(
        output_dir=a.out, per_device_train_batch_size=a.bs, per_device_eval_batch_size=1,  # dev pads to 40k tokens
        gradient_accumulation_steps=a.grad_accum, learning_rate=a.lr, num_train_epochs=a.epochs, **by_length,
        lr_scheduler_type="cosine", warmup_steps=max(1, int(0.03 * total_steps)), weight_decay=0.0,
        bf16=torch.cuda.is_available(),
        gradient_checkpointing=True, gradient_checkpointing_kwargs={"use_reentrant": False},
        logging_steps=10, eval_strategy="steps", eval_steps=a.save_steps, save_strategy="steps",
        save_steps=a.save_steps, save_total_limit=3, remove_unused_columns=False,
        report_to=a.report_to, deepspeed=a.deepspeed, ddp_find_unused_parameters=False)
    trainer = Trainer(model=model, args=args, train_dataset=train, eval_dataset=dev,
                      data_collator=DataCollatorForSeq2Seq(tok, padding=True, label_pad_token_id=-100))
    trainer.train(resume_from_checkpoint=True if a.resume else None)
    trainer.save_model(f"{a.out}/final")
    tok.save_pretrained(f"{a.out}/final")


if __name__ == "__main__":
    main()