metrale_model_weights/fast_weights/
load.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! `WeightLoader` for `FastSafetensorsLoader`: the shard walk that drives
4//! `load_shard_fast`.
5
6use super::*;
7
8impl WeightLoader for FastSafetensorsLoader {
9    fn load(
10        &self,
11        model_dir: &Path,
12        gpu: &dyn GpuBackend,
13        oom_reserve_bytes: usize,
14    ) -> Result<WeightStore> {
15        let skip_fn = |name: &str| self.should_skip_tensor(name);
16        let defer_fn =
17            |name: &str, dtype: crate::weights::WeightDtype| self.is_deferred(name, dtype);
18
19        // Resolve shard list (sharded index, single file, or unindexed shards).
20        let (shard_files, tensor_to_shard): (Vec<PathBuf>, Option<HashMap<String, String>>) =
21            resolve_shards(model_dir)?;
22
23        // Pre-flight OOM estimate (identical to SafetensorsLoader).
24        //
25        // Deferred tensors are DEFERRED further down — they are never
26        // uploaded, so counting them here refuses a model that fits. On
27        // LongCat-Flash-Lite the n-gram tables are 62.8 of the checkpoint's
28        // 138 GB, which is the difference between a 167 GB "peak" and a 98 GB
29        // one; on `nvidia/GLM-5.3-Flash-NVFP4` the MTP block's full-width
30        // routed experts are ~7.25 GB per EP=2 rank.
31        //
32        // 🔴 These two closures and the retain loop below MUST stay the same
33        // rule. A tensor counted here but deferred there inflates the peak;
34        // one deferred here but swept there is the OOM this pre-flight exists
35        // to prevent.
36        let preflight_skip = |name: &str| skip_fn(name) || crate::weights::is_ngram_table(name);
37        {
38            let estimated = estimate_load_bytes(&shard_files, &preflight_skip, &defer_fn)?;
39            let has_fp8 = estimate_has_fp8(&shard_files, &preflight_skip, &defer_fn)?;
40            let mult = self
41                .peak_memory_multiplier
42                .unwrap_or(if has_fp8 { 1.5 } else { 1.3 });
43            let peak = (estimated as f64 * mult) as usize;
44            let free = gpu.free_memory()?;
45            let gib = |b: usize| b as f64 / (1024.0 * 1024.0 * 1024.0);
46            tracing::info!(
47                "Fast-load pre-flight: {:.2} GB on-disk, {:.1}x overhead = {:.2} GB peak, \
48                 {:.2} GB free, {:.1} GB reserve (FP8: {})",
49                gib(estimated),
50                mult,
51                gib(peak),
52                gib(free),
53                gib(oom_reserve_bytes),
54                has_fp8,
55            );
56            metrale_telemetry::progress::preflight(gib(estimated), gib(free));
57            if peak + oom_reserve_bytes > free {
58                bail!(
59                    "OOM pre-flight: peak {:.2} GB + {:.2} GB reserve exceeds {:.2} GB free. \
60                     Use a smaller quantization or add more GPUs for EP.",
61                    gib(peak),
62                    gib(oom_reserve_bytes),
63                    gib(free),
64                );
65            }
66        }
67
68        // Load each shard. Loaded tensors filtered by EP rules upstream.
69        let mut weights: HashMap<String, WeightTensor> = HashMap::new();
70        // Locations of tensors deliberately NOT uploaded (the n-gram tables).
71        let mut deferred: HashMap<String, crate::weights::DeferredTensor> = HashMap::new();
72        let total_shards = shard_files.len();
73        let initial_free = gpu.free_memory()?;
74        let mut offload_logged = false;
75
76        for (i, shard_path) in shard_files.iter().enumerate() {
77            // When an index is present, only load the tensors it routes here;
78            // otherwise load everything in the shard. `None` means "load all".
79            let shard_name = shard_path
80                .file_name()
81                .and_then(|n| n.to_str())
82                .unwrap_or_default();
83            let tensor_filter: Option<Vec<String>> = tensor_to_shard.as_ref().map(|map| {
84                map.iter()
85                    .filter(|(_, s)| *s == shard_name)
86                    .map(|(t, _)| t.clone())
87                    .collect()
88            });
89
90            tracing::info!(
91                "Fast-loading shard {}/{}: {}{}",
92                i + 1,
93                total_shards,
94                shard_name,
95                tensor_filter
96                    .as_ref()
97                    .map(|v| format!(" ({} tensors)", v.len()))
98                    .unwrap_or_default(),
99            );
100            metrale_telemetry::progress::shard_start(i + 1, total_shards, shard_name);
101
102            load_shard_fast(
103                shard_path,
104                tensor_filter.as_deref(),
105                gpu,
106                &skip_fn,
107                &defer_fn,
108                self.try_direct_io,
109                self.direct_io_tensor_cap,
110                self.prefetch_shards,
111                &mut weights,
112                &mut deferred,
113                &mut offload_logged,
114            )?;
115
116            let free_now = gpu.free_memory().unwrap_or(0);
117            let used = initial_free.saturating_sub(free_now);
118            tracing::info!(
119                "  Shard {}/{} done — GPU memory: {:.2} GB used, {:.2} GB free",
120                i + 1,
121                total_shards,
122                used as f64 / (1024.0 * 1024.0 * 1024.0),
123                free_now as f64 / (1024.0 * 1024.0 * 1024.0),
124            );
125            metrale_telemetry::progress::shard_done(
126                i + 1,
127                total_shards,
128                used as f64 / (1024.0 * 1024.0 * 1024.0),
129                free_now as f64 / (1024.0 * 1024.0 * 1024.0),
130            );
131            if !offload_logged {
132                check_oom_guard(
133                    gpu,
134                    oom_reserve_bytes,
135                    &format!("fast weight loading (shard {}/{})", i + 1, total_shards),
136                )?;
137            }
138        }
139
140        // Extra weights (e.g. MTP grafted from another quantization).
141        let no_skip = |_: &str| false;
142        let no_defer = |_: &str, _: crate::weights::WeightDtype| false;
143        let extra = model_dir.join("extra_weights.safetensors");
144        if extra.exists() {
145            tracing::info!("Fast-loading extra_weights.safetensors");
146            let mut extra_offload = false;
147            load_shard_fast(
148                &extra,
149                None,
150                gpu,
151                &no_skip,
152                &no_defer,
153                self.try_direct_io,
154                self.direct_io_tensor_cap,
155                self.prefetch_shards,
156                &mut weights,
157                &mut deferred,
158                &mut extra_offload,
159            )?;
160        }
161
162        tracing::info!("Fast-loaded {} weight tensors", weights.len());
163        let mut store = WeightStore::from_map(weights);
164        for (name, d) in deferred {
165            store.defer(name, d);
166        }
167        Ok(store)
168    }
169}
metrale_model_weights/fast_weights/load.rs

metrale_model_weights/fast_weights/
load.rs
