fix(gpt2): 重构gpt2的cpu的单机推理

Author: onenewcode

Cargo.toml

@@ -35,7 +35,7 @@ itertools = "0.13"
 env_logger = "0.11"
 build-script-cfg = "0.0"
 
-operators = { git = "https://github.com/YdrMaster/operators-rs", rev = "7886d54", default-features = false }
+operators = { git = "https://github.com/YdrMaster/operators-rs", rev = "359b86a", default-features = false }
 
 search-cl-tools = { git = "https://github.com/InfiniTensor/clrt", rev = "f69b160" }
 search-infini-tools = { git = "https://github.com/InfiniTensor/infini-rt", rev = "e8362c3" }

models/gpt2/common-cpu/Cargo.toml

@@ -12,7 +12,7 @@ common.workspace = true
 operators = { workspace = true, features = ["common-cpu"] }
 
 [dev-dependencies]
-test-utils = { workspace = true, features = ["llama"] }
+test-utils = { workspace = true, features = ["gpt2"] }
 gguf.workspace = true
 regex.workspace = true
 

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

@@ -22,16 +22,6 @@ pub struct Weights<'w> {
     output_norm_b: &'w [u8],
     output: &'w [u8],
     pos_embd: &'w [u8],
-    // dt_embd: DigitLayout,
-    // dt_mat: DigitLayout,
-    // size_qkv_b: usize,
-    // size_qkv_w: usize,
-    // size_o_b: usize,
-    // size_o_w: usize,
-    // size_up_b: usize,
-    // size_up_w: usize,
-    // size_down_b: usize,
-    // size_down_w: usize,
 }
 
 macro_rules! op {

models/gpt2/common/src/storage.rs

@@ -291,4 +291,4 @@ impl Gpt2Meta {
         }
         own(ans.take())
     }
-}
+}

models/gpt2/cuda/Cargo.toml

@@ -18,6 +18,6 @@ search-cuda-tools.workspace = true
 search-corex-tools.workspace = true
 
 [dev-dependencies]
-test-utils.workspace = true
+test-utils = { workspace = true, features = ["gpt2"] }
 gguf.workspace = true
 regex.workspace = true

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

