support custom NCCL_HOME

build_utils/src/lib.rs

@@ -140,6 +140,10 @@ pub fn find_cuda_home() -> Option<String> {
     cuda_home
 }
 
+pub fn find_nccl_home() -> Option<String> {
+    get_env_var_with_rerun("NCCL_HOME").ok().or_else(|| get_env_var_with_rerun("NCCL_ROOT").ok())
+}
+
 /// Discover CUDA configuration including home, include dirs, and lib dirs
 pub fn discover_cuda_config() -> Result<CudaConfig, BuildError> {
     let cuda_home = find_cuda_home().ok_or(BuildError::CudaNotFound)?;
@@ -288,6 +292,14 @@ mod tests {
         assert_eq!(result, Some("/test/cuda".to_string()));
     }
 
+    #[test]
+    fn test_find_nccl_home_env_var() {
+        env::set_var("NCCL_HOME", "/test/nccl");
+        let result = find_nccl_home();
+        env::remove_var("NCCL_HOME");
+        assert_eq!(result, Some("/test/nccl".to_string()));
+    }
+
     #[test]
     fn test_python_scripts_constants() {
         assert!(PYTHON_PRINT_DIRS.contains("sysconfig"));

nccl-sys/build.rs

@@ -24,6 +24,10 @@ fn main() {
             "-I{}/include",
             build_utils::find_cuda_home().unwrap()
         ))
+        .clang_arg(format!(
+            "-I{}/include",
+            build_utils::find_nccl_home().unwrap()
+        ))
         .parse_callbacks(Box::new(bindgen::CargoCallbacks::new()))
         // Communicator creation and management
         .allowlist_function("ncclGetLastError")
@@ -113,6 +117,7 @@ fn main() {
         .expect("Couldn't write bindings!");
 
     println!("cargo::rustc-link-lib=nccl");
+    println!("cargo::rustc-link-search=native={}/lib", build_utils::find_nccl_home().unwrap());
     println!("cargo::rustc-cfg=cargo");
     println!("cargo::rustc-check-cfg=cfg(cargo)");
 }

torch-sys-cuda/build.rs

@@ -137,6 +137,11 @@ fn main() {
     // Add CUDA toolkit includes
     libtorch_include_dirs.push(format!("{}/include", cuda_home.display()).into());
 
+    // Add NCCL includes
+    if let Some(nccl_home) = build_utils::find_nccl_home() {
+        libtorch_include_dirs.push(format!("{}/include", nccl_home).into());
+    }
+
     // Prefix includes with `monarch` to maintain consistency with fbcode
     // folder structure
     CFG.include_prefix = "monarch/torch-sys-cuda";