zsfm_tirex/
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;
9
10use zsfm_gguf::{GGMLType, GGUFMetaValue, GGUFWriter};
11
12use crate::config::TiRexConfig;
13use crate::tensor_map::{map_tensor_name, needs_bias_permute};
14
15pub struct ConvertOptions {
16    pub output_dtype: GGMLType,
17}
18
19/// Convert a TiRex PyTorch Lightning checkpoint (.ckpt) to GGUF.
20pub fn convert(
21    ckpt_path: &Path,
22    config: &TiRexConfig,
23    opts: &ConvertOptions,
24    output_path: &Path,
25) -> anyhow::Result<()> {
26    let mut writer = GGUFWriter::new();
27    write_metadata(&mut writer, config);
28
29    println!("Reading checkpoint {} …", ckpt_path.display());
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    let nh = config.num_heads;
54    let dh = config.head_dim();
55    let ng = 4usize; // num_gates
56
57    for raw in &ckpt.tensors {
58        let ckpt_name = raw.name.clone();
59        pb.set_message(ckpt_name.clone());
60
61        let gguf_name = match map_tensor_name(&ckpt_name) {
62            Some(n) => n,
63            None => {
64                skipped.push(ckpt_name.clone());
65                pb.inc(1);
66                continue;
67            }
68        };
69
70        let py_shape: Vec<usize> = raw.shape.iter().map(|&d| d as usize).collect();
71        let n_elems: usize = py_shape.iter().product();
72        let innermost = py_shape.last().copied().unwrap_or(1);
73
74        let mut f32_vals: Vec<f32> = decode_to_f32(&raw.data, raw.dtype)
75            .with_context(|| format!("tensor {ckpt_name}: decode failed"))?;
76
77        // Permute sLSTM bias from [NH, NG, DH] storage → [NG, NH, DH]
78        if needs_bias_permute(&gguf_name) {
79            f32_vals = permute_bias(&f32_vals, nh, ng, dh);
80        }
81
82        let (dst_dtype, gguf_shape, tensor_data) =
83            if opts.output_dtype == GGMLType::Q8_0 && (innermost % 32 != 0 || n_elems % 32 != 0) {
84                fallback_count += 1;
85                let gs = py_shape.iter().rev().map(|&d| d as u64).collect();
86                (GGMLType::F32, gs, f32_to_bytes(&f32_vals))
87            } else {
88                let (dst, data) = apply_dtype(&f32_vals, opts.output_dtype)?;
89                let gs = py_shape.iter().rev().map(|&d| d as u64).collect();
90                (dst, gs, data)
91            };
92
93        writer.add_tensor(gguf_name, gguf_shape, dst_dtype, tensor_data);
94        mapped += 1;
95        pb.inc(1);
96    }
97
98    pb.finish_with_message("tensors processed");
99
100    if !skipped.is_empty() {
101        eprintln!("\nWarning: {} tensor(s) skipped:", skipped.len());
102        for name in &skipped {
103            eprintln!("  {name}");
104        }
105    }
106    if fallback_count > 0 {
107        eprintln!("\nNote: {fallback_count} tensor(s) fell back to F32.");
108    }
109
110    println!("Writing {mapped} tensors to {} …", output_path.display());
111    let out_file =
112        File::create(output_path).with_context(|| format!("create {}", output_path.display()))?;
113    let mut buf_writer = BufWriter::new(out_file);
114    writer.write_to(&mut buf_writer)?;
115    println!("Done.");
116    Ok(())
117}
118
119fn write_metadata(writer: &mut GGUFWriter, config: &TiRexConfig) {
120    writer.add_metadata(
121        "general.architecture",
122        GGUFMetaValue::String("tirex".into()),
123    );
124    writer.add_metadata("general.name", GGUFMetaValue::String("TiRex".into()));
125    writer.add_metadata(
126        "tirex.patch_size",
127        GGUFMetaValue::Uint32(config.patch_size as u32),
128    );
129    writer.add_metadata(
130        "tirex.num_blocks",
131        GGUFMetaValue::Uint32(config.num_blocks as u32),
132    );
133    writer.add_metadata(
134        "tirex.embedding_dim",
135        GGUFMetaValue::Uint32(config.embedding_dim as u32),
136    );
137    writer.add_metadata(
138        "tirex.num_heads",
139        GGUFMetaValue::Uint32(config.num_heads as u32),
140    );
141    writer.add_metadata(
142        "tirex.input_ff_dim",
143        GGUFMetaValue::Uint32(config.input_ff_dim as u32),
144    );
145    writer.add_metadata(
146        "tirex.ffn_up_dim",
147        GGUFMetaValue::Uint32(config.ffn_up_dim as u32),
148    );
149    writer.add_metadata(
150        "tirex.train_ctx_len",
151        GGUFMetaValue::Uint32(config.train_ctx_len as u32),
152    );
153    writer.add_metadata(
154        "tirex.num_quantiles",
155        GGUFMetaValue::Uint32(config.num_quantiles as u32),
156    );
157}
158
159/// Permute bias from stored [NH, NG, DH] order to [NG, NH, DH] order.
160/// The sLSTM cell stores bias as [NH*NG*DH] = [NH, NG, DH] row-major,
161/// but during forward pass it permutes to [NG, NH, DH] before adding to gate projections.
162fn permute_bias(vals: &[f32], nh: usize, ng: usize, dh: usize) -> Vec<f32> {
163    let mut out = vec![0.0f32; vals.len()];
164    for h in 0..nh {
165        for g in 0..ng {
166            for d in 0..dh {
167                let src = h * (ng * dh) + g * dh + d;
168                let dst = g * (nh * dh) + h * dh + d;
169                out[dst] = vals[src];
170            }
171        }
172    }
173    out
174}
175
176fn f32_to_bytes(vals: &[f32]) -> Vec<u8> {
177    vals.iter().flat_map(|v| v.to_le_bytes()).collect()
178}
179
180fn apply_dtype(f32_vals: &[f32], dst: GGMLType) -> anyhow::Result<(GGMLType, Vec<u8>)> {
181    match dst {
182        GGMLType::F32 => Ok((GGMLType::F32, f32_to_bytes(f32_vals))),
183        GGMLType::F16 => {
184            let bytes: Vec<u8> = f32_vals
185                .iter()
186                .flat_map(|&v| f32_to_f16_bits(v).to_le_bytes())
187                .collect();
188            Ok((GGMLType::F16, bytes))
189        }
190        GGMLType::BF16 => {
191            let bytes: Vec<u8> = f32_vals
192                .iter()
193                .flat_map(|&v| ((v.to_bits() >> 16) as u16).to_le_bytes())
194                .collect();
195            Ok((GGMLType::BF16, bytes))
196        }
197        GGMLType::Q8_0 => Ok((GGMLType::Q8_0, quantize_q8_0(f32_vals)?)),
198    }
199}
200
201fn quantize_q8_0(values: &[f32]) -> anyhow::Result<Vec<u8>> {
202    const BLOCK: usize = 32;
203    if values.len() % BLOCK != 0 {
204        anyhow::bail!(
205            "Q8_0 requires count divisible by {BLOCK}, got {}",
206            values.len()
207        );
208    }
209    let n_blocks = values.len() / BLOCK;
210    let mut out = vec![0u8; n_blocks * 34];
211    for b in 0..n_blocks {
212        let blk = &values[b * BLOCK..(b + 1) * BLOCK];
213        let amax = blk.iter().copied().map(f32::abs).fold(0.0f32, f32::max);
214        let d = if amax == 0.0 { 0.0f32 } else { amax / 127.0 };
215        let d_inv = if d == 0.0 { 0.0f32 } else { 1.0 / d };
216        let base = b * 34;
217        out[base..base + 2].copy_from_slice(&f32_to_f16_bits(d).to_le_bytes());
218        for i in 0..BLOCK {
219            out[base + 2 + i] = (blk[i] * d_inv).round().clamp(-127.0, 127.0) as i8 as u8;
220        }
221    }
222    Ok(out)
223}
224
225fn f32_to_f16_bits(v: f32) -> u16 {
226    let bits = v.to_bits();
227    let sign = ((bits >> 16) & 0x8000) as u16;
228    let exp = ((bits >> 23) & 0xFF) as i32;
229    let mantissa = bits & 0x007F_FFFF;
230    if exp == 0xFF {
231        return sign | 0x7C00 | if mantissa != 0 { 0x0200 } else { 0 };
232    }
233    let new_exp = exp - 127 + 15;
234    if new_exp >= 31 {
235        return sign | 0x7C00;
236    }
237    if new_exp <= 0 {
238        if new_exp < -10 {
239            return sign;
240        }
241        let m = (mantissa | 0x0080_0000) >> (1 - new_exp);
242        return sign | (m >> 13) as u16;
243    }
244    sign | ((new_exp as u16) << 10) | (mantissa >> 13) as u16
245}
246
247fn f16_to_f32(bits: u16) -> f32 {
248    let sign = ((bits & 0x8000) as u32) << 16;
249    let exp = ((bits >> 10) & 0x1F) as i32;
250    let mantissa = (bits & 0x03FF) as u32;
251    let f32_bits = if exp == 0 {
252        if mantissa == 0 {
253            sign
254        } else {
255            let mut m = mantissa;
256            let mut e = 0i32;
257            while m & 0x0400 == 0 {
258                m <<= 1;
259                e += 1;
260            }
261            sign | ((127 - 15 - e + 1) as u32) << 23 | (m & 0x03FF) << 13
262        }
263    } else if exp == 31 {
264        sign | 0x7F80_0000 | (mantissa << 13)
265    } else {
266        sign | ((exp + 127 - 15) as u32) << 23 | (mantissa << 13)
267    };
268    f32::from_bits(f32_bits)
269}
zsfm_tirex/convert.rs

zsfm_tirex/
convert.rs
