style(llama): clean up parallel inference

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

Author: YdrMaster

Cargo.toml

@@ -21,7 +21,7 @@ common.path = "common"
 gguf.path = "gguf"
 tensor.path = "tensor"
 causal-lm.path = "causal-lm"
-test-utils.path = "test-utils"
+test-utils = { path = "test-utils", default-features = false }
 
 ggus = "0.3"
 itertools = "0.13"

models/llama/common-cpu/Cargo.toml

@@ -11,6 +11,6 @@ llama.path = "../common"
 operators = { workspace = true, features = ["common-cpu"] }
 
 [dev-dependencies]
-test-utils.workspace = true
+test-utils = { workspace = true, features = ["llama"] }
 gguf.workspace = true
 regex.workspace = true

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

@@ -8,14 +8,8 @@ use operators::{
     Blob,
 };
 use regex::Regex;
-use std::{
-    iter::zip,
-    ptr::copy_nonoverlapping,
-    slice::from_raw_parts_mut,
-    sync::mpsc::{Receiver, Sender},
-    thread,
-};
-use test_utils::{Inference, TokenizerAndPrompt};
+use std::{iter::zip, ptr::copy_nonoverlapping, slice::from_raw_parts_mut, thread};
+use test_utils::{test_infer_paralle, Inference, Task, TokenizerAndPrompt, WorkerSeed};
 
 type Worker<'w> = LlamaWorker<Operators<InprocNode<usize>, AllReduce>, Weights<'w>>;
 
