| """LoRA training of the BT policy (EXPERIMENT_LOG.md 2026-10-17): Qwen3.5-9B reads one frame and BT_PROMPT with the |
| target name and writes the behavior code (one digit token per FSQ dimension). Runs with torchrun in .train-venv. |
| |
| Rows (jsonl): image (PNG path), name, code (digits); ppo rows also adv and vllm_logp, and memory (the two 0/1 strings |
| of the decision's prompt) when the rollouts ran with memory. A row is the chat template of [image, bt_prompt(name, |
| memory)] with the generation prompt (thinking off), as vLLM renders BTPolicy's request, followed by the code and |
| <|im_end|>; batches are right-padded and the scored tokens gathered per row. |
| - sft: cross entropy on the code tokens and <|im_end|>, new LoRA (r 16, alpha 32, train_iql.py's modules), AdamW, |
| 10-update warmup then cosine; per-digit and whole-code (all digits argmax) accuracy on --valid after each epoch. |
| Saves the adapter to --out (vLLM-loadable). |
| - ppo: continues --adapter (and its optimizer.pt if it has one, i.e. a PPO round's), loads --ref (the SFT adapter) as |
| a second adapter; old and ref log-probs of the code tokens computed once with no_grad, then --epochs of PPO: |
| loss = -min(ratio A, clip(ratio, 0.8, 1.2) A) + beta * sum over the code tokens of k3 = exp(ref - lp) - (ref - lp) - 1, |
| ratio = exp(sum lp - sum old). Saves the adapter, optimizer.pt and stats.json (pg_loss, kl, clip_frac of the last |
| epoch; logp_gap = mean of old - vllm_logp, logp_gap_abs) to --out. |
| Usage: CUDA_VISIBLE_DEVICES=6,7 .train-venv/bin/python -m torch.distributed.run --nproc_per_node=2 train/bt_train.py |
| sft --data TRAIN.jsonl --valid VALID.jsonl --out DIR [--epochs 1 --lr 1e-4 --max-updates N] |
| ... train/bt_train.py ppo --data ROUND/train.jsonl --adapter ADAPTER --ref SFT --out DIR |
| """ |
| import argparse |
| import json |
| import math |
| import os |
| import random |
| import time |
| from pathlib import Path |
|
|
| import torch |
| import torch.distributed as dist |
| from peft import LoraConfig, PeftModel, get_peft_model |
| from PIL import Image |
| from torch.nn.parallel import DistributedDataParallel |
| from torch.utils.data import DataLoader, DistributedSampler |
| from transformers import AutoProcessor, Qwen3_5ForConditionalGeneration |
|
|
| from ctrl_policy import bt_prompt |
|
|
| ROOT = Path(__file__).resolve().parents[1] |
| MODULES = (r"model.language_model.*\.(q_proj|k_proj|v_proj|o_proj|in_proj_qkv|in_proj_z|in_proj_b|in_proj_a|out_proj" |
| r"|gate_proj|up_proj|down_proj)$") |
| END = "<|im_end|>" |
| CLIP, BETA = 0.2, 0.02 |
|
|
|
|
| def log(message): |
| print(f"{time.strftime('%Y-%m-%d %H:%M:%S')} {message}", flush=True) |
|
|
|
|
| class Collate: |
| def __init__(self, processor): |
| self.processor = processor |
|
|
| def __call__(self, rows): |
| p = self.processor |
| texts = [p.apply_chat_template([{"role": "user", "content": [{"type": "image"}, {"type": "text", |
| "text": bt_prompt(r["name"], r.get("memory"))}]}], tokenize=False, |
| add_generation_prompt=True, enable_thinking=False) + r["code"] + END |
| for r in rows] |
| inputs = p(text=texts, images=[Image.open(r["image"]).convert("RGB") for r in rows], padding=True, |
| padding_side="right", return_tensors="pt") |
| n = len(rows[0]["code"]) |
| ends = inputs["attention_mask"].sum(1) |
| targets = torch.stack([inputs["input_ids"][i, e - n - 1:e] for i, e in enumerate(ends)]) |
| extra = {k: torch.tensor([r[k] for r in rows], dtype=torch.float32) for k in ("adv", "vllm_logp", "index") |
| if k in rows[0]} |
| return {"inputs": dict(inputs), "ends": ends, "targets": targets, **extra} |
|
|
|
|
| def token_logprobs(model, batch, device): |
| """(B, n + 1) log-probs of the code tokens and <|im_end|>, and the logits' argmax at those positions.""" |
| inputs = {k: v.to(device) for k, v in batch["inputs"].items()} |
| ends, k = batch["ends"], batch["targets"].shape[1] |
| lo, hi = int(ends.min()) - k - 1, int(ends.max()) - 1 |
| with torch.autocast("cuda", dtype=torch.bfloat16): |
| logits = model(**inputs, logits_to_keep=torch.arange(lo, hi, device=device), use_cache=False).logits |
| rows = torch.stack([logits[i, e - k - 1 - lo:e - 1 - lo] for i, e in enumerate(ends.tolist())]).float() |
| lp = rows.log_softmax(-1) |
| targets = batch["targets"].to(device) |
| return lp.gather(-1, targets[..., None])[..., 0], rows.argmax(-1) == targets |
|
|
|
|
| def load_model(local): |
| model = Qwen3_5ForConditionalGeneration.from_pretrained( |
| ROOT / "models/Qwen3.5-9B", dtype=torch.bfloat16, device_map={"": local}, attn_implementation="sdpa", |
| local_files_only=True, use_kernels=True) |
| model.requires_grad_(False) |
| return model |
|
|
|
|
| def train_setup(model): |
| model.enable_input_require_grads() |
| model.gradient_checkpointing_enable(gradient_checkpointing_kwargs={"use_reentrant": False}) |
| model.config.use_cache = False |
| model.config.text_config.use_cache = False |
|
|
|
|
| def loader(rows, collate, batch, rank, world, seed, shuffle=True): |
| sampler = DistributedSampler(rows, num_replicas=world, rank=rank, seed=seed, shuffle=shuffle) |
| return sampler, DataLoader(rows, batch_size=batch, sampler=sampler, collate_fn=collate, num_workers=4, |
| persistent_workers=True, prefetch_factor=2) |
|
|
|
|
| @torch.no_grad() |
| def accuracy(model, rows, collate, batch, rank, world, device): |
| """Per-digit and whole-code accuracy (teacher-forced argmax) of the code tokens on rows.""" |
| model.eval() |
| totals = torch.zeros(3, device=device) |
| for b in DataLoader(rows[rank::world], batch_size=batch, collate_fn=collate, num_workers=4): |
| _, right = token_logprobs(model, b, device) |
| right = right[:, :-1] |
| totals += torch.tensor([right.sum(), right.all(1).sum(), len(right)], device=device, dtype=torch.float32) |
| dist.all_reduce(totals) |
| model.train() |
| return {"digit_acc": round((totals[0] / totals[2] / right.shape[1]).item(), 4), |
| "code_acc": round((totals[1] / totals[2]).item(), 4), "rows": int(totals[2])} |
|
|
|
|
| def sft(args, rank, world, local, device): |
| rows = [json.loads(l) for l in args.data.read_text().splitlines()] |
| valid = [json.loads(l) for l in args.valid.read_text().splitlines()] |
| processor = AutoProcessor.from_pretrained(ROOT / "models/Qwen3.5-9B", local_files_only=True) |
| collate = Collate(processor) |
| torch.manual_seed(args.seed) |
| model = get_peft_model(load_model(local), LoraConfig(r=16, lora_alpha=32, lora_dropout=0.0, |
| task_type="CAUSAL_LM", target_modules=MODULES)) |
| train_setup(model) |
| ddp = DistributedDataParallel(model, device_ids=[local]) |
| sampler, data = loader(rows, collate, args.batch, rank, world, args.seed) |
| opt = torch.optim.AdamW((p for p in model.parameters() if p.requires_grad), lr=args.lr) |
| total = min(args.max_updates or 10**9, len(data) * args.epochs) |
| if rank == 0: |
| log(f"SFT rows {len(rows)} valid {len(valid)}, {world} GPUs x batch {args.batch}, updates {total}") |
| log("VALID before " + json.dumps(accuracy(model, valid, collate, args.batch, rank, world, device))) |
| else: |
| accuracy(model, valid, collate, args.batch, rank, world, device) |
| update, start = 0, time.monotonic() |
| for epoch in range(args.epochs): |
| sampler.set_epoch(epoch) |
| for b in data: |
| lp, _ = token_logprobs(ddp, b, device) |
| loss = -lp.mean() |
| opt.zero_grad(set_to_none=True) |
| loss.backward() |
| torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) |
| for g in opt.param_groups: |
| g["lr"] = args.lr * min(1, (update + 1) / 10) * 0.5 * (1 + math.cos(math.pi * update / total)) |
| opt.step() |
| update += 1 |
| if rank == 0 and (update <= 5 or update % 50 == 0): |
| s = time.monotonic() - start |
| log(f"SFT update {update}/{total} loss {loss.item():.4f} {s / update:.2f} s/update " |
| f"eta {(total - update) * s / update / 60:.0f} min peak {torch.cuda.max_memory_allocated() / 2**30:.1f} GiB") |
| if update >= total: |
| break |
| acc = accuracy(model, valid, collate, args.batch, rank, world, device) |
| if rank == 0: |
| log(f"VALID epoch {epoch + 1} " + json.dumps(acc)) |
| model.save_pretrained(args.out) |
| (args.out / "valid.json").write_text(json.dumps({"epoch": epoch + 1, "updates": update, **acc})) |
| if update >= total: |
| break |
|
|
|
|
| def ppo(args, rank, world, local, device): |
| rows = [{**json.loads(l), "index": i} for i, l in enumerate(args.data.read_text().splitlines())] |
| processor = AutoProcessor.from_pretrained(ROOT / "models/Qwen3.5-9B", local_files_only=True) |
| collate = Collate(processor) |
| torch.manual_seed(args.seed) |
| model = PeftModel.from_pretrained(load_model(local), args.adapter, is_trainable=True) |
| model.load_adapter(args.ref, adapter_name="ref") |
| n = len(rows[0]["code"]) |
| lp_of = {} |
| for name in ("ref", "default"): |
| model.set_adapter(name) |
| model.eval() |
| mine = {} |
| with torch.no_grad(): |
| for b in DataLoader(rows[rank::world], batch_size=2 * args.batch, collate_fn=collate, num_workers=4): |
| lp, _ = token_logprobs(model, b, device) |
| mine.update(zip(b["index"].long().tolist(), lp[:, :n].cpu())) |
| gathered = [None] * world |
| dist.all_gather_object(gathered, mine) |
| lp_of[name] = torch.stack([{k: v for g in gathered for k, v in g.items()}[i] for i in range(len(rows))]) |
| model.set_adapter("default") |
| for name, p in model.named_parameters(): |
| p.requires_grad = "lora_" in name and ".default." in name |
| model.train() |
| train_setup(model) |
| ddp = DistributedDataParallel(model, device_ids=[local]) |
| params = [p for p in model.parameters() if p.requires_grad] |
| opt = torch.optim.AdamW(params, lr=args.lr) |
| if (args.adapter / "optimizer.pt").exists(): |
| opt.load_state_dict(torch.load(args.adapter / "optimizer.pt", map_location=device)) |
| for g in opt.param_groups: |
| g["lr"] = args.lr |
| sampler, data = loader(rows, collate, args.batch, rank, world, args.seed) |
| old_all, ref_all = lp_of["default"].to(device), lp_of["ref"].to(device) |
| vllm = torch.tensor([r["vllm_logp"] for r in rows], device=device) |
| gap = old_all.sum(1) - vllm |
| start = time.monotonic() |
| for epoch in range(args.epochs): |
| sampler.set_epoch(epoch) |
| stats = [] |
| for step, b in enumerate(data): |
| idx = b["index"].long().to(device) |
| lp, _ = token_logprobs(ddp, b, device) |
| lp = lp[:, :n] |
| adv = b["adv"].to(device) |
| ratio = (lp.sum(1) - old_all[idx].sum(1)).exp() |
| pg = -torch.min(ratio * adv, ratio.clamp(1 - CLIP, 1 + CLIP) * adv) |
| d = ref_all[idx] - lp |
| kl = (d.exp() - d - 1).sum(1) |
| loss = (pg + BETA * kl).mean() / args.accumulation |
| sync = (step + 1) % args.accumulation == 0 or step + 1 == len(data) |
| if sync: |
| loss.backward() |
| torch.nn.utils.clip_grad_norm_(params, 1.0) |
| opt.step() |
| opt.zero_grad(set_to_none=True) |
| else: |
| with ddp.no_sync(): |
| loss.backward() |
| stats.append([pg.mean().item(), kl.mean().item(), ((ratio - 1).abs() > CLIP).float().mean().item()]) |
| if rank == 0: |
| log(f"PPO epoch {epoch + 1} pg {sum(s[0] for s in stats) / len(stats):+.4f} " |
| f"kl {sum(s[1] for s in stats) / len(stats):.4f} {time.monotonic() - start:.0f} s") |
| totals = torch.tensor(stats, device=device).mean(0) |
| dist.all_reduce(totals) |
| totals /= world |
| if rank == 0: |
| args.out.mkdir(parents=True, exist_ok=True) |
| model.save_pretrained(args.out, selected_adapters=["default"]) |
| torch.save(opt.state_dict(), args.out / "optimizer.pt") |
| result = {"pg_loss": round(totals[0].item(), 4), "kl": round(totals[1].item(), 4), |
| "clip_frac": round(totals[2].item(), 4), "logp_gap": round(gap.mean().item(), 4), |
| "logp_gap_abs": round(gap.abs().mean().item(), 4), "rows": len(rows), |
| "train_s": round(time.monotonic() - start)} |
| (args.out / "stats.json").write_text(json.dumps(result)) |
| log("PPO_DONE " + json.dumps(result)) |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument("mode", choices=["sft", "ppo"]) |
| ap.add_argument("--data", type=Path, required=True) |
| ap.add_argument("--valid", type=Path) |
| ap.add_argument("--adapter", type=Path) |
| ap.add_argument("--ref", type=Path) |
| ap.add_argument("--out", type=Path, required=True) |
| ap.add_argument("--epochs", type=int) |
| ap.add_argument("--lr", type=float) |
| ap.add_argument("--batch", type=int, default=16, help="rows per GPU per forward") |
| ap.add_argument("--accumulation", type=int, default=2, help="ppo: forwards per update (64 rows on 2 GPUs)") |
| ap.add_argument("--max-updates", type=int) |
| ap.add_argument("--seed", type=int, default=0) |
| args = ap.parse_args() |
| rank, world, local = int(os.environ["RANK"]), int(os.environ["WORLD_SIZE"]), int(os.environ["LOCAL_RANK"]) |
| torch.set_num_threads(4) |
| random.seed(args.seed) |
| torch.cuda.set_device(local) |
| device = torch.device("cuda", local) |
| dist.init_process_group("nccl", device_id=device) |
| if args.mode == "sft": |
| args.epochs, args.lr = args.epochs or 1, args.lr or 1e-4 |
| sft(args, rank, world, local, device) |
| else: |
| args.epochs, args.lr = args.epochs or 2, args.lr or 3e-6 |
| ppo(args, rank, world, local, device) |
| dist.barrier() |
| dist.destroy_process_group() |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
