Add NCCL auto-detection for InfiniBand and network interfaces

- Add get_nccl_env_vars() function to automatically detect network configuration
- Detects InfiniBand interfaces (ibp*, ib*) and enables them automatically
- Falls back to Ethernet (^lo) if no InfiniBand is available
- Respects user-set NCCL environment variables
- Sets NCCL_SOCKET_IFNAME and NCCL_IB_DISABLE based on actual cluster hardware
- Eliminates need for manual NCCL configuration

Author: Hossein Kavianihamedani

src/forge/controller/provisioner.py

@@ -14,6 +14,10 @@
 
 import torch
 
+from forge.controller.launcher import BaseLauncher, get_launcher
+from forge.env import all_env_vars, FORGE_DISABLE_METRICS
+from forge.types import ProcessConfig, ProvisionerConfig
+
 from monarch._src.actor.actor_mesh import ActorMesh
 from monarch._src.actor.shape import Extent
 
@@ -22,10 +26,6 @@
 from monarch.tools import commands
 from monarch.utils import setup_env_for_distributed
 
-from forge.controller.launcher import BaseLauncher, get_launcher
-from forge.env import all_env_vars, FORGE_DISABLE_METRICS
-from forge.types import ProcessConfig, ProvisionerConfig
-
 logger = logging.getLogger(__name__)
 logger.setLevel(logging.DEBUG)
 
@@ -134,6 +134,23 @@ async def set_environment(proc_mesh: ProcMesh, env_vars: dict[str, str]):
     await env_setter.set_env.call(env_vars)
 
 
+async def get_nccl_env_vars() -> dict[str, str]:
+    """Get NCCL environment variables by detecting network interfaces."""
+    if "NCCL_SOCKET_IFNAME" in os.environ and "NCCL_IB_DISABLE" in os.environ:
+        return {}
+
+    try:
+        interfaces = os.listdir("/sys/class/net/")
+        ib_interfaces = [i for i in interfaces if i.startswith("ib")]
+
+        return {
+            "NCCL_SOCKET_IFNAME": ",".join(ib_interfaces) if ib_interfaces else "^lo",
+            "NCCL_IB_DISABLE": "0" if ib_interfaces else "1",
+        }
+    except Exception:
+        return {"NCCL_SOCKET_IFNAME": "^lo", "NCCL_IB_DISABLE": "1"}
+
+
 class GpuManager:
     """Tracks and assigns GPU devices on a host.
 
@@ -347,11 +364,16 @@ async def get_proc_mesh(
             if with_gpus:
                 if not addr or not port:
                     addr, port = await get_remote_info(host_mesh)
-                gpu_ids = gpu_manager.get_gpus(num_procs)
+                gpu_ids: list[str] = gpu_manager.get_gpus(num_procs)
 
+                # Set PyTorch distributed environment variables
                 env_vars["MASTER_ADDR"] = addr
                 env_vars["MASTER_PORT"] = port
 
+                # Get NCCL-specific environment variables
+                nccl_vars = await get_nccl_env_vars()
+                env_vars.update(nccl_vars)
+
                 # Set the PTD world size
                 world_size = num_procs * (num_hosts or 1)
                 env_vars["WORLD_SIZE"] = str(world_size)