add min_gpu verification similart tu tensor_parallel_example.py example

Signed-off-by: jafraustro <[email protected]>

Author: jafraustro

distributed/ddp/example.py

@@ -11,6 +11,12 @@
 
 from torch.nn.parallel import DistributedDataParallel as DDP
 
+def verify_min_gpu_count(min_gpus: int = 2) -> bool:
+    """ verification that we have at least 2 gpus to run dist examples """
+    has_gpu = torch.accelerator.is_available()
+    gpu_count = torch.accelerator.device_count()
+    return has_gpu and gpu_count >= min_gpus
+
 class ToyModel(nn.Module):
     def __init__(self):
         super(ToyModel, self).__init__()
@@ -88,4 +94,8 @@ def main():
     dist.destroy_process_group()
 
 if __name__ == "__main__":
+    _min_gpu_count = 2
+    if not verify_min_gpu_count(min_gpus=_min_gpu_count):
+        print(f"Unable to locate sufficient {_min_gpu_count} gpus to run this example. Exiting.")
+        sys.exit()
     main()

distributed/ddp/run_example.sh

@@ -1,10 +1,10 @@
 # /bin/bash
 # bash run_example.sh {file_to_run.py} {num_gpus}
 # where file_to_run = example to run. Default = 'example.py'
-# num_gpus = num local gpus to use (must be at least 2). Default = 4
+# num_gpus = num local gpus to use (must be at least 2). Default = 2
 
 # samples to run include:
 # example.py
 
-echo "Launching ${1:-example.py} with ${2:-4} gpus"
-torchrun --nnodes=1 --nproc_per_node=${2:-4} --rdzv_id=101 --rdzv_endpoint="localhost:5972" ${1:-example.py}
+echo "Launching ${1:-example.py} with ${2:-2} gpus"
+torchrun --nnodes=1 --nproc_per_node=${2:-2} --rdzv_id=101 --rdzv_endpoint="localhost:5972" ${1:-example.py}