#!/usr/bin/env python3 """Resumable formal Robometer-Low-LoRA training on joint preference pairs.""" from __future__ import annotations import argparse import hashlib import json import math import os import random import sys import time from pathlib import Path # Avoid large reserved-but-unallocated CUDA segments on near-capacity batches. os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True") import numpy as np import torch import torch.nn.functional as F EXPERIMENT_ROOT = Path(__file__).resolve().parents[1] if str(EXPERIMENT_ROOT) not in sys.path: sys.path.insert(0, str(EXPERIMENT_ROOT)) from scripts.smoke_phase21a_robometer_low import ( # noqa: E402 SPECIAL_TOKENS, RobometerHeads, build_conversation, finite_grad_audit, group_grad_norm, make_targets, process_batch, ) def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser() parser.add_argument("--snapshot", type=Path, required=True) parser.add_argument("--manifest", type=Path, required=True) parser.add_argument("--processed-root", type=Path, required=True) parser.add_argument("--seed", type=int, required=True) parser.add_argument("--optimizer-steps", type=int, default=3000) parser.add_argument("--micro-batch-pairs", type=int, default=32) parser.add_argument("--gradient-accumulation", type=int, default=2) parser.add_argument("--lora-lr", type=float, default=2.0e-6) parser.add_argument("--head-lr", type=float, default=1.0e-4) parser.add_argument("--checkpoint-every", type=int, default=250) parser.add_argument("--eval-every", type=int, default=250) parser.add_argument("--eval-micro-batch-pairs", type=int, default=32) parser.add_argument("--early-stop-min-steps", type=int, default=1000) parser.add_argument("--early-stop-patience-evals", type=int, default=4) parser.add_argument("--early-stop-min-delta", type=float, default=5.0e-3) parser.add_argument("--output-dir", type=Path, required=True) parser.add_argument("--resume", action="store_true") return parser.parse_args() def read_rows(path: Path) -> list[dict]: rows = [json.loads(line) for line in path.read_text(encoding="utf-8").splitlines() if line] train = [row for row in rows if row["split"] == "train" and row["sample_type"] == "preference"] if len(train) != 2021: raise RuntimeError(f"expected 2021 train preference pairs, found {len(train)}") return train def read_dev_rows(path: Path) -> list[dict]: rows = [json.loads(line) for line in path.read_text(encoding="utf-8").splitlines() if line] dev = [row for row in rows if row["split"] == "dev" and row["sample_type"] == "preference"] if len(dev) != 108: raise RuntimeError(f"expected 108 ID-dev preference pairs, found {len(dev)}") return dev def global_success_pos_weight(rows: list[dict]) -> float: positive = 0.0 total = 0 for row in rows: for role in ("chosen", "rejected"): record = row[role] valid = np.asarray(record["frame_valid_mask"], dtype=bool) labels = np.asarray(record["success_label"], dtype=np.float32) positive += float(labels[valid].sum()) total += int(valid.sum()) if not 0 < positive < total: raise RuntimeError("success labels require both classes") return (total - positive) / positive def atomic_torch_save(payload: dict, path: Path) -> None: temporary = path.with_suffix(".tmp") torch.save(payload, temporary) temporary.replace(path) def reconcile_progress_log(path: Path, completed_steps: int) -> None: """Drop records newer than the checkpoint before a resumed run appends.""" if not path.is_file(): return records = [json.loads(line) for line in path.read_text(encoding="utf-8").splitlines() if line] retained = [record for record in records if int(record["step"]) <= completed_steps] temporary = path.with_suffix(".tmp") payload = "".join(json.dumps(record, sort_keys=True) + "\n" for record in retained) temporary.write_text(payload, encoding="utf-8") temporary.replace(path) def cpu_state_dict(state: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]: return {name: tensor.detach().cpu().clone() for name, tensor in state.items()} @torch.inference_mode() def evaluate_dev( *, model, processor, heads, rows: list[dict], processed_root: Path, token_ids: dict[str, int], hidden_size: int, pos_weight: torch.Tensor, micro_batch_pairs: int, ) -> dict: """Evaluate both A/B orientations on the frozen ID-dev split only.""" model.eval() heads.eval() presentations = [(row, chosen_is_a) for row in rows for chosen_is_a in (True, False)] sums = {"preference": 0.0, "progress": 0.0, "success": 0.0} counts = {"preference": 0, "progress": 0, "success": 0} preference_correct = 0 for start in range(0, len(presentations), micro_batch_pairs): batch = presentations[start : start + micro_batch_pairs] conversations, records_a, preference_labels = [], [], [] for row, chosen_is_a in batch: conversation, record_a, label = build_conversation(row, processed_root, chosen_is_a) conversations.append(conversation) records_a.append(record_a) preference_labels.append(label) inputs = process_batch(processor, conversations) inputs = {key: value.to("cuda") for key, value in inputs.items()} targets = make_targets(records_a, torch.device("cuda")) preference_target = torch.tensor(preference_labels, dtype=torch.float32, device="cuda") output = model.model(**inputs, output_hidden_states=False, use_cache=False) hidden = output.last_hidden_state prog_mask = inputs["input_ids"].eq(token_ids["<|prog_token|>"]) pref_mask = inputs["input_ids"].eq(token_ids["<|pref_token|>"]) batch_size = len(batch) if not torch.equal(prog_mask.sum(1), torch.full_like(prog_mask.sum(1), 16)): raise RuntimeError("ID-dev progress token extraction mismatch") if not torch.equal(pref_mask.sum(1), torch.ones_like(pref_mask.sum(1))): raise RuntimeError("ID-dev preference token extraction mismatch") prog_hidden = hidden[prog_mask].reshape(batch_size, 16, hidden_size) pref_hidden = hidden[pref_mask].reshape(batch_size, hidden_size) progress_logits = heads.progress_head(prog_hidden[:, :8]) success_logits = heads.success_head(prog_hidden[:, :8]).squeeze(-1) preference_logits = heads.preference_head(pref_hidden).squeeze(-1) progress_frames = -( targets["progress"] * F.log_softmax(progress_logits.float(), dim=-1) ).sum(-1) progress_values = progress_frames[targets["progress_mask"]] success_sum = F.binary_cross_entropy_with_logits( success_logits.float()[targets["success_mask"]], targets["success"][targets["success_mask"]], pos_weight=pos_weight, reduction="sum", ) preference_sum = F.binary_cross_entropy_with_logits( preference_logits.float(), preference_target, reduction="sum" ) sums["progress"] += float(progress_values.sum()) sums["success"] += float(success_sum) sums["preference"] += float(preference_sum) counts["progress"] += int(progress_values.numel()) counts["success"] += int(targets["success_mask"].sum()) counts["preference"] += batch_size preference_correct += int( ((preference_logits > 0) == (preference_target > 0.5)).sum() ) components = {name: sums[name] / counts[name] for name in sums} model.train() heads.train() return { "loss": sum(components.values()), "loss_components": components, "preference_accuracy": preference_correct / counts["preference"], "pair_count": len(rows), "orientation_presentations": len(presentations), "supervision_counts": counts, } def main() -> None: args = parse_args() if args.micro_batch_pairs * args.gradient_accumulation != 64: raise ValueError("formal protocol requires exact effective batch 64") if min(args.optimizer_steps, args.checkpoint_every, args.eval_every) <= 0: raise ValueError("steps, checkpoint interval, and eval interval must be positive") if args.eval_micro_batch_pairs <= 0 or args.early_stop_patience_evals <= 0: raise ValueError("eval batch and early-stop patience must be positive") if not 0 < args.early_stop_min_steps <= args.optimizer_steps: raise ValueError("early-stop minimum must be in (0, optimizer_steps]") if args.checkpoint_every != args.eval_every: raise ValueError("v2 requires checkpoint_every == eval_every") os.environ.setdefault("TOKENIZERS_PARALLELISM", "false") random.seed(args.seed) np.random.seed(args.seed) torch.manual_seed(args.seed) torch.cuda.manual_seed_all(args.seed) torch.backends.cuda.matmul.allow_tf32 = True from peft import LoraConfig, get_peft_model, get_peft_model_state_dict, set_peft_model_state_dict from transformers import AutoModelForImageTextToText, AutoProcessor args.output_dir.mkdir(parents=True, exist_ok=True) progress_path = args.output_dir / "progress.jsonl" checkpoint_path = args.output_dir / "latest_training_state.pt" best_model_path = args.output_dir / "best_model_state.pt" eval_path = args.output_dir / "id_dev_metrics.jsonl" manifest_sha256 = hashlib.sha256(args.manifest.read_bytes()).hexdigest() rows = read_rows(args.manifest) dev_rows = read_dev_rows(args.manifest) processor = AutoProcessor.from_pretrained(args.snapshot, local_files_only=True, padding_side="right") for token in SPECIAL_TOKENS: if token not in processor.tokenizer.get_vocab(): processor.tokenizer.add_special_tokens({"additional_special_tokens": [token]}) token_ids = {token: int(processor.tokenizer.convert_tokens_to_ids(token)) for token in SPECIAL_TOKENS} model = AutoModelForImageTextToText.from_pretrained( args.snapshot, local_files_only=True, dtype=torch.bfloat16, attn_implementation="flash_attention_2", device_map={"": 0}, ) model.resize_token_embeddings(len(processor.tokenizer)) for parameter in model.parameters(): parameter.requires_grad_(False) model.model.language_model = get_peft_model( model.model.language_model, LoraConfig( r=8, lora_alpha=16, lora_dropout=0.05, bias="none", target_modules=["q_proj", "k_proj", "v_proj", "o_proj"], trainable_token_indices=list(token_ids.values()), ), ) model.config.use_cache = False model.gradient_checkpointing_enable(gradient_checkpointing_kwargs={"use_reentrant": False}) model.enable_input_require_grads() hidden_size = int(model.config.text_config.hidden_size) heads = RobometerHeads(hidden_size).to(device="cuda", dtype=torch.bfloat16) backbone_parameters = [parameter for parameter in model.parameters() if parameter.requires_grad] head_parameters = list(heads.parameters()) optimizer = torch.optim.AdamW( [ {"params": backbone_parameters, "lr": args.lora_lr}, {"params": head_parameters, "lr": args.head_lr}, ], weight_decay=0.01, ) warmup_steps = max(1, round(args.optimizer_steps * 0.05)) def lr_lambda(step: int) -> float: if step < warmup_steps: return (step + 1) / warmup_steps progress = (step - warmup_steps) / max(1, args.optimizer_steps - warmup_steps) return 0.5 * (1.0 + math.cos(math.pi * progress)) scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda) generator = np.random.default_rng(args.seed) order = generator.permutation(len(rows)).tolist() cursor = 0 start_step = 0 step_seconds: list[float] = [] eval_history: list[dict] = [] best_eval_loss = math.inf best_step = 0 bad_eval_count = 0 pos_weight = torch.tensor(global_success_pos_weight(rows), dtype=torch.float32, device="cuda") if args.resume: if not checkpoint_path.is_file(): raise FileNotFoundError(checkpoint_path) checkpoint = torch.load(checkpoint_path, map_location="cpu", weights_only=False) if checkpoint.get("protocol_id") != "phase21a_robometer_low_lora_train_v2": raise RuntimeError("refusing to resume a checkpoint from a different protocol") if checkpoint["manifest_sha256"] != manifest_sha256 or checkpoint["seed"] != args.seed: raise RuntimeError("resume manifest or seed mismatch") set_peft_model_state_dict(model.model.language_model, checkpoint["adapter_state"]) heads.load_state_dict(checkpoint["heads_state"]) optimizer.load_state_dict(checkpoint["optimizer_state"]) scheduler.load_state_dict(checkpoint["scheduler_state"]) generator.bit_generator.state = checkpoint["numpy_generator_state"] order = checkpoint["order"] cursor = checkpoint["cursor"] start_step = checkpoint["completed_steps"] step_seconds = checkpoint["step_seconds"] eval_history = checkpoint.get("eval_history", []) best_eval_loss = float(checkpoint.get("best_eval_loss", math.inf)) best_step = int(checkpoint.get("best_step", 0)) bad_eval_count = int(checkpoint.get("bad_eval_count", 0)) torch.set_rng_state(checkpoint["torch_rng_state"]) torch.cuda.set_rng_state_all(checkpoint["cuda_rng_state_all"]) reconcile_progress_log(progress_path, start_step) reconcile_progress_log(eval_path, start_step) if start_step >= args.optimizer_steps: print(json.dumps({"status": "already_complete", "completed_steps": start_step}), flush=True) return model.train() heads.train() torch.cuda.empty_cache() torch.cuda.reset_peak_memory_stats() named_trainable = list(model.named_parameters()) + [ (f"heads.{name}", parameter) for name, parameter in heads.named_parameters() ] completed = start_step stop_reason = "max_steps" for step in range(start_step, args.optimizer_steps): if cursor + 64 > len(order): order = generator.permutation(len(rows)).tolist() cursor = 0 selected_indices = order[cursor : cursor + 64] cursor += 64 selected_rows = [rows[index] for index in selected_indices] optimizer.zero_grad(set_to_none=True) components = {"preference": 0.0, "progress": 0.0, "success": 0.0} sequence_min = 10**9 sequence_max = 0 torch.cuda.synchronize() step_started = time.perf_counter() for microstep in range(args.gradient_accumulation): batch_rows = selected_rows[ microstep * args.micro_batch_pairs : (microstep + 1) * args.micro_batch_pairs ] conversations = [] records_a = [] preference_labels = [] for offset, row in enumerate(batch_rows): global_presentation = step * 64 + microstep * args.micro_batch_pairs + offset chosen_is_a = global_presentation % 2 == 0 conversation, record_a, label = build_conversation(row, args.processed_root, chosen_is_a) conversations.append(conversation) records_a.append(record_a) preference_labels.append(label) inputs = process_batch(processor, conversations) lengths = inputs["attention_mask"].sum(1) sequence_min = min(sequence_min, int(lengths.min())) sequence_max = max(sequence_max, int(lengths.max())) inputs = {key: value.to("cuda") for key, value in inputs.items()} targets = make_targets(records_a, torch.device("cuda")) preference_target = torch.tensor(preference_labels, dtype=torch.float32, device="cuda") output = model.model(**inputs, output_hidden_states=False, use_cache=False) hidden = output.last_hidden_state prog_mask = inputs["input_ids"].eq(token_ids["<|prog_token|>"]) pref_mask = inputs["input_ids"].eq(token_ids["<|pref_token|>"]) if not torch.equal(prog_mask.sum(1), torch.full_like(prog_mask.sum(1), 16)): raise RuntimeError("progress token extraction mismatch") if not torch.equal(pref_mask.sum(1), torch.ones_like(pref_mask.sum(1))): raise RuntimeError("preference token extraction mismatch") prog_hidden = hidden[prog_mask].reshape(args.micro_batch_pairs, 16, hidden_size) pref_hidden = hidden[pref_mask].reshape(args.micro_batch_pairs, hidden_size) progress_logits = heads.progress_head(prog_hidden[:, :8]) success_logits = heads.success_head(prog_hidden[:, :8]).squeeze(-1) preference_logits = heads.preference_head(pref_hidden).squeeze(-1) if targets["progress_mask"].any(): progress_frames = -( targets["progress"] * F.log_softmax(progress_logits.float(), dim=-1) ).sum(-1) progress_loss = progress_frames[targets["progress_mask"]].mean() else: raise RuntimeError("formal micro-batch has empty progress supervision") success_loss = F.binary_cross_entropy_with_logits( success_logits.float()[targets["success_mask"]], targets["success"][targets["success_mask"]], pos_weight=pos_weight, ) preference_loss = F.binary_cross_entropy_with_logits( preference_logits.float(), preference_target ) total_loss = preference_loss + progress_loss + success_loss if not torch.isfinite(total_loss): raise FloatingPointError(f"non-finite loss at step {step + 1}") (total_loss / args.gradient_accumulation).backward() components["preference"] += float(preference_loss.detach()) / args.gradient_accumulation components["progress"] += float(progress_loss.detach()) / args.gradient_accumulation components["success"] += float(success_loss.detach()) / args.gradient_accumulation finite, grad_norm, grad_tensors = finite_grad_audit(named_trainable) if not finite: raise FloatingPointError(f"non-finite gradient at step {step + 1}") group_norms = { "lora": group_grad_norm(named_trainable, "lora_"), "token_rows": group_grad_norm(named_trainable, "trainable_tokens"), "progress_head": group_grad_norm(named_trainable, "heads.progress_head"), "success_head": group_grad_norm(named_trainable, "heads.success_head"), "preference_head": group_grad_norm(named_trainable, "heads.preference_head"), } if any(value <= 0 for value in group_norms.values()): raise RuntimeError(f"missing gradient group: {group_norms}") torch.nn.utils.clip_grad_norm_(backbone_parameters + head_parameters, 1.0) optimizer.step() scheduler.step() torch.cuda.synchronize() elapsed = time.perf_counter() - step_started step_seconds.append(elapsed) record = { "step": step + 1, "loss": sum(components.values()), "loss_components": components, "gradient_norm_preclip": grad_norm, "gradient_tensor_count": grad_tensors, "gradient_group_norms": group_norms, "step_seconds": elapsed, "lora_lr": optimizer.param_groups[0]["lr"], "head_lr": optimizer.param_groups[1]["lr"], "sequence_tokens_min": sequence_min, "sequence_tokens_max": sequence_max, "peak_allocated_gib": torch.cuda.max_memory_allocated() / 2**30, "peak_reserved_gib": torch.cuda.max_memory_reserved() / 2**30, } with progress_path.open("a", encoding="utf-8") as stream: stream.write(json.dumps(record, sort_keys=True) + "\n") print(json.dumps(record, sort_keys=True), flush=True) completed = step + 1 should_evaluate = completed % args.eval_every == 0 or completed == args.optimizer_steps if should_evaluate: metrics = evaluate_dev( model=model, processor=processor, heads=heads, rows=dev_rows, processed_root=args.processed_root, token_ids=token_ids, hidden_size=hidden_size, pos_weight=pos_weight, micro_batch_pairs=args.eval_micro_batch_pairs, ) metrics["step"] = completed eval_history.append(metrics) with eval_path.open("a", encoding="utf-8") as stream: stream.write(json.dumps(metrics, sort_keys=True) + "\n") print(json.dumps({"id_dev": metrics}, sort_keys=True), flush=True) if metrics["loss"] < best_eval_loss - args.early_stop_min_delta: best_eval_loss = float(metrics["loss"]) best_step = completed bad_eval_count = 0 atomic_torch_save( { "protocol_id": "phase21a_robometer_low_lora_train_v2", "selected_step": best_step, "id_dev_metrics": metrics, "adapter_state": cpu_state_dict( get_peft_model_state_dict(model.model.language_model) ), "heads_state": cpu_state_dict(heads.state_dict()), }, best_model_path, ) else: bad_eval_count += 1 should_checkpoint = completed % args.checkpoint_every == 0 or completed == args.optimizer_steps if should_checkpoint: atomic_torch_save( { "protocol_id": "phase21a_robometer_low_lora_train_v2", "completed_steps": completed, "target_steps": args.optimizer_steps, "seed": args.seed, "manifest_sha256": manifest_sha256, "adapter_state": get_peft_model_state_dict(model.model.language_model), "heads_state": heads.state_dict(), "optimizer_state": optimizer.state_dict(), "scheduler_state": scheduler.state_dict(), "numpy_generator_state": generator.bit_generator.state, "order": order, "cursor": cursor, "step_seconds": step_seconds, "eval_history": eval_history, "best_eval_loss": best_eval_loss, "best_step": best_step, "bad_eval_count": bad_eval_count, "torch_rng_state": torch.get_rng_state(), "cuda_rng_state_all": torch.cuda.get_rng_state_all(), }, checkpoint_path, ) if ( should_evaluate and completed >= args.early_stop_min_steps and bad_eval_count >= args.early_stop_patience_evals ): stop_reason = "id_dev_early_stop" break if not best_model_path.is_file(): raise RuntimeError("training completed without an ID-dev model selection checkpoint") best_model = torch.load(best_model_path, map_location="cpu", weights_only=False) set_peft_model_state_dict(model.model.language_model, best_model["adapter_state"]) heads.load_state_dict(best_model["heads_state"]) adapter_dir = args.output_dir / "adapter" model.model.language_model.save_pretrained(adapter_dir) torch.save(heads.state_dict(), args.output_dir / "reward_heads.pt") summary = { "status": "training_complete", "protocol_id": "phase21a_robometer_low_lora_train_v2", "seed": args.seed, "optimizer_steps_cap": args.optimizer_steps, "completed_steps": completed, "selected_step": best_step, "best_id_dev_loss": best_eval_loss, "stop_reason": stop_reason, "effective_batch_pairs": 64, "mean_step_seconds": float(np.mean(step_seconds)), "total_training_hours": float(np.sum(step_seconds) / 3600), "final_ood_used": False, } (args.output_dir / "summary.json").write_text(json.dumps(summary, indent=2) + "\n", encoding="utf-8") print(json.dumps(summary, indent=2), flush=True) if __name__ == "__main__": main()