test(distributed): 恢复分布式切分的测试

Signed-off-by: YdrMaster <[email protected]>
InfiniTensor · Apr 30, 2024 · 0171f99 · 0171f99
1 parent 165b3cd
commit 0171f99
Show file tree

Hide file tree

Showing 3 changed files with 76 additions and 54 deletions.
diff --git a/nvidia/distributed/src/distribute.rs b/nvidia/distributed/src/distribute.rs
@@ -204,26 +204,25 @@ impl DistributeScheme {
     }
 }
 
-// #[test]
-// fn test() {
-//     use super::Memory;
-//     use std::time::Instant;
-
-//     let Some(model_dir) = common::test_model::find() else {
-//         return;
-//     };
-//     println!("model_dir: {}", model_dir.display());
-
-//     let time = Instant::now();
-//     let model = Memory::load_safetensors(model_dir).unwrap();
-//     println!("mmap {:?}", time.elapsed());
-
-//     let distributer = Distributer::new(&model, 4, 512);
-//     let time = Instant::now();
-//     for layer in 0..model.num_hidden_layers() {
-//         for i in 0..4 {
-//             let _ = distributer.distribute(layer, i);
-//         }
-//     }
-//     println!("distribute {:?}", time.elapsed());
-// }
+#[test]
+fn test() {
+    use std::time::Instant;
+
+    let Some(model_dir) = common::test_model::find() else {
+        return;
+    };
+    println!("model_dir: {}", model_dir.display());
+
+    let time = Instant::now();
+    let model = llama::Storage::load_safetensors(model_dir).unwrap();
+    println!("mmap {:?}", time.elapsed());
+
+    let distributer = Distributer::new(&model, 4, 512);
+    let time = Instant::now();
+    for layer in 0..model.config.nlayers as usize {
+        for i in 0..4 {
+            let _ = distributer.distribute(layer, i);
+        }
+    }
+    println!("distribute {:?}", time.elapsed());
+}
diff --git a/nvidia/distributed/src/parameters.rs b/nvidia/distributed/src/parameters.rs
@@ -133,32 +133,32 @@ impl Layer<'_> {
     }
 }
 
-// #[test]
-// fn test_load() {
-//     use common_nv::cuda::{self, Device};
-//     use log::LevelFilter::Trace;
-//     use simple_logger::SimpleLogger;
-
-//     let Some(model_dir) = common_nv::test_model::find() else {
-//         return;
-//     };
-//     println!("model_dir: {}", model_dir.display());
-
-//     const N: usize = 1;
-
-//     cuda::init();
-//     if Device::count() < N {
-//         return;
-//     }
-
-//     SimpleLogger::new().with_level(Trace).init().unwrap();
-
-//     let time = Instant::now();
-//     let model = Memory::load_safetensors(model_dir).unwrap();
-//     info!("mmap {:?}", time.elapsed());
-
-//     let contexts = (0..N as _)
-//         .map(|i| Device::new(i).retain_primary())
-//         .collect::<Vec<_>>();
-//     unsafe { ParameterMatrix::load(&model, &contexts).kill(&contexts) };
-// }
+#[test]
+fn test_load() {
+    use common_nv::cuda::{self, Device};
+    use log::LevelFilter::Trace;
+    use simple_logger::SimpleLogger;
+
+    let Some(model_dir) = common::test_model::find() else {
+        return;
+    };
+    println!("model_dir: {}", model_dir.display());
+
+    const N: usize = 1;
+
+    cuda::init();
+    if Device::count() < N {
+        return;
+    }
+
+    SimpleLogger::new().with_level(Trace).init().unwrap();
+
+    let time = Instant::now();
+    let model = llama::Storage::load_safetensors(model_dir).unwrap();
+    info!("mmap {:?}", time.elapsed());
+
+    let contexts = (0..N as _)
+        .map(|i| Device::new(i).retain_primary())
+        .collect::<Vec<_>>();
+    unsafe { ParameterMatrix::load(&model, &contexts).kill(&contexts) };
+}
diff --git a/nvidia/transformer/src/lib.rs b/nvidia/transformer/src/lib.rs
@@ -32,10 +32,12 @@ pub struct Transformer {
     transfer: StreamSpore,
     compute: StreamSpore,
     kernels: NvidiaKernels,
+
     embed_tokens: Tensor<HostMemSpore>,
     layers: Vec<LayerStorage<HostMemSpore>>,
     lm_layernorm: Tensor<DevMemSpore>,
     lm_head: Tensor<DevMemSpore>,
+
     pool: Mutex<VecDeque<(LayerStorage<DevMemSpore>, EventSpore)>>,
 }
 
@@ -289,8 +291,29 @@ impl Drop for Transformer {
     #[inline]
     fn drop(&mut self) {
         self.context.apply(|ctx| unsafe {
-            self.transfer.kill(ctx);
-            self.compute.kill(ctx);
+            ctx.kill(&mut self.transfer);
+            ctx.kill(&mut self.compute);
+            ctx.kill(self.embed_tokens.physical_mut());
+            ctx.kill(self.lm_layernorm.physical_mut());
+            ctx.kill(self.lm_head.physical_mut());
+            for layer in self.layers.iter_mut() {
+                ctx.kill(layer.att_layernorm.physical_mut());
+                ctx.kill(layer.att_qkv.physical_mut());
+                ctx.kill(layer.att_o.physical_mut());
+                ctx.kill(layer.mlp_layernorm.physical_mut());
+                ctx.kill(layer.mlp_gate_up.physical_mut());
+                ctx.kill(layer.mlp_down.physical_mut());
+            }
+            let mut pool = self.pool.lock().unwrap();
+            while let Some((mut layer, mut event)) = pool.pop_front() {
+                ctx.kill(layer.att_layernorm.physical_mut());
+                ctx.kill(layer.att_qkv.physical_mut());
+                ctx.kill(layer.att_o.physical_mut());
+                ctx.kill(layer.mlp_layernorm.physical_mut());
+                ctx.kill(layer.mlp_gate_up.physical_mut());
+                ctx.kill(layer.mlp_down.physical_mut());
+                ctx.kill(&mut event);
+            }
             self.kernels.kill(ctx);
         });
     }