perf(nv): 加速大模型加载

Signed-off-by: YdrMaster <ydrml@hotmail.com>

Author: YdrMaster

Cargo.toml

@@ -31,7 +31,7 @@ itertools = "0.13"
 env_logger = "0.11"
 build-script-cfg = "0.0"
 
-operators = { git = "https://github.com/YdrMaster/operators-rs", rev = "44ad48", default-features = false }
+operators = { git = "https://github.com/YdrMaster/operators-rs", rev = "89ffbf1", default-features = false }
 search-cl-tools = { git = "https://github.com/InfiniTensor/clrt", rev = "9b6289d" }
-search-infini-tools = { git = "https://github.com/InfiniTensor/infini-rt", rev = "0e57976" }
-search-cuda-tools = { git = "https://github.com/YdrMaster/cuda-driver", rev = "5b9dbd9" }
+search-infini-tools = { git = "https://github.com/InfiniTensor/infini-rt", rev = "f40bcb5" }
+search-cuda-tools = { git = "https://github.com/YdrMaster/cuda-driver", rev = "041badf" }

common/src/lib.rs

@@ -1,4 +1,4 @@
-use std::ops::Deref;
+use std::{borrow::Borrow, collections::HashMap, hash::Hash, ops::Deref};
 
 pub enum Contiguous<'a, T> {
     Borrowed(&'a [u8]),
@@ -25,3 +25,30 @@ pub fn borrow<T>(t: &[u8]) -> Contiguous<'_, T> {
 pub fn own<'a, T>(t: T) -> Contiguous<'a, T> {
     Contiguous::Owned(t)
 }
+
+#[derive(Clone, Default, Debug)]
+#[repr(transparent)]
+pub struct Slab<K, V>(HashMap<K, Vec<V>>);
+
+impl<K, V> Slab<K, V> {
+    #[inline]
+    pub fn new() -> Self {
+        Self(HashMap::new())
+    }
+}
+
+impl<K: Eq + Hash, V> Slab<K, V> {
+    #[inline]
+    pub fn take<Q>(&mut self, key: &Q) -> Option<V>
+    where
+        K: Borrow<Q>,
+        Q: ?Sized + Hash + Eq,
+    {
+        self.0.get_mut(key).and_then(|pool| pool.pop())
+    }
+
+    #[inline]
+    pub fn put(&mut self, key: K, value: V) {
+        self.0.entry(key).or_default().push(value);
+    }
+}

models/llama/common-cpu/Cargo.toml

@@ -8,6 +8,7 @@ authors = ["YdrMaster <ydrml@hotmail.com>"]
 
 [dependencies]
 llama.path = "../common"
+common.workspace = true
 operators = { workspace = true, features = ["common-cpu"] }
 
 [dev-dependencies]

models/llama/common-cpu/src/lib.rs

@@ -1,6 +1,7 @@
+use common::Contiguous;
 use llama::{
     ext::ggml_quants::{self, digit_layout::DigitLayout, f16, DataBlock, QuantExt},
-    BlkWeight, Contiguous, LlamaBlkStorage, LlamaStorage, Tensor,
+    BlkWeight, LlamaBlkStorage, LlamaStorage, Tensor,
     TensorUsage::Computation,
     WeightLoader,
 };

models/llama/common/src/lib.rs

@@ -6,7 +6,6 @@ use gguf::ggml_quants::digit_layout::DigitLayout;
 use std::ops::{Range, RangeBounds};
 
 pub use args::{Args as LlamaArgs, Request as LlamaRequest};
-pub use common::Contiguous;
 pub use compute::{BlkWeight, LlamaWorker, Operators, WeightLoader};
 pub use storage::{BlkStorage as LlamaBlkStorage, Storage as LlamaStorage};
 pub use tensor::{RandomSample, Tensor};

models/llama/infini/Cargo.toml

@@ -8,6 +8,7 @@ authors = ["YdrMaster <ydrml@hotmail.com>"]
 
 [dependencies]
 llama.path = "../common"
+common.workspace = true
 operators = { workspace = true, features = ["infini"] }
 
 [build-dependencies]

models/llama/infini/src/lib.rs

@@ -1,6 +1,7 @@
 #![cfg(detected)]
 
-use llama::{BlkWeight, Contiguous, LlamaBlkStorage, LlamaStorage, Tensor, WeightLoader};
+use common::Contiguous;
+use llama::{BlkWeight, LlamaBlkStorage, LlamaStorage, Tensor, WeightLoader};
 use operators::{
     all_reduce::{infini::Operator as InfiniAllReduce, AllReduce},
     infini::{Device, InfiniNode},

models/llama/nvidia-gpu/Cargo.toml

@@ -8,6 +8,8 @@ authors = ["YdrMaster <ydrml@hotmail.com>"]
 
 [dependencies]
 llama.path = "../common"
+common.workspace = true
+log.workspace = true
 operators = { workspace = true, features = ["nvidia-gpu"] }
 
 [build-dependencies]
@@ -17,5 +19,4 @@ search-cuda-tools.workspace = true
 [dev-dependencies]
 test-utils = { workspace = true, features = ["llama"] }
 gguf.workspace = true
-log.workspace = true
 regex.workspace = true

models/llama/nvidia-gpu/src/lib.rs

@@ -1,9 +1,11 @@
 #![cfg(driver_detected)]
 
-use llama::{BlkWeight, Contiguous, LlamaBlkStorage, LlamaStorage, Tensor, WeightLoader};
+use common::{Contiguous, Slab};
+use llama::{BlkWeight, LlamaBlkStorage, LlamaStorage, Tensor, WeightLoader};
+use log::trace;
 use operators::{
     all_reduce::{AllReduce, NonAllReduce},
-    cuda::{memcpy_d2h, CurrentCtx, DevByte, DevMem, Event, HostMem, Stream},
+    cuda::{memcpy_d2h, AsRaw, CurrentCtx, DevByte, DevMem, Event, HostMem, Stream},
     nvidia_gpu::Gpu,
     random_sample::nvidia_gpu::Operator as RandomSampleGpu,
     rearrange::nvidia_gpu::Operator as Rearrange,
@@ -15,6 +17,7 @@ use std::{
     mem::replace,
     ops::{Deref, RangeBounds},
     rc::Rc,
+    time::Instant,
 };
 
 pub struct Operators<N = Gpu, R = NonAllReduce<Gpu, Rearrange>>(PhantomData<(N, R)>);
@@ -157,11 +160,14 @@ impl<'blk> Weights<'blk> {
     ) -> Self {
         assert!(pool_size > 0);
         let stream = Rc::new(ctx.stream());
+        let igpu = unsafe { ctx.dev().as_raw() };
+        let mut slab = Slab::new();
         let blks = if pool_size < model.meta.nblk {
             let mut blks_host = model.blocks[0]
                 .as_ref()
                 .map(|_| Vec::with_capacity(model.meta.nblk));
-            for blk in model.blocks.iter() {
+            for (iblk, blk) in model.blocks.iter().enumerate() {
+                let time = Instant::now();
                 let blk = blk
                     .distribute(&model.meta, range.clone(), count, |len| {
                         ctx.malloc_host::<u8>(len)
@@ -188,6 +194,7 @@ impl<'blk> Weights<'blk> {
                     ffn_gate_up
                     ffn_down
                 }
+                trace!("blk{iblk} loaded to gpu{igpu} in {:?}", time.elapsed())
             }
             blks_host.map(|vec| {
                 let roll_cache = vec
@@ -206,18 +213,26 @@ impl<'blk> Weights<'blk> {
             let mut blks_dev = model.blocks[0]
                 .as_ref()
                 .map(|_| Vec::with_capacity(model.meta.nblk));
-            for blk in &model.blocks {
-                let blk = blk.distribute(&model.meta, range.clone(), count, |len| {
-                    ctx.malloc_host::<u8>(len)
+            for (iblk, blk) in model.blocks.iter().enumerate() {
+                let blk = blk.distribute(&model.meta, range.clone(), count, |size| {
+                    slab.take(&size)
+                        .unwrap_or_else(|| ctx.malloc_host::<u8>(size))
                 });
                 let loader = loader
                     .get_or_insert_with(|| blk.as_ref().map(|s| H2DLoader::new(s.len(), &stream)));
 
                 macro_rules! load {
                     ($( $ident:ident )+ ) => {
-                        $({ blks_dev.$ident.push(loader.$ident.load(blk.$ident, &stream)); })+
+                        $(
+                            let (dev, host) = loader.$ident.load(blk.$ident, &stream);
+                            if let Some(host) = host {
+                                slab.put(host.len(), host)
+                            }
+                            blks_dev.$ident.push(dev);
+                        )+
                     };
                 }
+                let time = Instant::now();
                 load! {
                     attn_norm
                     attn_qkv
@@ -226,6 +241,7 @@ impl<'blk> Weights<'blk> {
                     ffn_gate_up
                     ffn_down
                 }
+                trace!("blk{iblk} loaded to gpu{igpu} in {:?}", time.elapsed())
             }
             blks_dev.map(|vec| Cache::Static(vec.into_boxed_slice()))
         };
@@ -253,15 +269,25 @@ impl<'ctx> H2DLoader<'ctx> {
         }
     }
 
-    fn load(&mut self, host: Contiguous<HostMem<'ctx>>, stream: &Stream<'ctx>) -> DevMem<'ctx> {
+    fn load(
+        &mut self,
+        host: Contiguous<HostMem<'ctx>>,
+        stream: &Stream<'ctx>,
+    ) -> (DevMem<'ctx>, Option<HostMem<'ctx>>) {
         self.event.synchronize();
-        match host {
-            Contiguous::Borrowed(host) => self.host.copy_from_slice(host),
-            Contiguous::Owned(host) => self.host = host,
+        let cache = match host {
+            Contiguous::Borrowed(host) => {
+                self.host.copy_from_slice(host);
+                None
+            }
+            Contiguous::Owned(host) => Some(replace(&mut self.host, host)),
         };
         stream.memcpy_h2d(&mut self.dev, &self.host);
         self.event = stream.record();
-        replace(&mut self.dev, stream.malloc::<u8>(self.host.len()))
+        (
+            replace(&mut self.dev, stream.malloc::<u8>(self.host.len())),
+            cache,
+        )
     }
 }
 

models/llama/nvidia-gpu/src/nccl_parallel.rs

@@ -14,7 +14,7 @@ use regex::Regex;
 use std::{
     iter::zip,
     slice::{from_raw_parts, from_raw_parts_mut},
-    thread,
+    thread, u64,
 };
 use test_utils::{test_infer_paralle, Inference, Task, TokenizerAndPrompt, WorkerSeed};
 
@@ -88,6 +88,11 @@ fn test_infer() {
                     let WorkerSeed { node, tasks } = seed;
                     node.processor().apply(|ctx| {
                         let stream = ctx.stream();
+                        let (free, _) = ctx.mem_info();
+
+                        ctx.dev().set_mempool_threshold(u64::MAX);
+                        let _ = stream.malloc::<u8>((free.0 >> 30).saturating_sub(1) << 30);
+
                         info!("worker[{id}] loading weights...");
                         let weights = Weights::new(model, range, count, usize::MAX, ctx);
                         let mut worker = Worker::new(id, &node, meta.clone(), weights, id == 0);