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