zsfm_lag_llama/
convert.rs

1use std::fs::File;
2use std::io::BufWriter;
3use std::path::Path;
4
5use anyhow::Context;
6use candle_core::{pickle, DType, Device};
7use indicatif::{ProgressBar, ProgressStyle};
8
9use zsfm_gguf::{GGMLType, GGUFMetaValue, GGUFWriter};
10
11use crate::config::LagLlamaConfig;
12use crate::tensor_map::map_tensor_name;
13
14pub struct ConvertOptions {
15    pub output_dtype: GGMLType,
16}
17
18/// Convert a Lag-Llama PyTorch Lightning checkpoint (.ckpt) to GGUF.
19pub fn convert(
20    ckpt_path: &Path,
21    config: &LagLlamaConfig,
22    opts: &ConvertOptions,
23    output_path: &Path,
24) -> anyhow::Result<()> {
25    let mut writer = GGUFWriter::new();
26    write_metadata(&mut writer, config);
27
28    println!("Reading checkpoint {} …", ckpt_path.display());
29    // PyTorch Lightning checkpoints store model weights under the "state_dict" key.
30    let tensors = pickle::read_all_with_key(ckpt_path, Some("state_dict"))
31        .with_context(|| format!("read checkpoint {}", ckpt_path.display()))?;
32
33    let total = tensors.len();
34    let pb = ProgressBar::new(total as u64);
35    pb.set_style(
36        ProgressStyle::with_template(
37            "{spinner:.green} [{elapsed_precise}] [{bar:40.cyan/blue}] {pos}/{len} {msg}",
38        )
39        .unwrap()
40        .progress_chars("=>-"),
41    );
42
43    let device = Device::Cpu;
44    let mut mapped = 0usize;
45    let mut skipped: Vec<String> = Vec::new();
46    let mut fallback_count = 0usize;
47
48    for (hf_name, tensor) in &tensors {
49        pb.set_message(hf_name.clone());
50
51        let gguf_name = match map_tensor_name(hf_name) {
52            Some(n) => n,
53            None => {
54                skipped.push(hf_name.clone());
55                pb.inc(1);
56                continue;
57            }
58        };
59
60        // Shape from candle is in Python order (e.g. [out, in]).
61        let py_shape: Vec<usize> = tensor.shape().dims().to_vec();
62        let n_elems: usize = py_shape.iter().product();
63        let innermost = py_shape.last().copied().unwrap_or(1);
64
65        // Decode to f32 using candle — handles BF16/F16/F32 transparently.
66        let f32_vals: Vec<f32> = tensor
67            .to_device(&device)?
68            .to_dtype(DType::F32)?
69            .flatten_all()?
70            .to_vec1()?;
71
72        let (dst_dtype, gguf_shape, tensor_data) =
73            if opts.output_dtype == GGMLType::Q8_0 && (innermost % 32 != 0 || n_elems % 32 != 0) {
74                fallback_count += 1;
75                let gs = py_shape.iter().rev().map(|&d| d as u64).collect();
76                (GGMLType::F32, gs, f32_to_bytes(&f32_vals))
77            } else {
78                let (dst, data) = apply_dtype(&f32_vals, opts.output_dtype)?;
79                let gs = py_shape.iter().rev().map(|&d| d as u64).collect();
80                (dst, gs, data)
81            };
82
83        writer.add_tensor(gguf_name, gguf_shape, dst_dtype, tensor_data);
84        mapped += 1;
85        pb.inc(1);
86    }
87
88    pb.finish_with_message("tensors processed");
89
90    if !skipped.is_empty() {
91        eprintln!("\nWarning: {} tensor(s) skipped (unrecognised names):", skipped.len());
92        for name in &skipped { eprintln!("  {name}"); }
93    }
94    if fallback_count > 0 {
95        eprintln!("\nNote: {fallback_count} tensor(s) fell back to F32.");
96    }
97
98    println!("Writing {mapped} tensors to {} …", output_path.display());
99    let out_file = File::create(output_path)
100        .with_context(|| format!("create {}", output_path.display()))?;
101    let mut buf_writer = BufWriter::new(out_file);
102    writer.write_to(&mut buf_writer)?;
103    println!("Done.");
104    Ok(())
105}
106
107fn write_metadata(writer: &mut GGUFWriter, config: &LagLlamaConfig) {
108    writer.add_metadata("general.architecture",   GGUFMetaValue::String("lag_llama".into()));
109    writer.add_metadata("general.name",           GGUFMetaValue::String("Lag-Llama".into()));
110    writer.add_metadata("lag_llama.n_layer",      GGUFMetaValue::Uint32(config.n_layer as u32));
111    writer.add_metadata("lag_llama.n_head",       GGUFMetaValue::Uint32(config.n_head as u32));
112    writer.add_metadata("lag_llama.n_embd_per_head", GGUFMetaValue::Uint32(config.n_embd_per_head as u32));
113    writer.add_metadata("lag_llama.n_embd",       GGUFMetaValue::Uint32(config.n_embd as u32));
114    writer.add_metadata("lag_llama.mlp_hidden",   GGUFMetaValue::Uint32(config.mlp_hidden as u32));
115    writer.add_metadata("lag_llama.feature_size", GGUFMetaValue::Uint32(config.feature_size as u32));
116    writer.add_metadata("lag_llama.n_lags",       GGUFMetaValue::Uint32(config.n_lags as u32));
117    writer.add_metadata("lag_llama.n_time_feat",  GGUFMetaValue::Uint32(config.n_time_feat as u32));
118    writer.add_metadata("lag_llama.max_context_length", GGUFMetaValue::Uint32(config.max_context_length as u32));
119}
120
121fn f32_to_bytes(vals: &[f32]) -> Vec<u8> {
122    vals.iter().flat_map(|v| v.to_le_bytes()).collect()
123}
124
125fn apply_dtype(f32_vals: &[f32], dst: GGMLType) -> anyhow::Result<(GGMLType, Vec<u8>)> {
126    match dst {
127        GGMLType::F32 => Ok((GGMLType::F32, f32_to_bytes(f32_vals))),
128        GGMLType::F16 => {
129            let bytes: Vec<u8> = f32_vals.iter()
130                .flat_map(|&v| f32_to_f16_bits(v).to_le_bytes())
131                .collect();
132            Ok((GGMLType::F16, bytes))
133        }
134        GGMLType::BF16 => {
135            let bytes: Vec<u8> = f32_vals.iter()
136                .flat_map(|&v| ((v.to_bits() >> 16) as u16).to_le_bytes())
137                .collect();
138            Ok((GGMLType::BF16, bytes))
139        }
140        GGMLType::Q8_0 => Ok((GGMLType::Q8_0, quantize_q8_0(f32_vals)?)),
141    }
142}
143
144fn quantize_q8_0(values: &[f32]) -> anyhow::Result<Vec<u8>> {
145    const BLOCK: usize = 32;
146    if values.len() % BLOCK != 0 {
147        anyhow::bail!("Q8_0 requires count divisible by {BLOCK}, got {}", values.len());
148    }
149    let n_blocks = values.len() / BLOCK;
150    let mut out = vec![0u8; n_blocks * 34];
151    for b in 0..n_blocks {
152        let blk = &values[b * BLOCK..(b + 1) * BLOCK];
153        let amax = blk.iter().copied().map(f32::abs).fold(0.0f32, f32::max);
154        let d = if amax == 0.0 { 0.0f32 } else { amax / 127.0 };
155        let d_inv = if d == 0.0 { 0.0f32 } else { 1.0 / d };
156        let base = b * 34;
157        out[base..base + 2].copy_from_slice(&f32_to_f16_bits(d).to_le_bytes());
158        for i in 0..BLOCK {
159            out[base + 2 + i] = (blk[i] * d_inv).round().clamp(-127.0, 127.0) as i8 as u8;
160        }
161    }
162    Ok(out)
163}
164
165fn f32_to_f16_bits(v: f32) -> u16 {
166    let bits = v.to_bits();
167    let sign = ((bits >> 16) & 0x8000) as u16;
168    let exp = ((bits >> 23) & 0xFF) as i32;
169    let mantissa = bits & 0x007F_FFFF;
170    if exp == 0xFF { return sign | 0x7C00 | if mantissa != 0 { 0x0200 } else { 0 }; }
171    let new_exp = exp - 127 + 15;
172    if new_exp >= 31 { return sign | 0x7C00; }
173    if new_exp <= 0 {
174        if new_exp < -10 { return sign; }
175        let m = (mantissa | 0x0080_0000) >> (1 - new_exp);
176        return sign | (m >> 13) as u16;
177    }
178    sign | ((new_exp as u16) << 10) | (mantissa >> 13) as u16
179}
zsfm_lag_llama/convert.rs

zsfm_lag_llama/
convert.rs
