update min gpu = 4 for fsdp+tp

pytorch · Nov 22, 2023 · 77fe3d8 · 77fe3d8
1 parent 2de0144
commit 77fe3d8
Show file tree

Hide file tree

Showing 3 changed files with 4 additions and 4 deletions.
diff --git a/distributed/tensor_parallelism/fsdp_tp_example.py b/distributed/tensor_parallelism/fsdp_tp_example.py
@@ -47,11 +47,11 @@
 https://docs.google.com/presentation/d/17g6WqrO00rP3MsxbRENsPpjrlSkwiA_QB4r93_eB5is/
 """
 
-_min_gpu_count = 2
+_min_gpu_count = 4
 
 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(0)
+    sys.exit()
 
 
 def find_multiple(n: int, k: int) -> int:

diff --git a/distributed/tensor_parallelism/sequence_parallel_example.py b/distributed/tensor_parallelism/sequence_parallel_example.py
@@ -34,7 +34,7 @@
 
 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(0)
+    sys.exit()
 
 
 class ToyModel(nn.Module):

diff --git a/distributed/tensor_parallelism/tensor_parallel_example.py b/distributed/tensor_parallelism/tensor_parallel_example.py
@@ -52,7 +52,7 @@
 
 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(0)
+    sys.exit()
 
 class ToyModel(nn.Module):
     """MLP based model"""