| """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): |
| b = None |
| torch.cuda.reset_peak_memory_stats() |
| try: |
| 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) |
| del b |
| torch.cuda.empty_cache() |
|
|
| if step("512") is None: |
| 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: |
| 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: |
| 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): |
| 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) |
| if not a.full_ft: |
| model.enable_input_require_grads() |
| 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 |
|
|
| |
| |
| 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 |
| |
| 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, |
| 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() |
|
|
