scripts/train_phase21a_robometer_low_lora.py

#!/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()