zsfm_lag_llama/
convert.rs

1use std::fs::File;
2use std::io::BufWriter;
3use std::path::Path;
4
5use anyhow::Context;
6use indicatif::{ProgressBar, ProgressStyle};
7
8use zsfm_checkpoint::cast::decode_to_f32;
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 ckpt = zsfm_checkpoint::load_checkpoint(
31        &[ckpt_path.to_path_buf()],
32        &zsfm_checkpoint::LoadOptions {
33            pickle_key: Some("state_dict".to_string()),
34            ..Default::default()
35        },
36    )
37    .with_context(|| format!("read checkpoint {}", ckpt_path.display()))?;
38
39    let total = ckpt.tensors.len();
40    let pb = ProgressBar::new(total as u64);
41    pb.set_style(
42        ProgressStyle::with_template(
43            "{spinner:.green} [{elapsed_precise}] [{bar:40.cyan/blue}] {pos}/{len} {msg}",
44        )
45        .unwrap()
46        .progress_chars("=>-"),
47    );
48
49    let mut mapped = 0usize;
50    let mut skipped: Vec<String> = Vec::new();
51    let mut fallback_count = 0usize;
52
53    for raw in &ckpt.tensors {
54        let hf_name = raw.name.clone();
55        pb.set_message(hf_name.clone());
56
57        let gguf_name = match map_tensor_name(&hf_name) {
58            Some(n) => n,
59            None => {
60                skipped.push(hf_name.clone());
61                pb.inc(1);
62                continue;
63            }
64        };
65
66        // Shape is in Python order (e.g. [out, in]).
67        let py_shape: Vec<usize> = raw.shape.iter().map(|&d| d as usize).collect();
68        let n_elems: usize = py_shape.iter().product();
69        let innermost = py_shape.last().copied().unwrap_or(1);
70
71        // Decode to f32 — handles BF16/F16/F32 transparently.
72        let f32_vals: Vec<f32> = decode_to_f32(&raw.data, raw.dtype)
73            .with_context(|| format!("tensor {hf_name}: decode failed"))?;
74
75        let (dst_dtype, gguf_shape, tensor_data) =
76            if opts.output_dtype == GGMLType::Q8_0 && (innermost % 32 != 0 || n_elems % 32 != 0) {
77                fallback_count += 1;
78                let gs = py_shape.iter().rev().map(|&d| d as u64).collect();
79                (GGMLType::F32, gs, f32_to_bytes(&f32_vals))
80            } else {
81                let (dst, data) = apply_dtype(&f32_vals, opts.output_dtype)?;
82                let gs = py_shape.iter().rev().map(|&d| d as u64).collect();
83                (dst, gs, data)
84            };
85
86        writer.add_tensor(gguf_name, gguf_shape, dst_dtype, tensor_data);
87        mapped += 1;
88        pb.inc(1);
89    }
90
91    pb.finish_with_message("tensors processed");
92
93    if !skipped.is_empty() {
94        eprintln!(
95            "\nWarning: {} tensor(s) skipped (unrecognised names):",
96            skipped.len()
97        );
98        for name in &skipped {
99            eprintln!("  {name}");
100        }
101    }
102    if fallback_count > 0 {
103        eprintln!("\nNote: {fallback_count} tensor(s) fell back to F32.");
104    }
105
106    println!("Writing {mapped} tensors to {} …", output_path.display());
107    let out_file =
108        File::create(output_path).with_context(|| format!("create {}", output_path.display()))?;
109    let mut buf_writer = BufWriter::new(out_file);
110    writer.write_to(&mut buf_writer)?;
111    println!("Done.");
112    Ok(())
113}
114
115fn write_metadata(writer: &mut GGUFWriter, config: &LagLlamaConfig) {
116    writer.add_metadata(
117        "general.architecture",
118        GGUFMetaValue::String("lag_llama".into()),
119    );
120    writer.add_metadata("general.name", GGUFMetaValue::String("Lag-Llama".into()));
121    writer.add_metadata(
122        "lag_llama.n_layer",
123        GGUFMetaValue::Uint32(config.n_layer as u32),
124    );
125    writer.add_metadata(
126        "lag_llama.n_head",
127        GGUFMetaValue::Uint32(config.n_head as u32),
128    );
129    writer.add_metadata(
130        "lag_llama.n_embd_per_head",
131        GGUFMetaValue::Uint32(config.n_embd_per_head as u32),
132    );
133    writer.add_metadata(
134        "lag_llama.n_embd",
135        GGUFMetaValue::Uint32(config.n_embd as u32),
136    );
137    writer.add_metadata(
138        "lag_llama.mlp_hidden",
139        GGUFMetaValue::Uint32(config.mlp_hidden as u32),
140    );
141    writer.add_metadata(
142        "lag_llama.feature_size",
143        GGUFMetaValue::Uint32(config.feature_size as u32),
144    );
145    writer.add_metadata(
146        "lag_llama.n_lags",
147        GGUFMetaValue::Uint32(config.n_lags as u32),
148    );
149    writer.add_metadata(
150        "lag_llama.n_time_feat",
151        GGUFMetaValue::Uint32(config.n_time_feat as u32),
152    );
153    writer.add_metadata(
154        "lag_llama.max_context_length",
155        GGUFMetaValue::Uint32(config.max_context_length as u32),
156    );
157}
158
159fn f32_to_bytes(vals: &[f32]) -> Vec<u8> {
160    vals.iter().flat_map(|v| v.to_le_bytes()).collect()
161}
162
163fn apply_dtype(f32_vals: &[f32], dst: GGMLType) -> anyhow::Result<(GGMLType, Vec<u8>)> {
164    match dst {
165        GGMLType::F32 => Ok((GGMLType::F32, f32_to_bytes(f32_vals))),
166        GGMLType::F16 => {
167            let bytes: Vec<u8> = f32_vals
168                .iter()
169                .flat_map(|&v| f32_to_f16_bits(v).to_le_bytes())
170                .collect();
171            Ok((GGMLType::F16, bytes))
172        }
173        GGMLType::BF16 => {
174            let bytes: Vec<u8> = f32_vals
175                .iter()
176                .flat_map(|&v| ((v.to_bits() >> 16) as u16).to_le_bytes())
177                .collect();
178            Ok((GGMLType::BF16, bytes))
179        }
180        GGMLType::Q8_0 => Ok((GGMLType::Q8_0, quantize_q8_0(f32_vals)?)),
181    }
182}
183
184fn quantize_q8_0(values: &[f32]) -> anyhow::Result<Vec<u8>> {
185    const BLOCK: usize = 32;
186    if values.len() % BLOCK != 0 {
187        anyhow::bail!(
188            "Q8_0 requires count divisible by {BLOCK}, got {}",
189            values.len()
190        );
191    }
192    let n_blocks = values.len() / BLOCK;
193    let mut out = vec![0u8; n_blocks * 34];
194    for b in 0..n_blocks {
195        let blk = &values[b * BLOCK..(b + 1) * BLOCK];
196        let amax = blk.iter().copied().map(f32::abs).fold(0.0f32, f32::max);
197        let d = if amax == 0.0 { 0.0f32 } else { amax / 127.0 };
198        let d_inv = if d == 0.0 { 0.0f32 } else { 1.0 / d };
199        let base = b * 34;
200        out[base..base + 2].copy_from_slice(&f32_to_f16_bits(d).to_le_bytes());
201        for i in 0..BLOCK {
202            out[base + 2 + i] = (blk[i] * d_inv).round().clamp(-127.0, 127.0) as i8 as u8;
203        }
204    }
205    Ok(out)
206}
207
208fn f32_to_f16_bits(v: f32) -> u16 {
209    let bits = v.to_bits();
210    let sign = ((bits >> 16) & 0x8000) as u16;
211    let exp = ((bits >> 23) & 0xFF) as i32;
212    let mantissa = bits & 0x007F_FFFF;
213    if exp == 0xFF {
214        return sign | 0x7C00 | if mantissa != 0 { 0x0200 } else { 0 };
215    }
216    let new_exp = exp - 127 + 15;
217    if new_exp >= 31 {
218        return sign | 0x7C00;
219    }
220    if new_exp <= 0 {
221        if new_exp < -10 {
222            return sign;
223        }
224        let m = (mantissa | 0x0080_0000) >> (1 - new_exp);
225        return sign | (m >> 13) as u16;
226    }
227    sign | ((new_exp as u16) << 10) | (mantissa >> 13) as u16
228}
zsfm_lag_llama/convert.rs

zsfm_lag_llama/
convert.rs
