1use std::fs::File;
2use std::io::BufWriter;
3use std::path::{Path, PathBuf};
4
5use anyhow::Context;
6use indicatif::{ProgressBar, ProgressStyle};
7use safetensors::SafeTensors;
8use safetensors::Dtype as StDtype;
9
10use zsfm_gguf::{GGMLType, GGUFMetaValue, GGUFWriter};
11
12use crate::config::MomentConfig;
13use crate::tensor_map::map_tensor_name;
14
15pub struct ConvertOptions {
16 pub output_dtype: GGMLType,
17}
18
19pub fn convert(
20 shard_paths: &[PathBuf],
21 config: &MomentConfig,
22 opts: &ConvertOptions,
23 output_path: &Path,
24) -> anyhow::Result<()> {
25 let mut writer = GGUFWriter::new();
26 write_metadata(&mut writer, config);
27
28 let shard_bytes: Vec<Vec<u8>> = shard_paths
29 .iter()
30 .map(|p| std::fs::read(p).with_context(|| format!("read {}", p.display())))
31 .collect::<anyhow::Result<_>>()?;
32 let shard_views: Vec<SafeTensors> = shard_bytes
33 .iter()
34 .map(|b| SafeTensors::deserialize(b).context("deserialize shard"))
35 .collect::<anyhow::Result<_>>()?;
36
37 let total: usize = shard_views.iter().map(|s| s.len()).sum();
38 let pb = ProgressBar::new(total as u64);
39 pb.set_style(
40 ProgressStyle::with_template(
41 "{spinner:.green} [{elapsed_precise}] [{bar:40.cyan/blue}] {pos}/{len} {msg}",
42 )
43 .unwrap()
44 .progress_chars("=>-"),
45 );
46
47 let mut mapped = 0usize;
48 let mut skipped: Vec<String> = Vec::new();
49 let mut fallback_count = 0usize;
50
51 for shard in &shard_views {
52 for (hf_name, tensor_view) in shard.tensors() {
53 pb.set_message(hf_name.to_string());
54
55 let gguf_name = match map_tensor_name(&hf_name) {
56 Some(n) => n,
57 None => {
58 skipped.push(hf_name.to_string());
59 pb.inc(1);
60 continue;
61 }
62 };
63
64 let src_dtype = ggml_type_from_st(tensor_view.dtype())
65 .with_context(|| format!("tensor {hf_name}: unsupported dtype {:?}", tensor_view.dtype()))?;
66
67 let raw_data = tensor_view.data();
68 let py_shape = tensor_view.shape();
69 let n_elems: usize = py_shape.iter().product();
70 let innermost = py_shape.last().copied().unwrap_or(1);
71 let _outermost = py_shape.first().copied().unwrap_or(1);
72
73 let (dst_dtype, gguf_shape, tensor_data) =
74 if opts.output_dtype == GGMLType::Q8_0 && (innermost % 32 != 0 || n_elems % 32 != 0) {
75 fallback_count += 1;
76 let data = cast_data(raw_data, src_dtype, GGMLType::F32)
77 .with_context(|| format!("tensor {hf_name}: cast failed"))?;
78 let gs = py_shape.iter().rev().map(|&d| d as u64).collect();
79 (GGMLType::F32, gs, data)
80 } else {
81 let dst = opts.output_dtype;
82 let data = cast_data(raw_data, src_dtype, dst)
83 .with_context(|| format!("tensor {hf_name}: cast failed"))?;
84 let gs = py_shape.iter().rev().map(|&d| d as u64).collect();
85 (dst, gs, data)
86 };
87
88 writer.add_tensor(gguf_name, gguf_shape, dst_dtype, tensor_data);
89 mapped += 1;
90 pb.inc(1);
91 }
92 } 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: &MomentConfig) {
114 writer.add_metadata("general.architecture", GGUFMetaValue::String("moment".into()));
115 writer.add_metadata("general.name", GGUFMetaValue::String("MOMENT-1-large".into()));
116 writer.add_metadata("moment.d_model", GGUFMetaValue::Uint32(config.d_model as u32));
117 writer.add_metadata("moment.n_layers", GGUFMetaValue::Uint32(config.n_layers as u32));
118 writer.add_metadata("moment.n_heads", GGUFMetaValue::Uint32(config.n_heads as u32));
119 writer.add_metadata("moment.head_dim", GGUFMetaValue::Uint32(config.head_dim as u32));
120 writer.add_metadata("moment.d_ff", GGUFMetaValue::Uint32(config.d_ff as u32));
121 writer.add_metadata("moment.seq_len", GGUFMetaValue::Uint32(config.seq_len as u32));
122 writer.add_metadata("moment.patch_len", GGUFMetaValue::Uint32(config.patch_len as u32));
123 writer.add_metadata("moment.patch_stride", GGUFMetaValue::Uint32(config.patch_stride as u32));
124 writer.add_metadata("moment.num_patches", GGUFMetaValue::Uint32(config.num_patches as u32));
125 writer.add_metadata("moment.layer_norm_eps", GGUFMetaValue::Float64(config.layer_norm_eps));
126}
127
128fn ggml_type_from_st(dtype: StDtype) -> anyhow::Result<GGMLType> {
129 match dtype {
130 StDtype::F32 => Ok(GGMLType::F32),
131 StDtype::F16 => Ok(GGMLType::F16),
132 StDtype::BF16 => Ok(GGMLType::BF16),
133 other => anyhow::bail!("unsupported safetensors dtype: {other:?}"),
134 }
135}
136
137fn cast_data(data: &[u8], src: GGMLType, dst: GGMLType) -> anyhow::Result<Vec<u8>> {
138 if src == dst { return Ok(data.to_vec()); }
139 if dst == GGMLType::Q8_0 {
140 let f32_values = decode_to_f32(data, src)?;
141 return quantize_q8_0(&f32_values);
142 }
143 match (src, dst) {
144 (GGMLType::F32, GGMLType::F16) => {
145 let vals = parse_f32_le(data)?;
146 let mut out = Vec::with_capacity(vals.len() * 2);
147 for v in vals { out.extend_from_slice(&f32_to_f16_bits(v).to_le_bytes()); }
148 Ok(out)
149 }
150 (GGMLType::F32, GGMLType::BF16) => {
151 let vals = parse_f32_le(data)?;
152 let mut out = Vec::with_capacity(vals.len() * 2);
153 for v in vals { out.extend_from_slice(&((v.to_bits() >> 16) as u16).to_le_bytes()); }
154 Ok(out)
155 }
156 (GGMLType::F16, GGMLType::BF16) => {
157 let mut out = Vec::with_capacity(data.len());
158 for c in data.chunks_exact(2) {
159 let f32_val = f16_to_f32(u16::from_le_bytes([c[0], c[1]]));
160 out.extend_from_slice(&((f32_val.to_bits() >> 16) as u16).to_le_bytes());
161 }
162 Ok(out)
163 }
164 (GGMLType::BF16, GGMLType::F32) => {
165 let mut out = Vec::with_capacity(data.len() * 2);
166 for c in data.chunks_exact(2) {
167 let bits = (u16::from_le_bytes([c[0], c[1]]) as u32) << 16;
168 out.extend_from_slice(&bits.to_le_bytes());
169 }
170 Ok(out)
171 }
172 (GGMLType::BF16, GGMLType::F16) => {
173 let mut out = Vec::with_capacity(data.len());
174 for c in data.chunks_exact(2) {
175 let bits = (u16::from_le_bytes([c[0], c[1]]) as u32) << 16;
176 out.extend_from_slice(&f32_to_f16_bits(f32::from_bits(bits)).to_le_bytes());
177 }
178 Ok(out)
179 }
180 (GGMLType::F16, GGMLType::F32) => {
181 let mut out = Vec::with_capacity(data.len() * 2);
182 for c in data.chunks_exact(2) {
183 out.extend_from_slice(&f16_to_f32(u16::from_le_bytes([c[0], c[1]])).to_bits().to_le_bytes());
184 }
185 Ok(out)
186 }
187 _ => anyhow::bail!("unsupported cast: {src:?} → {dst:?}"),
188 }
189}
190
191fn decode_to_f32(data: &[u8], src: GGMLType) -> anyhow::Result<Vec<f32>> {
192 match src {
193 GGMLType::F32 => parse_f32_le(data),
194 GGMLType::F16 => data.chunks_exact(2)
195 .map(|c| Ok(f16_to_f32(u16::from_le_bytes([c[0], c[1]]))))
196 .collect(),
197 GGMLType::BF16 => data.chunks_exact(2)
198 .map(|c| Ok(f32::from_bits((u16::from_le_bytes([c[0], c[1]]) as u32) << 16)))
199 .collect(),
200 GGMLType::Q8_0 => anyhow::bail!("Q8_0 as source not supported"),
201 }
202}
203
204fn quantize_q8_0(values: &[f32]) -> anyhow::Result<Vec<u8>> {
205 const BLOCK: usize = 32;
206 if values.len() % BLOCK != 0 {
207 anyhow::bail!("Q8_0 requires count divisible by {BLOCK}");
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 parse_f32_le(data: &[u8]) -> anyhow::Result<Vec<f32>> {
226 if data.len() % 4 != 0 { anyhow::bail!("f32 data length not divisible by 4"); }
227 Ok(data.chunks_exact(4).map(|c| f32::from_le_bytes([c[0], c[1], c[2], c[3]])).collect())
228}
229
230fn f32_to_f16_bits(v: f32) -> u16 {
231 let bits = v.to_bits();
232 let sign = ((bits >> 16) & 0x8000) as u16;
233 let exp = ((bits >> 23) & 0xFF) as i32;
234 let mantissa = bits & 0x007F_FFFF;
235 if exp == 0xFF { return sign | 0x7C00 | if mantissa != 0 { 0x0200 } else { 0 }; }
236 let new_exp = exp - 127 + 15;
237 if new_exp >= 31 { return sign | 0x7C00; }
238 if new_exp <= 0 {
239 if new_exp < -10 { return sign; }
240 let m = (mantissa | 0x0080_0000) >> (1 - new_exp);
241 return sign | (m >> 13) as u16;
242 }
243 sign | ((new_exp as u16) << 10) | (mantissa >> 13) as u16
244}
245
246fn f16_to_f32(bits: u16) -> f32 {
247 let sign = ((bits & 0x8000) as u32) << 16;
248 let exp = ((bits >> 10) & 0x1F) as i32;
249 let mantissa = (bits & 0x03FF) as u32;
250 let f32_bits = if exp == 0 {
251 if mantissa == 0 { sign }
252 else {
253 let mut m = mantissa; let mut e = 0i32;
254 while m & 0x0400 == 0 { m <<= 1; e += 1; }
255 sign | ((127 - 15 - e + 1) as u32) << 23 | (m & 0x03FF) << 13
256 }
257 } else if exp == 31 { sign | 0x7F80_0000 | (mantissa << 13) }
258 else { sign | ((exp + 127 - 15) as u32) << 23 | (mantissa << 13) };
259 f32::from_bits(f32_bits)
260}