spark_runtime/fast_weights/
mod.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2
3//! Fast safetensors loader (InstantTensor-style) — pure Rust.
4//!
5//! Two wins over the mmap-based loader in [`crate::weights`]:
6//!
7//! 1. **`O_DIRECT`** reads. Bypasses the OS page cache, so the bytes never
8//!    compete with GPU allocations on GB10 unified memory. The mmap path
9//!    already works around this with `POSIX_FADV_DONTNEED` post-load; here
10//!    we avoid the pollution in the first place.
11//! 2. **Pipelined read/copy**. One background reader thread fetches the
12//!    next tensor while the main thread does `copy_h2d` for the current
13//!    one. Overlaps disk I/O with the host→device memcpy.
14//!
15//! Behavioural parity with [`crate::weights::SafetensorsLoader`] is
16//! preserved — same EP filtering, same OOM pre-flight, same UVM fallback
17//! on GPU allocation failure, same extra-weights handling.
18
19use crate::gpu::GpuBackend;
20use crate::weights::{
21    WeightLoader, WeightStore, WeightTensor, check_oom_guard, estimate_has_fp8,
22    estimate_load_bytes, evict_page_cache, f16_to_bf16_bytes,
23};
24use anyhow::{Context, Result, bail};
25use std::collections::HashMap;
26use std::fs::File;
27use std::path::{Path, PathBuf};
28use std::sync::mpsc::sync_channel;
29
30mod direct_io;
31mod header;
32
33use header::{parse_header, resolve_shards};
34
35/// Pure-Rust InstantTensor-style loader. Same public shape as
36/// [`crate::weights::SafetensorsLoader`].
37pub struct FastSafetensorsLoader {
38    pub ep_rank: usize,
39    pub ep_world_size: usize,
40    pub num_experts: usize,
41    pub peak_memory_multiplier: Option<f64>,
42    /// Skip the W4A4 `*.input_scale` activation scales at load.
43    ///
44    /// ModelOpt NVFP4 checkpoints ship one 0-dim F32 scalar per quantized
45    /// projection. On a 512-expert model that is ~74k four-byte allocations,
46    /// each taking a full allocation granule — GBs of padding for values
47    /// Atlas never reads, because it serves w4a16 (BF16 activations) and the
48    /// NVFP4 loader already treats the key as optional.
49    ///
50    /// OPT-IN: `step3p7` reads this key on its own path, so it must stay off
51    /// unless the model's loader is known not to need it.
52    pub skip_activation_scales: bool,
53    /// Skip `mtp.*` tensors at load.
54    ///
55    /// For models whose loader deliberately does not build an MTP head,
56    /// uploading its weights is pure waste — on Qwen3.8-Flash-Next that is a
57    /// 1.49 GB expert shard plus the MTP backbone, held resident while the KV
58    /// cache goes without.
59    ///
60    /// OPT-IN: a model that DOES build an MTP head must keep them, so this is
61    /// set only where `load_mtp_weights` is known to return `None`.
62    pub skip_mtp: bool,
63    /// When true (default), attempt `O_DIRECT`; fall back to buffered reads if
64    /// the filesystem rejects it (tmpfs, overlayfs, some FUSE backends).
65    pub try_direct_io: bool,
66    /// Per-shard heuristic cap: if a shard's tensor count exceeds this,
67    /// we skip `O_DIRECT` for that shard and fall back to buffered +
68    /// pipelined reads even when [`Self::try_direct_io`] is `true`.
69    ///
70    /// Motivation: `O_DIRECT`'s 4 KiB-aligned per-tensor `pread` has a
71    /// fixed syscall + copy overhead that kernel readahead amortises for
72    /// free on the buffered path. Benchmarks on GB10 showed buffered wins
73    /// above ~5k tensors/shard; O_DIRECT wins below. Set to [`usize::MAX`]
74    /// to disable.
75    pub direct_io_tensor_cap: usize,
76    /// When true, advise the kernel to read a whole buffered shard
77    /// sequentially before the per-tensor copy loop starts. This helps NFS
78    /// mounts where many small tensor reads defeat normal readahead.
79    pub prefetch_shards: bool,
80    /// Skip a multimodal checkpoint's vision tower.
81    ///
82    /// Set by the caller from `ModelWeightLoader::binds_vision_encoder()`:
83    /// false by default, true only when the model's loader is a text-only
84    /// port that will never bind the tower. Reading it anyway costs the full
85    /// tower in unified memory (1.05 GiB/rank on GLM-5.3's checkpoint) from
86    /// load time until `build_model` frees it — which is after the inference
87    /// -buffer preflight has already refused the serve.
88    pub skip_vision: bool,
89    /// Tensors the MODEL's weight loader will read from disk itself, so this
90    /// loader records their location instead of uploading them.
91    ///
92    /// Set by the caller from `ModelWeightLoader::defer_predicate()`, the same
93    /// way `skip_vision` comes from `binds_vision_encoder()`. `None` for every
94    /// model that does not opt in. See [`crate::weights::DeferHook`] for the
95    /// contract — including the half that says the pre-flight must agree.
96    pub defer: Option<crate::weights::DeferHook>,
97}
98
99/// Is this tensor part of a multimodal checkpoint's vision tower?
100///
101/// Same three spellings `build_model`'s unbound-tower reclaim matches, kept
102/// here so the load-time skip and the post-bind free can never disagree.
103pub fn is_vision_tensor(name: &str) -> bool {
104    name.starts_with("model.visual.")
105        || name.starts_with("model.vision")
106        || name.starts_with("visual.")
107}
108
109/// Default tensor-count cap for per-shard `O_DIRECT`. Above this, the fast
110/// loader uses buffered reads even when `try_direct_io = true`. See the
111/// field doc on [`FastSafetensorsLoader::direct_io_tensor_cap`].
112pub const DEFAULT_DIRECT_IO_TENSOR_CAP: usize = 5000;
113
114impl Default for FastSafetensorsLoader {
115    fn default() -> Self {
116        Self::new()
117    }
118}
119
120#[path = "skip.rs"]
121mod skip;
122
123impl FastSafetensorsLoader {
124    pub fn new() -> Self {
125        Self {
126            ep_rank: 0,
127            ep_world_size: 1,
128            num_experts: 0,
129            peak_memory_multiplier: None,
130            skip_activation_scales: false,
131            skip_mtp: false,
132            try_direct_io: true,
133            direct_io_tensor_cap: DEFAULT_DIRECT_IO_TENSOR_CAP,
134            prefetch_shards: false,
135            skip_vision: false,
136            defer: None,
137        }
138    }
139
140    pub fn with_ep(ep_rank: usize, ep_world_size: usize, num_experts: usize) -> Self {
141        Self {
142            ep_rank,
143            ep_world_size,
144            num_experts,
145            peak_memory_multiplier: None,
146            skip_activation_scales: false,
147            skip_mtp: false,
148            try_direct_io: true,
149            direct_io_tensor_cap: DEFAULT_DIRECT_IO_TENSOR_CAP,
150            prefetch_shards: false,
151            skip_vision: false,
152            defer: None,
153        }
154    }
155}
156
157impl WeightLoader for FastSafetensorsLoader {
158    fn load(
159        &self,
160        model_dir: &Path,
161        gpu: &dyn GpuBackend,
162        oom_reserve_bytes: usize,
163    ) -> Result<WeightStore> {
164        let skip_fn = |name: &str| self.should_skip_tensor(name);
165        let defer_fn =
166            |name: &str, dtype: crate::weights::WeightDtype| self.is_deferred(name, dtype);
167
168        // Resolve shard list (sharded index, single file, or unindexed shards).
169        let (shard_files, tensor_to_shard): (Vec<PathBuf>, Option<HashMap<String, String>>) =
170            resolve_shards(model_dir)?;
171
172        // Pre-flight OOM estimate (identical to SafetensorsLoader).
173        //
174        // Deferred tensors are DEFERRED further down — they are never
175        // uploaded, so counting them here refuses a model that fits. On
176        // LongCat-Flash-Lite the n-gram tables are 62.8 of the checkpoint's
177        // 138 GB, which is the difference between a 167 GB "peak" and a 98 GB
178        // one; on `nvidia/GLM-5.3-Flash-NVFP4` the MTP block's full-width
179        // routed experts are ~7.25 GB per EP=2 rank.
180        //
181        // 🔴 These two closures and the retain loop below MUST stay the same
182        // rule. A tensor counted here but deferred there inflates the peak;
183        // one deferred here but swept there is the OOM this pre-flight exists
184        // to prevent.
185        let preflight_skip = |name: &str| skip_fn(name) || crate::weights::is_ngram_table(name);
186        {
187            let estimated = estimate_load_bytes(&shard_files, &preflight_skip, &defer_fn)?;
188            let has_fp8 = estimate_has_fp8(&shard_files, &preflight_skip, &defer_fn)?;
189            let mult = self
190                .peak_memory_multiplier
191                .unwrap_or(if has_fp8 { 1.5 } else { 1.3 });
192            let peak = (estimated as f64 * mult) as usize;
193            let free = gpu.free_memory()?;
194            let gib = |b: usize| b as f64 / (1024.0 * 1024.0 * 1024.0);
195            tracing::info!(
196                "Fast-load pre-flight: {:.2} GB on-disk, {:.1}x overhead = {:.2} GB peak, \
197                 {:.2} GB free, {:.1} GB reserve (FP8: {})",
198                gib(estimated),
199                mult,
200                gib(peak),
201                gib(free),
202                gib(oom_reserve_bytes),
203                has_fp8,
204            );
205            crate::progress::preflight(gib(estimated), gib(free));
206            if peak + oom_reserve_bytes > free {
207                bail!(
208                    "OOM pre-flight: peak {:.2} GB + {:.2} GB reserve exceeds {:.2} GB free. \
209                     Use a smaller quantization or add more GPUs for EP.",
210                    gib(peak),
211                    gib(oom_reserve_bytes),
212                    gib(free),
213                );
214            }
215        }
216
217        // Load each shard. Loaded tensors filtered by EP rules upstream.
218        let mut weights: HashMap<String, WeightTensor> = HashMap::new();
219        // Locations of tensors deliberately NOT uploaded (the n-gram tables).
220        let mut deferred: HashMap<String, crate::weights::DeferredTensor> = HashMap::new();
221        let total_shards = shard_files.len();
222        let initial_free = gpu.free_memory()?;
223        let mut offload_logged = false;
224
225        for (i, shard_path) in shard_files.iter().enumerate() {
226            // When an index is present, only load the tensors it routes here;
227            // otherwise load everything in the shard. `None` means "load all".
228            let shard_name = shard_path
229                .file_name()
230                .and_then(|n| n.to_str())
231                .unwrap_or_default();
232            let tensor_filter: Option<Vec<String>> = tensor_to_shard.as_ref().map(|map| {
233                map.iter()
234                    .filter(|(_, s)| *s == shard_name)
235                    .map(|(t, _)| t.clone())
236                    .collect()
237            });
238
239            tracing::info!(
240                "Fast-loading shard {}/{}: {}{}",
241                i + 1,
242                total_shards,
243                shard_name,
244                tensor_filter
245                    .as_ref()
246                    .map(|v| format!(" ({} tensors)", v.len()))
247                    .unwrap_or_default(),
248            );
249            crate::progress::shard_start(i + 1, total_shards, shard_name);
250
251            load_shard_fast(
252                shard_path,
253                tensor_filter.as_deref(),
254                gpu,
255                &skip_fn,
256                &defer_fn,
257                self.try_direct_io,
258                self.direct_io_tensor_cap,
259                self.prefetch_shards,
260                &mut weights,
261                &mut deferred,
262                &mut offload_logged,
263            )?;
264
265            let free_now = gpu.free_memory().unwrap_or(0);
266            let used = initial_free.saturating_sub(free_now);
267            tracing::info!(
268                "  Shard {}/{} done — GPU memory: {:.2} GB used, {:.2} GB free",
269                i + 1,
270                total_shards,
271                used as f64 / (1024.0 * 1024.0 * 1024.0),
272                free_now as f64 / (1024.0 * 1024.0 * 1024.0),
273            );
274            crate::progress::shard_done(
275                i + 1,
276                total_shards,
277                used as f64 / (1024.0 * 1024.0 * 1024.0),
278                free_now as f64 / (1024.0 * 1024.0 * 1024.0),
279            );
280            if !offload_logged {
281                check_oom_guard(
282                    gpu,
283                    oom_reserve_bytes,
284                    &format!("fast weight loading (shard {}/{})", i + 1, total_shards),
285                )?;
286            }
287        }
288
289        // Extra weights (e.g. MTP grafted from another quantization).
290        let no_skip = |_: &str| false;
291        let no_defer = |_: &str, _: crate::weights::WeightDtype| false;
292        let extra = model_dir.join("extra_weights.safetensors");
293        if extra.exists() {
294            tracing::info!("Fast-loading extra_weights.safetensors");
295            let mut extra_offload = false;
296            load_shard_fast(
297                &extra,
298                None,
299                gpu,
300                &no_skip,
301                &no_defer,
302                self.try_direct_io,
303                self.direct_io_tensor_cap,
304                self.prefetch_shards,
305                &mut weights,
306                &mut deferred,
307                &mut extra_offload,
308            )?;
309        }
310
311        tracing::info!("Fast-loaded {} weight tensors", weights.len());
312        let mut store = WeightStore::from_map(weights);
313        for (name, d) in deferred {
314            store.defer(name, d);
315        }
316        Ok(store)
317    }
318}
319
320/// Load a single shard with O_DIRECT + pipelined read/copy.
321///
322/// Pipeline:
323///   reader thread: pread tensor N into aligned buffer → sync_channel ──▶
324///   main thread:   recv → copy_h2d → store tensor
325///
326/// The channel has capacity 1, so at any time the reader is ≤1 tensor
327/// ahead of the copier. Memory overhead per shard: 2 × max_tensor_bytes
328/// (rounded up to O_DIRECT alignment).
329#[allow(clippy::too_many_arguments)]
330fn load_shard_fast(
331    shard_path: &Path,
332    tensor_filter: Option<&[String]>,
333    gpu: &dyn GpuBackend,
334    skip_fn: &dyn Fn(&str) -> bool,
335    defer_fn: &dyn Fn(&str, crate::weights::WeightDtype) -> bool,
336    try_direct_io: bool,
337    direct_io_tensor_cap: usize,
338    prefetch_shards: bool,
339    out: &mut HashMap<String, WeightTensor>,
340    deferred_out: &mut HashMap<String, crate::weights::DeferredTensor>,
341    offload_logged: &mut bool,
342) -> Result<()> {
343    // Header parsing uses a buffered fd — header is a few KB, cache pollution
344    // is negligible and buffered I/O handles short reads cleanly.
345    let mut meta_file = File::open(shard_path)
346        .with_context(|| format!("Failed to open {}", shard_path.display()))?;
347    let mut tensors = parse_header(&mut meta_file)?;
348    let file_size = meta_file.metadata()?.len();
349
350    // Filter down to tensors we actually want (index filter + EP filter).
351    if let Some(allow) = tensor_filter {
352        let allow_set: std::collections::HashSet<&str> = allow.iter().map(|s| s.as_str()).collect();
353        tensors.retain(|t| allow_set.contains(t.name.as_str()));
354    }
355    // The n-gram embedding TABLES are never uploaded with the checkpoint —
356    // 63 GB (LongCat-Lite) to ~102 GB (Flash-Next) of BF16 would exhaust a
357    // 121 GB unified box before any quantization could run, and the fallback
358    // on GB10 is managed memory, i.e. Linux swap, i.e. a kernel freeze. They
359    // are recorded with their on-disk location and served either by streaming
360    // per-table quantize-on-load or straight off NVMe by the row cache.
361    //
362    // The MODEL's loader can claim tensors too (`defer_fn`), for the other
363    // reason a tensor must not be swept: the bytes that reach the device are
364    // not the bytes on disk. `nvidia/GLM-5.3-Flash-NVFP4` ships the MTP
365    // block's routed experts full-width BF16 against a w4a16-only forward, so
366    // every one of them would be uploaded, read back, quantised and freed —
367    // ~7.25 GB per EP=2 rank resident across the whole 45-layer build.
368    let mut deferred_here: Vec<(String, crate::weights::DeferredTensor)> = Vec::new();
369    let mut ngram_count = 0usize;
370    #[allow(clippy::items_after_statements)]
371    tensors.retain(|t| {
372        let mut defer = |t: &header::TensorMeta| {
373            deferred_here.push((
374                t.name.clone(),
375                crate::weights::DeferredTensor {
376                    path: shard_path.to_path_buf(),
377                    offset: t.abs_offset,
378                    shape: t.shape.clone(),
379                    dtype: t.dtype,
380                },
381            ));
382        };
383        if crate::weights::is_ngram_table(&t.name) {
384            defer(t);
385            ngram_count += 1;
386            return false;
387        }
388        // 🪤 The model hook is asked only about tensors this rank KEEPS. Asked
389        // first, it would record a remote expert's on-disk location and invite
390        // a binder to read weights EP gave to another rank.
391        if skip_fn(&t.name) {
392            return false;
393        }
394        // 🪤 `from_f16` is not deferrable: the store's BF16 is a REWRITE of the
395        // disk bytes, so a (path, offset) locator would hand its reader F16.
396        // `estimate_load_bytes` refuses the same case, so the two agree.
397        if !t.from_f16 && defer_fn(&t.name, t.dtype) {
398            defer(t);
399            return false;
400        }
401        true
402    });
403    if !deferred_here.is_empty() {
404        tracing::info!(
405            "Deferred {} tensor(s) in {} ({} n-gram table(s), {} claimed by the model's \
406             weight loader) — served from disk, not uploaded",
407            deferred_here.len(),
408            shard_path.display(),
409            ngram_count,
410            deferred_here.len() - ngram_count,
411        );
412        deferred_out.extend(deferred_here);
413    }
414
415    // Per-shard heuristic: above `direct_io_tensor_cap` tensors, O_DIRECT's
416    // per-tensor syscall + 4 KiB alignment overhead costs more than kernel
417    // readahead on the buffered path saves. Skip the direct-open attempt
418    // entirely in that case — keeps the log clean and avoids a wasted fd.
419    let wants_direct = try_direct_io && tensors.len() <= direct_io_tensor_cap;
420    if try_direct_io && !wants_direct {
421        tracing::info!(
422            "  Shard has {} tensors (> {} cap) — using buffered+pipelined path",
423            tensors.len(),
424            direct_io_tensor_cap
425        );
426    }
427
428    // File for data reads. Try O_DIRECT; if it fails, fall through to buffered.
429    let (direct_file, using_direct) = match wants_direct
430        .then(|| direct_io::open_direct(shard_path))
431        .transpose()
432    {
433        Ok(Some(f)) => (Some(f), true),
434        Ok(None) => (None, false),
435        Err(e) => {
436            tracing::warn!(
437                "O_DIRECT open failed for {} ({e}); falling back to buffered reads",
438                shard_path.display()
439            );
440            (None, false)
441        }
442    };
443    let buffered_file = File::open(shard_path)?;
444    let data_fd = direct_file.as_ref().unwrap_or(&buffered_file);
445    if prefetch_shards && !using_direct {
446        advise_prefetch_shard(&buffered_file, shard_path, file_size);
447    }
448
449    // Pipelined reader: sends (tensor_index, aligned_buffer, slice_start) to main.
450    type ReadMsg = (usize, direct_io::AlignedBuffer, usize);
451    let (tx, rx) = sync_channel::<Result<ReadMsg>>(1);
452    let tensors_for_reader: Vec<(u64, usize)> =
453        tensors.iter().map(|t| (t.abs_offset, t.len)).collect();
454    let raw_fd = {
455        use std::os::unix::io::AsRawFd;
456        data_fd.as_raw_fd()
457    };
458
459    let _ = file_size; // retained for future use (tail-fragment buffered read)
460    let reader_handle = std::thread::spawn(move || {
461        for (idx, (abs_offset, len)) in tensors_for_reader.iter().enumerate() {
462            let msg = direct_io::read_tensor_aligned(raw_fd, *abs_offset, *len, using_direct)
463                .map(|(buf, slice_start)| (idx, buf, slice_start));
464            if tx.send(msg).is_err() {
465                break; // receiver dropped
466            }
467        }
468    });
469
470    // Copier: drains the channel, does gpu alloc + copy_h2d, inserts into the map.
471    for result in rx {
472        let (idx, buf, slice_start) = result?;
473        let meta = &tensors[idx];
474        let raw = &buf.as_slice()[slice_start..slice_start + meta.len];
475        // F16 shards: convert bytes to BF16 before upload (same length,
476        // different bit layout — meta.dtype is already staged as BF16).
477        let converted: Vec<u8>;
478        let src: &[u8] = if meta.from_f16 {
479            converted = f16_to_bf16_bytes(raw);
480            &converted
481        } else {
482            raw
483        };
484
485        let ptr = match gpu.alloc(meta.len) {
486            Ok(p) => {
487                gpu.copy_h2d(src, p)?;
488                p
489            }
490            Err(_) => {
491                if !*offload_logged {
492                    tracing::warn!(
493                        "GPU alloc failed for {} ({} bytes) — switching to managed (UVM) memory",
494                        meta.name,
495                        meta.len
496                    );
497                    *offload_logged = true;
498                }
499                let p = gpu.alloc_managed(meta.len)?;
500                unsafe {
501                    std::ptr::copy_nonoverlapping(src.as_ptr(), p.0 as *mut u8, meta.len);
502                }
503                p
504            }
505        };
506
507        out.insert(
508            meta.name.clone(),
509            WeightTensor {
510                ptr,
511                shape: meta.shape.clone(),
512                dtype: meta.dtype,
513            },
514        );
515    }
516
517    reader_handle
518        .join()
519        .map_err(|_| anyhow::anyhow!("reader thread panicked"))?;
520
521    // Release file handles, then advise the kernel to drop any pages we did
522    // end up caching on the buffered fallback path. O_DIRECT reads never hit
523    // the page cache, so the posix_fadvise is a no-op there but cheap.
524    drop(direct_file);
525    evict_page_cache(&buffered_file);
526    drop(buffered_file);
527    Ok(())
528}
529
530#[cfg(target_os = "linux")]
531fn advise_prefetch_shard(file: &File, shard_path: &Path, file_size: u64) {
532    use std::os::unix::io::AsRawFd;
533
534    let fd = file.as_raw_fd();
535    let seq_rc = unsafe { libc::posix_fadvise(fd, 0, 0, libc::POSIX_FADV_SEQUENTIAL) };
536    let willneed_rc = unsafe { libc::posix_fadvise(fd, 0, 0, libc::POSIX_FADV_WILLNEED) };
537    if seq_rc == 0 && willneed_rc == 0 {
538        tracing::info!(
539            "  NFS/shard prefetch requested for {} ({:.2} GB)",
540            shard_path.display(),
541            file_size as f64 / (1024.0 * 1024.0 * 1024.0)
542        );
543    } else {
544        tracing::warn!(
545            "  NFS/shard prefetch hint failed for {}: sequential_rc={}, willneed_rc={}",
546            shard_path.display(),
547            seq_rc,
548            willneed_rc
549        );
550    }
551}
552
553#[cfg(not(target_os = "linux"))]
554fn advise_prefetch_shard(_file: &File, _shard_path: &Path, _file_size: u64) {}
555
556#[cfg(test)]
557mod skip_vision_tests {
558    use super::{FastSafetensorsLoader, is_vision_tensor};
559
560    fn loader(skip_vision: bool, ep: usize) -> FastSafetensorsLoader {
561        let mut l = FastSafetensorsLoader::with_ep(0, ep, 288);
562        l.skip_vision = skip_vision;
563        l
564    }
565
566    #[test]
567    fn vision_names_are_recognised() {
568        assert!(is_vision_tensor("model.visual.blocks.0.attn.proj.weight"));
569        assert!(is_vision_tensor(
570            "model.vision_tower.encoder.layer.0.weight"
571        ));
572        assert!(is_vision_tensor("visual.merger.proj.weight"));
573        assert!(!is_vision_tensor(
574            "model.language_model.layers.45.eh_proj.weight"
575        ));
576        // The trap: a text tensor whose name merely CONTAINS "vision".
577        assert!(!is_vision_tensor(
578            "model.language_model.layers.3.mlp.revision.weight"
579        ));
580    }
581
582    #[test]
583    fn skip_vision_drops_only_the_tower() {
584        let l = loader(true, 2);
585        assert!(l.should_skip_tensor("model.visual.blocks.0.attn.proj.weight"));
586        assert!(!l.should_skip_tensor("model.language_model.layers.45.eh_proj.weight"));
587        assert!(!l.should_skip_tensor("lm_head.weight"));
588    }
589
590    #[test]
591    fn skip_vision_applies_without_ep() {
592        // The EP short-circuit must not swallow the vision rule at ep=1.
593        let l = loader(true, 1);
594        assert!(l.should_skip_tensor("model.visual.blocks.0.attn.proj.weight"));
595        assert!(!l.should_skip_tensor("model.layers.0.self_attn.q_proj.weight"));
596    }
597
598    #[test]
599    fn default_loader_keeps_the_tower() {
600        let l = loader(false, 2);
601        assert!(!l.should_skip_tensor("model.visual.blocks.0.attn.proj.weight"));
602        assert!(!FastSafetensorsLoader::new().skip_vision);
603    }
604
605    #[test]
606    fn ep_expert_filtering_is_unchanged_by_the_vision_rule() {
607        let l = loader(true, 2); // ep_rank 0 of 2, 288 experts -> keeps 0..143
608        assert!(
609            !l.should_skip_tensor("model.language_model.layers.4.mlp.experts.7.up_proj.weight")
610        );
611        assert!(
612            l.should_skip_tensor("model.language_model.layers.4.mlp.experts.200.up_proj.weight")
613        );
614    }
615}
spark_runtime/fast_weights/mod.rs

spark_runtime/fast_weights/
mod.rs
