1use super::*;
2
3pub(super) fn open(dir: &Path) -> Result<Checkpoint> {
4 let manifest = Manifest::load(dir)?;
5 let config_path = if dir.join("config.json").is_file() {
6 dir.join("config.json")
7 } else {
8 dir.parent()
9 .context("packed configuration missing")?
10 .join("config.json")
11 };
12 let config: serde_json::Value = serde_json::from_slice(&std::fs::read(config_path)?)?;
13 let paths = ["dense.bin", "experts.bin", "ngram.bin"].map(|name| dir.join(name));
14 let objects = MappedObjects::open(paths.iter().map(|path| path.as_path()))?;
15 let mut tensors = dense(&manifest)?;
16
17 experts(&manifest, &mut tensors)?;
18
19 let table_id = TensorId(tensors.len());
20 let table = ngram(&manifest)?;
21
22 tensors.push(table);
23
24 let text = config.get("text_config").unwrap_or(&config);
25 let n = &manifest.ngram;
26 let ngram = Some(NgramTable {
27 hashing: NgramHash::Qwen4Exp {
28 ngram_size: text["ngram_size"].as_u64().unwrap_or(3),
29 heads_per_ngram: text["heads_per_ngram"].as_u64().unwrap_or(8),
30 head_offsets: n.head_offsets.clone(),
31 head_vocab_sizes: n.head_vocab_sizes.clone(),
32 layer_multipliers: n.layer_multipliers.clone(),
33 },
34 shards: vec![TableShard {
35 first_row: 0,
36 rows: n.rows,
37 tensor: table_id,
38 }],
39 });
40 let description = ModelDescription {
41 schema_version: 1,
42 format: CheckpointFormat {
43 container: ContainerFormat::CherenkovPacked {
44 version: manifest.version,
45 },
46 conventions: CheckpointConventions::Qwen4ExpCherenkov,
47 },
48 architecture: architecture(&config),
49 metadata: serde_json::json!({"config": config, "manifest": manifest}),
50 tensors,
51 ngram,
52 };
53 let checkpoint = Checkpoint {
54 description,
55 objects,
56 };
57
58 for tensor in &checkpoint.description.tensors {
59 tensor.validate()?;
60
61 for data in tensor.encoding.data() {
62 checkpoint.map(data)?;
63 }
64 }
65
66 Ok(checkpoint)
67}
68
69fn dense(manifest: &Manifest) -> Result<Vec<Tensor>> {
70 let mut tensors = Vec::new();
71
72 for entry in &manifest.dense {
73 if entry.name.ends_with(".scales") || entry.name.ends_with(".biases") {
74 continue;
75 }
76
77 let codes = dense_part(entry)?;
78 let role = qwen_role(&entry.name, entry.shape.len());
79
80 if codes.dtype != Dtype::U32 || !entry.name.ends_with(".weight") {
81 tensors.push(Tensor::new(
82 entry.name.clone(),
83 role,
84 Some(codes.shape.clone()),
85 TensorEncoding::Dense { tensor: codes },
86 ));
87
88 continue;
89 }
90
91 let prefix = entry.name.trim_end_matches(".weight");
92 let scales = dense_part(manifest.dense(&format!("{prefix}.scales"))?)?;
93 let biases = dense_part(manifest.dense(&format!("{prefix}.biases"))?)?;
94
95 tensors.push(affine(entry.name.clone(), role, codes, scales, biases)?);
96 }
97
98 Ok(tensors)
99}
100
101fn dense_part(entry: &crate::qwen4_exp::DenseEntry) -> Result<StoredTensor> {
102 let dtype = match entry.dtype.as_str() {
103 "U32" => Dtype::U32,
104 "I64" => Dtype::I64,
105 "F32" => Dtype::F32,
106 "F16" => Dtype::F16,
107 "BF16" => Dtype::Bf16,
108 other => anyhow::bail!("unsupported packed dtype {other}"),
109 };
110
111 StoredTensor::contiguous(
112 dtype,
113 entry.shape.iter().map(|&v| v as u64).collect(),
114 DataSpan {
115 object: ObjectId(0),
116 offset: entry.offset,
117 length: entry.nbytes,
118 },
119 )
120}
121
122fn experts(manifest: &Manifest, tensors: &mut Vec<Tensor>) -> Result<()> {
123 let l = &manifest.experts;
124
125 for (layer, prefix) in l.layer_prefixes.iter().enumerate() {
126 let base = (layer as u64)
127 .checked_mul(l.experts as u64)
128 .and_then(|v| v.checked_mul(l.record_stride))
129 .context("expert offset overflow")?;
130 let records = Records {
131 object: ObjectId(1),
132 base,
133 count: l.experts as u64,
134 stride: l.record_stride,
135 };
136
137 for (name, rows, width, offsets) in [
138 (
139 "gate_proj",
140 l.inter,
141 l.hidden,
142 [l.gate_w, l.gate_s, l.gate_b],
143 ),
144 ("up_proj", l.inter, l.hidden, [l.up_w, l.up_s, l.up_b]),
145 (
146 "down_proj",
147 l.hidden,
148 l.inter,
149 [l.down_w, l.down_s, l.down_b],
150 ),
151 ] {
152 ensure!(
153 l.group > 0 && width.is_multiple_of(l.group) && width.is_multiple_of(8),
154 "invalid expert quantization group"
155 );
156
157 let codes = records.part(Dtype::U32, rows as u64, (width / 8) as u64, offsets[0])?;
158 let scales = records.part(
159 Dtype::Bf16,
160 rows as u64,
161 (width / l.group) as u64,
162 offsets[1],
163 )?;
164 let biases = records.part(
165 Dtype::Bf16,
166 rows as u64,
167 (width / l.group) as u64,
168 offsets[2],
169 )?;
170
171 tensors.push(affine(
172 format!("{prefix}.{name}.weight"),
173 TensorRole::Expert,
174 codes,
175 scales,
176 biases,
177 )?);
178 }
179 }
180
181 Ok(())
182}
183
184pub(super) fn ngram(manifest: &Manifest) -> Result<Tensor> {
185 let n = &manifest.ngram;
186
187 ensure!(
188 n.group > 0 && n.dim.is_multiple_of(n.group) && n.dim.is_multiple_of(8),
189 "invalid n-gram quantization group"
190 );
191 ensure!(
192 n.weight_bytes == n.dim as u64 / 2 && n.scale_bytes == (n.dim / n.group * 2) as u64,
193 "n-gram component sizes disagree with dimensions"
194 );
195 ensure!(
196 n.row_bytes >= n.weight_bytes + 2 * n.scale_bytes,
197 "n-gram record is too small"
198 );
199
200 let records = Records {
201 object: ObjectId(2),
202 base: 0,
203 count: n.rows,
204 stride: n.row_bytes,
205 };
206 let row_part = |dtype, width, offset| -> Result<StoredTensor> {
207 let mut part = records.part(dtype, 1, width, offset)?;
208
209 part.shape.remove(1);
210 part.byte_strides.remove(1);
211
212 Ok(part)
213 };
214 let codes = row_part(Dtype::U32, n.dim as u64 / 8, 0)?;
215 let scales = row_part(Dtype::Bf16, (n.dim / n.group) as u64, n.weight_bytes)?;
216 let biases = row_part(
217 Dtype::Bf16,
218 (n.dim / n.group) as u64,
219 n.weight_bytes + n.scale_bytes,
220 )?;
221
222 affine(
223 "ngram_embedding".into(),
224 TensorRole::NgramEmbedding,
225 codes,
226 scales,
227 biases,
228 )
229}
230
231struct Records {
232 object: ObjectId,
233 base: u64,
234 count: u64,
235 stride: u64,
236}
237
238impl Records {
239 fn part(&self, dtype: Dtype, rows: u64, width: u64, offset: u64) -> Result<StoredTensor> {
240 ensure!(
241 self.count > 0 && rows > 0 && width > 0,
242 "empty packed record component"
243 );
244
245 let row_bytes = width
246 .checked_mul(dtype.bytes())
247 .context("row size overflow")?;
248 let part_bytes = rows
249 .checked_mul(row_bytes)
250 .context("record size overflow")?;
251
252 ensure!(
253 offset
254 .checked_add(part_bytes)
255 .is_some_and(|end| end <= self.stride),
256 "component exceeds record stride"
257 );
258
259 let length = (self.count - 1)
260 .checked_mul(self.stride)
261 .and_then(|v| v.checked_add(part_bytes))
262 .context("record span overflow")?;
263 let offset = self
264 .base
265 .checked_add(offset)
266 .context("record offset overflow")?;
267 let part = StoredTensor {
268 dtype,
269 shape: vec![self.count, rows, width],
270 byte_strides: vec![self.stride, row_bytes, dtype.bytes()],
271 data: DataSpan {
272 object: self.object,
273 offset,
274 length,
275 },
276 };
277
278 part.validate()?;
279
280 Ok(part)
281 }
282}
283
284fn affine(
285 name: String,
286 role: TensorRole,
287 codes: StoredTensor,
288 scales: StoredTensor,
289 biases: StoredTensor,
290) -> Result<Tensor> {
291 let axis = codes
292 .shape
293 .len()
294 .checked_sub(1)
295 .context("scalar packed codes")?;
296 let mut shape = codes.shape.clone();
297 shape[axis] = shape[axis]
298 .checked_mul(8)
299 .context("packed width overflow")?;
300
301 ensure!(
302 scales.shape.len() == shape.len() && biases.shape == scales.shape,
303 "affine component rank mismatch"
304 );
305 ensure!(
306 scales.shape[..axis] == shape[..axis],
307 "affine component row mismatch"
308 );
309
310 let groups = scales.shape[axis];
311
312 ensure!(
313 groups > 0 && shape[axis].is_multiple_of(groups),
314 "invalid affine group count"
315 );
316
317 Ok(Tensor::new(
318 name,
319 role,
320 Some(shape.clone()),
321 TensorEncoding::Affine {
322 bits: 4,
323 group_size: shape[axis] / groups,
324 group_axis: axis,
325 packing: BitPacking::LowFirstU32,
326 codes,
327 scales,
328 offset: AffineOffset::Bias { tensor: biases },
329 },
330 ))
331}