[cuda] Switching CI to CUDA 12.1 timing out linux-bionic-cuda12.1-py3.10-gcc7 / test (distributed, 3, 3, linux.8xlarge.nvidia.gpu)
Labels: module: cuda, module: ci, triaged

### 🐛 Describe the bug

Following PR, switching CI from 11.8 to 12.1 CUDA : https://github.com/pytorch/pytorch/pull/102178

Failure of  linux-bionic-cuda12.1-py3.10-gcc7 / test (distributed, 3, 3, linux.8xlarge.nvidia.gpu) : https://github.com/pytorch/pytorch/actions/runs/5071677184/jobs/9113227787

```
023-05-24T23:05:58.3284344Z distributed/optim/test_zero_redundancy_optimizer.py::TestZeroRedundancyOptimizerDistributed::test_zero_join_gpu [2023-05-24 22:25:58,524] torch.testing._internal.common_distributed: [ERROR] Process 0 timed out with traceback:
2023-05-24T23:05:58.3284800Z 
2023-05-24T23:05:58.3284961Z Thread 0x00007fe2aaccc700 (most recent call first):
2023-05-24T23:05:58.3285257Z   <no Python frame>
2023-05-24T23:05:58.3285416Z 
2023-05-24T23:05:58.3285571Z Thread 0x00007fe2b371e700 (most recent call first):
2023-05-24T23:05:58.3285931Z   <no Python frame>
2023-05-24T23:05:58.3286098Z 
2023-05-24T23:05:58.3286273Z Current thread 0x00007fe2d8b0a700 (most recent call first):
2023-05-24T23:05:58.3286879Z   File "/opt/conda/envs/py_3.10/lib/python3.10/site-packages/torch/testing/_internal/common_distributed.py", line 620 in _event_listener
2023-05-24T23:05:58.3287377Z   File "/opt/conda/envs/py_3.10/lib/python3.10/threading.py", line 953 in run
2023-05-24T23:05:58.3287791Z   File "/opt/conda/envs/py_3.10/lib/python3.10/threading.py", line 1016 in _bootstrap_inner
2023-05-24T23:05:58.3288226Z   File "/opt/conda/envs/py_3.10/lib/python3.10/threading.py", line 973 in _bootstrap
2023-05-24T23:05:58.3288465Z 
2023-05-24T23:05:58.3288623Z Thread 0x00007fe3825f9080 (most recent call first):
2023-05-24T23:05:58.3289167Z   File "/opt/conda/envs/py_3.10/lib/python3.10/site-packages/torch/nn/parallel/distributed.py", line 1401 in _pre_forward
2023-05-24T23:05:58.3289808Z   File "/opt/conda/envs/py_3.10/lib/python3.10/site-packages/torch/nn/parallel/distributed.py", line 1531 in forward
2023-05-24T23:05:58.3290429Z   File "/opt/conda/envs/py_3.10/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1511 in _call_impl
2023-05-24T23:05:58.3291053Z   File "/opt/conda/envs/py_3.10/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1502 in _wrapped_call_impl
2023-05-24T23:05:58.3291673Z   File "/var/lib/jenkins/workspace/test/distributed/optim/test_zero_redundancy_optimizer.py", line 1016 in _test_zero_join
2023-05-24T23:05:58.3292215Z   File "/var/lib/jenkins/workspace/test/distributed/optim/test_zero_redundancy_optimizer.py", line 1122 in test_zero_join_gpu
2023-05-24T23:05:58.3292886Z   File "/opt/conda/envs/py_3.10/lib/python3.10/site-packages/torch/testing/_internal/common_distributed.py", line 109 in wrapper
2023-05-24T23:05:58.3293553Z   File "/opt/conda/envs/py_3.10/lib/python3.10/site-packages/torch/testing/_internal/common_distributed.py", line 543 in wrapper
2023-05-24T23:05:58.3294211Z   File "/opt/conda/envs/py_3.10/lib/python3.10/site-packages/torch/testing/_internal/common_distributed.py", line 657 in run_test
2023-05-24T23:05:58.3294872Z   File "/opt/conda/envs/py_3.10/lib/python3.10/site-packages/torch/testing/_internal/common_distributed.py", line 636 in _run
2023-05-24T23:05:58.3295374Z   File "/opt/conda/envs/py_3.10/lib/python3.10/multiprocessing/process.py", line 108 in run
2023-05-24T23:05:58.3295840Z   File "/opt/conda/envs/py_3.10/lib/python3.10/multiprocessing/process.py", line 314 in _bootstrap
2023-05-24T23:05:58.3296278Z   File "/opt/conda/envs/py_3.10/lib/python3.10/multiprocessing/spawn.py", line 129 in _main
2023-05-24T23:05:58.3296737Z   File "/opt/conda/envs/py_3.10/lib/python3.10/multiprocessing/spawn.py", line 116 in spawn_main
2023-05-24T23:05:58.3297107Z   File "<string>", line 1 in <module>
2023-05-24T23:05:58.3297286Z 
2023-05-24T23:05:58.3297632Z [2023-05-24 22:25:58,524] torch.testing._internal.common_distributed: [ERROR] Process 1 timed out with traceback:
2023-05-24T23:05:58.3297913Z 
2023-05-24T23:05:58.3298093Z Current thread 0x00007fb1a6d38700 (most recent call first):
2023-05-24T23:05:58.3298690Z   File "/opt/conda/envs/py_3.10/lib/python3.10/site-packages/torch/testing/_internal/common_distributed.py", line 620 in _event_listener
2023-05-24T23:05:58.3299184Z   File "/opt/conda/envs/py_3.10/lib/python3.10/threading.py", line 953 in run
2023-05-24T23:05:58.3299594Z   File "/opt/conda/envs/py_3.10/lib/python3.10/threading.py", line 1016 in _bootstrap_inner
2023-05-24T23:05:58.3300028Z   File "/opt/conda/envs/py_3.10/lib/python3.10/threading.py", line 973 in _bootstrap
2023-05-24T23:05:58.3300271Z 
2023-05-24T23:05:58.3300426Z Thread 0x00007fb250796080 (most recent call first):
2023-05-24T23:05:58.3300956Z   File "/opt/conda/envs/py_3.10/lib/python3.10/site-packages/torch/nn/modules/linear.py", line 114 in forward
2023-05-24T23:05:58.3301962Z   File "/opt/conda/envs/py_3.10/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1511 in _call_impl
2023-05-24T23:05:58.3302727Z   File "/opt/conda/envs/py_3.10/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1502 in _wrapped_call_impl
2023-05-24T23:05:58.3303371Z   File "/opt/conda/envs/py_3.10/lib/python3.10/site-packages/torch/nn/modules/container.py", line 217 in forward
2023-05-24T23:05:58.3303991Z   File "/opt/conda/envs/py_3.10/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1511 in _call_impl
2023-05-24T23:05:58.3304603Z   File "/opt/conda/envs/py_3.10/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1502 in _wrapped_call_impl
2023-05-24T23:05:58.3305250Z   File "/opt/conda/envs/py_3.10/lib/python3.10/site-packages/torch/nn/parallel/distributed.py", line 1372 in _run_ddp_forward
2023-05-24T23:05:58.3305895Z   File "/opt/conda/envs/py_3.10/lib/python3.10/site-packages/torch/nn/parallel/distributed.py", line 1535 in forward
2023-05-24T23:05:58.3306519Z   File "/opt/conda/envs/py_3.10/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1511 in _call_impl
2023-05-24T23:05:58.3307140Z   File "/opt/conda/envs/py_3.10/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1502 in _wrapped_call_impl
2023-05-24T23:05:58.3307673Z   File "/var/lib/jenkins/workspace/test/distributed/optim/test_zero_redundancy_optimizer.py", line 1016 in _test_zero_join
2023-05-24T23:05:58.3308347Z   File "/var/lib/jenkins/workspace/test/distributed/optim/test_zero_redundancy_optimizer.py", line 1122 in test_zero_join_gpu
2023-05-24T23:05:58.3309017Z   File "/opt/conda/envs/py_3.10/lib/python3.10/site-packages/torch/testing/_internal/common_distributed.py", line 109 in wrapper
2023-05-24T23:05:58.3309666Z   File "/opt/conda/envs/py_3.10/lib/python3.10/site-packages/torch/testing/_internal/common_distributed.py", line 543 in wrapper
2023-05-24T23:05:58.3310329Z   File "/opt/conda/envs/py_3.10/lib/python3.10/site-packages/torch/testing/_internal/common_distributed.py", line 657 in run_test
2023-05-24T23:05:58.3310991Z   File "/opt/conda/envs/py_3.10/lib/python3.10/site-packages/torch/testing/_internal/common_distributed.py", line 636 in _run
2023-05-24T23:05:58.3311469Z   File "/opt/conda/envs/py_3.10/lib/python3.10/multiprocessing/process.py", line 108 in run
2023-05-24T23:05:58.3311933Z   File "/opt/conda/envs/py_3.10/lib/python3.10/multiprocessing/process.py", line 314 in _bootstrap
2023-05-24T23:05:58.3312399Z   File "/opt/conda/envs/py_3.10/lib/python3.10/multiprocessing/spawn.py", line 129 in _main
2023-05-24T23:05:58.3312859Z   File "/opt/conda/envs/py_3.10/lib/python3.10/multiprocessing/spawn.py", line 116 in spawn_main
2023-05-24T23:05:58.3313208Z   File "<string>", line 1 in <module>
2023-05-24T23:05:58.3313390Z 
2023-05-24T23:05:58.3313504Z FAILED [300.0435s] [ 95%]
2023-05-24T23:05:58.3314080Z distributed/optim/test_zero_redundancy_optimizer.py::TestZeroRedundancyOptimizerDistributed::test_zero_model_parallel_parameters_as_bucket_view_False SKIPPED [8.1162s] (Need at least 4 CUDA devices) [ 97%]
2023-05-24T23:05:58.3314917Z distributed/optim/test_zero_redundancy_optimizer.py::TestZeroRedundancyOptimizerDistributed::test_zero_model_parallel_parameters_as_bucket_view_True SKIPPED [4.8112s] (Need at least 4 CUDA devices) [100%]
```

### Versions

2.1.0 nightly

cc @ngimel @seemethere @malfet @pytorch/pytorch-dev-infra @ptrblck