pytorch / pytorch/pytorch

Inductor crash with ops that return lists of tensors in separate partitions under CUDA graphs.

Open
#193,953 0 comments 0 reactions 0 assignees View on GitHub
bot-triaged module: cuda graphs module: inductor oncall: pt2 triaged
Dominant language
Python
Stars
103k
Forks
29.5k
PR merge metrics
PR metrics pending

Description

### 🐛 Describe the bug

Hi. Below is some code that I believe ought to work, but crashes.

```
import torch
import torch._inductor.config as inductor_config
# Keep _foreach_abs an opaque FallbackKernel that returns list[Tensor]
inductor_config.fallback_by_default = True

def fn(a: torch.Tensor, b: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
ys = torch._foreach_abs([a, b])
return torch.unique(ys[0]).sum(), ys[1] + 1.0

a = torch.ones(4, device="cuda")
b = torch.ones(4, device="cuda")
compiled = torch.compile(fn, mode="reduce-overhead", fullgraph=True)
compiled(a, b)
```

The error is something like this:
```
RuntimeError Traceback (most recent call last)
[/tmp/ipykernel_1210/248242094.py](https://localhost:8080/#) in ()
11 b = torch.ones(4, device="cuda")
12 compiled = torch.compile(fn, mode="reduce-overhead", fullgraph=True)
---> 13 compiled(a, b)

19 frames
[/usr/local/lib/python3.12/dist-packages/torch/_inductor/cudagraph_trees.py](https://localhost:8080/#) in check_memory_pool(device, pool_id, live_storages_ptrs)
1948 f"Allocations:\n{formatted_s}"
1949 )
-> 1950 raise RuntimeError(msg)
1951
1952

RuntimeError: Detected 1 tensor(s) in the cudagraph pool not tracked as outputs. All live allocations must be tracked for correctness.
Debugging:
- Set torch._inductor.config.triton.cudagraph_trees_history_recording = True for allocation origins
- Objgraph backrefs disabled; set torch._inductor.config.triton.cudagraph_trees_objgraph = True (refs_live_tensor_{index}.svg)
- Search gc.get_objects() for tensors with data_ptr() in OrderedSet([133885696410112])
- Use refcycle to find what is preventing cleanup
Allocations:
Data Pointer: 133885696410112, history:
```

The issue seems to be that when inductor splits a graph into partitions and a list-returning op (e.g. _foreach_abs) lands in a different partition from the getitem that unpacks it, the producing partition returns the Python list as a partition output. But `CUDAWarmupNode.run` skips any output that isn't a torch.Tensor, so the tensor inside the list is never registered as a tracked cudagraph output, and then this causes check_memory_pool to throw a "Detected N tensor(s) in the cudagraph pool not tracked as outputs".

I verified with 2.11 in google colab and 2.12 on my machine. I also read through the current main branch and the relevant code seems to be unchanged on main.

### Versions

Linux. PyTorch 2.12

cc @mcarilli @ezyang @eellison @penguinwu @BoyuanFeng @chauhang @voznesenskym @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @ipiszy @kadeng @muchulee8 @amjames @aakhundov @coconutruben @jataylo

Contributor guide

Open the contributing guide

Research direction

Reproduce the crash with the Python example using torch.compile(..., mode="reduce-overhead", fullgraph=True), then inspect torch/_inductor/cudagraph_trees.py, especially check_memory_pool and CUDAWarmupNode.run. Trace how partition outputs containing lists of tensors are registered, and verify that the repro completes without the untracked-allocation error while preserving CUDA graph tracking.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
compilers
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Active
Clarity
Mostly clear
Newbie friendliness
55/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.