Blame for scripts/smoke_phase21a_robometer_low_full.py

#!/usr/bin/env python3
"""Single-GPU full-language-decoder Robometer-Low smoke on real pairs."""

from __future__ import annotations

import argparse
import json
import math
import os
import random
import sys
import time
from pathlib import Path

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,
    read_train_rows,
)


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("--micro-batch-pairs", type=int, required=True)
    parser.add_argument("--gradient-accumulation", type=int, required=True)
    parser.add_argument("--optimizer-steps", type=int, default=1)
    parser.add_argument("--pair-pool-size", type=int, default=8)
    parser.add_argument("--optimizer", choices=("adamw_torch", "paged_adamw_8bit"), required=True)
    parser.add_argument("--learning-rate", type=float, default=2.0e-5)
    parser.add_argument("--seed", type=int, default=7)
    parser.add_argument("--require-all-gradient-groups", action="store_true")
    parser.add_argument("--output", type=Path, required=True)
    return parser.parse_args()


def tensor_bytes(value) -> tuple[int, int]:
    """Return CUDA and CPU tensor bytes recursively for optimizer state."""
    cuda_bytes = 0
    cpu_bytes = 0
    if torch.is_tensor(value):
        size = value.numel() * value.element_size()
        if value.device.type == "cuda":
            cuda_bytes += size
        elif value.device.type == "cpu":
            cpu_bytes += size
    elif isinstance(value, dict):
        for item in value.values():
            cuda, cpu = tensor_bytes(item)
            cuda_bytes += cuda
            cpu_bytes += cpu
    elif isinstance(value, (list, tuple)):
        for item in value:
            cuda, cpu = tensor_bytes(item)
            cuda_bytes += cuda
            cpu_bytes += cpu
    return cuda_bytes, cpu_bytes


def make_optimizer(name: str, parameters: list[torch.nn.Parameter], learning_rate: float):
    if name == "adamw_torch":
        return torch.optim.AdamW(
            parameters,
            lr=learning_rate,
            weight_decay=0.01,
            foreach=False,
        )
    import bitsandbytes as bnb

    return bnb.optim.PagedAdamW8bit(
        parameters,
        lr=learning_rate,
        weight_decay=0.01,
        min_8bit_size=4096,
    )


