probes/p0_gpu_probe.py

# Phase 0 GPU probe. Runs on the 2x Tesla T4 shape -- the only accelerator §2 permits.
# Purpose: prove the GPU path end to end and settle facts the main run's schedule depends on:
#   (1) the real accelerator(s) handed out, (2) Turing capability envelope (fp16 yes / bf16 no),
#   (3) whether NCCL all-reduce actually works across both cards, (4) torch.compile and
#       flash-attention viability on cc 7.5, (5) a first tokens/sec reading, and
#       (6) wall-clock, so the quota delta around this job reveals the burn rate per hour.
# No credentials here. Nothing is installed at runtime beyond what the image already has.

import hashlib
import json
import os
import platform
import socket
import subprocess
import sys
import time

R = {}
MONO0 = time.monotonic()
EPOCH0 = time.time()


def step(name, fn):
    a = time.monotonic()
    try:
        R[name] = fn()
    except Exception as e:
        R[name] = {"error": f"{type(e).__name__}: {e}"}
    R.setdefault("_timing", {})[name] = round(time.monotonic() - a, 2)


def run(cmd, timeout=180):
    p = subprocess.run(cmd, shell=True, capture_output=True, text=True, timeout=timeout)
    return {"rc": p.returncode, "out": p.stdout[-2500:], "err": p.stderr[-1500:]}


def gpu_identity():
    import shutil

    d = {"CUDA_VISIBLE_DEVICES": os.environ.get("CUDA_VISIBLE_DEVICES"),
         "KAGGLE_DOCKER_IMAGE": os.environ.get("KAGGLE_DOCKER_IMAGE"),
         "python": sys.version.split()[0]}
    for p in ["/kaggle/working", "/tmp"]:
        u = shutil.disk_usage(p)
        d[f"free_GB_{p.strip('/')}"] = round(u.free / 1024**3, 1)
    d["smi_L"] = run("nvidia-smi -L", 60)
    d["smi_query"] = run(
        'nvidia-smi --query-gpu=index,name,driver_version,memory.total,compute_cap,'
        'multiprocessor_count --format=csv', 60)
    return d


def torch_env():
    import torch

    out = {
        "torch": torch.__version__, "cuda_built": torch.version.cuda,
        "cudnn": torch.backends.cudnn.version(),
        "nccl_built": torch.cuda.nccl.version() if hasattr(torch.cuda, "nccl") else None,
        "arch_list": list(torch.cuda.get_arch_list()),
        "device_count": torch.cuda.device_count(),
        "tf32_default": torch.backends.cuda.matmul.allow_tf32,
    }
    devs = []
    for i in range(torch.cuda.device_count()):
        p = torch.cuda.get_device_properties(i)
        devs.append({"i": i, "name": p.name, "cc": f"{p.major}.{p.minor}", "sm": p.multi_processor_count,
                     "mem_GB": round(p.total_memory / 1024**3, 2)})
    out["devices"] = devs
    out["all_turing_T4"] = all(d["cc"] == "7.5" and "T4" in d["name"] for d in devs)
    out["exactly_two"] = len(devs) == 2
    for extra in ("triton", "flash_attn", "xformers", "transformers", "datasets", "accelerate", "trl"):
        try:
            from importlib.metadata import version

            out[f"pkg_{extra}"] = version(extra if extra != "flash_attn" else "flash-attn")
        except Exception:
            out[f"pkg_{extra}"] = "absent"
    return out


def precision_and_gemm():
    import torch

    out = {}
    for dt in (torch.float32, torch.float16, torch.bfloat16):
        for tf32 in ([False] if dt != torch.float32 else [False, True]):
            torch.backends.cuda.matmul.allow_tf32 = tf32
            key = f"{str(dt).replace('torch.', '')}{'_tf32' if tf32 else ''}"
            try:
                t0 = time.monotonic()
                a = torch.randn(2048, 2048, device="cuda", dtype=dt)
                b = torch.randn(2048, 2048, device="cuda", dtype=dt)
                for _ in range(10):
                    c = a @ b
                torch.cuda.synchronize()
                dt_s = (time.monotonic() - t0) / 10
                out[key] = {"ok": True, "ms_per_gemm": round(dt_s * 1000, 2),
                            "tflops": round(2 * 2048**3 / dt_s / 1e12, 2),
                            "checksum": round(float(c.float().abs().mean()), 4)}
            except Exception as e:
                out[key] = {"ok": False, "err": f"{type(e).__name__}: {e}"[:180]}
    torch.backends.cuda.matmul.allow_tf32 = False
    return out


def sdpa_backends():
    import torch

    out = {}
    for name in ("flash", "mem_efficient", "math"):
        try:
            fn = getattr(torch.backends.cuda, f"{name}_sdpa_enabled")
            out[name + "_enabled"] = fn()
        except Exception as e:
            out[name + "_enabled"] = str(e)[:80]
    from torch.nn.attention import SDPBackend, sdpa_kernel
    q = torch.randn(2, 4, 128, 64, device="cuda", dtype=torch.float16)
    for name in ("FLASH_ATTENTION", "EFFICIENT_ATTENTION", "MATH"):
        try:
            be = getattr(SDPBackend, name)
            with sdpa_kernel(be):
                o = torch.nn.functional.scaled_dot_product_attention(q, q, q)
            torch.cuda.synchronize()
            out[f"sdpa_{name}"] = {"ok": True, "shape": list(o.shape)}
        except Exception as e:
            out[f"sdpa_{name}"] = {"ok": False, "err": f"{type(e).__name__}: {e}"[:180]}
    return out


