Fix handling of device when compiled for but disabled nccl (#1227)

Author: EricLBuehler

mistralrs-bench/src/main.rs

@@ -370,7 +370,7 @@ fn main() -> anyhow::Result<()> {
     #[cfg(feature = "metal")]
     let device = Device::new_metal(0)?;
     #[cfg(not(feature = "metal"))]
-    let device = if cfg!(feature = "nccl") {
+    let device = if mistralrs_core::distributed::use_nccl() {
         Device::Cpu
     } else {
         Device::cuda_if_available(0)?
@@ -433,7 +433,7 @@ fn main() -> anyhow::Result<()> {
         DeviceMapSetting::Auto(auto_device_map_params)
     };
 
-    let no_paged_attn = if device.is_cuda() || cfg!(feature = "nccl") {
+    let no_paged_attn = if device.is_cuda() || mistralrs_core::distributed::use_nccl() {
         args.no_paged_attn
     } else if device.is_metal() {
         !args.paged_attn

mistralrs-core/src/distributed.rs

@@ -4,6 +4,7 @@ use core::ffi::c_char;
 use interprocess::local_socket::traits::{Listener, Stream};
 use interprocess::local_socket::{GenericNamespaced, Name, ToNsName};
 use interprocess::local_socket::{ListenerOptions, Stream as LocalStream};
+pub use mistralrs_quant::distributed::use_nccl;
 use mistralrs_quant::{ShardedSafeTensors, ShardedVarBuilder};
 use serde::{Deserialize, Serialize};
 use serde_big_array::BigArray;

mistralrs-pyo3/src/lib.rs

@@ -44,7 +44,7 @@ static DEVICE: OnceLock<Result<Device>> = OnceLock::new();
 #[cfg(not(feature = "metal"))]
 fn get_device(seed: Option<u64>) -> &'static Result<Device> {
     DEVICE.get_or_init(|| {
-        let device = if cfg!(feature = "nccl") {
+        let device = if mistralrs_core::distributed::use_nccl() {
             Device::Cpu
         } else {
             Device::cuda_if_available(0)?
@@ -652,7 +652,7 @@ impl Runner {
             None => DeviceMapSetting::Auto(auto_map_params),
         };
 
-        let no_paged_attn = if device.is_cuda() || cfg!(feature = "nccl") {
+        let no_paged_attn = if device.is_cuda() || mistralrs_core::distributed::use_nccl() {
             no_paged_attn
         } else if device.is_metal() {
             !paged_attn

mistralrs-server/src/main.rs

@@ -317,7 +317,7 @@ async fn main() -> Result<()> {
     let device = if args.cpu {
         args.no_paged_attn = true;
         Device::Cpu
-    } else if cfg!(feature = "nccl") {
+    } else if mistralrs_core::distributed::use_nccl() {
         Device::Cpu
     } else {
         Device::cuda_if_available(0)?
@@ -379,7 +379,7 @@ async fn main() -> Result<()> {
         DeviceMapSetting::Auto(auto_device_map_params)
     };
 
-    let no_paged_attn = if device.is_cuda() || cfg!(feature = "nccl") {
+    let no_paged_attn = if device.is_cuda() || mistralrs_core::distributed::use_nccl() {
         args.no_paged_attn
     } else if device.is_metal() {
         !args.paged_attn