Switch torch dependency from ~=2.9.1 to ~=2.10.0 (silent bfloat16 memory regression)

pyproject.toml

@@ -28,7 +28,7 @@ requires-python = ">=3.11"
 dependencies = [
     "numpy~=2.3.3",
     "scikit-learn>=1.7.2",
-    "torch~=2.9.1",
+    "torch~=2.10.0",
     "torchvision>=0.24.0",
     "lightning~= 2.5.6",
     "wandb>=0.12.10",

scripts/benchmark_conv3d_memory.py

@@ -0,0 +1,170 @@
+"""Standalone nn.Conv3d peak-memory benchmark across dtypes.
+
+Fixed tensor (consistent across all torch versions):
+    Input:  (N=1, C_in=32, D=64, H=64, W=64)
+    Conv3d: in=32, out=32, kernel=5, padding=2, padding_mode='zeros'
+
+Collects:
+    - torch / CUDA / cuDNN versions
+    - GPU name and total VRAM
+    - Forward and backward peak GPU memory (MB) for float32 and bfloat16
+    - Forward and backward wall-clock time (ms)
+
+Usage:
+    python benchmark_conv3d_memory.py [--output results.json]
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import sys
+import time
+from pathlib import Path
+
+import torch
+import torch.nn as nn
+
+# ── Fixed benchmark configuration ─────────────────────────────────────────────
+N, CIN, COUT = 1, 32, 32
+D = H = W = 64
+K = 5
+PADDING = K // 2
+WARMUP_ITERS = 3
+BENCH_ITERS = 5  # average over multiple runs for stable timing
+
+
+# ── Memory helpers ────────────────────────────────────────────────────────────
+
+
+def reset_mem():
+    torch.cuda.empty_cache()
+    torch.cuda.reset_peak_memory_stats()
+
+
+def peak_mb() -> float:
+    return torch.cuda.max_memory_allocated() / 1024**2
+
+
+# ── Per-dtype measurement ─────────────────────────────────────────────────────
+
+
+def bench_dtype(dtype_str: str, device: torch.device) -> dict:
+    dtype = {"float32": torch.float32, "bfloat16": torch.bfloat16}[dtype_str]
+
+    torch.manual_seed(42)
+    conv = nn.Conv3d(CIN, COUT, K, padding=PADDING).to(device=device, dtype=dtype)
+    x_base = torch.randn(N, CIN, D, H, W, device=device, dtype=dtype)
+
+    # ── Warmup (no grad, no memory tracking) ──────────────────────────────────
+    for _ in range(WARMUP_ITERS):
+        with torch.no_grad():
+            _ = conv(x_base)
+    torch.cuda.synchronize()
+
+    # ── Forward pass ──────────────────────────────────────────────────────────
+    fwd_peaks, fwd_times = [], []
+    for _ in range(BENCH_ITERS):
+        x = x_base.detach().requires_grad_(True)
+        reset_mem()
+        torch.cuda.synchronize()
+        t0 = time.perf_counter()
+        out = conv(x)
+        torch.cuda.synchronize()
+        fwd_times.append(time.perf_counter() - t0)
+        fwd_peaks.append(peak_mb())
+        del out
+
+    # ── Backward pass ─────────────────────────────────────────────────────────
+    # Note: reset_peak_memory_stats() zeros the peak counter but does not free
+    # live tensors, so bwd_peak_mb reflects the peak *during* backward, which
+    # includes the fresh forward activations still held by the autograd graph —
+    # not purely the backward-specific allocation. The methodology is identical
+    # across torch versions, so relative comparisons remain valid.
+    bwd_peaks, bwd_times = [], []
+    for _ in range(BENCH_ITERS):
+        x = x_base.detach().requires_grad_(True)
+        out = conv(x)  # fresh forward to build the computation graph
+        reset_mem()
+        torch.cuda.synchronize()
+        t0 = time.perf_counter()
+        out.sum().backward()
+        torch.cuda.synchronize()
+        bwd_times.append(time.perf_counter() - t0)
+        bwd_peaks.append(peak_mb())
+
+    def avg(lst: list) -> float:
+        return round(sum(lst) / len(lst), 2)
+
+    return {
+        "fwd_peak_mb": avg(fwd_peaks),
+        "bwd_peak_mb": avg(bwd_peaks),
+        "fwd_time_ms": round(avg(fwd_times) * 1e3, 3),
+        "bwd_time_ms": round(avg(bwd_times) * 1e3, 3),
+    }
+
+
+# ── Entry point ───────────────────────────────────────────────────────────────
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Conv3d memory benchmark")
+    parser.add_argument(
+        "--output",
+        default="results.json",
+        help="Path to write JSON results (default: results.json)",
+    )
+    args = parser.parse_args()
+
+    if not torch.cuda.is_available():
+        result = {"error": "CUDA not available"}
+        with Path.open(args.output, "w") as f:
+            json.dump(result, f, indent=2)
+        sys.exit(1)
+
+    device = torch.device("cuda:0")
+    gpu_props = torch.cuda.get_device_properties(0)
+
+    # Decode cuDNN version integer.
+    # cuDNN < 9 : MAJOR*1000 + MINOR*100 + PATCH  (e.g. 8904 → "8.9.4")
+    # cuDNN 9+  : MAJOR*10000 + MINOR*1000 + PATCH*100 + BUILD (e.g. 91002 → "9.1.0.2")
+    raw_cudnn = torch.backends.cudnn.version()
+    if not isinstance(raw_cudnn, int):
+        cudnn_str = str(raw_cudnn)
+    elif raw_cudnn >= 10000:
+        major = raw_cudnn // 10000
+        minor = (raw_cudnn % 10000) // 1000
+        patch = (raw_cudnn % 1000) // 100
+        build = raw_cudnn % 100
+        cudnn_str = f"{major}.{minor}.{patch}.{build}"
+    else:
+        major = raw_cudnn // 1000
+        minor = (raw_cudnn % 1000) // 100
+        patch = raw_cudnn % 100
+        cudnn_str = f"{major}.{minor}.{patch}"
+
+    results = {
+        "torch_version": torch.__version__,
+        "cuda_runtime": torch.version.cuda,
+        "cudnn_version": cudnn_str,
+        "gpu_name": gpu_props.name,
+        "gpu_total_memory_gb": round(gpu_props.total_memory / 1024**3, 1),
+        "gpu_sm_count": gpu_props.multi_processor_count,
+        "tensor_shape": [N, CIN, D, H, W],
+        "conv_config": f"Conv3d(in={CIN}, out={COUT}, k={K}, padding={PADDING})",
+        "bench_iters": BENCH_ITERS,
+        "measurements": {},
+    }
+
+    for dtype in ("float32", "bfloat16"):
+        try:
+            results["measurements"][dtype] = bench_dtype(dtype, device)
+        except Exception as exc:
+            results["measurements"][dtype] = {"error": str(exc)}
+
+    with Path.open(args.output, "w") as f:
+        json.dump(results, f, indent=2)
+
+
+if __name__ == "__main__":
+    main()

scripts/benchmark_gpus.py

@@ -0,0 +1,149 @@
+"""
+Compare per-sample benchmark results across two GPUs (e.g. A100 vs H200).
+
+Reads two JSON files produced by benchmark_precision.py and prints
+a markdown report section with a combined table (memory, time, and ratios).
+
+Usage:
+    uv run python scripts/benchmark_gpus.py \
+        --gpu1 benchmark_results/per_task_precision.json \
+        --gpu2 benchmark_results/per_task_precision_h200.json \
+        --out  benchmark_results/gpu_comparison.md
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import sys
+from pathlib import Path
+
+
+def load(path: Path) -> tuple[dict, list[dict]]:
+    d = json.loads(path.read_text())
+    return d["gpu"], d["results"]
+
+
+def by_key(results: list[dict]) -> dict[tuple[str, str], dict]:
+    return {(r["task_id"], r["precision"]): r for r in results}
+
+
+def ratio(a, b):
+    """Return a/b ratio string."""
+    if a is None or b is None:
+        return "—"
+    return f"{a / b:.2f}x"
+
+
+def status(r):
+    if r is None:
+        return "—"
+    return "❌ OOM" if r["oom"] else "✅"
+
+
+def peak_gb(r):
+    if r is None or r["oom"] or r["peak_mem_mb"] is None:
+        return None
+    return round(r["peak_mem_mb"] / 1024, 2)
+
+
+def epoch_s(r):
+    if r is None or r["oom"]:
+        return None
+    return r["avg_epoch_s"]
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--gpu1", type=Path, required=True, help="First GPU results JSON (reference)"
+    )
+    parser.add_argument(
+        "--gpu2", type=Path, required=True, help="Second GPU results JSON"
+    )
+    parser.add_argument(
+        "--out", type=Path, default=None, help="Output markdown file (default: stdout)"
+    )
+    args = parser.parse_args()
+
+    gpu1_info, res1 = load(args.gpu1)
+    gpu2_info, res2 = load(args.gpu2)
+
+    idx1 = by_key(res1)
+    idx2 = by_key(res2)
+
+    gpu1_name = gpu1_info["name"]
+    gpu2_name = gpu2_info["name"]
+
+    # grid shape lookup
+    grid_shapes = {}
+    for r in res1 + res2:
+        if r.get("grid_shape"):
+            grid_shapes[r["task_id"]] = r["grid_shape"]
+
+    def shape_str(tid):
+        s = grid_shapes.get(tid)
+        return f"{s[0]} x {s[1]} x {s[2]}" if s else "—"
+
+    def voxels(tid):
+        s = grid_shapes.get(tid)
+        return f"{s[0] * s[1] * s[2] / 1e6:.1f} M" if s else "—"
+
+    all_task_ids = sorted({k[0] for k in set(idx1) | set(idx2)})
+
+    lines = []
+    a = lines.append
+
+    a("| GPU | Model | VRAM |")
+    a("|-----|-------|:----:|")
+    a(f"| GPU 1 (reference) | {gpu1_name} | {gpu1_info['total_mem_gb']} GB |")
+    a(f"| GPU 2 | {gpu2_name} | {gpu2_info['total_mem_gb']} GB |")
+    a("")
+
+    for prec in ["f32", "bf16-mixed"]:
+        a(f"### {prec}\n")
+        a(
+            f"| Task ID | Grid shape | Voxels "
+            f"| {gpu1_name} status | {gpu1_name} peak (GB) | {gpu1_name} epoch (s) "
+            f"| {gpu2_name} status | {gpu2_name} peak (GB) | {gpu2_name} epoch (s) "
+            f"| Peak mem ratio (GPU1/GPU2) | Epoch time ratio (GPU1/GPU2) |"
+        )
+        a(
+            "|---------|-----------|:------:"
+            "|:---------:|:-------------------:|:--------------------:"
+            "|:---------:|:-------------------:|:--------------------:"
+            "|:-------------------------:|:----------------------------:|"
+        )
+
+        for tid in all_task_ids:
+            key = (tid, prec)
+            r1 = idx1.get(key)
+            r2 = idx2.get(key)
+
+            p1, p2 = peak_gb(r1), peak_gb(r2)
+            e1, e2 = epoch_s(r1), epoch_s(r2)
+
+            p1_str = "—" if p1 is None else f"{p1:.2f}"
+            p2_str = "—" if p2 is None else f"{p2:.2f}"
+            e1_str = "—" if e1 is None else f"{e1:.2f}"
+            e2_str = "—" if e2 is None else f"{e2:.2f}"
+
+            a(
+                f"| {tid} | {shape_str(tid)} | {voxels(tid)} "
+                f"| {status(r1)} | {p1_str} | {e1_str} "
+                f"| {status(r2)} | {p2_str} | {e2_str} "
+                f"| {ratio(p1, p2)} | {ratio(e1, e2)} |"
+            )
+        a("")
+
+    output = "\n".join(lines)
+
+    if args.out:
+        args.out.parent.mkdir(parents=True, exist_ok=True)
+        args.out.write_text(output)
+    else:
+        sys.stdout.write(output + "\n")
+
+
+if __name__ == "__main__":
+    main()