zsfm_chronos/
convert.rs

1use std::fs::File;
2use std::io::BufWriter;
3use std::path::Path;
4
5use anyhow::Context;
6use indicatif::{ProgressBar, ProgressStyle};
7use safetensors::SafeTensors;
8use safetensors::Dtype as StDtype;
9
10use zsfm_gguf::{GGMLType, GGUFMetaValue, GGUFWriter};
11use zsfm_hub::ModelFiles;
12
13use crate::config::Chronos2Config;
14use crate::tensor_map::map_tensor_name;
15
16/// Options that control how tensors are stored in the output GGUF.
17pub struct ConvertOptions {
18    pub output_dtype: GGMLType,
19}
20
21/// Read `files`, convert every tensor, and write a GGUF file to `output_path`.
22pub fn convert(
23    model_id: &str,
24    files: &ModelFiles,
25    config: &Chronos2Config,
26    opts: &ConvertOptions,
27    output_path: &Path,
28) -> anyhow::Result<()> {
29    let mut writer = GGUFWriter::new();
30
31    write_metadata(&mut writer, model_id, config);
32
33    let shard_bytes = load_shard_bytes(&files.safetensors_shards)?;
34    let shard_views: Vec<SafeTensors> = shard_bytes
35        .iter()
36        .map(|b| SafeTensors::deserialize(b).context("deserialize shard"))
37        .collect::<anyhow::Result<_>>()?;
38
39    let total_tensors: usize = shard_views.iter().map(|s| s.len()).sum();
40    println!(
41        "Found {} tensors across {} shard(s).",
42        total_tensors,
43        shard_views.len()
44    );
45
46    let pb = ProgressBar::new(total_tensors as u64);
47    pb.set_style(
48        ProgressStyle::with_template(
49            "{spinner:.green} [{elapsed_precise}] [{bar:40.cyan/blue}] {pos}/{len} {msg}",
50        )
51        .unwrap()
52        .progress_chars("=>-"),
53    );
54
55    let mut mapped = 0usize;
56    let mut skipped: Vec<String> = Vec::new();
57    let mut fallback_count = 0usize;
58
59    for shard in &shard_views {
60        for (hf_name, tensor_view) in shard.tensors() {
61            pb.set_message(hf_name.to_string());
62
63            let gguf_name = match map_tensor_name(&hf_name) {
64                Some(n) => n,
65                None => {
66                    skipped.push(hf_name.to_string());
67                    pb.inc(1);
68                    continue;
69                }
70            };
71
72            let src_dtype = ggml_type_from_st(tensor_view.dtype())
73                .with_context(|| format!("tensor {hf_name}: unsupported dtype {:?}", tensor_view.dtype()))?;
74
75            let raw_data = tensor_view.data();
76            let py_shape = tensor_view.shape();
77            let n_elems: usize = py_shape.iter().product();
78            let innermost = py_shape.last().copied().unwrap_or(1);
79            let outermost = py_shape.first().copied().unwrap_or(1);
80
81            // Q8_0: candle requires the innermost dim (last Python dim, first GGUF dim) to be
82            // divisible by 32. When it isn't for a 2D weight whose outermost dim is aligned,
83            // transpose data and store GGUF shape in Python order (not reversed). Candle
84            // reverses the shape at load time, placing the aligned dim innermost. The inference
85            // loader detects the transposition by checking dim(0) against the expected d_out.
86            let (dst_dtype, gguf_shape, tensor_data) =
87                if opts.output_dtype == GGMLType::Q8_0 && (innermost % 32 != 0 || n_elems % 32 != 0) {
88                    fallback_count += 1;
89                    let data = cast_data(raw_data, src_dtype, GGMLType::F32)
90                        .with_context(|| format!("tensor {hf_name}: cast failed"))?;
91                    let gs = py_shape.iter().rev().map(|&d| d as u64).collect();
92                    (GGMLType::F32, gs, data)
93                } else {
94                    let dst = opts.output_dtype;
95                    let data = cast_data(raw_data, src_dtype, dst)
96                        .with_context(|| format!("tensor {hf_name}: cast failed"))?;
97                    let gs = py_shape.iter().rev().map(|&d| d as u64).collect();
98                    (dst, gs, data)
99                };
100
101            writer.add_tensor(gguf_name, gguf_shape, dst_dtype, tensor_data);
102            mapped += 1;
103            pb.inc(1);
104        }
105    }
106
107    pb.finish_with_message("tensors processed");
108
109    if !skipped.is_empty() {
110        eprintln!(
111            "\nWarning: {} tensor(s) had unrecognised names and were skipped:",
112            skipped.len()
113        );
114        for name in &skipped {
115            eprintln!("  {name}");
116        }
117        eprintln!("Update tensor_map.rs to include these if needed.");
118    }
119
120    if fallback_count > 0 {
121        eprintln!(
122            "\nNote: {fallback_count} tensor(s) fell back to F32 (scalars/biases too \
123             small for Q8_0 blocks)."
124        );
125    }
126    println!("Writing {mapped} tensors to {} …", output_path.display());
127    let out_file = File::create(output_path)
128        .with_context(|| format!("create output file {}", output_path.display()))?;
129    let mut buf_writer = BufWriter::new(out_file);
130    writer.write_to(&mut buf_writer)?;
131    println!("Done.");
132
133    Ok(())
134}
135
136fn ggml_type_from_st(dtype: StDtype) -> anyhow::Result<GGMLType> {
137    match dtype {
138        StDtype::F32 => Ok(GGMLType::F32),
139        StDtype::F16 => Ok(GGMLType::F16),
140        StDtype::BF16 => Ok(GGMLType::BF16),
141        other => anyhow::bail!("unsupported safetensors dtype: {other:?}"),
142    }
143}
144
145/// Build GGUF metadata section from parsed config.
146fn write_metadata(writer: &mut GGUFWriter, model_id: &str, config: &Chronos2Config) {
147    let cc = &config.chronos_config;
148
149    writer.add_metadata("general.architecture", GGUFMetaValue::String("chronos2".into()));
150    writer.add_metadata("general.name", GGUFMetaValue::String(model_id.into()));
151    writer.add_metadata("chronos2.block_count",        GGUFMetaValue::Uint32(config.num_layers));
152    writer.add_metadata("chronos2.embedding_length",   GGUFMetaValue::Uint32(config.d_model));
153    writer.add_metadata("chronos2.feed_forward_length",GGUFMetaValue::Uint32(config.d_ff));
154    writer.add_metadata("chronos2.attention.head_count", GGUFMetaValue::Uint32(config.num_heads));
155    writer.add_metadata("chronos2.attention.head_dim",   GGUFMetaValue::Uint32(config.d_kv));
156    writer.add_metadata("chronos2.rope_theta",         GGUFMetaValue::Float64(config.rope_theta));
157    writer.add_metadata("chronos2.layer_norm_epsilon", GGUFMetaValue::Float64(config.layer_norm_epsilon));
158    writer.add_metadata("chronos2.context_length",     GGUFMetaValue::Uint32(cc.context_length));
159    writer.add_metadata("chronos2.patch_size",         GGUFMetaValue::Uint32(cc.input_patch_size));
160    writer.add_metadata("chronos2.patch_stride",       GGUFMetaValue::Uint32(cc.input_patch_stride));
161    writer.add_metadata("chronos2.quantile_count",     GGUFMetaValue::Uint32(cc.quantiles.len() as u32));
162    writer.add_metadata("chronos2.quantiles",          GGUFMetaValue::ArrayFloat32(cc.quantiles.clone()));
163    writer.add_metadata("chronos2.use_reg_token",      GGUFMetaValue::Bool(cc.use_reg_token));
164    writer.add_metadata("chronos2.use_arcsinh",        GGUFMetaValue::Bool(cc.use_arcsinh));
165    writer.add_metadata("chronos2.time_encoding_scale",GGUFMetaValue::Uint32(config.time_encoding_scale()));
166    writer.add_metadata("chronos2.dense_act_fn",       GGUFMetaValue::String(config.dense_act_fn().into()));
167}
168
169fn load_shard_bytes(shards: &[std::path::PathBuf]) -> anyhow::Result<Vec<Vec<u8>>> {
170    shards
171        .iter()
172        .map(|p| std::fs::read(p).with_context(|| format!("read shard {}", p.display())))
173        .collect()
174}
175
176/// Cast tensor bytes from `src_dtype` to `dst_dtype`.
177fn cast_data(data: &[u8], src: GGMLType, dst: GGMLType) -> anyhow::Result<Vec<u8>> {
178    if src == dst {
179        return Ok(data.to_vec());
180    }
181    if dst == GGMLType::Q8_0 {
182        let f32_values = decode_to_f32(data, src)?;
183        return quantize_q8_0(&f32_values);
184    }
185    match (src, dst) {
186        (GGMLType::F32, GGMLType::F16) => {
187            let f32_values = parse_f32_le(data)?;
188            let mut out = Vec::with_capacity(f32_values.len() * 2);
189            for v in f32_values {
190                let bits = f32_to_f16_bits(v);
191                out.extend_from_slice(&bits.to_le_bytes());
192            }
193            Ok(out)
194        }
195        (GGMLType::F32, GGMLType::BF16) => {
196            let f32_values = parse_f32_le(data)?;
197            let mut out = Vec::with_capacity(f32_values.len() * 2);
198            for v in f32_values {
199                let bits = (v.to_bits() >> 16) as u16;
200                out.extend_from_slice(&bits.to_le_bytes());
201            }
202            Ok(out)
203        }
204        (GGMLType::F16, GGMLType::BF16) => {
205            let mut out = Vec::with_capacity(data.len());
206            for chunk in data.chunks_exact(2) {
207                let f16_bits = u16::from_le_bytes([chunk[0], chunk[1]]);
208                let f32_val = f16_to_f32(f16_bits);
209                let bits = (f32_val.to_bits() >> 16) as u16;
210                out.extend_from_slice(&bits.to_le_bytes());
211            }
212            Ok(out)
213        }
214        (GGMLType::BF16, GGMLType::F32) => {
215            let mut out = Vec::with_capacity(data.len() * 2);
216            for chunk in data.chunks_exact(2) {
217                let bf16_bits = u16::from_le_bytes([chunk[0], chunk[1]]);
218                let f32_bits = (bf16_bits as u32) << 16;
219                out.extend_from_slice(&f32_bits.to_le_bytes());
220            }
221            Ok(out)
222        }
223        (GGMLType::BF16, GGMLType::F16) => {
224            let mut out = Vec::with_capacity(data.len());
225            for chunk in data.chunks_exact(2) {
226                let bf16_bits = u16::from_le_bytes([chunk[0], chunk[1]]);
227                let f32_bits = (bf16_bits as u32) << 16;
228                let f32_val = f32::from_bits(f32_bits);
229                let f16_bits = f32_to_f16_bits(f32_val);
230                out.extend_from_slice(&f16_bits.to_le_bytes());
231            }
232            Ok(out)
233        }
234        (GGMLType::F16, GGMLType::F32) => {
235            let mut out = Vec::with_capacity(data.len() * 2);
236            for chunk in data.chunks_exact(2) {
237                let f16_bits = u16::from_le_bytes([chunk[0], chunk[1]]);
238                let f32_val = f16_to_f32(f16_bits);
239                out.extend_from_slice(&f32_val.to_bits().to_le_bytes());
240            }
241            Ok(out)
242        }
243        _ => anyhow::bail!("unsupported cast: {src:?} → {dst:?}"),
244    }
245}
246
247fn decode_to_f32(data: &[u8], src: GGMLType) -> anyhow::Result<Vec<f32>> {
248    match src {
249        GGMLType::F32 => parse_f32_le(data),
250        GGMLType::F16 => data
251            .chunks_exact(2)
252            .map(|c| Ok(f16_to_f32(u16::from_le_bytes([c[0], c[1]]))))
253            .collect(),
254        GGMLType::BF16 => data
255            .chunks_exact(2)
256            .map(|c| {
257                let bf16_bits = u16::from_le_bytes([c[0], c[1]]);
258                Ok(f32::from_bits((bf16_bits as u32) << 16))
259            })
260            .collect(),
261        GGMLType::Q8_0 => anyhow::bail!("Q8_0 → Q8_0 re-quantization not supported as source"),
262    }
263}
264
265fn transpose_f32(data: &[f32], n_rows: usize, n_cols: usize) -> Vec<f32> {
266    let mut out = vec![0.0f32; n_rows * n_cols];
267    for r in 0..n_rows {
268        for c in 0..n_cols {
269            out[c * n_rows + r] = data[r * n_cols + c];
270        }
271    }
272    out
273}
274
275/// Quantize a slice of f32 values to GGML Q8_0 block format.
276/// Each block: `[f16 scale (2 bytes)][32 × i8 (32 bytes)]` = 34 bytes.
277fn quantize_q8_0(values: &[f32]) -> anyhow::Result<Vec<u8>> {
278    const BLOCK: usize = 32;
279    if values.len() % BLOCK != 0 {
280        anyhow::bail!(
281            "Q8_0 requires element count divisible by {BLOCK}, got {}",
282            values.len()
283        );
284    }
285    let n_blocks = values.len() / BLOCK;
286    let mut out = vec![0u8; n_blocks * 34];
287
288    for b in 0..n_blocks {
289        let blk = &values[b * BLOCK..(b + 1) * BLOCK];
290        let amax = blk.iter().copied().map(f32::abs).fold(0.0f32, f32::max);
291        let d = if amax == 0.0 { 0.0f32 } else { amax / 127.0 };
292        let d_inv = if d == 0.0 { 0.0f32 } else { 1.0 / d };
293
294        let base = b * 34;
295        let d_f16 = f32_to_f16_bits(d);
296        out[base..base + 2].copy_from_slice(&d_f16.to_le_bytes());
297        for i in 0..BLOCK {
298            let q = (blk[i] * d_inv).round().clamp(-127.0, 127.0) as i8;
299            out[base + 2 + i] = q as u8;
300        }
301    }
302    Ok(out)
303}
304
305fn parse_f32_le(data: &[u8]) -> anyhow::Result<Vec<f32>> {
306    if data.len() % 4 != 0 {
307        anyhow::bail!("f32 data length not divisible by 4");
308    }
309    Ok(data
310        .chunks_exact(4)
311        .map(|c| f32::from_le_bytes([c[0], c[1], c[2], c[3]]))
312        .collect())
313}
314
315fn f32_to_f16_bits(v: f32) -> u16 {
316    let bits = v.to_bits();
317    let sign = ((bits >> 16) & 0x8000) as u16;
318    let exp = ((bits >> 23) & 0xFF) as i32;
319    let mantissa = bits & 0x007F_FFFF;
320
321    if exp == 0xFF {
322        return sign | 0x7C00 | if mantissa != 0 { 0x0200 } else { 0 };
323    }
324    let new_exp = exp - 127 + 15;
325    if new_exp >= 31 {
326        return sign | 0x7C00;
327    }
328    if new_exp <= 0 {
329        if new_exp < -10 {
330            return sign;
331        }
332        let m = (mantissa | 0x0080_0000) >> (1 - new_exp);
333        return sign | (m >> 13) as u16;
334    }
335    sign | ((new_exp as u16) << 10) | (mantissa >> 13) as u16
336}
337
338fn f16_to_f32(bits: u16) -> f32 {
339    let sign = ((bits & 0x8000) as u32) << 16;
340    let exp = ((bits >> 10) & 0x1F) as i32;
341    let mantissa = (bits & 0x03FF) as u32;
342
343    let f32_bits = if exp == 0 {
344        if mantissa == 0 {
345            sign
346        } else {
347            let mut m = mantissa;
348            let mut e = 0i32;
349            while m & 0x0400 == 0 {
350                m <<= 1;
351                e += 1;
352            }
353            sign | ((127 - 15 - e + 1) as u32) << 23 | (m & 0x03FF) << 13
354        }
355    } else if exp == 31 {
356        sign | 0x7F80_0000 | (mantissa << 13)
357    } else {
358        sign | ((exp + 127 - 15) as u32) << 23 | (mantissa << 13)
359    };
360    f32::from_bits(f32_bits)
361}
zsfm_chronos/convert.rs

zsfm_chronos/
convert.rs
