spikes/002a-mini-gpu-smoke/run_gpu_smoke.py

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