Repository: pytorch
Issue: Profiler Hangs on Non-Blocking H2D Transfer in Non-Default Stream
Labels: oncall: profiler

### 🐛 Describe the bug

Running the following script hangs at:
```
STAGE:2022-09-11 17:08:41 59397:59397 ActivityProfilerController.cpp:294] Completed Stage: Warm Up
```

```
import torch

from torch.profiler import profile, ProfilerActivity
from torch.testing._internal.common_utils import get_cycles_per_ms

stream = torch.cuda.Stream()
sleep_duration_ms = 25
sleep_duration_cycles = int(sleep_duration_ms * get_cycles_per_ms())
tensor_numel = 1 * 1024 * 1024
NON_BLOCKING = True
USE_SEPARATE_STREAM = True

cpu_tensor = torch.ones((tensor_numel,))
cpu_tensor = cpu_tensor.pin_memory()

stream_context = (
    torch.cuda.stream(stream) if USE_SEPARATE_STREAM
    else torch.cuda.stream(torch.cuda.current_stream())
)

with profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA]) as prof:
    with stream_context:
        torch.cuda._sleep(sleep_duration_cycles)
        gpu_tensor = cpu_tensor.to(torch.cuda.current_device(), non_blocking=NON_BLOCKING)

print(f"non-blocking={NON_BLOCKING}")
print(f"use separate stream={USE_SEPARATE_STREAM}")
prof.export_chrome_trace(f"trace_{NON_BLOCKING}_{USE_SEPARATE_STREAM}.json")
```

- If we use `USE_SEPARATE_STREAM = False` instead, then the profiler does not hang.
- If we do not include the H2D transfer (`gpu_tensor = ...`), then the profiler does not hang.
- If we do not include the `torch.cuda._sleep(sleep_duration_cycles)`, then the profiler does not hang.
- If we run the `torch.cuda._sleep(sleep_duration_cycles)` in the default stream but the H2D copy in the separate stream, then the profiler does hang.
- If we use `non_blocking=False`, then the profiler does not hang.
- If we do not use the profiler, then the code does not hang.


### Versions

```
PyTorch version: 1.13.0a0+gita879a8a
Is debug build: False
CUDA used to build PyTorch: 11.1
ROCM used to build PyTorch: N/A

OS: Ubuntu 18.04.6 LTS (x86_64)
GCC version: (Ubuntu 7.5.0-3ubuntu1~18.04) 7.5.0
Clang version: 6.0.0-1ubuntu2 (tags/RELEASE_600/final)
CMake version: version 3.21.0
Libc version: glibc-2.27

Python version: 3.8.10 (default, May 19 2021, 18:05:58)  [GCC 7.3.0] (64-bit runtime)
Python platform: Linux-5.4.0-1069-aws-x86_64-with-glibc2.10
Is CUDA available: True
CUDA runtime version: 11.1.105
GPU models and configuration: GPU 0: NVIDIA A100-SXM4-40GB
Nvidia driver version: 510.47.03
cuDNN version: Probably one of the following:
/usr/local/cuda-11.0/targets/x86_64-linux/lib/libcudnn.so.8.0.5
/usr/local/cuda-11.0/targets/x86_64-linux/lib/libcudnn_adv_infer.so.8.0.5
/usr/local/cuda-11.0/targets/x86_64-linux/lib/libcudnn_adv_train.so.8.0.5
/usr/local/cuda-11.0/targets/x86_64-linux/lib/libcudnn_cnn_infer.so.8.0.5
/usr/local/cuda-11.0/targets/x86_64-linux/lib/libcudnn_cnn_train.so.8.0.5
/usr/local/cuda-11.0/targets/x86_64-linux/lib/libcudnn_ops_infer.so.8.0.5
/usr/local/cuda-11.0/targets/x86_64-linux/lib/libcudnn_ops_train.so.8.0.5
/usr/local/cuda-11.1/targets/x86_64-linux/lib/libcudnn.so.8.0.5
/usr/local/cuda-11.1/targets/x86_64-linux/lib/libcudnn_adv_infer.so.8.0.5
/usr/local/cuda-11.1/targets/x86_64-linux/lib/libcudnn_adv_train.so.8.0.5
/usr/local/cuda-11.1/targets/x86_64-linux/lib/libcudnn_cnn_infer.so.8.0.5
/usr/local/cuda-11.1/targets/x86_64-linux/lib/libcudnn_cnn_train.so.8.0.5
/usr/local/cuda-11.1/targets/x86_64-linux/lib/libcudnn_ops_infer.so.8.0.5
/usr/local/cuda-11.1/targets/x86_64-linux/lib/libcudnn_ops_train.so.8.0.5
/usr/local/cuda-11.2/targets/x86_64-linux/lib/libcudnn.so.8.1.1
/usr/local/cuda-11.2/targets/x86_64-linux/lib/libcudnn_adv_infer.so.8.1.1
/usr/local/cuda-11.2/targets/x86_64-linux/lib/libcudnn_adv_train.so.8.1.1
/usr/local/cuda-11.2/targets/x86_64-linux/lib/libcudnn_cnn_infer.so.8.1.1
/usr/local/cuda-11.2/targets/x86_64-linux/lib/libcudnn_cnn_train.so.8.1.1
/usr/local/cuda-11.2/targets/x86_64-linux/lib/libcudnn_ops_infer.so.8.1.1
/usr/local/cuda-11.2/targets/x86_64-linux/lib/libcudnn_ops_train.so.8.1.1
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

Versions of relevant libraries:
[pip3] mypy==0.950
[pip3] mypy-extensions==0.4.3
[pip3] numpy==1.20.3
[pip3] torch==1.10.0a0+gitaa378da
[pip3] torchaudio==0.9.0
[pip3] torchdistx==0.3.0.dev0+cu111
[pip3] torchmultimodal==0.1.0a0
[pip3] torchtext==0.10.0a0+27695bb
[pip3] torchvision==0.11.0a0+31d5336
[conda] blas                      1.0                         mkl  
[conda] magma-cuda110             2.5.2                         1    pytorch
[conda] magma-cuda111             2.5.2                         1    pytorch
[conda] mkl                       2021.3.0           h06a4308_520  
[conda] mkl-include               2021.3.0           h06a4308_520  
[conda] mkl-service               2.4.0                    pypi_0    pypi
[conda] mkl_fft                   1.3.0            py38h42c9631_2  
[conda] mkl_random                1.2.1            py38ha9443f7_2  
[conda] numpy                     1.19.0                   pypi_0    pypi
[conda] numpy-base                1.20.3           py38h74d4b33_0  
[conda] pytorch-sphinx-theme      0.0.24                   pypi_0    pypi
[conda] torch                     1.9.0                    pypi_0    pypi
[conda] torchaudio                0.9.0                    pypi_0    pypi
[conda] torchdistx                0.3.0.dev0+cu111          pypi_0    pypi
[conda] torchmultimodal           0.1.0a0                   dev_0    <develop>
[conda] torchtext                 0.10.0a0+27695bb          pypi_0    pypi
[conda] torchvision               0.11.0a0+31d5336          pypi_0    pypi
```
My PyTorch build is from source with some unrelated commits on top. I think it can be effectively treated as `master` or `viable/strict`.

cc @robieta @chaekit @aaronenyeshi @ngimel @nbcsm @guotuofeng @guyang3532 @gaoteng-git @tiffzhaofb