Repository: pytorch
Issue: Tracing GPU usage of `torch.vmap` vectorized operation using `torch.profiler` cause CUDA OOM
Labels: triaged,oncall: profiler,module: vmap

### 🐛 Describe the bug

Hi there! I'm using **vmap** to vectorize my customized **rankdata** operation for batched 1D tensor. I want to trace the GPU memory usage to determine if the implementation could be more efficient.

If I wrap the **vmap** with `torch.profiler.profile`
```python
import torch
from torch.profiler import profile, record_function, ProfilerActivity

def rank1d(data: torch.Tensor):
    indices = torch.arange(1, data.numel() + 1, dtype=data.dtype, device=data.device)
    data.scatter_(dim=0, index=torch.argsort(data), src=indices)
    return data

device = torch.device('cuda:0')
input_data = torch.randn((8, 2048, 1024)).to(device)

with profile(activities=[
        ProfilerActivity.CPU, ProfilerActivity.CUDA], record_shapes=True) as prof:
    with record_function("model_inference"):
        with torch.no_grad():
            for i in range(10):
                v_func = torch.vmap(torch.vmap(rank1d, in_dims=1, out_dims=-1), in_dims=0, out_dims=0)
                output = v_func(input_data)

                torch.cuda.synchronize(device=device)
                torch.cuda.empty_cache()

                print(torch.cuda.memory_allocated() / 1024 / 1024 / 1024, torch.cuda.memory_reserved() / 1024 / 1024 / 1024)
```
The usage of GPU memory will grow linearly as following output shows
```
STAGE:2024-01-03 20:07:43 24961:24961 ActivityProfilerController.cpp:312] Completed Stage: Warm Up
/home/test.py:6: UserWarning: There is a performance drop because we have not yet implemented the batching rule for aten::scatter_.src. Please file us an issue on GitHub so that we can prioritize its implementation. (Triggered internally at ../aten/src/ATen/functorch/BatchedFallback.cpp:82.)
  data.scatter_(dim=0, index=torch.argsort(data), src=indices)
0.1875 0.1875
0.3125 0.3125
0.4375 0.4375
0.5625 0.5625
0.6875 0.6875
0.8125 0.8125
0.9375 0.9375
1.0625 1.0625
1.1875 1.1875
1.3125 1.3125
STAGE:2024-01-03 20:07:49 24961:24961 ActivityProfilerController.cpp:318] Completed Stage: Collection
STAGE:2024-01-03 20:07:50 24961:24961 ActivityProfilerController.cpp:322] Completed Stage: Post Processing
```
However, if I remove the outside wrap of pytorch profiler as follows
```python
import torch

def rank1d(data: torch.Tensor):
    indices = torch.arange(1, data.numel() + 1, dtype=data.dtype, device=data.device)
    data.scatter_(dim=0, index=torch.argsort(data), src=indices)
    return data

device = torch.device('cuda:0')
input_data = torch.randn((8, 2048, 1024)).to(device)


with torch.no_grad():
    for i in range(10):
        v_func = torch.vmap(torch.vmap(rank1d, in_dims=1, out_dims=-1), in_dims=0, out_dims=0)
        output = v_func(input_data)

        torch.cuda.synchronize(device=device)
        torch.cuda.empty_cache()

        print(torch.cuda.memory_allocated() / 1024 / 1024 / 1024, torch.cuda.memory_reserved() / 1024 / 1024 / 1024)
```
everything works as expected and I get a constant memory usage.
```
/home/test.py:5: UserWarning: There is a performance drop because we have not yet implemented the batching rule for aten::scatter_.src. Please file us an issue on GitHub so that we can prioritize its implementation. (Triggered internally at ../aten/src/ATen/functorch/BatchedFallback.cpp:82.)
  data.scatter_(dim=0, index=torch.argsort(data), src=indices)
0.0625 0.0625
0.0625 0.0625
0.0625 0.0625
0.0625 0.0625
0.0625 0.0625
0.0625 0.0625
0.0625 0.0625
0.0625 0.0625
0.0625 0.0625
0.0625 0.0625
```
In practice, it will loop over 1000 times and finally cause OOM. Besides, if I profile `rank1d` individually, everything goes fine.

**My problem is: Am I using `torch.profiler` correctly when profiling **vmap**? Or this is an unexpected interaction between `torch.profiler` and `torch.vmap`?**

### Versions

PyTorch version: 2.1.0+cu121
Is debug build: False
CUDA used to build PyTorch: 12.1
ROCM used to build PyTorch: N/A

OS: Ubuntu 20.04.6 LTS (x86_64)
GCC version: (Ubuntu 9.4.0-1ubuntu1~20.04.2) 9.4.0
Clang version: Could not collect
CMake version: version 3.16.3
Libc version: glibc-2.31

Python version: 3.8.8 (default, Apr 13 2021, 19:58:26)  [GCC 7.3.0] (64-bit runtime)
Python platform: Linux-5.4.0-146-generic-x86_64-with-glibc2.10
Is CUDA available: True
CUDA runtime version: Could not collect
CUDA_MODULE_LOADING set to: LAZY
GPU models and configuration: GPU 0: NVIDIA A100-PCIE-40GB
Nvidia driver version: 535.86.10
cuDNN version: Probably one of the following:
/usr/lib/x86_64-linux-gnu/libcudnn.so.8.9.6
/usr/lib/x86_64-linux-gnu/libcudnn_adv_infer.so.8.9.6
/usr/lib/x86_64-linux-gnu/libcudnn_adv_train.so.8.9.6
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_infer.so.8.9.6
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_train.so.8.9.6
/usr/lib/x86_64-linux-gnu/libcudnn_ops_infer.so.8.9.6
/usr/lib/x86_64-linux-gnu/libcudnn_ops_train.so.8.9.6
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

Versions of relevant libraries:
[pip3] flake8==3.9.0
[pip3] mypy-extensions==0.4.3
[pip3] numpy==1.24.4
[pip3] numpydoc==1.1.0
[pip3] onnx==1.15.0
[pip3] onnxmltools==1.12.0
[pip3] torch==2.1.0+cu121
[pip3] triton==2.1.0
[conda] blas                      1.0                         mkl  
[conda] mkl                       2021.2.0           h06a4308_296  
[conda] mkl-service               2.3.0            py38h27cfd23_1  
[conda] mkl_fft                   1.3.0            py38h42c9631_2  
[conda] mkl_random                1.2.1            py38ha9443f7_2  
[conda] numpy                     1.24.4                   pypi_0    pypi
[conda] numpydoc                  1.1.0              pyhd3eb1b0_1  
[conda] torch                     2.1.0+cu121              pypi_0    pypi
[conda] triton                    2.1.0                    pypi_0    pypi

cc @robieta @chaekit @aaronenyeshi @nbcsm @guotuofeng @guyang3532 @gaoteng-git @tiffzhaofb @dzhulgakov @davidberard98 @zou3519