Switch torch dependency from ~=2.9.1 to ~=2.10.0 (silent bfloat16 memory regression)#118
Conversation
|
@claude review this PR 👀 👀 👀 |
|
Claude finished @hanaol's task in 2m 33s —— View job PR Review: Switch torch ~=2.9.1 → ~=2.10.0
Summary: The fix is correct and well-motivated. The bfloat16 Conv3d regression in torch 2.9.x is real (pytorch/pytorch#166643), and 2.10.0 cleanly resolves it. The benchmark methodology is solid. A few things worth flagging below. pyproject.tomlThe one-line change is correct: -"torch~=2.9.1",
+"torch~=2.10.0",
benchmark script (
|
…semantics, drop stale ref - Make `--output` default to `results.json` so the script runs standalone without args. - Add comment explaining that `bwd_peak_mb` reflects peak *during* backward (includes fresh-forward activations held by the autograd graph); relative comparisons across torch versions remain valid. - `del out` after the forward loop so the stale reference doesn't linger into the backward loop. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
forklady42
left a comment
forklady42
left a comment
There was a problem hiding this comment.
I addressed the minor comments about the benchmark script and reviewed myself
Claude's point about commit history is reasonable, but I'm not concerned as they'll be squashed when merging to main. Could be helpful to keep in mind in the future.
Good to merge ✅
## Summary This PR adds a benchmark script that trains one sample (task ID) at a time for 3 epochs under both f32 and bf16-mixed precision on a single A100 GPU, recording peak GPU memory, forward/backward times, and OOM status for 10 large-grid samples. ## The problem The 10 task IDs selected for this benchmark are Materials Project entries with relatively large charge-density grids, spanning 3.4 M – 46.7 M voxels across a variety of shapes and aspect ratios. It was not previously known which samples would fit on an A100 (79.3 GB) under f32 vs bf16-mixed, or how training time scales with grid volume. This script establishes those baselines. ## Benchmark results (A100-SXM4-80GB, CUDA 12.8) **Config**: `n_channels=32`, `n_residual_blocks=1`, `kernel_size=5`, `depth=2`, `batch_size=1`, single GPU, 3 epochs per experiment. | Task ID | Grid shape | Voxels | f32 status | f32 peak (GB) | f32 epoch (s) | bf16 status | bf16 peak (GB) | bf16 epoch (s) | Mem ratio | |---------|-----------|:------:|:----------:|:-------------:|:-------------:|:-----------:|:--------------:|:--------------:|:---------:| | mp-1890579 | 56 × 56 × 1080 | 3.4 M | ✅ | 10.5 | 1.65 | ✅ | 5.5 | 1.07 | 1.91× | | mp-1849767 | 60 × 60 × 1120 | 4.0 M | ✅ | 12.4 | 1.95 | ✅ | 6.5 | 1.27 | 1.91× | | mp-1851604 | 60 × 60 × 1120 | 4.0 M | ✅ | 12.4 | 1.93 | ✅ | 6.5 | 1.25 | 1.91× | | mp-1862536 | 80 × 80 × 1024 | 6.6 M | ✅ | 19.9 | 3.17 | ✅ | 10.4 | 1.94 | 1.91× | | mp-1847208 | 1120 × 84 × 84 | 7.9 M | ✅ | 23.9 | 3.87 | ✅ | 12.5 | 2.46 | 1.91× | | mp-1936557 | 80 × 756 × 216 | 13.1 M | ✅ | 39.1 | 6.47 | ✅ | 20.4 | 3.73 | 1.91× | | mp-1850168 | 972 × 240 × 128 | 29.9 M | ❌ OOM | 70.2 | — | ✅ | 46.3 | 8.77 | — | | mp-1887804 | 320 × 320 × 320 | 32.8 M | ❌ OOM | 68.8 | — | ✅ | 50.6 | 169.91 | — | | mp-1889246 | 540 × 144 × 432 | 33.6 M | ❌ OOM | 70.6 | — | ✅ | 51.9 | 208.94 | — | | mp-1871122 | 360 × 360 × 360 | 46.7 M | ❌ OOM | 71.4 | — | ❌ OOM | 71.9 | — | — | ### Summary | Precision | ✅ Completed | ❌ OOM | |-----------|:-----------:|:------:| | f32 | 6 / 10 | 4 / 10 | | bf16-mixed | 9 / 10 | 1 / 10 (mp-1871122, 360³) | ### Key findings - **Memory**: bf16-mixed yields a consistent **1.91× reduction** in peak GPU memory for all grids that fit under both precisions. - **Speed**: bf16-mixed is **1.5–1.7× faster** per epoch on grids ≤ 13 M voxels; the backward pass benefits the most (~2.2×). - **OOM threshold on A100**: ~13 M voxels under f32; ~33 M voxels under bf16-mixed. - **Largest grid (360³, 46.7 M voxels)**: exceeds A100 capacity even with bf16-mixed. - **Very large grids (320³, 540×144×432)**: fit under bf16-mixed but epoch times are 170–209 s, dominated by the backward pass (~144–176 s through large skip connections). ## Files - `scripts/benchmark_precision.py` -- Runs 3-epoch single-GPU training for each (task_id, precision) pair, recording CUDA-event forward/backward times and peak GPU memory --------- 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>
2 participants
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:
F.conv3dwithbfloat16Inputs inPyTorch 2.9.0pytorch/pytorch#166643 (issue)Benchmark results (A100-SXM4-80GB, CUDA 12.8, input shape [1, 32, 64, 64, 64])
Peak GPU memory —
float32Peak GPU memory —
bfloat16The 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