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