Skip to content

Commit a2be897

Browse files
hanaolHananeh Oliaeiforklady42claude
authored
Switch torch dependency from ~=2.9.1 to ~=2.10.0 (silent bfloat16 memory regression) (#118)
## Summary This PR updates the torch dependency from ~=2.9.1 to ~=2.10.0 to fix a silent bfloat16 memory regression introduced in torch 2.9.0. ## The problem torch 2.9.0 and 2.9.1 contain a cuDNN regression that inflates the nn.Conv3d bfloat16 forward-pass workspace by 26x -- from ~77 MB to ~2,053 MB -- relative to both the preceding (2.8.0) and following (2.10.0) releases. These numbers were measured on a fixed tensor of shape [1, 32, 64, 64, 64] with a Conv3d(in=32, out=32, k=5, padding=2) layer. float32 memory is completely unaffected (stable at ~123 MB across all versions), confirming the bug is specific to the bfloat16 cuDNN kernel selection path. This matters because we use (or plan to use) bf16-mixed precision training. This regression would silently consume an extra ~2 GB per Conv3d layer, directly undermining the memory savings that bf16 is supposed to provide -- without any crash or warning. This issue has been raised in the PyTorch community: - pytorch/pytorch#166643 (issue) - pytorch/pytorch#166480 (fix) ## Benchmark results (A100-SXM4-80GB, CUDA 12.8, input shape [1, 32, 64, 64, 64]) ### Peak GPU memory — `float32` | torch | cuDNN | Fwd peak (MB) | Bwd peak (MB) | |-------|-------|:-------------:|:-------------:| | 2.8.0 | 9.1.0.2 | 123 | 212 | | 2.9.0 | 9.1.0.2 | 123 | 212 | | 2.9.1 | 9.1.0.2 | 123 | 212 | | 2.10.0 | 9.1.0.2 | 123 | 212 | | 2.11.0 | 9.1.9.0 | 123 | 209 | ### Peak GPU memory — `bfloat16` | torch | cuDNN | Fwd peak (MB) | Bwd peak (MB) | |-------|-------|---------------|---------------| | 2.8.0 | 9.1.0.2 | 77 | 111 | | **2.9.0** | 9.1.0.2 | **2053** | **2081** | | **2.9.1** | 9.1.0.2 | **2053** | **2081** | | 2.10.0 | 9.1.0.2 | 77 | 111 | | 2.11.0 | 9.1.9.0 | 77 | 111 | The benchmark script is included at scripts/benchmark_conv3d_memory.py and can be run standalone on any CUDA node. ## Decision: 2.10.0 vs 2.11.0 Both 2.10.0 and 2.11.0 are clean. This PR pins to 2.10.0 for now. Upgrading to 2.11.0 is possible but introduces a CUDA 13.0 dependency (vs 12.8 for all prior versions), which pulls in a new set of nvidia-*-cu13 libraries and we have not tested it against our full stack (lightning, etc.). Once our ecosystem catches up to CUDA 13.0, bumping to ~=2.11.0 is an option worth revisiting. ## Files changed - pyproject.toml -- torch~=2.9.1 to torch~=2.10.0 - uv.lock -- regenerated - scripts/benchmark_conv3d_memory.py -- standalone benchmark used to produce the results above --------- Co-authored-by: Hananeh Oliaei <ho0950@della-vis2.princeton.edu> Co-authored-by: Betsy Cannon <betsy@openathena.ai> Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com> Co-authored-by: Betsy Cannon <forklady42@users.noreply.github.com>
1 parent ffb70db commit a2be897

5 files changed

Lines changed: 666 additions & 61 deletions

File tree

pyproject.toml

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -28,7 +28,7 @@ requires-python = ">=3.11"
2828
dependencies = [
2929
"numpy~=2.3.3",
3030
"scikit-learn>=1.7.2",
31-
"torch~=2.9.1",
31+
"torch~=2.10.0",
3232
"torchvision>=0.24.0",
3333
"lightning~= 2.5.6",
3434
"wandb>=0.12.10",

scripts/benchmark_conv3d_memory.py

Lines changed: 170 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,170 @@
1+
"""Standalone nn.Conv3d peak-memory benchmark across dtypes.
2+
3+
Fixed tensor (consistent across all torch versions):
4+
Input: (N=1, C_in=32, D=64, H=64, W=64)
5+
Conv3d: in=32, out=32, kernel=5, padding=2, padding_mode='zeros'
6+
7+
Collects:
8+
- torch / CUDA / cuDNN versions
9+
- GPU name and total VRAM
10+
- Forward and backward peak GPU memory (MB) for float32 and bfloat16
11+
- Forward and backward wall-clock time (ms)
12+
13+
Usage:
14+
python benchmark_conv3d_memory.py [--output results.json]
15+
"""
16+
17+
from __future__ import annotations
18+
19+
import argparse
20+
import json
21+
import sys
22+
import time
23+
from pathlib import Path
24+
25+
import torch
26+
import torch.nn as nn
27+
28+
# ── Fixed benchmark configuration ─────────────────────────────────────────────
29+
N, CIN, COUT = 1, 32, 32
30+
D = H = W = 64
31+
K = 5
32+
PADDING = K // 2
33+
WARMUP_ITERS = 3
34+
BENCH_ITERS = 5 # average over multiple runs for stable timing
35+
36+
37+
# ── Memory helpers ────────────────────────────────────────────────────────────
38+
39+
40+
def reset_mem():
41+
torch.cuda.empty_cache()
42+
torch.cuda.reset_peak_memory_stats()
43+
44+
45+
def peak_mb() -> float:
46+
return torch.cuda.max_memory_allocated() / 1024**2
47+
48+
49+
# ── Per-dtype measurement ─────────────────────────────────────────────────────
50+
51+
52+
def bench_dtype(dtype_str: str, device: torch.device) -> dict:
53+
dtype = {"float32": torch.float32, "bfloat16": torch.bfloat16}[dtype_str]
54+
55+
torch.manual_seed(42)
56+
conv = nn.Conv3d(CIN, COUT, K, padding=PADDING).to(device=device, dtype=dtype)
57+
x_base = torch.randn(N, CIN, D, H, W, device=device, dtype=dtype)
58+
59+
# ── Warmup (no grad, no memory tracking) ──────────────────────────────────
60+
for _ in range(WARMUP_ITERS):
61+
with torch.no_grad():
62+
_ = conv(x_base)
63+
torch.cuda.synchronize()
64+
65+
# ── Forward pass ──────────────────────────────────────────────────────────
66+
fwd_peaks, fwd_times = [], []
67+
for _ in range(BENCH_ITERS):
68+
x = x_base.detach().requires_grad_(True)
69+
reset_mem()
70+
torch.cuda.synchronize()
71+
t0 = time.perf_counter()
72+
out = conv(x)
73+
torch.cuda.synchronize()
74+
fwd_times.append(time.perf_counter() - t0)
75+
fwd_peaks.append(peak_mb())
76+
del out
77+
78+
# ── Backward pass ─────────────────────────────────────────────────────────
79+
# Note: reset_peak_memory_stats() zeros the peak counter but does not free
80+
# live tensors, so bwd_peak_mb reflects the peak *during* backward, which
81+
# includes the fresh forward activations still held by the autograd graph —
82+
# not purely the backward-specific allocation. The methodology is identical
83+
# across torch versions, so relative comparisons remain valid.
84+
bwd_peaks, bwd_times = [], []
85+
for _ in range(BENCH_ITERS):
86+
x = x_base.detach().requires_grad_(True)
87+
out = conv(x) # fresh forward to build the computation graph
88+
reset_mem()
89+
torch.cuda.synchronize()
90+
t0 = time.perf_counter()
91+
out.sum().backward()
92+
torch.cuda.synchronize()
93+
bwd_times.append(time.perf_counter() - t0)
94+
bwd_peaks.append(peak_mb())
95+
96+
def avg(lst: list) -> float:
97+
return round(sum(lst) / len(lst), 2)
98+
99+
return {
100+
"fwd_peak_mb": avg(fwd_peaks),
101+
"bwd_peak_mb": avg(bwd_peaks),
102+
"fwd_time_ms": round(avg(fwd_times) * 1e3, 3),
103+
"bwd_time_ms": round(avg(bwd_times) * 1e3, 3),
104+
}
105+
106+
107+
# ── Entry point ───────────────────────────────────────────────────────────────
108+
109+
110+
def main():
111+
parser = argparse.ArgumentParser(description="Conv3d memory benchmark")
112+
parser.add_argument(
113+
"--output",
114+
default="results.json",
115+
help="Path to write JSON results (default: results.json)",
116+
)
117+
args = parser.parse_args()
118+
119+
if not torch.cuda.is_available():
120+
result = {"error": "CUDA not available"}
121+
with Path.open(args.output, "w") as f:
122+
json.dump(result, f, indent=2)
123+
sys.exit(1)
124+
125+
device = torch.device("cuda:0")
126+
gpu_props = torch.cuda.get_device_properties(0)
127+
128+
# Decode cuDNN version integer.
129+
# cuDNN < 9 : MAJOR*1000 + MINOR*100 + PATCH (e.g. 8904 → "8.9.4")
130+
# cuDNN 9+ : MAJOR*10000 + MINOR*1000 + PATCH*100 + BUILD (e.g. 91002 → "9.1.0.2")
131+
raw_cudnn = torch.backends.cudnn.version()
132+
if not isinstance(raw_cudnn, int):
133+
cudnn_str = str(raw_cudnn)
134+
elif raw_cudnn >= 10000:
135+
major = raw_cudnn // 10000
136+
minor = (raw_cudnn % 10000) // 1000
137+
patch = (raw_cudnn % 1000) // 100
138+
build = raw_cudnn % 100
139+
cudnn_str = f"{major}.{minor}.{patch}.{build}"
140+
else:
141+
major = raw_cudnn // 1000
142+
minor = (raw_cudnn % 1000) // 100
143+
patch = raw_cudnn % 100
144+
cudnn_str = f"{major}.{minor}.{patch}"
145+
146+
results = {
147+
"torch_version": torch.__version__,
148+
"cuda_runtime": torch.version.cuda,
149+
"cudnn_version": cudnn_str,
150+
"gpu_name": gpu_props.name,
151+
"gpu_total_memory_gb": round(gpu_props.total_memory / 1024**3, 1),
152+
"gpu_sm_count": gpu_props.multi_processor_count,
153+
"tensor_shape": [N, CIN, D, H, W],
154+
"conv_config": f"Conv3d(in={CIN}, out={COUT}, k={K}, padding={PADDING})",
155+
"bench_iters": BENCH_ITERS,
156+
"measurements": {},
157+
}
158+
159+
for dtype in ("float32", "bfloat16"):
160+
try:
161+
results["measurements"][dtype] = bench_dtype(dtype, device)
162+
except Exception as exc:
163+
results["measurements"][dtype] = {"error": str(exc)}
164+
165+
with Path.open(args.output, "w") as f:
166+
json.dump(results, f, indent=2)
167+
168+
169+
if __name__ == "__main__":
170+
main()

scripts/benchmark_gpus.py

Lines changed: 149 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,149 @@
1+
"""
2+
Compare per-sample benchmark results across two GPUs (e.g. A100 vs H200).
3+
4+
Reads two JSON files produced by benchmark_precision.py and prints
5+
a markdown report section with a combined table (memory, time, and ratios).
6+
7+
Usage:
8+
uv run python scripts/benchmark_gpus.py \
9+
--gpu1 benchmark_results/per_task_precision.json \
10+
--gpu2 benchmark_results/per_task_precision_h200.json \
11+
--out benchmark_results/gpu_comparison.md
12+
"""
13+
14+
from __future__ import annotations
15+
16+
import argparse
17+
import json
18+
import sys
19+
from pathlib import Path
20+
21+
22+
def load(path: Path) -> tuple[dict, list[dict]]:
23+
d = json.loads(path.read_text())
24+
return d["gpu"], d["results"]
25+
26+
27+
def by_key(results: list[dict]) -> dict[tuple[str, str], dict]:
28+
return {(r["task_id"], r["precision"]): r for r in results}
29+
30+
31+
def ratio(a, b):
32+
"""Return a/b ratio string."""
33+
if a is None or b is None:
34+
return "—"
35+
return f"{a / b:.2f}x"
36+
37+
38+
def status(r):
39+
if r is None:
40+
return "—"
41+
return "❌ OOM" if r["oom"] else "✅"
42+
43+
44+
def peak_gb(r):
45+
if r is None or r["oom"] or r["peak_mem_mb"] is None:
46+
return None
47+
return round(r["peak_mem_mb"] / 1024, 2)
48+
49+
50+
def epoch_s(r):
51+
if r is None or r["oom"]:
52+
return None
53+
return r["avg_epoch_s"]
54+
55+
56+
def main():
57+
parser = argparse.ArgumentParser()
58+
parser.add_argument(
59+
"--gpu1", type=Path, required=True, help="First GPU results JSON (reference)"
60+
)
61+
parser.add_argument(
62+
"--gpu2", type=Path, required=True, help="Second GPU results JSON"
63+
)
64+
parser.add_argument(
65+
"--out", type=Path, default=None, help="Output markdown file (default: stdout)"
66+
)
67+
args = parser.parse_args()
68+
69+
gpu1_info, res1 = load(args.gpu1)
70+
gpu2_info, res2 = load(args.gpu2)
71+
72+
idx1 = by_key(res1)
73+
idx2 = by_key(res2)
74+
75+
gpu1_name = gpu1_info["name"]
76+
gpu2_name = gpu2_info["name"]
77+
78+
# grid shape lookup
79+
grid_shapes = {}
80+
for r in res1 + res2:
81+
if r.get("grid_shape"):
82+
grid_shapes[r["task_id"]] = r["grid_shape"]
83+
84+
def shape_str(tid):
85+
s = grid_shapes.get(tid)
86+
return f"{s[0]} x {s[1]} x {s[2]}" if s else "—"
87+
88+
def voxels(tid):
89+
s = grid_shapes.get(tid)
90+
return f"{s[0] * s[1] * s[2] / 1e6:.1f} M" if s else "—"
91+
92+
all_task_ids = sorted({k[0] for k in set(idx1) | set(idx2)})
93+
94+
lines = []
95+
a = lines.append
96+
97+
a("| GPU | Model | VRAM |")
98+
a("|-----|-------|:----:|")
99+
a(f"| GPU 1 (reference) | {gpu1_name} | {gpu1_info['total_mem_gb']} GB |")
100+
a(f"| GPU 2 | {gpu2_name} | {gpu2_info['total_mem_gb']} GB |")
101+
a("")
102+
103+
for prec in ["f32", "bf16-mixed"]:
104+
a(f"### {prec}\n")
105+
a(
106+
f"| Task ID | Grid shape | Voxels "
107+
f"| {gpu1_name} status | {gpu1_name} peak (GB) | {gpu1_name} epoch (s) "
108+
f"| {gpu2_name} status | {gpu2_name} peak (GB) | {gpu2_name} epoch (s) "
109+
f"| Peak mem ratio (GPU1/GPU2) | Epoch time ratio (GPU1/GPU2) |"
110+
)
111+
a(
112+
"|---------|-----------|:------:"
113+
"|:---------:|:-------------------:|:--------------------:"
114+
"|:---------:|:-------------------:|:--------------------:"
115+
"|:-------------------------:|:----------------------------:|"
116+
)
117+
118+
for tid in all_task_ids:
119+
key = (tid, prec)
120+
r1 = idx1.get(key)
121+
r2 = idx2.get(key)
122+
123+
p1, p2 = peak_gb(r1), peak_gb(r2)
124+
e1, e2 = epoch_s(r1), epoch_s(r2)
125+
126+
p1_str = "—" if p1 is None else f"{p1:.2f}"
127+
p2_str = "—" if p2 is None else f"{p2:.2f}"
128+
e1_str = "—" if e1 is None else f"{e1:.2f}"
129+
e2_str = "—" if e2 is None else f"{e2:.2f}"
130+
131+
a(
132+
f"| {tid} | {shape_str(tid)} | {voxels(tid)} "
133+
f"| {status(r1)} | {p1_str} | {e1_str} "
134+
f"| {status(r2)} | {p2_str} | {e2_str} "
135+
f"| {ratio(p1, p2)} | {ratio(e1, e2)} |"
136+
)
137+
a("")
138+
139+
output = "\n".join(lines)
140+
141+
if args.out:
142+
args.out.parent.mkdir(parents=True, exist_ok=True)
143+
args.out.write_text(output)
144+
else:
145+
sys.stdout.write(output + "\n")
146+
147+
148+
if __name__ == "__main__":
149+
main()

0 commit comments

Comments
 (0)