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        // The n-gram tables are deferred, never uploaded, so they must not
81        // count toward the peak — see the note in `fast_weights`.
82        let preflight_skip = |name: &str| skip_fn(name) || super::is_ngram_table(name);
83        {
84            let estimated = estimate_load_bytes(&shard_files, &preflight_skip)?;
85            let has_fp8 = estimate_has_fp8(&shard_files, &preflight_skip)?;
86            let overhead_multiplier: f64 =
87                self.peak_memory_multiplier
88                    .unwrap_or(if has_fp8 { 1.5 } else { 1.3 });
89            let peak_estimated = (estimated as f64 * overhead_multiplier) as usize;
90            let free = gpu.free_memory()?;
91            let free_gb = free as f64 / (1024.0 * 1024.0 * 1024.0);
92            let est_gb = estimated as f64 / (1024.0 * 1024.0 * 1024.0);
93            let peak_gb = peak_estimated as f64 / (1024.0 * 1024.0 * 1024.0);
94            let reserve_gb = oom_reserve_bytes as f64 / (1024.0 * 1024.0 * 1024.0);
95            tracing::info!(
96                "Pre-flight estimate: {:.2} GB on-disk weights, {:.1}x overhead = {:.2} GB peak, \
97                 {:.2} GB free, {:.1} GB reserve (FP8: {})",
98                est_gb,
99                overhead_multiplier,
100                peak_gb,
101                free_gb,
102                reserve_gb,
103                has_fp8,
104            );
105            if peak_estimated + oom_reserve_bytes > free {
106                bail!(
107                    "OOM pre-flight: model peak memory ({:.2} GB = {:.2} GB weights × {:.1}x \
108                     model-building overhead) + {:.1} GB reserve = {:.2} GB, \
109                     but only {:.2} GB GPU memory is available. \
110                     This model is too large. Use a smaller quantization (NVFP4 instead of FP8) \
111                     or add more GPUs for expert parallelism.",
112                    peak_gb,
113                    est_gb,
114                    overhead_multiplier,
115                    reserve_gb,
116                    peak_gb + reserve_gb,
117                    free_gb,
118                );
119            }
120        }
121
122        // Locations of tensors deliberately NOT uploaded (the n-gram tables).
123        let mut deferred: HashMap<String, crate::weights::DeferredTensor> = HashMap::new();
124        let mut weight_map = if use_index {
125            load_sharded(
126                model_dir,
127                &actual_index_path,
128                gpu,
129                oom_reserve_bytes,
130                &skip_fn,
131                self.peak_memory_multiplier,
132                &mut deferred,
133            )?
134        } else if shard_files.len() == 1 {
135            load_single(&shard_files[0], gpu, oom_reserve_bytes, &skip_fn)?
136        } else {
137            tracing::info!("Loading {} unindexed safetensor shards", shard_files.len());
138            let initial_free = gpu.free_memory()?;
139            let mut combined = HashMap::new();
140            for (i, shard) in shard_files.iter().enumerate() {
141                let map = load_single(shard, gpu, oom_reserve_bytes, &skip_fn)?;
142                let free_now = gpu.free_memory().unwrap_or(0);
143                let used = initial_free.saturating_sub(free_now);
144                tracing::info!(
145                    "  Shard {}/{} done — GPU memory: {:.2} GB used, {:.2} GB free",
146                    i + 1,
147                    shard_files.len(),
148                    used as f64 / (1024.0 * 1024.0 * 1024.0),
149                    free_now as f64 / (1024.0 * 1024.0 * 1024.0),
150                );
151                check_oom_guard(
152                    gpu,
153                    oom_reserve_bytes,
154                    &format!("weight loading (shard {}/{})", i + 1, shard_files.len()),
155                )?;
156                combined.extend(map);
157            }
158            combined
159        };
160
161        // Load extra weight files (e.g. MTP weights grafted from another quantization).
162        // Extra weights (MTP) are always fully loaded — they have their own expert lists.
163        let no_skip = |_: &str| false;
164        let extra = model_dir.join("extra_weights.safetensors");
165        if extra.exists() {
166            let extra_weights = load_single(&extra, gpu, oom_reserve_bytes, &no_skip)?;
167            tracing::info!(
168                "Loaded {} extra weight tensors from extra_weights.safetensors",
169                extra_weights.len()
170            );
171            weight_map.extend(extra_weights);
172        }
173
174        let mut store = WeightStore::from_map(weight_map);
175        for (name, d) in deferred {
176            store.defer(name, d);
177        }
178        Ok(store)
179    }
180}
181
182/// Index file format: { "weight_map": { "tensor_name": "shard_filename" } }
183#[derive(serde::Deserialize)]
184struct SafetensorsIndex {
185    weight_map: HashMap<String, String>,
186}
187
188/// Read only the safetensor header from a file (no mmap, no GPU memory impact).
189/// The header is typically a few KB of JSON — safe to read on GB10 unified memory
190/// without consuming GPU pages.
191pub(crate) fn read_safetensor_header(
192    path: &Path,
193) -> Result<Vec<(String, Vec<usize>, safetensors::Dtype)>> {
194    use std::io::Read;
195    let mut file = std::fs::File::open(path)
196        .with_context(|| format!("Pre-flight: failed to open {}", path.display()))?;
197
198    // Safetensors format: 8-byte LE header size, then JSON header, then data.
199    let mut size_buf = [0u8; 8];
200    file.read_exact(&mut size_buf)?;
201    let header_size = u64::from_le_bytes(size_buf) as usize;
202
203    // Sanity check: header shouldn't exceed 64 MB.
204    if header_size > 64 * 1024 * 1024 {
205        bail!(
206            "Safetensor header too large ({} bytes) in {}",
207            header_size,
208            path.display()
209        );
210    }
211
212    let mut header_buf = vec![0u8; header_size];
213    file.read_exact(&mut header_buf)?;
214
215    // Parse the JSON header manually to extract tensor metadata.
216    let header: serde_json::Value = serde_json::from_slice(&header_buf)?;
217    let obj = header.as_object().context("Invalid safetensor header")?;
218
219    let mut tensors = Vec::new();
220    for (name, info) in obj {
221        if name == "__metadata__" {
222            continue;
223        }
224        let dtype_str = info["dtype"].as_str().unwrap_or("BF16");
225        let dtype = match dtype_str {
226            "F32" => safetensors::Dtype::F32,
227            "F16" => safetensors::Dtype::F16,
228            "BF16" => safetensors::Dtype::BF16,
229            "I32" => safetensors::Dtype::I32,
230            "I16" => safetensors::Dtype::I16,
231            "I8" => safetensors::Dtype::I8,
232            "U8" => safetensors::Dtype::U8,
233            "F8_E4M3" => safetensors::Dtype::F8_E4M3,
234            "F8_E5M2" => safetensors::Dtype::F8_E5M2,
235            _ => safetensors::Dtype::BF16,
236        };
237        let shape: Vec<usize> = info["shape"]
238            .as_array()
239            .map(|a| {
240                a.iter()
241                    .filter_map(|v| v.as_u64().map(|n| n as usize))
242                    .collect()
243            })
244            .unwrap_or_default();
245        tensors.push((name.clone(), shape, dtype));
246    }
247    Ok(tensors)
248}
249
250/// Scan safetensor file headers (metadata only, no data loaded) to estimate
251/// total GPU bytes this rank will load. Reads only the JSON header from each
252/// file — does NOT mmap, so it's safe on GB10 unified memory.
253pub(crate) fn estimate_load_bytes(
254    files: &[std::path::PathBuf],
255    skip_fn: &dyn Fn(&str) -> bool,
256) -> Result<usize> {
257    let mut total = 0usize;
258    for path in files {
259        for (name, shape, dtype) in read_safetensor_header(path)? {
260            if skip_fn(&name) {
261                continue;
262            }
263            let numel: usize = shape.iter().product();
264            let elem_bytes = match dtype {
265                safetensors::Dtype::F32 | safetensors::Dtype::I32 | safetensors::Dtype::U32 => 4,
266                safetensors::Dtype::F16
267                | safetensors::Dtype::BF16
268                | safetensors::Dtype::I16
269                | safetensors::Dtype::U16 => 2,
270                safetensors::Dtype::I8
271                | safetensors::Dtype::U8
272                | safetensors::Dtype::F8_E4M3
273                | safetensors::Dtype::F8_E5M2 => 1,
274                _ => 2,
275            };
276            total += numel * elem_bytes;
277        }
278    }
279    Ok(total)
280}
281
282/// Check if the model is predominantly FP8 (>50% of weight bytes are FP8).
283/// Sehyo NVFP4 models have a few FP8 scale tensors but the bulk is uint8 (NVFP4 packed).
284/// True FP8 checkpoints (e.g. Qwen/Qwen3.5-122B-A10B-FP8) have most bytes as FP8.
285pub(crate) fn estimate_has_fp8(
286    files: &[std::path::PathBuf],
287    skip_fn: &dyn Fn(&str) -> bool,
288) -> Result<bool> {
289    let mut fp8_bytes = 0usize;
290    let mut total_bytes = 0usize;
291    for path in files {
292        for (name, shape, dtype) in read_safetensor_header(path)? {
293            if skip_fn(&name) {
294                continue;
295            }
296            let numel: usize = shape.iter().product();
297            let elem_bytes = match dtype {
298                safetensors::Dtype::F32 | safetensors::Dtype::I32 | safetensors::Dtype::U32 => 4,
299                safetensors::Dtype::F16
300                | safetensors::Dtype::BF16
301                | safetensors::Dtype::I16
302                | safetensors::Dtype::U16 => 2,
303                safetensors::Dtype::I8
304                | safetensors::Dtype::U8
305                | safetensors::Dtype::F8_E4M3
306                | safetensors::Dtype::F8_E5M2 => 1,
307                _ => 2,
308            };
309            let bytes = numel * elem_bytes;
310            total_bytes += bytes;
311            if matches!(
312                dtype,
313                safetensors::Dtype::F8_E4M3 | safetensors::Dtype::F8_E5M2
314            ) {
315                fp8_bytes += bytes;
316            }
317        }
318    }
319    let fp8_frac = if total_bytes > 0 {
320        fp8_bytes as f64 / total_bytes as f64
321    } else {
322        0.0
323    };
324    tracing::debug!(
325        "FP8 fraction: {:.1}% ({} / {} bytes)",
326        fp8_frac * 100.0,
327        fp8_bytes,
328        total_bytes
329    );
330    Ok(fp8_frac > 0.5)
331}
332
333pub(crate) fn check_oom_guard(
334    gpu: &dyn GpuBackend,
335    reserve_bytes: usize,
336    phase: &str,
337) -> Result<()> {
338    let free = gpu.free_memory()?;
339    if free < reserve_bytes {
340        let free_gb = free as f64 / (1024.0 * 1024.0 * 1024.0);
341        let reserve_gb = reserve_bytes as f64 / (1024.0 * 1024.0 * 1024.0);
342        bail!(
343            "OOM guard: aborting during {phase}. \
344             Free GPU memory ({free_gb:.2} GB) is below the {reserve_gb:.1} GB safety reserve. \
345             This model is too large for available GPU memory. \
346             Reduce --max-seq-len, increase --oom-guard-mb, or use a smaller model."
347        );
348    }
349    Ok(())
350}
351mod load_fns;
352use load_fns::{load_sharded, load_single};
spark_runtime/weights/loader.rs

spark_runtime/weights/
loader.rs