def main() -> None:
    args = parse_args()
    if args.micro_batch_pairs <= 0 or args.gradient_accumulation <= 0:
        raise ValueError("batch and accumulation must be positive")
    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 transformers import AutoModelForImageTextToText, AutoProcessor

    started = time.perf_counter()
    processor = AutoProcessor.from_pretrained(
        args.snapshot, local_files_only=True, padding_side="right"
    )
    original_vocab = len(processor.tokenizer)
    for token in SPECIAL_TOKENS:
        if token not in processor.tokenizer.get_vocab():
            processor.tokenizer.add_special_tokens({"additional_special_tokens": [token]})
    special_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)
    for parameter in model.model.language_model.parameters():
        parameter.requires_grad_(True)
    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)
    language_parameters = [
        parameter for parameter in model.parameters() if parameter.requires_grad
    ]
    head_parameters = list(heads.parameters())
    trainable_parameters = language_parameters + head_parameters
    vision_trainable = sum(
        parameter.numel()
        for name, parameter in model.named_parameters()
        if parameter.requires_grad and ".visual." in name
    )
    if vision_trainable:
        raise RuntimeError(f"vision boundary violation: {vision_trainable} trainable parameters")
    optimizer = make_optimizer(args.optimizer, trainable_parameters, args.learning_rate)
    rows = read_train_rows(args.manifest, args.pair_pool_size)

    model.train()
    heads.train()
    torch.cuda.empty_cache()
    torch.cuda.reset_peak_memory_stats()
    load_seconds = time.perf_counter() - started
    run_started = time.perf_counter()
    component_history = []
    sequence_lengths = []
    optimizer.zero_grad(set_to_none=True)
    named_trainable = list(model.named_parameters()) + [
        (f"heads.{name}", parameter) for name, parameter in heads.named_parameters()
    ]
    last_gradient_audit = None

    for optimizer_step in range(args.optimizer_steps):
        for accumulation_index in range(args.gradient_accumulation):
            global_microstep = optimizer_step * args.gradient_accumulation + accumulation_index
            batch_rows = [
                rows[(global_microstep * args.micro_batch_pairs + offset) % len(rows)]
                for offset in range(args.micro_batch_pairs)
            ]
            conversations = []
            records_a = []
            preference_labels = []
            for offset, row in enumerate(batch_rows):
                chosen_is_a = ((global_microstep * args.micro_batch_pairs + offset) % 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)
            sequence_lengths.extend(inputs["attention_mask"].sum(1).tolist())
            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(special_token_ids["<|prog_token|>"])
            pref_mask = inputs["input_ids"].eq(special_token_ids["<|pref_token|>"])
            prog_counts = prog_mask.sum(1)
            pref_counts = pref_mask.sum(1)
            if not torch.equal(prog_counts, torch.full_like(prog_counts, 16)):
                raise RuntimeError(f"progress token counts: {prog_counts.tolist()}")
            if not torch.equal(pref_counts, torch.ones_like(pref_counts)):
                raise RuntimeError(f"preference token counts: {pref_counts.tolist()}")
            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():
                per_frame_progress = -(
                    targets["progress"] * F.log_softmax(progress_logits.float(), dim=-1)
                ).sum(-1)
                progress_loss = per_frame_progress[targets["progress_mask"]].mean()
            else:
                progress_loss = progress_logits.float().sum() * 0.0
            valid_success = targets["success"][targets["success_mask"]]
            valid_success_logits = success_logits.float()[targets["success_mask"]]
            positives = valid_success.sum()
            negatives = valid_success.numel() - positives
            pos_weight = (negatives / positives.clamp_min(1.0)).clamp_min(1.0)
            success_loss = F.binary_cross_entropy_with_logits(
                valid_success_logits, valid_success, pos_weight=pos_weight
            )
            preference_loss = F.binary_cross_entropy_with_logits(
                preference_logits.float(), preference_target
            )
            total_loss = progress_loss + success_loss + preference_loss
            if not torch.isfinite(total_loss):
                raise FloatingPointError(f"non-finite loss: {float(total_loss)}")
            (total_loss / args.gradient_accumulation).backward()
            component_history.append(
                {
                    "total": float(total_loss.detach()),
                    "preference": float(preference_loss.detach()),
                    "progress": float(progress_loss.detach()),
                    "success": float(success_loss.detach()),
                    "progress_supervised_frames": int(targets["progress_mask"].sum()),
                    "success_supervised_frames": int(targets["success_mask"].sum()),
                }
            )

        finite_gradients, preclip_grad_norm, gradient_tensor_count = finite_grad_audit(
            named_trainable
        )
        if not finite_gradients:
            raise FloatingPointError("non-finite trainable gradient")
        embedding_gradient = model.get_input_embeddings().weight.grad
        if embedding_gradient is None:
            special_token_grad_norm = 0.0
        else:
            special_token_grad_norm = float(
                torch.linalg.vector_norm(
                    embedding_gradient[list(special_token_ids.values())].float()
                )
            )
        grad_groups = {
            "language_decoder": group_grad_norm(named_trainable, "language_model"),
            "special_token_rows": special_token_grad_norm,
            "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"),
        }
        missing_grad_groups = [name for name, value in grad_groups.items() if not value > 0]
        if args.require_all_gradient_groups and missing_grad_groups:
            raise RuntimeError(f"missing gradient groups: {missing_grad_groups}")
        last_gradient_audit = {
            "finite": finite_gradients,
            "preclip_global_norm": preclip_grad_norm,
            "tensor_count": gradient_tensor_count,
            "group_norms": grad_groups,
            "missing_groups": missing_grad_groups,
        }
        torch.nn.utils.clip_grad_norm_(trainable_parameters, 10.0)
        optimizer.step()
        optimizer.zero_grad(set_to_none=True)

    torch.cuda.synchronize()
    run_seconds = time.perf_counter() - run_started
    device = torch.cuda.get_device_properties(0)
    peak_allocated = torch.cuda.max_memory_allocated()
    peak_reserved = torch.cuda.max_memory_reserved()
    optimizer_cuda_bytes, optimizer_cpu_bytes = tensor_bytes(optimizer.state)
    effective_batch = args.micro_batch_pairs * args.gradient_accumulation
    language_trainable_count = sum(parameter.numel() for parameter in language_parameters)
    head_trainable_count = sum(parameter.numel() for parameter in head_parameters)
    payload = {
        "status": "pass",
        "protocol_id": "phase21a_robometer_low_full_smoke_v1",
        "scope": {
            "smoke_only": True,
            "formal_training_started": False,
            "cosmos3_edge_modified": False,
        },
        "micro_batch_pairs": args.micro_batch_pairs,
        "gradient_accumulation": args.gradient_accumulation,
        "effective_batch_pairs": effective_batch,
        "optimizer_steps": args.optimizer_steps,
        "microsteps": args.optimizer_steps * args.gradient_accumulation,
        "optimizer": args.optimizer,
        "optimizer_official_default": args.optimizer == "adamw_torch",
        "real_pair_pool_size": len(rows),
        "pair_pool_selection": "longest_prompt_characters_in_train_split",
        "pair_pool_prompt_character_range": [
            min(len(row["prompt_task"]) for row in rows),
            max(len(row["prompt_task"]) for row in rows),
        ],
        "frames_per_pair": 16,
        "precision": "BF16",
        "gradient_checkpointing": True,
        "architecture": {
            "joint_a_b_forward": True,
            "vision_frozen": True,
            "language_decoder_full_tune": True,
            "peft_or_lora": False,
            "direct_preference_token_head": True,
            "per_frame_progress_tokens": True,
            "progress_bins": 10,
            "special_tokens": special_token_ids,
            "original_vocab_size": original_vocab,
            "resized_vocab_size": len(processor.tokenizer),
        },
        "parameters": {
            "model_total": sum(parameter.numel() for parameter in model.parameters()),
            "language_decoder_trainable": language_trainable_count,
            "heads_trainable": head_trainable_count,
            "all_trainable": language_trainable_count + head_trainable_count,
            "vision_trainable": vision_trainable,
            "trainable_fraction_of_model_and_heads": (
                (language_trainable_count + head_trainable_count)
                / (sum(parameter.numel() for parameter in model.parameters()) + head_trainable_count)
            ),
        },
        "optimizer_state": {
            "cuda_gib": optimizer_cuda_bytes / 2**30,
            "cpu_gib": optimizer_cpu_bytes / 2**30,
        },
        "loss_last": component_history[-1],
        "loss_all_finite": all(
            math.isfinite(value)
            for item in component_history
            for key, value in item.items()
            if key in {"total", "preference", "progress", "success"}
        ),
        "gradient_audit": last_gradient_audit,
        "sequence_tokens": {
            "min": int(min(sequence_lengths)),
            "max": int(max(sequence_lengths)),
            "mean": float(np.mean(sequence_lengths)),
        },
        "timing": {
            "model_load_seconds": load_seconds,
            "run_seconds": run_seconds,
            "seconds_per_microstep": run_seconds
            / (args.optimizer_steps * args.gradient_accumulation),
            "pairs_per_second": effective_batch * args.optimizer_steps / run_seconds,
        },
        "gpu": {
            "name": device.name,
            "total_gib": device.total_memory / 2**30,
            "peak_allocated_gib": peak_allocated / 2**30,
            "peak_reserved_gib": peak_reserved / 2**30,
            "peak_reserved_fraction": peak_reserved / device.total_memory,
        },
    }
    args.output.parent.mkdir(parents=True, exist_ok=True)
    args.output.write_text(json.dumps(payload, indent=2) + "\n", encoding="utf-8")
    print(json.dumps(payload, indent=2), flush=True)


if __name__ == "__main__":
    main()