From 259e5db7bf0bfb60a0c7621691282fdb74790327 Mon Sep 17 00:00:00 2001 From: qinyiqun Date: Mon, 22 Jul 2024 11:42:49 +0800 Subject: [PATCH] update style and ban mutiple split --- ggus/src/metadata/mod.rs | 12 ++++++- xtask/src/gguf_file.rs | 6 ++-- xtask/src/merge.rs | 6 ++-- xtask/src/split.rs | 68 +++++++++++++++------------------------- 4 files changed, 42 insertions(+), 50 deletions(-) diff --git a/ggus/src/metadata/mod.rs b/ggus/src/metadata/mod.rs index b8502fc..76e3697 100644 --- a/ggus/src/metadata/mod.rs +++ b/ggus/src/metadata/mod.rs @@ -1,4 +1,4 @@ -//! See . +//! See . mod general; mod llm; @@ -141,6 +141,16 @@ impl<'a> GGufMetaKVPairs<'a> { .map(|(&key, &len)| GGufMetaKV { key, len }) } + pub fn remove(&mut self, _key: &str) -> bool { + match self.indices.swap_remove_entry(_key) { + Some((_, v)) => { + self.nbytes -= v; + true + } + None => false, + } + } + fn get_typed( &self, name: impl AsRef, diff --git a/xtask/src/gguf_file.rs b/xtask/src/gguf_file.rs index 85eb130..ef1c86d 100644 --- a/xtask/src/gguf_file.rs +++ b/xtask/src/gguf_file.rs @@ -51,15 +51,15 @@ impl<'a> GGufFile<'a> { }) } - pub(crate) fn get_header(&self) -> &GGufFileHeader { + pub fn header(&self) -> &GGufFileHeader { &self.header } - pub(crate) fn get_meta_kvs(&self) -> &GGufMetaKVPairs<'a> { + pub fn meta_kvs(&self) -> &GGufMetaKVPairs<'a> { &self.meta_kvs } - pub(crate) fn get_tensors_as_indexmap(&self) -> IndexMap { + pub fn tensors_as_indexmap(&self) -> IndexMap { self.tensors .iter() .map(move |t| (t, self.data)) diff --git a/xtask/src/merge.rs b/xtask/src/merge.rs index a8b0e03..db39245 100644 --- a/xtask/src/merge.rs +++ b/xtask/src/merge.rs @@ -45,7 +45,7 @@ impl MergeArgs { let kvs = files .iter() - .flat_map(|file| file.get_meta_kvs().kvs()) + .flat_map(|file| file.meta_kvs().kvs()) .filter(|kv| { let key = kv.key(); !key.starts_with("split.") && key != "general.alignment" @@ -53,7 +53,7 @@ impl MergeArgs { .collect::>(); let tensors = files .iter() - .flat_map(|file| file.get_tensors_as_indexmap()) + .flat_map(|file| file.tensors_as_indexmap()) .collect::>(); let out = File::create(shards.single_file()).unwrap(); @@ -62,7 +62,7 @@ impl MergeArgs { let align = files .iter() - .map(|file| file.get_meta_kvs().alignment()) + .map(|file| file.meta_kvs().alignment()) .max() .unwrap(); diff --git a/xtask/src/split.rs b/xtask/src/split.rs index c01760b..fc55bc1 100644 --- a/xtask/src/split.rs +++ b/xtask/src/split.rs @@ -1,4 +1,7 @@ -use crate::gguf_file::{pad, GGufError, GGufFile}; +use crate::{ + gguf_file::{pad, GGufError, GGufFile}, + loose_shards::LooseShards, +}; use ggus::{GGufFileHeader, GGufMetaDataValueType, GGufMetaKVPairs, GGufWriter}; use std::{fs::File, path::PathBuf}; @@ -50,6 +53,11 @@ impl<'a> GGufFileInfo<'a> { impl SplitArgs { pub fn split(self) { + let shards = LooseShards::from(&*self.input); + if shards.count() > 1 { + println!("Model has already been splited"); + return; + } let file = File::open(&self.input) .map_err(|e| { println!("Failed to open"); @@ -60,10 +68,10 @@ impl SplitArgs { let mmap = unsafe { memmap2::Mmap::map(&file).unwrap() }; let ctx_gguf: GGufFile = GGufFile::new(&mmap).unwrap(); - let align = ctx_gguf.get_meta_kvs().alignment(); + let align = ctx_gguf.meta_kvs().alignment(); let ggufs = self.split_strategy(ctx_gguf.clone()).unwrap(); - let tensors = ctx_gguf.get_tensors_as_indexmap(); + let tensors = ctx_gguf.tensors_as_indexmap(); let mut tensor_iter: indexmap::map::Iter = tensors.iter(); for gguf in ggufs { @@ -75,12 +83,6 @@ impl SplitArgs { let kvs = gguf.meta_kvs.kvs(); for kv in kvs { - if kv.key() == LLM_KV_SPLIT_TENSORS_COUNT - || kv.key() == LLM_KV_SPLIT_COUNT - || kv.key() == LLM_KV_SPLIT_NO - { - continue; - } writer .write_meta_kv(kv.key(), kv.ty(), kv.value_bytes()) .unwrap(); @@ -197,10 +199,10 @@ impl SplitArgs { }, } - let tensors = ctx_gguf.get_tensors_as_indexmap(); + let tensors = ctx_gguf.tensors_as_indexmap(); let mut ggufs: Vec = Vec::new(); - let n_tensors: u64 = ctx_gguf.get_header().tensor_count; + let n_tensors: u64 = ctx_gguf.header().tensor_count; let setup_gguf_file = |i_split: u64, n_tensors: u64| { let mut gguf_file = GGufFileInfo::new_empty(); @@ -221,20 +223,20 @@ impl SplitArgs { let mut i_split: u64 = 1; let mut gguf_file = setup_gguf_file(i_split, n_tensors); - gguf_file.meta_kvs = ctx_gguf.get_meta_kvs().clone(); + gguf_file.meta_kvs = ctx_gguf.meta_kvs().clone(); + + gguf_file.header.metadata_kv_count += ctx_gguf.header().metadata_kv_count; - if gguf_file.meta_kvs.get(LLM_KV_SPLIT_NO).is_some() { + if gguf_file.meta_kvs.remove(LLM_KV_SPLIT_NO) { gguf_file.header.metadata_kv_count -= 1; } - if gguf_file.meta_kvs.get(LLM_KV_SPLIT_COUNT).is_some() { + if gguf_file.meta_kvs.remove(LLM_KV_SPLIT_COUNT) { gguf_file.header.metadata_kv_count -= 1; } - if gguf_file.meta_kvs.get(LLM_KV_SPLIT_TENSORS_COUNT).is_some() { + if gguf_file.meta_kvs.remove(LLM_KV_SPLIT_TENSORS_COUNT) { gguf_file.header.metadata_kv_count -= 1; } - gguf_file.header.metadata_kv_count += ctx_gguf.get_header().metadata_kv_count; - if self.no_tensor_first_split { ggufs.push(gguf_file); gguf_file = setup_gguf_file(i_split, n_tensors); @@ -267,7 +269,12 @@ impl SplitArgs { let output_path = &self.output.unwrap(); let mut index = 0; while index < ggufs.len() { - ggufs[index].output_path = split_path(output_path, index + 1, ggufs.len()); + ggufs[index].output_path = format!( + "{}-{:05}-of-{:05}.gguf", + output_path, + index + 1, + ggufs.len() + ); ggufs[index].new_kv_tuples[1] = (LLM_KV_SPLIT_COUNT.to_string(), ty::U16, tensor_count); index += 1; } @@ -283,28 +290,3 @@ impl SplitArgs { } } } - -fn split_path(path_prefix: &String, split_no: usize, split_count: usize) -> String { - format!("{}-{:05}-of-{:05}.gguf", path_prefix, split_no, split_count) -} - -#[cfg(test)] -mod tests { - use super::*; - #[test] - fn it_works() { - let input = PathBuf::from("/home/qinyiqun/gguf/xtask/src/test/rust/rust_ori.gguf"); - let output = Some("/home/qinyiqun/gguf/xtask/src/test/rust/rust_ori_oi".to_string()); - let split_args = SplitArgs { - input, - output, - split_max_tensors: None, - split_max_size: Some("300M".to_string()), - no_tensor_first_split: true, - n_bytes_split: 0, - n_split_tensors: 0, - }; - - split_args.split(); - } -}