@@ -48,17 +42,19 @@ fn test_infer() {
     let sample_args = SampleArgs::new(temperature, top_p, top_k).expect("invalid sample args");
     println!("{sample_args:?}");
 
-    let lens = match devices {
-        Some(devices) => Regex::new(r"\d+")
-            .unwrap()
-            .find_iter(&devices)
-            .map(|c| c.as_str().parse::<usize>().unwrap())
-            .collect::<Vec<_>>(),
-        None => vec![1],
-    };
-    println!("distribution: {lens:?}");
+    let lens = devices
+        .map(|devices| {
+            Regex::new(r"\d+")
+                .unwrap()
+                .find_iter(&devices)
+                .map(|c| c.as_str().parse().unwrap())
+                .collect()
+        })
+        .unwrap_or_else(|| vec![1]);
     let count = lens.iter().sum();
-    let (seeds, senders) = WorkerSeed::new(lens.len());
+    println!("distribution: {lens:?}");
+
+    let (seeds, senders) = WorkerSeed::new(InprocNode::new(lens.len()));
     thread::scope(|s| {
         let _workers = zip(lens, seeds)
             .enumerate()
@@ -70,7 +66,6 @@ fn test_infer() {
                 meta.distribute(range.clone(), count);
 
                 let model = &model;
-
                 Some(s.spawn(move || {
                     let WorkerSeed { node, tasks } = seed;
                     let weights = Weights::new(model, range, count);
@@ -141,63 +136,7 @@ fn test_infer() {
             })
             .collect::<Vec<_>>();
 
-        let (next, next_recv) = std::sync::mpsc::channel();
-        test_utils::test_infer(eos, tokenizer, &prompt, max_steps, |input, pos| {
-            let mut embd = model.meta.embd(input.len()).map(Blob::new);
-
-            let d = embd.get().len() / input.len();
-            for (i, &tok) in input.iter().enumerate() {
-                embd.get_mut()[i * d..][..d]
-                    .copy_from_slice(&model.token_embd[tok as usize * d..][..d]);
-            }
-            let embd = embd.take();
-
-            for sender in &senders {
-                sender
-                    .send(Task {
-                        nt: input.len(),
-                        pos,
-                        embd: embd.as_ptr(),
-                        next: next.clone(),
-                    })
-                    .unwrap();
-            }
-            next_recv.recv().unwrap()
-        });
-
-        drop(senders)
+        let senders = senders.into_boxed_slice();
+        test_infer_paralle(&model, senders, eos, tokenizer, &prompt, max_steps)
     })
 }
-
-struct Task {
-    nt: usize,
-    pos: usize,
-    embd: *const u8,
-    next: Sender<u32>,
-}
-
-unsafe impl Send for Task {}
-
-struct WorkerSeed {
-    tasks: Receiver<Task>,
-    node: InprocNode<usize>,
-}
-
-impl WorkerSeed {
-    fn new(n: usize) -> (Vec<Self>, Vec<Sender<Task>>) {
-        let mut tasks = Vec::with_capacity(n);
-        let mut senders = Vec::with_capacity(n);
-        let nodes = InprocNode::new(n);
-        for _ in 0..n {
-            let (sender, receiver) = std::sync::mpsc::channel();
-            tasks.push(receiver);
-            senders.push(sender);
-        }
-        (
-            zip(nodes, tasks)
-                .map(|(node, tasks)| Self { node, tasks })
-                .collect(),
-            senders,
-        )
-    }
-}

models/llama/nvidia-gpu/Cargo.toml

@@ -15,6 +15,6 @@ build-script-cfg.workspace = true
 search-cuda-tools.workspace = true
 
 [dev-dependencies]
-test-utils.workspace = true
+test-utils = { workspace = true, features = ["llama"] }
 gguf.workspace = true
 regex.workspace = true

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

@@ -7,16 +7,15 @@ use operators::{
     nccl::CommunicatorGroup,
     nvidia_gpu::NcclNode,
     random_sample::{KVPair, SampleArgs},
-    Blob, TopoNode,
+    TopoNode,
 };
 use regex::Regex;
 use std::{
     iter::zip,
     slice::{from_raw_parts, from_raw_parts_mut},
-    sync::mpsc::{Receiver, Sender},
-    thread, usize,
+    thread,
 };
-use test_utils::{Inference, TokenizerAndPrompt};
+use test_utils::{test_infer_paralle, Inference, Task, TokenizerAndPrompt, WorkerSeed};
 
 type Worker<'w> = LlamaWorker<Operators<NcclNode, AllReduce>, Weights<'w>>;
 
@@ -49,21 +48,27 @@ fn test_infer() {
     let sample_args = SampleArgs::new(temperature, top_p, top_k).expect("invalid sample args");
     println!("{sample_args:?}");
 
-    let devices = match devices {
-        Some(devices) => Regex::new(r"\d+")
-            .unwrap()
-            .find_iter(&devices)
-            .map(|c| c.as_str().parse().unwrap())
-            .collect::<Vec<_>>(),
-        None => vec![0],
-    };
-    println!("distribution: {devices:?}");
-
+    let devices = devices
+        .map(|devices| {
+            Regex::new(r"\d+")
+                .unwrap()
+                .find_iter(&devices)
+                .map(|c| c.as_str().parse().unwrap())
+                .collect()
+        })
+        .unwrap_or_else(|| vec![1]);
     let lens = vec![1; devices.len()];
     let count = devices.len();
+    println!("distribution: {devices:?}");
 
     let (seeds, senders) = match cuda::init() {
-        Ok(()) => WorkerSeed::new(&devices),
+        Ok(()) => WorkerSeed::new(
+            CommunicatorGroup::new(&devices)
+                .into_vec()
+                .into_iter()
+                .map(|comm| NcclNode::new(comm, Default::default()))
+                .collect(),
+        ),
         Err(NoDevice) => return,
     };
     thread::scope(|s| {
@@ -77,7 +82,6 @@ fn test_infer() {
                 meta.distribute(range.clone(), count);
 
                 let model = &model;
-
                 Some(s.spawn(move || {
                     let WorkerSeed { node, tasks } = seed;
                     node.processor().apply(|ctx| {
@@ -163,68 +167,7 @@ fn test_infer() {
             })
             .collect::<Vec<_>>();
 
-        let (next, next_recv) = std::sync::mpsc::channel();
-        test_utils::test_infer(eos, tokenizer, &prompt, max_steps, |input, pos| {
-            let mut embd = model.meta.embd(input.len()).map(Blob::new);
-
-            let d = embd.get().len() / input.len();
-            for (i, &tok) in input.iter().enumerate() {
-                embd.get_mut()[i * d..][..d]
-                    .copy_from_slice(&model.token_embd[tok as usize * d..][..d]);
-            }
-            let embd = embd.take();
-
-            for sender in &senders {
-                sender
-                    .send(Task {
-                        nt: input.len(),
-                        pos,
-                        embd: embd.as_ptr(),
-                        next: next.clone(),
-                    })
-                    .unwrap();
-            }
-            next_recv.recv().unwrap()
-        });
-
-        drop(senders)
+        let senders = senders.into_boxed_slice();
+        test_infer_paralle(&model, senders, eos, tokenizer, &prompt, max_steps)
     })
 }
-
-struct Task {
-    nt: usize,
-    pos: usize,
-    embd: *const u8,
-    next: Sender<u32>,
-}
-
-unsafe impl Send for Task {}
-
-struct WorkerSeed {
-    tasks: Receiver<Task>,
-    node: NcclNode,
-}
-
-impl WorkerSeed {
-    fn new(devices: &[i32]) -> (Vec<Self>, Vec<Sender<Task>>) {
-        let nodes = CommunicatorGroup::new(devices)
-            .into_vec()
-            .into_iter()
-            .map(|comm| NcclNode::new(comm, Default::default()))
-            .collect::<Vec<_>>();
-        let n = nodes.len();
-        let mut tasks = Vec::with_capacity(n);
-        let mut senders = Vec::with_capacity(n);
-        for _ in 0..n {
-            let (sender, receiver) = std::sync::mpsc::channel();
-            tasks.push(receiver);
-            senders.push(sender);
-        }
-        (
-            zip(nodes, tasks)
-                .map(|(node, tasks)| Self { node, tasks })
-                .collect(),
-            senders,
-        )
-    }
-}

test-utils/Cargo.toml

@@ -4,7 +4,12 @@ version = "0.0.0"
 edition = "2021"
 authors = ["YdrMaster <ydrml@hotmail.com>"]
 
+[features]
+default = ["llama"]
+
 [dependencies]
+llama = { path = "../models/llama/common", optional = true }
 gguf.workspace = true
+tensor.workspace = true
 env_logger.workspace = true
 cli-table = "0.4.9"

test-utils/src/lib.rs

@@ -5,11 +5,14 @@ use gguf::{
 use std::{
     env::{var, var_os},
     fmt,
+    iter::zip,
     path::{Path, PathBuf},
     str::FromStr,
-    sync::Once,
+    sync::{
+        mpsc::{self, Sender},
+        Once,
+    },
     time::{Duration, Instant},
-    vec,
 };
 
 pub struct Inference {
@@ -176,3 +179,71 @@ pub fn test_infer(
         ]
     }
 }
+
+#[cfg(feature = "llama")]
+pub fn test_infer_paralle<'w>(
+    model: &llama::LlamaStorage<&'w [u8]>,
+    senders: Box<[mpsc::Sender<Task>]>,
+    eos: utok,
+    tokenizer: Tokenizer,
+    prompt: &str,
+    max_steps: usize,
+) {
+    use tensor::Blob;
+
+    let (next, next_recv) = mpsc::channel();
+    test_infer(eos, tokenizer, prompt, max_steps, |input, pos| {
+        let mut embd = model.meta.embd(input.len()).map(Blob::new).take();
+
+        let d = embd.len() / input.len();
+        for (i, &tok) in input.iter().enumerate() {
+            embd[i * d..][..d].copy_from_slice(&model.token_embd[tok as usize * d..][..d]);
+        }
+
+        for sender in &senders {
+            sender
+                .send(Task {
+                    nt: input.len(),
+                    pos,
+                    embd: embd.as_ptr(),
+                    next: next.clone(),
+                })
+                .unwrap()
+        }
+        next_recv.recv().unwrap()
+    });
+}
+
+pub struct Task {
+    pub nt: usize,
+    pub pos: usize,
+    pub embd: *const u8,
+    pub next: mpsc::Sender<utok>,
+}
+
+unsafe impl Send for Task {}
+
+pub struct WorkerSeed<N> {
+    pub tasks: mpsc::Receiver<Task>,
+    pub node: N,
+}
+
+impl<N> WorkerSeed<N> {
+    pub fn new(nodes: Vec<N>) -> (Vec<Self>, Vec<Sender<Task>>) {
+        let n = nodes.len();
+
+        let mut tasks = Vec::with_capacity(n);
+        let mut senders = Vec::with_capacity(n);
+        for _ in 0..n {
+            let (sender, receiver) = std::sync::mpsc::channel();
+            tasks.push(receiver);
+            senders.push(sender);
+        }
+        (
+            zip(nodes, tasks)
+                .map(|(node, tasks)| Self { node, tasks })
+                .collect(),
+            senders,
+        )
+    }
+}