refactor(transformer): 简化模型加载

Signed-off-by: YdrMaster <[email protected]>
InfiniTensor · Apr 18, 2024 · 75c4f29 · 75c4f29
1 parent 076b05c
commit 75c4f29
Show file tree

Hide file tree

Showing 11 changed files with 23 additions and 39 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/nvidia/distributed/src/lib.rs b/nvidia/distributed/src/lib.rs
@@ -422,7 +422,7 @@ impl transformer::Transformer for Transformer {
 impl Transformer {
     pub fn new(model_dir: impl AsRef<Path>, dev: &[Device]) -> Self {
         let time = Instant::now();
-        let host = Memory::load_safetensors_from_dir(model_dir).unwrap();
+        let host = Memory::load_safetensors(model_dir).unwrap();
         info!("load host: {:?}", time.elapsed());
 
         let block_size = dev.iter().map(|dev| dev.max_block_dims().0).min().unwrap();

diff --git a/nvidia/distributed/src/parameters.rs b/nvidia/distributed/src/parameters.rs
@@ -165,7 +165,7 @@ fn test_load() {
     SimpleLogger::new().with_level(Trace).init().unwrap();
 
     let time = Instant::now();
-    let safetensors = Memory::load_safetensors_from_dir(model_dir);
+    let safetensors = Memory::load_safetensors(model_dir);
     info!("mmap {:?}", time.elapsed());
 
     let model = match safetensors {

diff --git a/nvidia/transformer/Cargo.toml b/nvidia/transformer/Cargo.toml
@@ -9,7 +9,6 @@ authors = ["YdrMaster <[email protected]>"]
 [dependencies]
 transformer = { path = "../../transformer" }
 common-nv = { path = "../common" }
-log.workspace = true
 half.workspace = true
 
 [dev-dependencies]

diff --git a/nvidia/transformer/src/lib.rs b/nvidia/transformer/src/lib.rs
@@ -2,25 +2,19 @@
 
 mod parameters;
 
-#[macro_use]
-extern crate log;
-
 pub use common_nv::cuda;
 
 use ::half::f16;
 use common_nv::{
     cuda::{memcpy_d2h, DevMem, DevMemSpore},
-    slice, udim, utok, DataType, LocalSplitable, NvidiaKernels, NvidiaKernelsPtx, SafeTensors,
-    Tensor,
+    slice, udim, utok, DataType, LocalSplitable, NvidiaKernels, NvidiaKernelsPtx, Tensor,
 };
 use cuda::{Context, ContextResource, ContextSpore, Device, Stream, StreamSpore};
 use parameters::{LayerParameter, LayersParameters, ModelParameters};
 use std::{
-    fs::File,
     path::Path,
     slice::from_raw_parts,
     sync::{Arc, Mutex},
-    time::Instant,
 };
 use transformer::{pos, Kernels, LayerBuffer, LayerCache, Llama2, Memory, Request, SampleArgs};
 
@@ -140,15 +134,9 @@ type Splitable<'ctx> = LocalSplitable<DevMem<'ctx>>;
 
 impl Transformer {
     pub fn new(model_dir: impl AsRef<Path>, preload_layers: usize, dev: Device) -> Self {
-        let time = Instant::now();
-        let config = File::open(model_dir.as_ref().join("config.json")).unwrap();
-        let model = SafeTensors::load_from_dir(model_dir).unwrap();
-        info!("open file {:?}", time.elapsed());
-
         let context = Arc::new(dev.retain_primary());
-        let host = Memory::load_safetensors(
-            config,
-            model,
+        let host = Memory::load_safetensors_realloc(
+            model_dir,
             Some(|l| context.apply(|ctx| ctx.malloc_host::<u8>(l).sporulate())),
         )
         .unwrap();

diff --git a/service/src/cpu.rs b/service/src/cpu.rs
@@ -4,7 +4,7 @@ use transformer_cpu::Transformer;
 
 pub fn transformer(model_dir: impl AsRef<Path>) -> Transformer {
     let time = Instant::now();
-    let model = Memory::load_safetensors_from_dir(model_dir).unwrap();
+    let model = Memory::load_safetensors(model_dir).unwrap();
     info!("load model ... {:?}", time.elapsed());
 
     let time = Instant::now();

diff --git a/transformer-cpu/src/lib.rs b/transformer-cpu/src/lib.rs
@@ -337,7 +337,7 @@ fn test_build() {
     println!("model_dir: {}", model_dir.display());
 
     let t0 = Instant::now();
-    let safetensors = Memory::load_safetensors_from_dir(model_dir);
+    let safetensors = Memory::load_safetensors(model_dir);
     let t1 = Instant::now();
     println!("mmap {:?}", t1 - t0);
 

diff --git a/transformer/src/parameters/distribute.rs b/transformer/src/parameters/distribute.rs
@@ -293,7 +293,7 @@ fn test() {
     println!("model_dir: {}", model_dir.display());
 
     let time = Instant::now();
-    let safetensors = Memory::load_safetensors_from_dir(model_dir);
+    let safetensors = Memory::load_safetensors(model_dir);
     println!("mmap {:?}", time.elapsed());
 
     let model = match safetensors {

diff --git a/transformer/src/parameters/memory.rs b/transformer/src/parameters/memory.rs
@@ -184,7 +184,7 @@ fn test_load() {
     println!("model_dir: {}", model_dir.display());
 
     let t0 = Instant::now();
-    let safetensors = Memory::load_safetensors_from_dir(model_dir);
+    let safetensors = Memory::load_safetensors(model_dir);
     let t1 = Instant::now();
     println!("mmap {:?}", t1 - t0);
 

diff --git a/transformer/src/parameters/safe_tensors.rs b/transformer/src/parameters/safe_tensors.rs
@@ -1,29 +1,27 @@
 use super::{memory::Layer, storage::HostMem, ConfigJson, Memory, Storage};
 use common::{
-    safe_tensors::{Dtype, SafeTensors, SafeTensorsError},
+    safe_tensors::{
+        Dtype, SafeTensors,
+        SafeTensorsError::{self, Io, Json},
+    },
     Blob,
 };
-use std::{fs::File, io::Read, ops::DerefMut, path::Path, sync::Arc};
+use std::{fs::File, ops::DerefMut, path::Path, sync::Arc};
 use tensor::{udim, DataType, Shape, Tensor};
 
 impl Memory {
-    pub fn load_safetensors_from_dir(
-        model_dir: impl AsRef<Path>,
-    ) -> Result<Self, SafeTensorsError> {
-        let model_dir = model_dir.as_ref();
-        let config = File::open(model_dir.join("config.json")).map_err(SafeTensorsError::Io)?;
-        let model = SafeTensors::load_from_dir(model_dir)?;
-        Self::load_safetensors(config, model, Some(Blob::new)).map_err(SafeTensorsError::Json)
+    pub fn load_safetensors(model_dir: impl AsRef<Path>) -> Result<Self, SafeTensorsError> {
+        Self::load_safetensors_realloc(model_dir, Some(Blob::new))
     }
 
-    pub fn load_safetensors<T: HostMem + DerefMut<Target = [u8]>>(
-        config: impl Read,
-        model: SafeTensors,
+    pub fn load_safetensors_realloc<T: HostMem + DerefMut<Target = [u8]>>(
+        model_dir: impl AsRef<Path>,
         mut realloc: Option<impl FnMut(usize) -> T>,
-    ) -> Result<Self, serde_json::Error> {
-        let config: ConfigJson = serde_json::from_reader(config)?;
+    ) -> Result<Self, SafeTensorsError> {
+        let config = File::open(model_dir.as_ref().join("config.json")).map_err(Io)?;
+        let config: ConfigJson = serde_json::from_reader(&config).map_err(Json)?;
+        let model = SafeTensors::load_from_dir(model_dir)?.share();
 
-        let model = model.share();
         let tensor = |name: &str| {
             let shared = model
                 .share_tensor(name)

diff --git a/xtask/src/cast.rs b/xtask/src/cast.rs
@@ -26,7 +26,7 @@ impl CastArgs {
         let model_dir = PathBuf::from(self.model);
 
         let time = Instant::now();
-        let model = Memory::load_safetensors_from_dir(&model_dir).unwrap();
+        let model = Memory::load_safetensors(&model_dir).unwrap();
         println!("load model ... {:?}", time.elapsed());
 
         let target = self.target.map(PathBuf::from).unwrap_or_else(|| {