Add NCCL exports

Author: amorehead

project/lit_autoencoder.py

@@ -127,7 +127,7 @@ def cli_main():
     # plugins
     # ------------
     trainer.plugins = [
-        'deepspeed_stage_2'
+        # 'deepspeed_stage_2'
     ]
 
     # ------------

project/lit_image_classifier.py

@@ -193,7 +193,7 @@ def cli_main():
     # plugins
     # ------------
     trainer.plugins = [
-        'deepspeed_stage_2'
+        # 'deepspeed_stage_2'
     ]
 
     # ------------

project/lit_mnist.py

@@ -184,7 +184,7 @@ def cli_main():
     # plugins
     # ------------
     trainer.plugins = [
-        'deepspeed_stage_2'
+        # 'deepspeed_stage_2'
     ]
 
     # ------------

train.bash

@@ -29,6 +29,10 @@ export https_proxy=https://proxy.ccs.ornl.gov:3128/
 export no_proxy='localhost,127.0.0.0/8,.ccs.ornl.gov,.ncrc.gov'
 export LC_ALL=en_US.utf8
 
+# Set NCCL settings for multi-node DDP
+export NCCL_NSOCKS_PERTHREAD=4
+export NCCL_SOCKET_NTHREADS=2
+
 # Run training script
 cd "$PROJDIR"/project || exit
 

train_nvme.bash

@@ -29,6 +29,10 @@ export https_proxy=https://proxy.ccs.ornl.gov:3128/
 export no_proxy='localhost,127.0.0.0/8,.ccs.ornl.gov,.ncrc.gov'
 export LC_ALL=en_US.utf8
 
+# Set NCCL settings for multi-node DDP
+export NCCL_NSOCKS_PERTHREAD=4
+export NCCL_SOCKET_NTHREADS=2
+
 # Run training script
 cd "$PROJDIR"/project || exit