build: update deps

Signed-off-by: YdrMaster <[email protected]>
InfiniTensor · Apr 3, 2024 · c14108b · c14108b
1 parent 245022d
commit c14108b
Show file tree

Hide file tree

Showing 3 changed files with 22 additions and 21 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/Cargo.toml b/Cargo.toml
@@ -21,6 +21,6 @@ serde = "1.0"
 log = "0.4"
 tokio = { version = "1.37", features = ["rt", "rt-multi-thread", "sync"] }
 
-cuda = { git = "https://github.com/YdrMaster/cuda-driver", rev = "659fc8" }
-cublas = { git = "https://github.com/YdrMaster/cuda-driver", rev = "659fc8" }
-search-cuda-tools = { git = "https://github.com/YdrMaster/cuda-driver", rev = "659fc8" }
+cuda = { git = "https://github.com/YdrMaster/cuda-driver", rev = "431d02" }
+cublas = { git = "https://github.com/YdrMaster/cuda-driver", rev = "431d02" }
+search-cuda-tools = { git = "https://github.com/YdrMaster/cuda-driver", rev = "431d02" }
diff --git a/nvidia/distributed/src/lib.rs b/nvidia/distributed/src/lib.rs
@@ -17,6 +17,7 @@ fn test_load() {
         .iter()
         .map(Device::retain_primary)
         .collect::<Vec<_>>();
+    let align = devices.iter().map(Device::alignment).max().unwrap();
 
     let time = Instant::now();
     let safetensors = Memory::load_safetensors_from_dir("../../../TinyLlama-1.1B-Chat-v1.0");
@@ -31,7 +32,7 @@ fn test_load() {
     let nlayers = model.num_hidden_layers();
     let mut matrix = Vec::with_capacity(contexts.len() * nlayers);
 
-    let distributer = Distributer::new(&model, contexts.len(), 512);
+    let distributer = Distributer::new(&model, contexts.len(), align);
     let time = Instant::now();
     for (i, context) in contexts.iter().enumerate() {
         context.apply(|ctx| {