pytorch - 💡(How to fix) Fix PYTORCH_ALLOC_CONF=expandable_segments:True silently ignored for CUDA in 2.9.1 (only deprecated PYTORCH_CUDA_ALLOC_CONF alias works) [3 pull requests]

🐛 Describe the bug

On PyTorch 2.9.1+cu128, setting PYTORCH_ALLOC_CONF=expandable_segments:True (the new unified env var name) is silently ignored for CUDA — expandable_segments stays False. Only the deprecated PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True alias actually enables the feature.

The runtime warning that fires when using the alias tells users to switch to the new name, but the new name doesn't work for expandable_segments on CUDA in 2.9.1.

Minimal repro

Test 1: new unified env var

PYTORCH_ALLOC_CONF=expandable_segments:True python -c "
import torch
torch.zeros(1, device='cuda')
print('expandable_segments:', torch.cuda.memory._snapshot()['allocator_settings']['expandable_segments'])
"
# → expandable_segments: False  ❌

Test 2: deprecated alias

PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True python -c "
import torch
torch.zeros(1, device='cuda')
print('expandable_segments:', torch.cuda.memory._snapshot()['allocator_settings']['expandable_segments'])
"
# → [W AllocatorConfig.cpp:28] Warning: PYTORCH_CUDA_ALLOC_CONF is deprecated, use PYTORCH_ALLOC_CONF instead
# → expandable_segments: True  ✓

Expected behavior

Both env vars should enable expandable_segments for CUDA, since the deprecation warning explicitly tells users to migrate to the new name. Per the docs, PYTORCH_ALLOC_CONF is the primary env var; PYTORCH_CUDA_ALLOC_CONF is the backward-compat alias. PyTorch's own header c10/core/AllocatorConfig.h documents the same:

The primary environment variable for configuration is PYTORCH_ALLOC_CONF. For backward compatibility, PYTORCH_CUDA_ALLOC_CONF is also supported

Practical impact

Code that follows the deprecation warning's guidance and switches to PYTORCH_ALLOC_CONF silently loses the fragmentation reduction it was asking for. Measured on an RTX 3090 (24 GB), this leaves ~2 GiB stranded as "reserved but unallocated" memory — enough to push training configs that should fit into OOM.

Encountered in karpathy/autoresearch's train.py:8, which sets the new name (forward-looking) and silently does nothing. The visible deprecation warning then misleads users into using the broken path.

Versions

PyTorch version: 2.9.1+cu128
CUDA runtime: 13.0
GPU: NVIDIA GeForce RTX 3090
NVIDIA driver: 580.126.09
OS: Linux x86_64 (Ubuntu 22.04+)
Python: 3.10