spark_runtime/weights/
loader.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2
3//! `impl WeightLoader for SafetensorsLoader` + sharded/single loader helpers.
4
5use anyhow::{Context, Result, bail};
6use std::collections::HashMap;
7use std::path::Path;
8
9use super::{SafetensorsLoader, WeightLoader, WeightStore};
10use crate::gpu::GpuBackend;
11
12impl WeightLoader for SafetensorsLoader {
13    fn load(
14        &self,
15        model_dir: &Path,
16        gpu: &dyn GpuBackend,
17        oom_reserve_bytes: usize,
18    ) -> Result<WeightStore> {
19        let skip_fn = |name: &str| self.should_skip_tensor(name);
20
21        // Collect all safetensor files (indexed, single, or unindexed shards).
22        // Supports both HuggingFace standard (model.safetensors*) and Mistral
23        // consolidated format (consolidated.safetensors*).
24        let index_path = model_dir.join("model.safetensors.index.json");
25        let consolidated_index = model_dir.join("consolidated.safetensors.index.json");
26        let shard_files: Vec<std::path::PathBuf>;
27        let use_index;
28        let actual_index_path;
29
30        if index_path.exists() {
31            use_index = true;
32            actual_index_path = index_path;
33            shard_files = vec![];
34        } else if consolidated_index.exists() {
35            use_index = true;
36            actual_index_path = consolidated_index;
37            shard_files = vec![];
38        } else {
39            use_index = false;
40            actual_index_path = index_path; // unused
41            let single = model_dir.join("model.safetensors");
42            if single.exists() {
43                shard_files = vec![single];
44            } else {
45                // Try both model.safetensors-* and consolidated-* shard patterns
46                let mut shards: Vec<_> = std::fs::read_dir(model_dir)?
47                    .filter_map(|e| e.ok())
48                    .map(|e| e.path())
49                    .filter(|p| {
50                        p.file_name().and_then(|n| n.to_str()).is_some_and(|n| {
51                            (n.starts_with("model.safetensors-") || n.starts_with("consolidated-"))
52                                && n.ends_with(".safetensors")
53                        })
54                    })
55                    .collect();
56                shards.sort();
57                if shards.is_empty() {
58                    bail!(
59                        "No safetensor files found in {}. Expected model.safetensors*, \
60                         consolidated.safetensors*, or consolidated-*-of-*.safetensors",
61                        model_dir.display()
62                    );
63                }
64                shard_files = shards;
65            }
66        }
67
68        // Pre-flight OOM estimate: scan safetensor headers (no data) to compute
69        // total bytes this rank will load, then apply a model-building overhead
70        // multiplier and abort early if the model won't fit.
71        //
72        // Model building creates additional GPU allocations on top of the raw
73        // weight store: transposed weight copies for prefill GEMM, predequanted
74        // FP8 copies, NVFP4 quantized copies (for FP8 checkpoints), and transient
75        // BF16 intermediates during FP8→NVFP4 conversion.
76        //
77        // Empirical overhead multipliers (peak memory / on-disk weight bytes):
78        //   NVFP4 (Sehyo): ~2.0x  (store aliased + transposed/predequant copies)
79        //   FP8 native:    ~1.5x  (store stays FP8, only attention prefill gets NVFP4 copies)
80        // Deferred tensors are never uploaded, so they must not count toward
81        // the peak — see the note in `fast_weights` and the contract in
82        // `weights::deferred`. Two rules produce them: the n-gram tables
83        // (always) and whatever the model's own loader claimed (opt-in).
84        let preflight_skip = |name: &str| skip_fn(name) || super::is_ngram_table(name);
85        let defer_fn = |name: &str, dtype: super::WeightDtype| self.is_deferred(name, dtype);
86        {
87            let estimated = estimate_load_bytes(&shard_files, &preflight_skip, &defer_fn)?;
88            let has_fp8 = estimate_has_fp8(&shard_files, &preflight_skip, &defer_fn)?;
89            let overhead_multiplier: f64 =
90                self.peak_memory_multiplier
91                    .unwrap_or(if has_fp8 { 1.5 } else { 1.3 });
92            let peak_estimated = (estimated as f64 * overhead_multiplier) as usize;
93            let free = gpu.free_memory()?;
94            let free_gb = free as f64 / (1024.0 * 1024.0 * 1024.0);
95            let est_gb = estimated as f64 / (1024.0 * 1024.0 * 1024.0);
96            let peak_gb = peak_estimated as f64 / (1024.0 * 1024.0 * 1024.0);
97            let reserve_gb = oom_reserve_bytes as f64 / (1024.0 * 1024.0 * 1024.0);
98            tracing::info!(
99                "Pre-flight estimate: {:.2} GB on-disk weights, {:.1}x overhead = {:.2} GB peak, \
100                 {:.2} GB free, {:.1} GB reserve (FP8: {})",
101                est_gb,
102                overhead_multiplier,
103                peak_gb,
104                free_gb,
105                reserve_gb,
106                has_fp8,
107            );
108            if peak_estimated + oom_reserve_bytes > free {
109                bail!(
110                    "OOM pre-flight: model peak memory ({:.2} GB = {:.2} GB weights × {:.1}x \
111                     model-building overhead) + {:.1} GB reserve = {:.2} GB, \
112                     but only {:.2} GB GPU memory is available. \
113                     This model is too large. Use a smaller quantization (NVFP4 instead of FP8) \
114                     or add more GPUs for expert parallelism.",
115                    peak_gb,
116                    est_gb,
117                    overhead_multiplier,
118                    reserve_gb,
119                    peak_gb + reserve_gb,
120                    free_gb,
121                );
122            }
123        }
124
125        // Locations of tensors deliberately NOT uploaded (the n-gram tables,
126        // plus whatever the model's loader claimed through `defer`).
127        let mut deferred: HashMap<String, crate::weights::DeferredTensor> = HashMap::new();
128        let rules = load_fns::TensorRules {
129            skip: &skip_fn,
130            defer: &defer_fn,
131        };
132        let mut weight_map = if use_index {
133            load_sharded(
134                model_dir,
135                &actual_index_path,
136                gpu,
137                oom_reserve_bytes,
138                &rules,
139                self.peak_memory_multiplier,
140                &mut deferred,
141            )?
142        } else if shard_files.len() == 1 {
143            load_single(
144                &shard_files[0],
145                gpu,
146                oom_reserve_bytes,
147                &rules,
148                &mut deferred,
149            )?
150        } else {
151            tracing::info!("Loading {} unindexed safetensor shards", shard_files.len());
152            let initial_free = gpu.free_memory()?;
153            let mut combined = HashMap::new();
154            for (i, shard) in shard_files.iter().enumerate() {
155                let map = load_single(shard, gpu, oom_reserve_bytes, &rules, &mut deferred)?;
156                let free_now = gpu.free_memory().unwrap_or(0);
157                let used = initial_free.saturating_sub(free_now);
158                tracing::info!(
159                    "  Shard {}/{} done — GPU memory: {:.2} GB used, {:.2} GB free",
160                    i + 1,
161                    shard_files.len(),
162                    used as f64 / (1024.0 * 1024.0 * 1024.0),
163                    free_now as f64 / (1024.0 * 1024.0 * 1024.0),
164                );
165                check_oom_guard(
166                    gpu,
167                    oom_reserve_bytes,
168                    &format!("weight loading (shard {}/{})", i + 1, shard_files.len()),
169                )?;
170                combined.extend(map);
171            }
172            combined
173        };
174
175        // Load extra weight files (e.g. MTP weights grafted from another quantization).
176        // Extra weights (MTP) are always fully loaded — they have their own expert lists.
177        let no_skip = |_: &str| false;
178        let no_defer = |_: &str, _: super::WeightDtype| false;
179        let extra = model_dir.join("extra_weights.safetensors");
180        if extra.exists() {
181            let mut extra_deferred = HashMap::new();
182            let extra_weights = load_single(
183                &extra,
184                gpu,
185                oom_reserve_bytes,
186                &load_fns::TensorRules {
187                    skip: &no_skip,
188                    defer: &no_defer,
189                },
190                &mut extra_deferred,
191            )?;
192            tracing::info!(
193                "Loaded {} extra weight tensors from extra_weights.safetensors",
194                extra_weights.len()
195            );
196            weight_map.extend(extra_weights);
197        }
198
199        let mut store = WeightStore::from_map(weight_map);
200        for (name, d) in deferred {
201            store.defer(name, d);
202        }
203        Ok(store)
204    }
205}
206
207/// Index file format: { "weight_map": { "tensor_name": "shard_filename" } }
208#[derive(serde::Deserialize)]
209struct SafetensorsIndex {
210    weight_map: HashMap<String, String>,
211}
212
213/// Read only the safetensor header from a file (no mmap, no GPU memory impact).
214/// The header is typically a few KB of JSON — safe to read on GB10 unified memory
215/// without consuming GPU pages.
216pub(crate) fn read_safetensor_header(
217    path: &Path,
218) -> Result<Vec<(String, Vec<usize>, safetensors::Dtype)>> {
219    use std::io::Read;
220    let mut file = std::fs::File::open(path)
221        .with_context(|| format!("Pre-flight: failed to open {}", path.display()))?;
222
223    // Safetensors format: 8-byte LE header size, then JSON header, then data.
224    let mut size_buf = [0u8; 8];
225    file.read_exact(&mut size_buf)?;
226    let header_size = u64::from_le_bytes(size_buf) as usize;
227
228    // Sanity check: header shouldn't exceed 64 MB.
229    if header_size > 64 * 1024 * 1024 {
230        bail!(
231            "Safetensor header too large ({} bytes) in {}",
232            header_size,
233            path.display()
234        );
235    }
236
237    let mut header_buf = vec![0u8; header_size];
238    file.read_exact(&mut header_buf)?;
239
240    // Parse the JSON header manually to extract tensor metadata.
241    let header: serde_json::Value = serde_json::from_slice(&header_buf)?;
242    let obj = header.as_object().context("Invalid safetensor header")?;
243
244    let mut tensors = Vec::new();
245    for (name, info) in obj {
246        if name == "__metadata__" {
247            continue;
248        }
249        let dtype_str = info["dtype"].as_str().unwrap_or("BF16");
250        let dtype = match dtype_str {
251            "F32" => safetensors::Dtype::F32,
252            "F16" => safetensors::Dtype::F16,
253            "BF16" => safetensors::Dtype::BF16,
254            "I32" => safetensors::Dtype::I32,
255            "I16" => safetensors::Dtype::I16,
256            "I8" => safetensors::Dtype::I8,
257            "U8" => safetensors::Dtype::U8,
258            "F8_E4M3" => safetensors::Dtype::F8_E4M3,
259            "F8_E5M2" => safetensors::Dtype::F8_E5M2,
260            _ => safetensors::Dtype::BF16,
261        };
262        let shape: Vec<usize> = info["shape"]
263            .as_array()
264            .map(|a| {
265                a.iter()
266                    .filter_map(|v| v.as_u64().map(|n| n as usize))
267                    .collect()
268            })
269            .unwrap_or_default();
270        tensors.push((name.clone(), shape, dtype));
271    }
272    Ok(tensors)
273}
274
275/// Would the model's loader defer this header entry?
276///
277/// 🪤 A dtype [`WeightDtype`] cannot hold is NEVER deferred, and that is the
278/// whole reason this is an `is_ok_and` rather than a fallback: a deferred
279/// tensor is served from its ON-DISK bytes, and F16 is the one width the disk
280/// loaders rewrite on the way to the store (`f16_to_bf16_bytes`). Recording it
281/// as BF16 would describe bytes that are not in the file, so an F16 tensor
282/// takes the ordinary upload path and is counted here — which is exactly what
283/// the loaders then do with it.
284fn defers(
285    name: &str,
286    dtype: safetensors::Dtype,
287    defer_fn: &dyn Fn(&str, super::WeightDtype) -> bool,
288) -> bool {
289    super::WeightDtype::from_safetensors(dtype).is_ok_and(|d| defer_fn(name, d))
290}
291
292/// Scan safetensor file headers (metadata only, no data loaded) to estimate
293/// total GPU bytes this rank will load. Reads only the JSON header from each
294/// file — does NOT mmap, so it's safe on GB10 unified memory.
295///
296/// 🔴 `defer_fn` is separate from `skip_fn` and not folded into it: a deferred
297/// tensor is withheld on `(name, dtype)`, and counting it at its on-disk size
298/// would refuse a load that fits. See [`super::DeferHook`].
299pub(crate) fn estimate_load_bytes(
300    files: &[std::path::PathBuf],
301    skip_fn: &dyn Fn(&str) -> bool,
302    defer_fn: &dyn Fn(&str, super::WeightDtype) -> bool,
303) -> Result<usize> {
304    let mut total = 0usize;
305    for path in files {
306        for (name, shape, dtype) in read_safetensor_header(path)? {
307            if skip_fn(&name) || defers(&name, dtype, defer_fn) {
308                continue;
309            }
310            let numel: usize = shape.iter().product();
311            let elem_bytes = match dtype {
312                safetensors::Dtype::F32 | safetensors::Dtype::I32 | safetensors::Dtype::U32 => 4,
313                safetensors::Dtype::F16
314                | safetensors::Dtype::BF16
315                | safetensors::Dtype::I16
316                | safetensors::Dtype::U16 => 2,
317                safetensors::Dtype::I8
318                | safetensors::Dtype::U8
319                | safetensors::Dtype::F8_E4M3
320                | safetensors::Dtype::F8_E5M2 => 1,
321                _ => 2,
322            };
323            total += numel * elem_bytes;
324        }
325    }
326    Ok(total)
327}
328
329/// Check if the model is predominantly FP8 (>50% of weight bytes are FP8).
330/// Sehyo NVFP4 models have a few FP8 scale tensors but the bulk is uint8 (NVFP4 packed).
331/// True FP8 checkpoints (e.g. Qwen/Qwen3.5-122B-A10B-FP8) have most bytes as FP8.
332pub(crate) fn estimate_has_fp8(
333    files: &[std::path::PathBuf],
334    skip_fn: &dyn Fn(&str) -> bool,
335    defer_fn: &dyn Fn(&str, super::WeightDtype) -> bool,
336) -> Result<bool> {
337    let mut fp8_bytes = 0usize;
338    let mut total_bytes = 0usize;
339    for path in files {
340        for (name, shape, dtype) in read_safetensor_header(path)? {
341            if skip_fn(&name) || defers(&name, dtype, defer_fn) {
342                continue;
343            }
344            let numel: usize = shape.iter().product();
345            let elem_bytes = match dtype {
346                safetensors::Dtype::F32 | safetensors::Dtype::I32 | safetensors::Dtype::U32 => 4,
347                safetensors::Dtype::F16
348                | safetensors::Dtype::BF16
349                | safetensors::Dtype::I16
350                | safetensors::Dtype::U16 => 2,
351                safetensors::Dtype::I8
352                | safetensors::Dtype::U8
353                | safetensors::Dtype::F8_E4M3
354                | safetensors::Dtype::F8_E5M2 => 1,
355                _ => 2,
356            };
357            let bytes = numel * elem_bytes;
358            total_bytes += bytes;
359            if matches!(
360                dtype,
361                safetensors::Dtype::F8_E4M3 | safetensors::Dtype::F8_E5M2
362            ) {
363                fp8_bytes += bytes;
364            }
365        }
366    }
367    let fp8_frac = if total_bytes > 0 {
368        fp8_bytes as f64 / total_bytes as f64
369    } else {
370        0.0
371    };
372    tracing::debug!(
373        "FP8 fraction: {:.1}% ({} / {} bytes)",
374        fp8_frac * 100.0,
375        fp8_bytes,
376        total_bytes
377    );
378    Ok(fp8_frac > 0.5)
379}
380
381pub(crate) fn check_oom_guard(
382    gpu: &dyn GpuBackend,
383    reserve_bytes: usize,
384    phase: &str,
385) -> Result<()> {
386    let free = gpu.free_memory()?;
387    if free < reserve_bytes {
388        let free_gb = free as f64 / (1024.0 * 1024.0 * 1024.0);
389        let reserve_gb = reserve_bytes as f64 / (1024.0 * 1024.0 * 1024.0);
390        bail!(
391            "OOM guard: aborting during {phase}. \
392             Free GPU memory ({free_gb:.2} GB) is below the {reserve_gb:.1} GB safety reserve. \
393             This model is too large for available GPU memory. \
394             Reduce --max-seq-len, increase --oom-guard-mb, or use a smaller model."
395        );
396    }
397    Ok(())
398}
399mod load_fns;
400use load_fns::{load_sharded, load_single};
spark_runtime/weights/loader.rs

spark_runtime/weights/
loader.rs