def ddp_smoke():
    # The main run is 2 GPUs. If NCCL across these two T4s does not work, every plan that
    # assumes data-parallel throughput is wrong -- prove it here on a toy, not at launch.
    worker = r'''
import os, json, time, torch, torch.distributed as dist
dist.init_process_group("nccl")
r = dist.get_rank()
torch.cuda.set_device(r)
x = torch.ones(1024, 1024, device=f"cuda:{r}") * (r + 1)
t0 = time.monotonic()
for _ in range(20):
    y = x.clone()
    dist.all_reduce(y)
torch.cuda.synchronize()
dt = (time.monotonic() - t0) / 20
print("WORKER_JSON " + json.dumps({"rank": r, "world": dist.get_world_size(),
      "allreduce_mean": float(y.mean()), "ms_per_allreduce": round(dt*1000, 3),
      "backend": dist.get_backend(), "device": torch.cuda.get_device_name(r)}), flush=True)
dist.destroy_process_group()
'''
    with open("ddp_worker.py", "w") as f:
        f.write(worker)
    t = time.monotonic()
    p = subprocess.run(
        [sys.executable, "-m", "torch.distributed.run", "--nproc_per_node=2", "--master_port=29511",
         "ddp_worker.py"], capture_output=True, text=True, timeout=300,
        env={**os.environ, "OMP_NUM_THREADS": "1"})
    lines = [l for l in (p.stdout + p.stderr).splitlines() if l.startswith("WORKER_JSON ")]
    return {"rc": p.returncode, "seconds": round(time.monotonic() - t, 1),
            "workers": [json.loads(l[len("WORKER_JSON "):]) for l in lines],
            "err_tail": p.stderr[-1200:]}


def tiny_lm_steps():
    # A few hundred steps of a real 10M-param causal LM, so tokens/sec is a measurement and not
    # an extrapolation. Not a model-building exercise: config is deliberately small.
    import torch
    from transformers import LlamaConfig, LlamaForCausalLM

    cfg = LlamaConfig(hidden_size=384, num_hidden_layers=6, num_attention_heads=6,
                      num_key_value_heads=6, intermediate_size=1024, vocab_size=32000,
                      max_position_embeddings=1024)
    model = LlamaForCausalLM(cfg).cuda().to(torch.float16)
    n = sum(p.numel() for p in model.parameters())
    opt = torch.optim.AdamW(model.parameters(), lr=3e-4)
    B, L = 8, 1024
    ids = torch.randint(0, 32000, (B, L), device="cuda")
    tg = torch.randint(0, 32000, (B, L), device="cuda")
    model.train()
    losses = []
    warm = 5
    t = time.monotonic()
    for s in range(warm + 60):
        opt.zero_grad(set_to_none=True)
        with torch.autocast("cuda", dtype=torch.float16):
            out = model(input_ids=ids, labels=tg)
        loss = out.loss
        loss.backward()
        opt.step()
        if s >= warm:
            losses.append(round(float(loss), 4))
    torch.cuda.synchronize()
    el = time.monotonic() - t
    tps = (60 * B * L) / el
    return {"params": n, "loss_first": losses[0], "loss_last": losses[-1],
            "losses_sampled": losses[::15], "seconds": round(el, 1),
            "tokens_per_sec": round(tps), "peak_mem_GB": round(torch.cuda.max_memory_allocated() / 1024**3, 2),
            "note": "single GPU, fp16 autocast, 8x1024 batch"}


def compile_smoke():
    import torch

    try:
        m = torch.nn.Sequential(torch.nn.Linear(512, 512), torch.nn.GELU(), torch.nn.Linear(512, 512)).cuda()
        cm = torch.compile(m)
        x = torch.randn(256, 512, device="cuda")
        t0 = time.monotonic()
        for _ in range(20):
            y = cm(x)
        torch.cuda.synchronize()
        compile_s = time.monotonic() - t0
        return {"ok": True, "warm_plus_20_steps_s": round(compile_s, 2),
                "out_shape": list(y.shape)}
    except Exception as e:
        return {"ok": False, "err": f"{type(e).__name__}: {e}"[:300]}


def artifact():
    payload = json.dumps({"probe": "p0-gpu", "epoch0": EPOCH0, "host": socket.gethostname()}).encode()
    open("ounce100m_p0_gpu_artifact.json", "wb").write(payload)
    return {"size": len(payload), "sha256": hashlib.sha256(payload).hexdigest()}


step("gpu_identity", gpu_identity)
step("torch_env", torch_env)
step("precision_gemm", precision_and_gemm)
step("sdpa", sdpa_backends)
step("ddp_nccl", ddp_smoke)
step("compile", compile_smoke)
step("tiny_lm", tiny_lm_steps)
step("smi_after", lambda: run("nvidia-smi --query-gpu=index,name,utilization.gpu,memory.used "
                              "--format=csv", 60))
step("artifact", artifact)

R["_timing"]["total"] = round(time.monotonic() - MONO0, 1)
R["_session"] = {"epoch_start": EPOCH0, "epoch_end": time.time(),
                 "wall_seconds": round(time.time() - EPOCH0, 1),
                 "hostname": socket.gethostname(), "platform": platform.platform()}
print("PROBE_JSON_BEGIN")
print(json.dumps(R, indent=1, default=str))
print("PROBE_JSON_END")