| |
| """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 ( |
| 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() |
|
|
