zsfm_tirex/
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::TiRexConfig;
12use crate::tensor_map::{map_tensor_name, needs_bias_permute};
13
14pub struct ConvertOptions {
15    pub output_dtype: GGMLType,
16}
17
18/// Convert a TiRex PyTorch Lightning checkpoint (.ckpt) to GGUF.
19pub fn convert(
20    ckpt_path: &Path,
21    config: &TiRexConfig,
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    let tensors = pickle::read_all_with_key(ckpt_path, Some("state_dict"))
30        .with_context(|| format!("read checkpoint {}", ckpt_path.display()))?;
31
32    let total = tensors.len();
33    let pb = ProgressBar::new(total as u64);
34    pb.set_style(
35        ProgressStyle::with_template(
36            "{spinner:.green} [{elapsed_precise}] [{bar:40.cyan/blue}] {pos}/{len} {msg}",
37        )
38        .unwrap()
39        .progress_chars("=>-"),
40    );
41
42    let device = Device::Cpu;
43    let mut mapped = 0usize;
44    let mut skipped: Vec<String> = Vec::new();
45    let mut fallback_count = 0usize;
46
47    let nh = config.num_heads;
48    let dh = config.head_dim();
49    let ng = 4usize; // num_gates
50
51    for (ckpt_name, tensor) in &tensors {
52        pb.set_message(ckpt_name.clone());
53
54        let gguf_name = match map_tensor_name(ckpt_name) {
55            Some(n) => n,
56            None => {
57                skipped.push(ckpt_name.clone());
58                pb.inc(1);
59                continue;
60            }
61        };
62
63        let py_shape: Vec<usize> = tensor.shape().dims().to_vec();
64        let n_elems: usize = py_shape.iter().product();
65        let innermost = py_shape.last().copied().unwrap_or(1);
66
67        let mut f32_vals: Vec<f32> = tensor
68            .to_device(&device)?
69            .to_dtype(DType::F32)?
70            .flatten_all()?
71            .to_vec1()?;
72
73        // Permute sLSTM bias from [NH, NG, DH] storage → [NG, NH, DH]
74        if needs_bias_permute(&gguf_name) {
75            f32_vals = permute_bias(&f32_vals, nh, ng, dh);
76        }
77
78        let (dst_dtype, gguf_shape, tensor_data) =
79            if opts.output_dtype == GGMLType::Q8_0 && (innermost % 32 != 0 || n_elems % 32 != 0) {
80                fallback_count += 1;
81                let gs = py_shape.iter().rev().map(|&d| d as u64).collect();
82                (GGMLType::F32, gs, f32_to_bytes(&f32_vals))
83            } else {
84                let (dst, data) = apply_dtype(&f32_vals, opts.output_dtype)?;
85                let gs = py_shape.iter().rev().map(|&d| d as u64).collect();
86                (dst, gs, data)
87            };
88
89        writer.add_tensor(gguf_name, gguf_shape, dst_dtype, tensor_data);
90        mapped += 1;
91        pb.inc(1);
92    }
93
94    pb.finish_with_message("tensors processed");
95
96    if !skipped.is_empty() {
97        eprintln!("\nWarning: {} tensor(s) skipped:", skipped.len());
98        for name in &skipped { eprintln!("  {name}"); }
99    }
100    if fallback_count > 0 {
101        eprintln!("\nNote: {fallback_count} tensor(s) fell back to F32.");
102    }
103
104    println!("Writing {mapped} tensors to {} …", output_path.display());
105    let out_file = File::create(output_path)
106        .with_context(|| format!("create {}", output_path.display()))?;
107    let mut buf_writer = BufWriter::new(out_file);
108    writer.write_to(&mut buf_writer)?;
109    println!("Done.");
110    Ok(())
111}
112
113fn write_metadata(writer: &mut GGUFWriter, config: &TiRexConfig) {
114    writer.add_metadata("general.architecture",   GGUFMetaValue::String("tirex".into()));
115    writer.add_metadata("general.name",           GGUFMetaValue::String("TiRex".into()));
116    writer.add_metadata("tirex.patch_size",       GGUFMetaValue::Uint32(config.patch_size as u32));
117    writer.add_metadata("tirex.num_blocks",       GGUFMetaValue::Uint32(config.num_blocks as u32));
118    writer.add_metadata("tirex.embedding_dim",    GGUFMetaValue::Uint32(config.embedding_dim as u32));
119    writer.add_metadata("tirex.num_heads",        GGUFMetaValue::Uint32(config.num_heads as u32));
120    writer.add_metadata("tirex.input_ff_dim",     GGUFMetaValue::Uint32(config.input_ff_dim as u32));
121    writer.add_metadata("tirex.ffn_up_dim",       GGUFMetaValue::Uint32(config.ffn_up_dim as u32));
122    writer.add_metadata("tirex.train_ctx_len",    GGUFMetaValue::Uint32(config.train_ctx_len as u32));
123    writer.add_metadata("tirex.num_quantiles",    GGUFMetaValue::Uint32(config.num_quantiles as u32));
124}
125
126/// Permute bias from stored [NH, NG, DH] order to [NG, NH, DH] order.
127/// The sLSTM cell stores bias as [NH*NG*DH] = [NH, NG, DH] row-major,
128/// but during forward pass it permutes to [NG, NH, DH] before adding to gate projections.
129fn permute_bias(vals: &[f32], nh: usize, ng: usize, dh: usize) -> Vec<f32> {
130    let mut out = vec![0.0f32; vals.len()];
131    for h in 0..nh {
132        for g in 0..ng {
133            for d in 0..dh {
134                let src = h * (ng * dh) + g * dh + d;
135                let dst = g * (nh * dh) + h * dh + d;
136                out[dst] = vals[src];
137            }
138        }
139    }
140    out
141}
142
143fn f32_to_bytes(vals: &[f32]) -> Vec<u8> {
144    vals.iter().flat_map(|v| v.to_le_bytes()).collect()
145}
146
147fn apply_dtype(f32_vals: &[f32], dst: GGMLType) -> anyhow::Result<(GGMLType, Vec<u8>)> {
148    match dst {
149        GGMLType::F32 => Ok((GGMLType::F32, f32_to_bytes(f32_vals))),
150        GGMLType::F16 => {
151            let bytes: Vec<u8> = f32_vals.iter()
152                .flat_map(|&v| f32_to_f16_bits(v).to_le_bytes())
153                .collect();
154            Ok((GGMLType::F16, bytes))
155        }
156        GGMLType::BF16 => {
157            let bytes: Vec<u8> = f32_vals.iter()
158                .flat_map(|&v| ((v.to_bits() >> 16) as u16).to_le_bytes())
159                .collect();
160            Ok((GGMLType::BF16, bytes))
161        }
162        GGMLType::Q8_0 => Ok((GGMLType::Q8_0, quantize_q8_0(f32_vals)?)),
163    }
164}
165
166fn quantize_q8_0(values: &[f32]) -> anyhow::Result<Vec<u8>> {
167    const BLOCK: usize = 32;
168    if values.len() % BLOCK != 0 {
169        anyhow::bail!("Q8_0 requires count divisible by {BLOCK}, got {}", values.len());
170    }
171    let n_blocks = values.len() / BLOCK;
172    let mut out = vec![0u8; n_blocks * 34];
173    for b in 0..n_blocks {
174        let blk = &values[b * BLOCK..(b + 1) * BLOCK];
175        let amax = blk.iter().copied().map(f32::abs).fold(0.0f32, f32::max);
176        let d = if amax == 0.0 { 0.0f32 } else { amax / 127.0 };
177        let d_inv = if d == 0.0 { 0.0f32 } else { 1.0 / d };
178        let base = b * 34;
179        out[base..base + 2].copy_from_slice(&f32_to_f16_bits(d).to_le_bytes());
180        for i in 0..BLOCK {
181            out[base + 2 + i] = (blk[i] * d_inv).round().clamp(-127.0, 127.0) as i8 as u8;
182        }
183    }
184    Ok(out)
185}
186
187fn f32_to_f16_bits(v: f32) -> u16 {
188    let bits = v.to_bits();
189    let sign = ((bits >> 16) & 0x8000) as u16;
190    let exp = ((bits >> 23) & 0xFF) as i32;
191    let mantissa = bits & 0x007F_FFFF;
192    if exp == 0xFF { return sign | 0x7C00 | if mantissa != 0 { 0x0200 } else { 0 }; }
193    let new_exp = exp - 127 + 15;
194    if new_exp >= 31 { return sign | 0x7C00; }
195    if new_exp <= 0 {
196        if new_exp < -10 { return sign; }
197        let m = (mantissa | 0x0080_0000) >> (1 - new_exp);
198        return sign | (m >> 13) as u16;
199    }
200    sign | ((new_exp as u16) << 10) | (mantissa >> 13) as u16
201}
202
203fn f16_to_f32(bits: u16) -> f32 {
204    let sign = ((bits & 0x8000) as u32) << 16;
205    let exp = ((bits >> 10) & 0x1F) as i32;
206    let mantissa = (bits & 0x03FF) as u32;
207    let f32_bits = if exp == 0 {
208        if mantissa == 0 { sign }
209        else {
210            let mut m = mantissa; let mut e = 0i32;
211            while m & 0x0400 == 0 { m <<= 1; e += 1; }
212            sign | ((127 - 15 - e + 1) as u32) << 23 | (m & 0x03FF) << 13
213        }
214    } else if exp == 31 {
215        sign | 0x7F80_0000 | (mantissa << 13)
216    } else {
217        sign | ((exp + 127 - 15) as u32) << 23 | (mantissa << 13)
218    };
219    f32::from_bits(f32_bits)
220}
zsfm_tirex/convert.rs

zsfm_tirex/
convert.rs
