chore(代码变更): 1. 清理无用代码 2. 更新依赖库 3. 优化项目结构 4. 修复小型bug 5. 增加注释以提高可读性

2026-07-08 10:44:32 +08:00
parent 4b342ef62f
commit d88afb5fe7
22 changed files with 1068 additions and 654372 deletions
@@ -2,7 +2,7 @@ use minicpm_core::config::LlamaConfig;
 use minicpm_core::model::LlamaForCausalLM;
 use burn::module::{Module, Param};
 use burn::nn::{EmbeddingRecord, LinearRecord};
-use burn::record::{FullPrecisionSettings, NamedMpkFileRecorder};
+use burn::record::{FullPrecisionSettings, HalfPrecisionSettings, NamedMpkFileRecorder};
 use burn::tensor::backend::Backend;
 use burn::tensor::{Float, Shape, Tensor, TensorData};
 use memmap2::Mmap;
@@ -40,6 +40,39 @@ pub fn export_model<B: Backend>(
    println!("模型已成功导出到: {:?}", output_dir);
    Ok(())
 }
+
+/// 导出模型为半精度（f16），输出文件约为全精度的一半大小。
+pub fn export_model_half<B: Backend>(
+    safetensors_path: &Path,
+    config_path: &Path,
+    tokenizer_path: &Path,
+    output_dir: &Path,
+    device: &B::Device,
+) -> anyhow::Result<()> {
+    println!("开始转换 MiniCPM 模型为 Burn 格式（半精度 f16）...");
+
+    println!("加载配置文件: {:?}", config_path);
+    let config = LlamaConfig::from_json(config_path.to_str().unwrap())?;
+
+    println!("创建模型结构...");
+    let model = LlamaForCausalLM::<B>::new(config, device);
+
+    println!("加载 safetensors 权重...");
+    let model = load_safetensors(model, safetensors_path, device)?;
+
+    println!("保存为半精度 MPK 格式...");
+    std::fs::create_dir_all(output_dir)?;
+    let output_path = output_dir.join("model");
+    let recorder = NamedMpkFileRecorder::<HalfPrecisionSettings>::new();
+    model.save_file(&output_path, &recorder)?;
+
+    println!("拷贝配置文件和 tokenizer...");
+    std::fs::copy(config_path, output_dir.join("config.json"))?;
+    std::fs::copy(tokenizer_path, output_dir.join("tokenizer.json"))?;
+
+    println!("半精度模型已成功导出到: {:?}", output_dir);
+    Ok(())
+}
 pub fn load_safetensors<B: Backend>(
    model: LlamaForCausalLM<B>,
    path: &Path,
@@ -256,3 +289,223 @@ fn f16_to_f32(f16: u16) -> f32 {
    let f32_bits = ((sign as u32) << 31) | (f32_exp << 23) | ((mant as u32) << 13);
    f32::from_bits(f32_bits)
 }
+
+// ==================== Q8 量化导出 ====================
+
+/// Q8 量化单个 f32 切片，返回 (scale, i8_data)
+fn quantize_q8(f32_data: &[f32]) -> (f32, Vec<i8>) {
+    let max_abs = f32_data
+        .iter()
+        .map(|&x| x.abs())
+        .fold(0.0f32, |a, b| a.max(b));
+    let scale = if max_abs > 0.0 { max_abs / 127.0 } else { 1.0 };
+    let q8: Vec<i8> = f32_data
+        .iter()
+        .map(|&x| (x / scale).round().clamp(-127.0, 127.0) as i8)
+        .collect();
+    (scale, q8)
+}
+
+/// Q8 导出元数据中的 tensor 条目
+#[derive(serde::Serialize, serde::Deserialize)]
+struct Q8TensorMeta {
+    shape: Vec<usize>,
+    scale: f32,
+    offset: u64,
+    numel: usize,
+}
+
+/// Q8 导出元数据
+#[derive(serde::Serialize, serde::Deserialize)]
+struct Q8Metadata {
+    tensors: std::collections::HashMap<String, Q8TensorMeta>,
+}
+
+/// 导出模型为 Q8 量化格式（INT8），输出 model.q8.json + model.q8.bin。
+/// 2D 权重（Linear/Embedding）量化为 INT8；1D norm 权重存为 F16。
+pub fn export_model_q8(
+    safetensors_path: &Path,
+    config_path: &Path,
+    tokenizer_path: &Path,
+    output_dir: &Path,
+) -> anyhow::Result<()> {
+    println!("开始转换 MiniCPM 模型为 Q8 量化格式...");
+
+    println!("解析 safetensors 文件: {:?}", safetensors_path);
+    let file = std::fs::File::open(safetensors_path)?;
+    let mmap = unsafe { Mmap::map(&file)? };
+
+    let header_len = u64::from_le_bytes(mmap[..8].try_into().unwrap()) as usize;
+    let header_bytes = &mmap[8..8 + header_len];
+    let header: serde_json::Value = serde_json::from_slice(header_bytes)?;
+
+    let data_offset = 8 + header_len;
+
+    // 收集 tensor 信息
+    let mut tensor_infos: Vec<(String, Vec<usize>, String, usize, usize)> = Vec::new();
+    if let Some(obj) = header.as_object() {
+        for (name, info) in obj {
+            if name == "__metadata__" {
+                continue;
+            }
+            if let Some(info_obj) = info.as_object() {
+                let mut shape = Vec::new();
+                if let Some(s) = info_obj.get("shape") {
+                    if let Some(arr) = s.as_array() {
+                        shape = arr.iter().map(|v| v.as_u64().unwrap() as usize).collect();
+                    }
+                }
+                let mut dtype = String::new();
+                if let Some(d) = info_obj.get("dtype") {
+                    dtype = d.as_str().unwrap_or("").to_string();
+                }
+                let mut start = 0;
+                let mut end = 0;
+                if let Some(offsets) = info_obj.get("data_offsets") {
+                    if let Some(arr) = offsets.as_array() {
+                        start = arr[0].as_u64().unwrap() as usize;
+                        end = arr[1].as_u64().unwrap() as usize;
+                    }
+                }
+                tensor_infos.push((name.clone(), shape, dtype, start, end));
+            }
+        }
+    }
+
+    // 定义需要量化的 2D 权重名称（与 load_safetensors 中的列表对应）
+    let mut q8_tensors: Vec<String> = Vec::new();
+    for i in 0..32 {
+        let p = format!("model.layers.{i}");
+        q8_tensors.push(format!("{p}.self_attn.q_proj.weight"));
+        q8_tensors.push(format!("{p}.self_attn.k_proj.weight"));
+        q8_tensors.push(format!("{p}.self_attn.v_proj.weight"));
+        q8_tensors.push(format!("{p}.self_attn.o_proj.weight"));
+        q8_tensors.push(format!("{p}.mlp.gate_proj.weight"));
+        q8_tensors.push(format!("{p}.mlp.up_proj.weight"));
+        q8_tensors.push(format!("{p}.mlp.down_proj.weight"));
+    }
+    q8_tensors.push("model.embed_tokens.weight".to_string());
+    q8_tensors.push("lm_head.weight".to_string());
+
+    let q8_set: std::collections::HashSet<String> = q8_tensors.into_iter().collect();
+
+    std::fs::create_dir_all(output_dir)?;
+    let json_path = output_dir.join("model.q8.json");
+    let bin_path = output_dir.join("model.q8.bin");
+
+    let mut metadata: std::collections::HashMap<String, Q8TensorMeta> = std::collections::HashMap::new();
+    let mut bin_file = std::fs::File::create(&bin_path)?;
+    let mut current_offset: u64 = 0;
+
+    // 先处理 2D 权重：量化为 INT8
+    for (name, shape, dtype, start, end) in &tensor_infos {
+        if shape.len() != 2 {
+            continue;
+        }
+        if !q8_set.contains(name.as_str()) {
+            continue;
+        }
+
+        println!("量化 2D 权重: {name}  shape={shape:?}  dtype={dtype}");
+        let raw = &mmap[data_offset + *start..data_offset + *end];
+        let f32_data = convert_to_f32(raw, dtype);
+        let numel = f32_data.len();
+        let (scale, q8_data) = quantize_q8(&f32_data);
+
+        // 写入 INT8 数据
+        let bytes: &[u8] = unsafe {
+            std::slice::from_raw_parts(q8_data.as_ptr() as *const u8, q8_data.len())
+        };
+        use std::io::Write;
+        bin_file.write_all(bytes)?;
+
+        metadata.insert(
+            name.clone(),
+            Q8TensorMeta {
+                shape: shape.clone(),
+                scale,
+                offset: current_offset,
+                numel,
+            },
+        );
+        current_offset += q8_data.len() as u64;
+    }
+
+    // 再处理 1D norm 权重：存为 F16
+    for (name, shape, dtype, start, end) in &tensor_infos {
+        if shape.len() != 1 {
+            continue;
+        }
+        let raw = &mmap[data_offset + *start..data_offset + *end];
+        let f32_data = convert_to_f32(raw, dtype);
+        let numel = f32_data.len();
+
+        // 转 F16 存储
+        let mut f16_bytes = Vec::with_capacity(numel * 2);
+        for &v in &f32_data {
+            let f16_bits = f32_to_f16_bits(v);
+            f16_bytes.extend_from_slice(&f16_bits.to_le_bytes());
+        }
+
+        use std::io::Write;
+        bin_file.write_all(&f16_bytes)?;
+
+        // 用负数 scale 表示 F16 类型（加载时 scale < 0 表示 F16）
+        metadata.insert(
+            name.clone(),
+            Q8TensorMeta {
+                shape: shape.clone(),
+                scale: -1.0, // 标记为 F16
+                offset: current_offset,
+                numel,
+            },
+        );
+        current_offset += f16_bytes.len() as u64;
+    }
+
+    // 写 JSON 元数据
+    let meta = Q8Metadata { tensors: metadata };
+    let json_str = serde_json::to_string_pretty(&meta)?;
+    std::fs::write(&json_path, json_str)?;
+
+    println!("拷贝配置文件和 tokenizer...");
+    std::fs::copy(config_path, output_dir.join("config.json"))?;
+    std::fs::copy(tokenizer_path, output_dir.join("tokenizer.json"))?;
+
+    println!("Q8 量化模型已导出到: {:?}", output_dir);
+    println!("  - model.q8.json (元数据)",);
+    println!("  - model.q8.bin (INT8 + F16 权重)",);
+    Ok(())
+}
+
+/// f32 → f16 位表示（返回 u16）
+fn f32_to_f16_bits(v: f32) -> u16 {
+    let bits = v.to_bits();
+    let sign = bits >> 31;
+    let exp = (bits >> 23) & 0xFF;
+    let mant = bits & 0x7FFFFF;
+
+    let res: u32 = if exp == 0xFF {
+        // NaN / Inf
+        let f16_exp = 0x1F;
+        let f16_mant = if mant == 0 { 0 } else { 0x3FF };
+        (sign << 15) | (f16_exp << 10) | f16_mant
+    } else if exp < 103 {
+        // 下溢，返回 0
+        0
+    } else {
+        let new_exp = (exp as i32) - 127 + 15;
+        if new_exp >= 31 {
+            // 上溢
+            (sign << 15) | (0x1F << 10)
+        } else if new_exp <= 0 {
+            // 次正规数
+            let mant_new = (mant | 0x800000) >> (113 - new_exp);
+            (sign << 15) | (mant_new & 0x3FF)
+        } else {
+            let f16_mant = mant >> 13;
+            (sign << 15) | ((new_exp as u32) << 10) | f16_mant
+        }
+    };
+    res as u16
+}