torch.compile mode='reduce-overhead' incompatible with PyTorch >=2.1 (CUDA Graphs tensor lifecycle)

[TheHiddenObserver]

Bug Description

torch.compile(fn, mode='reduce-overhead') fails at runtime on PyTorch 2.1+ (confirmed on 2.8.0+cu128) with a CUDA Graphs tensor lifecycle error. The code was written for PyTorch 2.0 where mode='reduce-overhead' used a different internal mechanism.

Error

RuntimeError: Error: accessing tensor output of CUDAGraphs that has been overwritten
by a subsequent run. Stack trace: File ".../statgpu/linear_model/_penalized.py", line 1534,
in _fista_elementwise
    y = c + _beta * (c - _coef_old). To prevent overwriting, clone the tensor outside of
torch.compile() or call torch.compiler.cudagraph_mark_step_begin() before each model invocation.

Root Cause

In PyTorch 2.0, mode='reduce-overhead' primarily used torch.cuda.make_graphed_callables. In PyTorch 2.1+, the Inductor backend with CUDA Graphs (cudagraph_trees) became the default for reduce-overhead mode. CUDA Graphs require that input/output tensors are not overwritten between replay calls — but the FISTA iteration loop reuses tensors like y_k, coef_old in-place across iterations.

Affected Code Locations (11 total)

All uses of mode='reduce-overhead':

File	Line	Function
linear_model/_penalized.py	1498	_fista_elementwise_l2 (ElasticNet L2 path)
linear_model/_penalized.py	1537	_fista_elementwise (Lasso/ElasticNet path)
linear_model/_elasticnet.py	194	_elastic_net_proximal_torch
glm_core/_solver.py	1551	_fused_update (squared_error+penalty)
penalties/_l1.py	27	L1 proximal
penalties/_adaptive_l1.py	32	Adaptive L1 proximal
penalties/_scad.py	48	SCAD proximal
penalties/_mcp.py	47	MCP proximal
penalties/_group_lasso.py	36	Group Lasso proximal
penalties/_group_mcp.py	42	Group MCP proximal
penalties/_group_scad.py	45	Group SCAD proximal

Note: _irls.py and _solver.py (FISTA/Newton steps) use dynamic=True, fullgraph=False without mode='reduce-overhead', so they are not affected.

Reproduction

import torch  # PyTorch 2.1+
import numpy as np
from statgpu.linear_model import Lasso

X = np.random.randn(2000, 200).astype(np.float64)
y = X @ np.concatenate([np.ones(20), np.zeros(180)]) + np.random.randn(2000) * 0.5

model = Lasso(alpha=0.1, device='torch')
model.fit(X, y)  # RuntimeError

Environment

• PyTorch: 2.8.0+cu128 (also affects 2.1+)
• GPU: NVIDIA GeForce RTX 4090 (CUDA cap 8.9)
• CUDA: 12.8
• Python: 3.13.5
• OS: Ubuntu (Docker)

Suggested Fixes

Option A: Replace mode='reduce-overhead' with mode='default':

torch.compile(fn, mode='default')  # no CUDA Graphs, still JIT-compiles

Option B: Add torch.compiler.cudagraph_mark_step_begin() before each compiled call in the iteration loop.

Option C: Make mode configurable via environment variable, defaulting to 'default' on PyTorch >= 2.1:

mode = os.environ.get('STATGPU_COMPILE_MODE', 'default')
torch.compile(fn, mode=mode)

Workaround

Users on PyTorch >= 2.1 can set the environment variable before importing statgpu:

import os
os.environ['TORCH_COMPILE_MODE'] = 'default'

(This requires code changes to respect the env var.)

Notes

• On PyTorch 2.0 + P100 (CUDA cap 6.0), torch.compile silently falls back to eager mode, so this bug is not triggered.
• On PyTorch 2.0 + GPU cap >= 7.0, reduce-overhead uses make_graphed_callables which handles tensor reuse differently — this bug may not manifest.
• The CuPy backend is unaffected.

[TheHiddenObserver]

Resolved by PR #87 and included in the 0.2.4 release tree.

Closure evidence:

• all statgpu-owned Torch compilation is routed through the centralized policy in statgpu/backends/_torch_compile.py;
• iterative kernels remain eager when STATGPU_TORCH_COMPILE_MODE is unset, auto, or disable;
• default and reduce-overhead are explicit opt-ins;
• the known CUDA Graph overwritten-output lifecycle failure produces an observable warning and permanently falls back to eager execution for that callable;
• unrelated runtime failures continue to propagate;
• regression coverage rejects unscoped reduce-overhead call sites;
• exact release-head RTX 4090 / PyTorch 2.8.0+cu128 acceptance passed, including the compile and CUDA Graph matrix: 9 passed.

Implementation merge: PR #87, commit 249456610c9d8ab62e9f2749e1fc7a769250b08a.
Release merge: PR #88, commit 0aeeb95b60e3e274053b8f1b6427ae50c8eec015.

The original PyTorch >=2.1 failure path is therefore fixed and covered by maintained regression and physical-GPU acceptance.