Repository: pytorch
Issue: Error during DDP, torch.compile, and cudagraph_trees
Labels: oncall: distributed,triaged,module: ddp,module: cuda graphs,oncall: pt2,module: inductor,pt2d-triage-nov2024

### 🐛 Describe the bug

I'm running DDP + torch.compile with cudagraph_trees but I'm running into errors related to cuda graphs. This is without DDPOptimizer; I set `torch._dynamo.config.optimize_ddp = False`.

For the 8 rank job, 2 of the ranks error due to failing the assert [on this line](https://github.com/pytorch/pytorch/blob/v2.1.0/c10/cuda/CUDACachingAllocator.cpp#L2111). 

The other ranks fail due to:
```
[4/8][2023-11-15 10:55:31,972]   File "/home/user/.local/lib/python3.10/site-packages/torch/_tensor.py", line 508, in backward
[4/8][2023-11-15 10:55:31,972]     torch.autograd.backward(
[4/8][2023-11-15 10:55:31,972]   File "/home/user/.local/lib/python3.10/site-packages/torch/autograd/__init__.py", line 251, in backward
[4/8][2023-11-15 10:55:31,972]     Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass
[4/8][2023-11-15 10:55:31,972]   File "/home/user/.local/lib/python3.10/site-packages/torch/autograd/function.py", line 288, in apply
[4/8][2023-11-15 10:55:31,972]     return user_fn(self, *args)
[4/8][2023-11-15 10:55:31,972]   File "/home/user/.local/lib/python3.10/site-packages/torch/_functorch/aot_autograd.py", line 3233, in backward
[4/8][2023-11-15 10:55:31,972]     out = call_compiled_backward()
[4/8][2023-11-15 10:55:31,972]   File "/home/user/.local/lib/python3.10/site-packages/torch/_functorch/aot_autograd.py", line 3205, in call_compiled_backward
[4/8][2023-11-15 10:55:31,972]     out = call_func_with_args(
[4/8][2023-11-15 10:55:31,972]   File "/home/user/.local/lib/python3.10/site-packages/torch/_functorch/aot_autograd.py", line 1506, in call_func_with_args
[4/8][2023-11-15 10:55:31,972]     out = normalize_as_list(f(args))
[4/8][2023-11-15 10:55:31,972]   File "/home/user/.local/lib/python3.10/site-packages/torch/_dynamo/eval_frame.py", line 328, in _fn
[4/8][2023-11-15 10:55:31,972]     return fn(*args, **kwargs)
[4/8][2023-11-15 10:55:31,972]   File "/home/user/.local/lib/python3.10/site-packages/torch/_dynamo/external_utils.py", line 17, in inner
[4/8][2023-11-15 10:55:31,972]     return fn(*args, **kwargs)
[4/8][2023-11-15 10:55:31,972]   File "/home/user/.local/lib/python3.10/site-packages/torch/_inductor/codecache.py", line 656, in __call__
[4/8][2023-11-15 10:55:31,972]     return self.get_current_callable()(inputs)
[4/8][2023-11-15 10:55:31,972]   File "/home/user/.local/lib/python3.10/site-packages/torch/_inductor/compile_fx.py", line 685, in run
[4/8][2023-11-15 10:55:31,972]     return compiled_fn(new_inputs)
[4/8][2023-11-15 10:55:31,972]   File "/home/user/.local/lib/python3.10/site-packages/torch/_inductor/cudagraph_trees.py", line 368, in deferred_cudagraphify
[4/8][2023-11-15 10:55:31,972]     return fn(inputs)
[4/8][2023-11-15 10:55:31,972]   File "/home/user/.local/lib/python3.10/site-packages/torch/_inductor/compile_fx.py", line 633, in run
[4/8][2023-11-15 10:55:31,972]     return model(new_inputs)
[4/8][2023-11-15 10:55:31,972]   File "/home/user/.local/lib/python3.10/site-packages/torch/_inductor/cudagraph_trees.py", line 1741, in run
[4/8][2023-11-15 10:55:31,972]     out = self._run(new_inputs, function_id)
[4/8][2023-11-15 10:55:31,972]   File "/home/user/.local/lib/python3.10/site-packages/torch/_inductor/cudagraph_trees.py", line 1815, in _run
[4/8][2023-11-15 10:55:31,972]     return self.record_function(new_inputs, function_id)
[4/8][2023-11-15 10:55:31,972]   File "/home/user/.local/lib/python3.10/site-packages/torch/_inductor/cudagraph_trees.py", line 1846, in record_function
[4/8][2023-11-15 10:55:31,972]     node = CUDAGraphNode(
[4/8][2023-11-15 10:55:31,972]   File "/home/user/.local/lib/python3.10/site-packages/torch/_inductor/cudagraph_trees.py", line 892, in __init__
[4/8][2023-11-15 10:55:31,972]     ] = self._record(wrapped_function.model, recording_inputs)
[4/8][2023-11-15 10:55:31,972]   File "/home/user/.local/lib/python3.10/site-packages/torch/_inductor/cudagraph_trees.py", line 1074, in _record
[4/8][2023-11-15 10:55:31,972]     check_memory_pool(self.device, self.cuda_graphs_pool, memory)
[4/8][2023-11-15 10:55:31,972]   File "/home/user/.local/lib/python3.10/site-packages/torch/_inductor/cudagraph_trees.py", line 1607, in check_memory_pool
[4/8][2023-11-15 10:55:31,972]     raise RuntimeError(msg)
[4/8][2023-11-15 10:55:31,972] RuntimeError: These live storage data ptrs are in the cudagraph pool but not accounted for as an output of cudagraph trees:
[4/8][2023-11-15 10:55:31,972]
[4/8][2023-11-15 10:55:31,972] Data Pointer: 23184050983424, history:
```

I am able to run torch.compile with cudagraph_trees on a single GPU, which suggests some interaction with DDP.

I have not yet been able to find a minimal example to reproduce.
Would greatly appreciate pointers on how to debug further.

cc: @eellison 

### Versions

Ubuntu 22.04.1 LTS, Pytorch 2.1.0 built from source, H100 GPU

cc @H-Huang @awgu @kwen2501 @wanchaol @fegin @fduwjj @wz337 @wconstab @d4l3k @c-p-i-o @mcarilli @ezyang @eellison @penguinwu @chauhang @voznesenskym @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @ipiszy @yf225 @chenyang78 @kadeng @muchulee8 @ColinPeppler @amjames @desertfire @aakhundov @mrshenli @pritamdamania87 @zhaojuanmao @satgera @rohan-varma @gqchen @aazzolini @osalpekar @jiayisuse @XilunWu @msaroufim @bdhirsh @anijain2305 @zou3519 @peterbell10