Set cuda device max connections based on cuda capability

Signed-off-by: Guyue Huang <[email protected]>

launcher_scripts/nemo_launcher/collections/conditional_cfgs.py

@@ -52,11 +52,23 @@ def get_ag_overlap(cfg):
     else:
         print(1)
 
+@hydra.main(version_base=None, config_path="conf", config_name="get_cuda_device_max_connections")
+def get_cuda_device_max_connections(cfg):
+    """
+    Set CUDA_DEVICE_MAX_CONNECTIONS to 32 for blackwell, to 1 for hopper and earlier generations
+    """
+    global cuda_capability
+    if cuda_capability >= 10:
+        print(32)
+    else:
+        print(1)
 
 if __name__ == "__main__":
     if sys.argv[1] == "name=get_ln_sm_margin":
         get_ln_sm_margin()
     elif sys.argv[1] == "name=get_ag_overlap":
         get_ag_overlap()
+    elif sys.argv[1] == "name=get_cuda_device_max_connections":
+        get_cuda_device_max_connections()
     else:
         raise ValueError("The provided conditional config function does not exist.")

launcher_scripts/nemo_launcher/core/stages.py

@@ -574,6 +574,17 @@ def _cuda_device_max_connections(self) -> str:
         tensor_model_parallel_size = model_cfg.get("tensor_model_parallel_size", 1)
         context_parallel_size = model_cfg.get("context_parallel_size", 1)
         fsdp = model_cfg.get("fsdp", False)
+        if (
+            (tensor_model_parallel_size > 1 or context_parallel_size > 1)
+            and not fsdp
+        ):
+            get_cuda_device_max_connections_command= (
+                f"python3 {self._launcher_scripts_path / 'nemo_launcher/collections/conditional_cfgs.py'} "
+                f"name=get_cuda_device_max_connections"
+            )
+            return f"CUDA_DEVICE_MAX_CONNECTIONS=\$({get_cuda_device_max_connections_command})"
+        return ""
+
         return (
             "CUDA_DEVICE_MAX_CONNECTIONS=1"
             if (