Repository: pytorch
Issue: MultiProcess tests fail when run on nodes with 1 GPU
Labels: oncall: distributed,triaged

### 🐛 Describe the bug

We are building PyTorch on HPC nodes, some of the (build) nodes only have 1 GPU.

Running the test then fails, e.g. with 

```
Running distributed/fsdp/test_distributed_checkpoint ... [2022-11-24 11:32:28.143789]
Executing ['/rds/projects/2017/branfosj-rse/easybuild/EL8-ice/software/Python/3.9.5-GCCcore-10.3.0/bin/python', 'distributed/fsdp/test_distributed_checkpoint.py', '-v'] ... [2022-11-24 11:32:28.143928]
test_distributed_checkpoint_state_dict_type_StateDictType_LOCAL_STATE_DICT (__main__.TestDistributedCheckpoint) ... INFO:torch.testing._internal.common_distributed:Started process 0 with pid 733462
INFO:torch.testing._internal.common_distributed:Started process 1 with pid 733463
dist init r=0, world=2
INFO:torch.distributed.distributed_c10d:Added key: store_based_barrier_key:1 to store for rank: 0
dist init r=1, world=2
INFO:torch.distributed.distributed_c10d:Added key: store_based_barrier_key:1 to store for rank: 1
INFO:torch.distributed.distributed_c10d:Rank 1: Completed store-based barrier for key:store_based_barrier_key:1 with 2 nodes.
INFO:torch.distributed.distributed_c10d:Rank 0: Completed store-based barrier for key:store_based_barrier_key:1 with 2 nodes.
Process process 0:
Traceback (most recent call last):
  File "/rds/projects/2017/branfosj-rse/easybuild/EL8-ice/software/Python/3.9.5-GCCcore-10.3.0/lib/python3.9/multiprocessing/process.py", line 315, in _bootstrap
    self.run()
  File "/rds/projects/2017/branfosj-rse/easybuild/EL8-ice/software/Python/3.9.5-GCCcore-10.3.0/lib/python3.9/multiprocessing/process.py", line 108, in run
    self._target(*self._args, **self._kwargs)
  File "/dev/shm/branfosj/tmp-up-EL8/eb-usbph1a0/tmp_jxl19_q/lib/python3.9/site-packages/torch/testing/_internal/common_fsdp.py", line 427, in _run
    dist.barrier()
  File "/dev/shm/branfosj/tmp-up-EL8/eb-usbph1a0/tmp_jxl19_q/lib/python3.9/site-packages/torch/distributed/distributed_c10d.py", line 2784, in barrier
    work = default_pg.barrier(opts=opts)
RuntimeError: NCCL error in: ../torch/csrc/distributed/c10d/ProcessGroupNCCL.cpp:1191, invalid usage, NCCL version 2.10.3
ncclInvalidUsage: This usually reflects invalid usage of NCCL library (such as too many async ops, too many collectives at once, mixing streams in a group, etc).
Process process 1:
Traceback (most recent call last):
  File "/rds/projects/2017/branfosj-rse/easybuild/EL8-ice/software/Python/3.9.5-GCCcore-10.3.0/lib/python3.9/multiprocessing/process.py", line 315, in _bootstrap
    self.run()
  File "/rds/projects/2017/branfosj-rse/easybuild/EL8-ice/software/Python/3.9.5-GCCcore-10.3.0/lib/python3.9/multiprocessing/process.py", line 108, in run
    self._target(*self._args, **self._kwargs)
  File "/dev/shm/branfosj/tmp-up-EL8/eb-usbph1a0/tmp_jxl19_q/lib/python3.9/site-packages/torch/testing/_internal/common_fsdp.py", line 427, in _run
    dist.barrier()
  File "/dev/shm/branfosj/tmp-up-EL8/eb-usbph1a0/tmp_jxl19_q/lib/python3.9/site-packages/torch/distributed/distributed_c10d.py", line 2784, in barrier
    work = default_pg.barrier(opts=opts)
RuntimeError: NCCL error in: ../torch/csrc/distributed/c10d/ProcessGroupNCCL.cpp:1191, invalid usage, NCCL version 2.10.3
ncclInvalidUsage: This usually reflects invalid usage of NCCL library (such as too many async ops, too many collectives at once, mixing streams in a group, etc).
FAIL
```

With a bit of debugging I found the cause is the [implementation of `skip_if_lt_x_gpu`](https://github.com/pytorch/pytorch/blob/master/torch/testing/_internal/common_distributed.py#L161):

It is implemented to call `sys.exit` when the test function is called and there are not enough GPUs. However for the distributed tests this is wrong:   
Before calling the test function the test class `setup` function is called which calls `spawn_processes`: https://github.com/pytorch/pytorch/blob/master/torch/testing/_internal/common_fsdp.py#L699

The ProcessGroup entry function `_bootstrap` will then make the processes run into a barrier which seems to be implemented using NCCL leading to the bug when there are not enough GPUs available, i.e. world_size is larger than the number of the GPUs.

The solution is to not try and implement a custom skip-decorator but rather use the one from `unittest`. E.g.:

```
def skip_if_lt_x_gpu(x):
  return unittest.skipIf(not torch.cuda.is_available() or torch.cuda.device_count() < x)
```


### Versions

PyTorch 1.12.1 and master

cc @mrshenli @pritamdamania87 @zhaojuanmao @satgera @rohan-varma @gqchen @aazzolini @osalpekar @jiayisuse @H-Huang @kwen2501 @awgu