"""run_gpu_smoke.py — real GPU smoke for the Composer Replication Framework. Runs the 3-channel loss composition on a real HuggingFace model on GPU, capturing memory + step-time + bf16 numerical sanity in addition to the loss curve. This is the verification that the framework's design choices (mixed-precision compatibility, GPU dtype casts, etc) work end-to-end on real hardware, NOT just CPU. Per docs/adrs/ADR-001-gpu-venue.md: target hardware is the local 5090 (sm_120, 32GB VRAM). Modal evaluated and rejected for this smoke phase (10x iteration penalty for verification work). Acceptance: 1. Model loads via AutoModelForCausalLM, bf16, device='cuda' 2. 50 steps run end-to-end with no nan/inf 3. Loss decreases meaningfully (final < 50% of initial) 4. Peak VRAM stays under 8 GB on 0.5B model (headroom check) 5. Step time stable (no thermal throttling, no swap thrashing) 6. CPU and GPU runs produce numerically equivalent results modulo bf16 quantization noise (numerical-equivalence test in tests/) """ from __future__ import annotations import argparse import csv import json import sys import time from pathlib import Path import torch HERE = Path(__file__).resolve().parent sys.path.insert(0, str(HERE.parent / "006-real-hf-model-smoke")) from compose_loss import compose_loss from real_batch import build_batch MODEL_REPO = "Qwen/Qwen2.5-0.5B-Instruct" DEFAULT_STEPS = 50 DEFAULT_LR = 1e-5 def main() -> int: parser = argparse.ArgumentParser() parser.add_argument("--steps", type=int, default=DEFAULT_STEPS) parser.add_argument("--lr", type=float, default=DEFAULT_LR) parser.add_argument("--alpha-sdpo", type=float, default=0.1) parser.add_argument("--beta-replay", type=float, default=0.05) parser.add_argument("--dtype", choices=["bf16", "fp32"], default="bf16") parser.add_argument("--results-dir", default=str(HERE / "results")) args = parser.parse_args() if not torch.cuda.is_available(): print("[gpu-smoke] CUDA not available — skipping (run on a host with a GPU)") return 1 results_dir = Path(args.results_dir) results_dir.mkdir(parents=True, exist_ok=True) dev_name = torch.cuda.get_device_name(0) cap = torch.cuda.get_device_capability(0) print(f"[gpu-smoke] device: {dev_name} (sm_{cap[0]}{cap[1]})") print(f"[gpu-smoke] dtype={args.dtype}, steps={args.steps}, lr={args.lr}, " f"alpha={args.alpha_sdpo}, beta={args.beta_replay}") torch_dtype = torch.bfloat16 if args.dtype == "bf16" else torch.float32 t_load_start = time.perf_counter() from transformers import AutoModelForCausalLM, AutoTokenizer print(f"[gpu-smoke] loading {MODEL_REPO} ...") tokenizer = AutoTokenizer.from_pretrained(MODEL_REPO) model = AutoModelForCausalLM.from_pretrained(MODEL_REPO, torch_dtype=torch_dtype) model = model.to("cuda") model.train() t_load_s = time.perf_counter() - t_load_start n_params = sum(p.numel() for p in model.parameters()) print(f"[gpu-smoke] model loaded in {t_load_s:.1f}s, {n_params / 1e9:.3f}B params") print(f"[gpu-smoke] VRAM after load: {torch.cuda.memory_allocated() / 1e9:.2f} GB") print("[gpu-smoke] building batch ...") batch = build_batch(tokenizer, device="cuda") optimizer = torch.optim.AdamW(model.parameters(), lr=args.lr) # Warmup CUDA graphs / kernel JIT print("[gpu-smoke] warmup pass ...") optimizer.zero_grad() _ = compose_loss(model, batch, alpha_sdpo=args.alpha_sdpo, beta_replay=args.beta_replay) torch.cuda.synchronize() optimizer.zero_grad() torch.cuda.reset_peak_memory_stats() rows: list[dict] = [] for step in range(args.steps): torch.cuda.synchronize() t0 = time.perf_counter() optimizer.zero_grad() components = compose_loss( model, batch, alpha_sdpo=args.alpha_sdpo, beta_replay=args.beta_replay, ) components.total.backward() finite_grads = all( (p.grad is None or torch.isfinite(p.grad).all().item()) for p in model.parameters() ) sq = sum( float((p.grad.detach() ** 2).sum()) for p in model.parameters() if p.grad is not None ) grad_norm = sq ** 0.5 optimizer.step() torch.cuda.synchronize() dt = time.perf_counter() - t0 c = components.detached() peak_mem_gb = torch.cuda.max_memory_allocated() / 1e9 row = { "step": step, "wall_s": dt, "lm_ce": c["lm_ce"], "sdpo_jsd": c["sdpo_jsd"], "trace_replay_dpo": c["trace_replay_dpo"], "total": c["total"], "grad_norm": grad_norm, "finite_grads": finite_grads, "peak_mem_gb": peak_mem_gb, } rows.append(row) if step % 5 == 0 or step == args.steps - 1: print(f"[step {step:3d}] total={c['total']:.4f} lm_ce={c['lm_ce']:.4f} " f"sdpo={c['sdpo_jsd']:.4f} dpo={c['trace_replay_dpo']:.4f} " f"|g|={grad_norm:.4f} dt={dt*1000:.1f}ms mem={peak_mem_gb:.2f}GB " f"finite={finite_grads}") losses = [r["total"] for r in rows] initial = losses[0] final = losses[-1] half = initial * 0.5 median_step_ms = sorted(r["wall_s"] for r in rows)[len(rows) // 2] * 1000 verdict = { "device": dev_name, "compute_capability": f"sm_{cap[0]}{cap[1]}", "dtype": args.dtype, "model": MODEL_REPO, "steps": args.steps, "model_load_s": t_load_s, "initial_loss": initial, "final_loss": final, "loss_decrease_pct": (1 - final / initial) * 100 if initial > 0 else 0, "all_grads_finite": all(r["finite_grads"] for r in rows), "loss_decreased_to_below_half": final < half, "peak_mem_gb": max(r["peak_mem_gb"] for r in rows), "median_step_ms": median_step_ms, "no_nan": all(not (l != l) for l in losses), # noqa: E741 "no_inf": all(abs(l) != float("inf") for l in losses), "passed": ( all(r["finite_grads"] for r in rows) and final < half and all(not (l != l) for l in losses) and all(abs(l) != float("inf") for l in losses) and max(r["peak_mem_gb"] for r in rows) < 8.0 ), } csv_path = results_dir / "gpu_loss_curve.csv" with csv_path.open("w", newline="") as f: writer = csv.DictWriter(f, fieldnames=list(rows[0].keys())) writer.writeheader() writer.writerows(rows) verdict_path = results_dir / "gpu_verdict.json" verdict_path.write_text(json.dumps(verdict, indent=2)) print() print("=" * 64) print(" GPU SMOKE VERDICT") print("=" * 64) for k, v in verdict.items(): print(f" {k:.<28} {v}") print("=" * 64) return 0 if verdict["passed"] else 1 if __name__ == "__main__": sys.exit(main())