zsfm_checkpoint/
cast.rs

1//! Canonical dtype casting for GGUF conversion: {F32, F16, BF16} sources to any
2//! writable [`GGMLType`], including Q8_0 block quantization. Bit-for-bit the same
3//! math as the per-model converters (which predate this crate and keep their own
4//! verified copies).
5
6use zsfm_gguf::GGMLType;
7
8/// Source dtype of raw checkpoint tensor bytes.
9#[derive(Debug, Clone, Copy, PartialEq, Eq)]
10pub enum SrcDtype {
11    F32,
12    F16,
13    BF16,
14}
15
16impl SrcDtype {
17    pub fn bytes_per_elem(self) -> usize {
18        match self {
19            SrcDtype::F32 => 4,
20            SrcDtype::F16 | SrcDtype::BF16 => 2,
21        }
22    }
23
24    pub fn name(self) -> &'static str {
25        match self {
26            SrcDtype::F32 => "F32",
27            SrcDtype::F16 => "F16",
28            SrcDtype::BF16 => "BF16",
29        }
30    }
31}
32
33/// Cast raw little-endian tensor bytes from `src` to `dst`.
34///
35/// Same-dtype casts are byte-for-byte passthrough. F32→BF16 truncates the
36/// mantissa (matching the existing converters); everything routed through F32
37/// uses exact widening.
38pub fn cast_data(data: &[u8], src: SrcDtype, dst: GGMLType) -> anyhow::Result<Vec<u8>> {
39    // Passthrough when the representation already matches.
40    match (src, dst) {
41        (SrcDtype::F32, GGMLType::F32)
42        | (SrcDtype::F16, GGMLType::F16)
43        | (SrcDtype::BF16, GGMLType::BF16) => return Ok(data.to_vec()),
44        _ => {}
45    }
46
47    let f32_values = decode_to_f32(data, src)?;
48    match dst {
49        GGMLType::F32 => Ok(f32_to_bytes(&f32_values)),
50        GGMLType::F16 => Ok(f32_values
51            .iter()
52            .flat_map(|&v| f32_to_f16_bits(v).to_le_bytes())
53            .collect()),
54        GGMLType::BF16 => Ok(f32_values
55            .iter()
56            .flat_map(|&v| ((v.to_bits() >> 16) as u16).to_le_bytes())
57            .collect()),
58        GGMLType::Q8_0 => quantize_q8_0(&f32_values),
59    }
60}
61
62/// Decode raw bytes of `src` dtype into f32 values (exact for all three sources).
63pub fn decode_to_f32(data: &[u8], src: SrcDtype) -> anyhow::Result<Vec<f32>> {
64    match src {
65        SrcDtype::F32 => parse_f32_le(data),
66        SrcDtype::F16 => {
67            if data.len() % 2 != 0 {
68                anyhow::bail!("f16 data length not divisible by 2");
69            }
70            Ok(data
71                .chunks_exact(2)
72                .map(|c| f16_to_f32(u16::from_le_bytes([c[0], c[1]])))
73                .collect())
74        }
75        SrcDtype::BF16 => {
76            if data.len() % 2 != 0 {
77                anyhow::bail!("bf16 data length not divisible by 2");
78            }
79            Ok(data
80                .chunks_exact(2)
81                .map(|c| f32::from_bits((u16::from_le_bytes([c[0], c[1]]) as u32) << 16))
82                .collect())
83        }
84    }
85}
86
87pub fn f32_to_bytes(vals: &[f32]) -> Vec<u8> {
88    vals.iter().flat_map(|v| v.to_le_bytes()).collect()
89}
90
91pub fn quantize_q8_0(values: &[f32]) -> anyhow::Result<Vec<u8>> {
92    const BLOCK: usize = 32;
93    if values.len() % BLOCK != 0 {
94        anyhow::bail!(
95            "Q8_0 requires element count divisible by {BLOCK}, got {}",
96            values.len()
97        );
98    }
99    let n_blocks = values.len() / BLOCK;
100    let mut out = vec![0u8; n_blocks * 34];
101    for b in 0..n_blocks {
102        let blk = &values[b * BLOCK..(b + 1) * BLOCK];
103        let amax = blk.iter().copied().map(f32::abs).fold(0.0f32, f32::max);
104        let d = if amax == 0.0 { 0.0f32 } else { amax / 127.0 };
105        let d_inv = if d == 0.0 { 0.0f32 } else { 1.0 / d };
106        let base = b * 34;
107        out[base..base + 2].copy_from_slice(&f32_to_f16_bits(d).to_le_bytes());
108        for i in 0..BLOCK {
109            out[base + 2 + i] = (blk[i] * d_inv).round().clamp(-127.0, 127.0) as i8 as u8;
110        }
111    }
112    Ok(out)
113}
114
115fn parse_f32_le(data: &[u8]) -> anyhow::Result<Vec<f32>> {
116    if data.len() % 4 != 0 {
117        anyhow::bail!("f32 data length not divisible by 4");
118    }
119    Ok(data
120        .chunks_exact(4)
121        .map(|c| f32::from_le_bytes([c[0], c[1], c[2], c[3]]))
122        .collect())
123}
124
125pub fn f32_to_f16_bits(v: f32) -> u16 {
126    let bits = v.to_bits();
127    let sign = ((bits >> 16) & 0x8000) as u16;
128    let exp = ((bits >> 23) & 0xFF) as i32;
129    let mantissa = bits & 0x007F_FFFF;
130    if exp == 0xFF {
131        return sign | 0x7C00 | if mantissa != 0 { 0x0200 } else { 0 };
132    }
133    let new_exp = exp - 127 + 15;
134    if new_exp >= 31 {
135        return sign | 0x7C00;
136    }
137    if new_exp <= 0 {
138        if new_exp < -10 {
139            return sign;
140        }
141        let m = (mantissa | 0x0080_0000) >> (1 - new_exp);
142        return sign | (m >> 13) as u16;
143    }
144    sign | ((new_exp as u16) << 10) | (mantissa >> 13) as u16
145}
146
147pub fn f16_to_f32(bits: u16) -> f32 {
148    let sign = ((bits & 0x8000) as u32) << 16;
149    let exp = ((bits >> 10) & 0x1F) as i32;
150    let mantissa = (bits & 0x03FF) as u32;
151    let f32_bits = if exp == 0 {
152        if mantissa == 0 {
153            sign
154        } else {
155            let mut m = mantissa;
156            let mut e = 0i32;
157            while m & 0x0400 == 0 {
158                m <<= 1;
159                e += 1;
160            }
161            sign | ((127 - 15 - e + 1) as u32) << 23 | (m & 0x03FF) << 13
162        }
163    } else if exp == 31 {
164        sign | 0x7F80_0000 | (mantissa << 13)
165    } else {
166        sign | ((exp + 127 - 15) as u32) << 23 | (mantissa << 13)
167    };
168    f32::from_bits(f32_bits)
169}
170
171#[cfg(test)]
172mod tests {
173    use super::*;
174
175    #[test]
176    fn f16_roundtrip_exact_values() {
177        for v in [0.0f32, 1.0, -1.0, 0.5, 65504.0, -65504.0, 0.099975586] {
178            assert_eq!(f16_to_f32(f32_to_f16_bits(v)), v, "roundtrip {v}");
179        }
180    }
181
182    #[test]
183    fn passthrough_same_dtype() {
184        let bytes: Vec<u8> = (0..16).collect();
185        assert_eq!(cast_data(&bytes, SrcDtype::F32, GGMLType::F32).unwrap(), bytes);
186        assert_eq!(cast_data(&bytes, SrcDtype::F16, GGMLType::F16).unwrap(), bytes);
187        assert_eq!(cast_data(&bytes, SrcDtype::BF16, GGMLType::BF16).unwrap(), bytes);
188    }
189
190    #[test]
191    fn f32_to_bf16_truncates() {
192        let v = 1.2345678f32;
193        let out = cast_data(&v.to_le_bytes(), SrcDtype::F32, GGMLType::BF16).unwrap();
194        let bits = u16::from_le_bytes([out[0], out[1]]);
195        assert_eq!(bits, (v.to_bits() >> 16) as u16);
196    }
197
198    #[test]
199    fn bf16_to_f32_exact_widening() {
200        let bits: u16 = 0x3FA0; // 1.25 in bf16
201        let out = cast_data(&bits.to_le_bytes(), SrcDtype::BF16, GGMLType::F32).unwrap();
202        let v = f32::from_le_bytes([out[0], out[1], out[2], out[3]]);
203        assert_eq!(v, 1.25);
204    }
205
206    #[test]
207    fn q8_0_block_layout() {
208        let vals: Vec<f32> = (0..32).map(|i| i as f32).collect();
209        let out = quantize_q8_0(&vals).unwrap();
210        assert_eq!(out.len(), 34);
211        let d = f16_to_f32(u16::from_le_bytes([out[0], out[1]]));
212        assert!((d - 31.0 / 127.0).abs() < 1e-3);
213        assert_eq!(out[2] as i8, 0); // 0.0 quantizes to 0
214        assert_eq!(out[33] as i8, 127); // amax quantizes to 127
215    }
216
217    #[test]
218    fn q8_0_rejects_partial_block() {
219        assert!(quantize_q8_0(&[1.0f32; 31]).is_err());
220    }
221}
zsfm_checkpoint/cast.rs

zsfm_checkpoint/
cast.rs
