Skip hierarchical partitioning ZeRO (hpZ) for single node (microsoft#717

)
LeetJoe · Sep 15, 2023 · 92d9b88 · 92d9b88
1 parent ffb8fb4
commit 92d9b88
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/applications/DeepSpeed-Chat/training/utils/ds_utils.py b/applications/DeepSpeed-Chat/training/utils/ds_utils.py
@@ -4,6 +4,7 @@
 # DeepSpeed Team
 
 import torch
+import deepspeed.comm as dist
 
 GLOBAL_BATCH_SIZE = 32
 MICRO_BATCH_SIZE = 4
@@ -38,7 +39,9 @@ def get_train_ds_config(offload,
     }
     if enable_mixed_precision_lora:
         zero_opt_dict["zero_quantized_nontrainable_weights"] = True
-        zero_opt_dict["zero_hpz_partition_size"] = torch.cuda.device_count()
+        if dist.get_world_size() != torch.cuda.device_count():
+            zero_opt_dict["zero_hpz_partition_size"] = torch.cuda.device_count(
+            )
     return {
         "train_batch_size": GLOBAL_BATCH_SIZE,
         "train_micro_batch_size_per_gpu": MICRO_BATCH_SIZE,