MiniCPM5-1B-rust/crates/minicpm-convert/src/lib.rs

use minicpm_core::config::LlamaConfig;
use minicpm_core::model::LlamaForCausalLM;
use burn::module::{Module, Param};
use burn::nn::{EmbeddingRecord, LinearRecord};
use burn::record::{FullPrecisionSettings, HalfPrecisionSettings, NamedMpkFileRecorder};
use burn::tensor::backend::Backend;
use burn::tensor::{Float, Shape, Tensor, TensorData};
use memmap2::Mmap;
use std::collections::HashMap;
use std::path::Path;

pub fn export_model<B: Backend>(
    safetensors_path: &Path,
    config_path: &Path,
    tokenizer_path: &Path,
    output_dir: &Path,
    device: &B::Device,
) -> anyhow::Result<()> {
    println!("开始转换 MiniCPM 模型为 Burn 格式...");

    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::<FullPrecisionSettings>::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(())
}

/// 导出模型为半精度（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,
    device: &B::Device,
) -> anyhow::Result<LlamaForCausalLM<B>> {
    let file = std::fs::File::open(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 mut tensors = HashMap::new();
    let mut shapes = HashMap::new();
    let mut dtypes = HashMap::new();

    let data_offset = 8 + header_len;

    if let Some(obj) = header.as_object() {
        for (name, info) in obj {
            if name == "__metadata__" {
                continue;
            }
            if let Some(info_obj) = info.as_object() {
                if let Some(offsets) = info_obj.get("data_offsets") {
                    if let Some(arr) = offsets.as_array() {
                        let start = arr[0].as_u64().unwrap() as usize;
                        let end = arr[1].as_u64().unwrap() as usize;
                        let data = &mmap[data_offset + start..data_offset + end];
                        tensors.insert(name.clone(), data.to_vec());
                    }
                }
                if let Some(shape) = info_obj.get("shape") {
                    if let Some(arr) = shape.as_array() {
                        let shape_vec: Vec<usize> = arr
                            .iter()
                            .map(|v| v.as_u64().unwrap() as usize)
                            .collect();
                        shapes.insert(name.clone(), shape_vec);
                    }
                }
                if let Some(dtype) = info_obj.get("dtype") {
                    if let Some(s) = dtype.as_str() {
                        dtypes.insert(name.clone(), s.to_string());
                    }
                }
            }
        }
    }

    let mut model = model;

    if let Some(weight) = load_tensor_2d::<B>(&tensors, &shapes, &dtypes, "model.embed_tokens.weight", device) {
        let record = EmbeddingRecord {
            weight: Param::from_tensor(weight),
        };
        model.model.embed_tokens = model.model.embed_tokens.clone().load_record(record);
    }

    for i in 0..model.config.num_hidden_layers {
        let prefix = format!("model.layers.{i}");

        if let Some(w) = load_tensor_2d::<B>(&tensors, &shapes, &dtypes, &format!("{prefix}.self_attn.q_proj.weight"), device) {
            let record = LinearRecord { weight: Param::from_tensor(w.transpose()), bias: None };
            model.model.layers[i].self_attn.q_proj = model.model.layers[i].self_attn.q_proj.clone().load_record(record);
        }
        if let Some(w) = load_tensor_2d::<B>(&tensors, &shapes, &dtypes, &format!("{prefix}.self_attn.k_proj.weight"), device) {
            let record = LinearRecord { weight: Param::from_tensor(w.transpose()), bias: None };
            model.model.layers[i].self_attn.k_proj = model.model.layers[i].self_attn.k_proj.clone().load_record(record);
        }
        if let Some(w) = load_tensor_2d::<B>(&tensors, &shapes, &dtypes, &format!("{prefix}.self_attn.v_proj.weight"), device) {
            let record = LinearRecord { weight: Param::from_tensor(w.transpose()), bias: None };
            model.model.layers[i].self_attn.v_proj = model.model.layers[i].self_attn.v_proj.clone().load_record(record);
        }
        if let Some(w) = load_tensor_2d::<B>(&tensors, &shapes, &dtypes, &format!("{prefix}.self_attn.o_proj.weight"), device) {
            let record = LinearRecord { weight: Param::from_tensor(w.transpose()), bias: None };
            model.model.layers[i].self_attn.o_proj = model.model.layers[i].self_attn.o_proj.clone().load_record(record);
        }

        if let Some(w) = load_tensor_2d::<B>(&tensors, &shapes, &dtypes, &format!("{prefix}.mlp.gate_proj.weight"), device) {
            let record = LinearRecord { weight: Param::from_tensor(w.transpose()), bias: None };
            model.model.layers[i].mlp.gate_proj = model.model.layers[i].mlp.gate_proj.clone().load_record(record);
        }
        if let Some(w) = load_tensor_2d::<B>(&tensors, &shapes, &dtypes, &format!("{prefix}.mlp.up_proj.weight"), device) {
            let record = LinearRecord { weight: Param::from_tensor(w.transpose()), bias: None };
            model.model.layers[i].mlp.up_proj = model.model.layers[i].mlp.up_proj.clone().load_record(record);
        }
        if let Some(w) = load_tensor_2d::<B>(&tensors, &shapes, &dtypes, &format!("{prefix}.mlp.down_proj.weight"), device) {
            let record = LinearRecord { weight: Param::from_tensor(w.transpose()), bias: None };
            model.model.layers[i].mlp.down_proj = model.model.layers[i].mlp.down_proj.clone().load_record(record);
        }

        if let Some(w) = load_tensor_1d::<B>(&tensors, &shapes, &dtypes, &format!("{prefix}.input_layernorm.weight"), device) {
            model.model.layers[i].input_layernorm.weight = Param::from_tensor(w);
        }
        if let Some(w) = load_tensor_1d::<B>(&tensors, &shapes, &dtypes, &format!("{prefix}.post_attention_layernorm.weight"), device) {
            model.model.layers[i].post_attention_layernorm.weight = Param::from_tensor(w);
        }
    }

    if let Some(w) = load_tensor_1d::<B>(&tensors, &shapes, &dtypes, "model.norm.weight", device) {
        model.model.norm.weight = Param::from_tensor(w);
    }

    if let Some(w) = load_tensor_2d::<B>(&tensors, &shapes, &dtypes, "lm_head.weight", device) {
        let record = LinearRecord { weight: Param::from_tensor(w.transpose()), bias: None };
        model.lm_head = model.lm_head.clone().load_record(record);
    } else if model.config.tie_word_embeddings {
        let embed_weight = model.model.embed_tokens.weight.val();
        let record = LinearRecord { weight: Param::from_tensor(embed_weight.transpose()), bias: None };
        model.lm_head = model.lm_head.clone().load_record(record);
    }

    Ok(model)
}

fn load_tensor_2d<B: Backend>(
    tensors: &HashMap<String, Vec<u8>>,
    shapes: &HashMap<String, Vec<usize>>,
    dtypes: &HashMap<String, String>,
    name: &str,
    device: &B::Device,
) -> Option<Tensor<B, 2, Float>> {
    let data = tensors.get(name)?;
    let shape = shapes.get(name)?;
    let dtype = dtypes.get(name)?;

    assert_eq!(shape.len(), 2, "Expected 2D tensor for {name}");

    let f32_data = convert_to_f32(data, dtype);
    let shape_arr: [usize; 2] = [shape[0], shape[1]];
    let tensor_data = TensorData::new(f32_data, Shape::new(shape_arr));
    Some(Tensor::from_data(tensor_data, device))
}

fn load_tensor_1d<B: Backend>(
    tensors: &HashMap<String, Vec<u8>>,
    shapes: &HashMap<String, Vec<usize>>,
    dtypes: &HashMap<String, String>,
    name: &str,
    device: &B::Device,
) -> Option<Tensor<B, 1, Float>> {
    let data = tensors.get(name)?;
    let shape = shapes.get(name)?;
    let dtype = dtypes.get(name)?;

    assert_eq!(shape.len(), 1, "Expected 1D tensor for {name}");

    let f32_data = convert_to_f32(data, dtype);
    let shape_arr: [usize; 1] = [shape[0]];
    let tensor_data = TensorData::new(f32_data, Shape::new(shape_arr));
    Some(Tensor::from_data(tensor_data, device))
}

fn convert_to_f32(data: &[u8], dtype: &str) -> Vec<f32> {
    match dtype {
        "BF16" => data
            .chunks_exact(2)
            .map(|chunk| {
                let bytes: [u8; 2] = chunk.try_into().unwrap();
                bf16_to_f32(u16::from_le_bytes(bytes))
            })
            .collect(),
        "F16" => data
            .chunks_exact(2)
            .map(|chunk| {
                let bytes: [u8; 2] = chunk.try_into().unwrap();
                f16_to_f32(u16::from_le_bytes(bytes))
            })
            .collect(),
        "F32" => data
            .chunks_exact(4)
            .map(|chunk| {
                let bytes: [u8; 4] = chunk.try_into().unwrap();
                f32::from_le_bytes(bytes)
            })
            .collect(),
        _ => panic!("Unsupported dtype: {dtype}"),
    }
}

fn bf16_to_f32(bf16: u16) -> f32 {
    let f32_bits = (bf16 as u32) << 16;
    f32::from_bits(f32_bits)
}

fn f16_to_f32(f16: u16) -> f32 {
    let sign = (f16 >> 15) & 1;
    let exp = (f16 >> 10) & 0x1F;
    let mant = f16 & 0x3FF;

    if exp == 0 {
        if mant == 0 {
            return if sign == 0 { 0.0 } else { -0.0 };
        }
        let mut m = mant as u32;
        let mut e = 0u32;
        while (m & 0x400) == 0 {
            m <<= 1;
            e += 1;
        }
        m &= 0x3FF;
        let f32_exp = 127 - 14 - e;
        let f32_bits = ((sign as u32) << 31) | (f32_exp << 23) | ((m as u32) << 13);
        return f32::from_bits(f32_bits);
    }

    if exp == 0x1F {
        let f32_bits = ((sign as u32) << 31) | (0xFFu32 << 23) | ((mant as u32) << 13);
        return f32::from_bits(f32_bits);
    }

    let f32_exp = (exp as u32) + 127 - 15;
    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
}