@@ -22,7 +22,6 @@ type Worker<'w> = LlamaWorker<Operators<InprocNode<usize>, AllReduce>, Weights<'
 
 #[test]
 fn test_infer() {
-    std::env::set_var("TEST_MODEL", r"F:\TinyLlama-1.1B-Chat-v1.0-F16.gguf");
     let Some(Inference {
         model,
         devices,

models/llama/opencl/src/lib.rs

@@ -1,6 +1,6 @@
 #![cfg(detected)]
 
-use common::{Distribution, WeightMemCalculator};
+use common::Distribution;
 use llama::{LlamaBlkStorage, LlamaBlkWeight, LlamaStorage, Tensor, WeightLoader};
 use operators::{
     all_reduce::{AllReduce, NonAllReduce},
@@ -10,12 +10,7 @@ use operators::{
     rearrange::opencl::Operator as Rearrange,
     Blob, ByteOf, QueueOf, TopoNode,
 };
-use std::{
-    iter::zip,
-    marker::PhantomData,
-    ops::{Deref, Range},
-    ptr::copy_nonoverlapping,
-};
+use std::{marker::PhantomData, ops::Deref, ptr::copy_nonoverlapping};
 
 pub struct Operators<N = ClDevice, R = NonAllReduce<ClDevice, Rearrange>>(PhantomData<(N, R)>);
 
@@ -65,10 +60,9 @@ where
 
 pub struct Weights {
     nexp: usize,
-    mem: SvmBlob,
-    blks: Box<[LlamaBlkStorage<Range<usize>>]>,
-    output_norm: Range<usize>,
-    output: Range<usize>,
+    blks: Box<[LlamaBlkStorage<SvmBlob>]>,
+    output_norm: SvmBlob,
+    output: SvmBlob,
 }
 
 impl Weights {
@@ -81,52 +75,40 @@ impl Weights {
             ..
         } = model;
 
-        let mut calculator = WeightMemCalculator::new(size_of::<usize>());
-        let meta_dist = meta.distribute(dist);
-        let blk_size = meta_dist.blk();
-        let off_blks = (0..meta_dist.nblk)
-            .map(|_| {
-                blk_size
-                    .clone()
+        let meta = meta.distribute(dist);
+        let queue = ctx.queue();
+        let blks = blocks
+            .iter()
+            .map(|blk| {
+                blk.clone()
                     .into_vec()
                     .into_iter()
-                    .map(|(which, size)| (which, calculator.push(size)))
+                    .map(|(which, data)| {
+                        let blob = meta.distribute_data(which, data, dist, Blob::new);
+                        let mut svm = ctx.malloc::<u8>(blob.len());
+                        let mut map = queue.map_mut(&mut svm, false);
+                        map.copy_from_slice(&blob);
+                        queue.unmap(map);
+                        (which, svm)
+                    })
                     .collect::<LlamaBlkStorage<_>>()
             })
             .collect::<Vec<_>>();
-        let off_output_norm = calculator.push(output_norm.len());
-        let off_output = calculator.push(output.len());
 
-        let mut mem = ctx.malloc::<u8>(calculator.size());
-        let queue = ctx.queue();
-
-        for (blk, off) in zip(blocks, off_blks.clone()) {
-            let blk = blk.clone().into_vec();
-            let off = off.into_vec();
-            for ((which, data), (which_, off)) in zip(blk, off) {
-                assert_eq!(which, which_);
-                if off.is_empty() {
-                    continue;
-                }
-                let data = meta.distribute_data(which, data, dist, Blob::new);
-                let mut map = queue.map_mut(&mut mem[off], false);
-                map.copy_from_slice(&data);
-                queue.unmap(map)
-            }
-        }
-        let mut map = queue.map_mut(&mut mem[off_output_norm.clone()], false);
-        map.copy_from_slice(output_norm);
-        queue.unmap(map);
-        let mut map = queue.map_mut(&mut mem[off_output.clone()], false);
-        map.copy_from_slice(output);
-        queue.unmap(map);
+        let mut output_norm_svm = ctx.malloc::<u8>(output_norm.len());
+        let mut output_svm = ctx.malloc::<u8>(output.len());
+        let mut output_norm_map = queue.map_mut(&mut output_norm_svm, false);
+        let mut output_map = queue.map_mut(&mut output_svm, false);
+        output_norm_map.copy_from_slice(output_norm);
+        output_map.copy_from_slice(output);
+        queue.unmap(output_norm_map);
+        queue.unmap(output_map);
 
         Self {
             nexp: meta.nexp,
-            mem,
-            blks: off_blks.into_boxed_slice(),
-            output_norm: off_output_norm,
-            output: off_output,
+            blks: blks.into_boxed_slice(),
+            output_norm: output_norm_svm,
+            output: output_svm,
         }
     }
 }
@@ -158,7 +140,7 @@ impl WeightLoader for Weights {
 
         use LlamaBlkWeight as W;
         #[rustfmt::skip]
-        let range = match which {
+        let ans = match which {
             W::AttnNorm    => attn_norm    ,
             W::AttnQKV     => attn_qkv     ,
             W::AttnQKVBias => attn_qkv_bias,
@@ -168,7 +150,7 @@ impl WeightLoader for Weights {
             W::FfnGateUp   => ffn_gate_up  ,
             W::FfnDown     => ffn_down     ,
         };
-        &self.mem[range.clone()]
+        ans
     }
 
     fn load_moe<'a>(
@@ -184,26 +166,25 @@ impl WeightLoader for Weights {
             ..
         } = &self.blks[iblk];
 
-        let range = match which {
+        let w = match which {
             LlamaBlkWeight::FfnGateUp => ffn_gate_up,
             LlamaBlkWeight::FfnDown => ffn_down,
             _ => unreachable!(),
         };
-        let w = &self.mem[range.clone()];
         let one = w.len() / self.nexp;
         &w[iexp * one..][..one]
     }
 
     #[inline]
     fn output_norm(&self, _queue: &QueueOf<Self::Hardware>) -> Self::Weight<'_> {
-        &self.mem[self.output_norm.clone()]
+        &self.output_norm
     }
 
     #[inline]
     fn output(&self, _queue: &QueueOf<Self::Hardware>) -> Self::Weight<'_> {
-        &self.mem[self.output.clone()]
+        &self.output
     }
 }
 
 #[cfg(test)]
-mod infer;
+mod infer;

test-utils/src/lib.rs

@@ -5,17 +5,17 @@ use gguf::{
 use std::{
     env::{var, var_os},
     fmt,
+    iter::zip,
     path::{Path, PathBuf},
     str::FromStr,
-    sync::{Once,mpsc},
-    iter::zip,
+    sync::{mpsc, Once},
     time::{Duration, Instant},
 };
 #[cfg(feature = "llama")]
-mod llama{
+mod llama {
+    use crate::InferStorage;
     use llama::LlamaStorage;
     use tensor::Tensor;
-    use crate::InferStorage;
 
     impl InferStorage for &LlamaStorage<&[u8]> {
         fn embd(&self, nt: usize) -> Tensor<usize> {
@@ -27,10 +27,10 @@ mod llama{
     }
 }
 #[cfg(feature = "llama")]
-mod gpt2{
+mod gpt2 {
+    use crate::InferStorage;
     use gpt2::GPT2Storage;
     use tensor::Tensor;
-    use crate::InferStorage;
 
     impl InferStorage for &GPT2Storage<&[u8]> {
         fn embd(&self, nt: usize) -> Tensor<usize> {
@@ -42,7 +42,6 @@ mod gpt2{
     }
 }
 
-
 pub trait InferStorage {
     fn embd(&self, nt: usize) -> Tensor<usize>;
     fn token_embd(&self) -> &[u8];