spark_model/quant_format/
modelopt.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2
3//! NVIDIA TensorRT ModelOpt NVFP4 serialization.
4//!
5//! ModelOpt exports are emitted by NVIDIA's `nvidia-modelopt` toolkit and
6//! appear on HuggingFace under `nvidia/*` and downstream community
7//! re-quants (`lukealonso/*`, `saricles/*`, `NVIDIA/*`). The tensor-name
8//! convention is distinct from compressed-tensors:
9//!
10//! | field                | tensor name          | dtype        |
11//! | -------------------- | -------------------- | ------------ |
12//! | packed FP4 payload   | `.weight`            | uint8 packed |
13//! | per-group FP8 scales | `.weight_scale`      | float8_e4m3  |
14//! | per-tensor scalar    | `.weight_scale_2`    | f32 scalar   |
15//! | activation scale     | `.input_scale`       | f32 scalar   |
16//!
17//! Unquantized modules (typically `lm_head`, embeddings, and the full
18//! attention tower on smaller re-quants) are declared in the top-level
19//! `ignore` array of `hf_quant_config.json` / `config.json`'s
20//! `quantization_config` block. Those modules ship as plain BF16 and
21//! must NOT be run through the NVFP4 loader — reading uint8-packed FP4
22//! at BF16 stride is a 4× byte overrun that lands as
23//! `CUDA_ERROR_ILLEGAL_ADDRESS` later on, which is the bug this module
24//! was written to eliminate.
25//!
26//! This maps to the existing [`Nvfp4Variant::Standard`] dispatch in
27//! `weight_map.rs` (which was always the ModelOpt path — the misleading
28//! name predates support for the compressed-tensors split).
29
30use crate::quant_format::{QuantFormat, module_matches_pattern};
31use crate::weight_map::Nvfp4Variant;
32
33/// ModelOpt-style NVFP4 checkpoint.
34#[derive(Debug)]
35pub struct ModeloptFormat {
36    /// `quant_algo` declared in config (`"NVFP4"`, `"FP8"`, …). Used for
37    /// diagnostic logging — the actual dispatch uses `base_variant`.
38    pub algo: String,
39    /// Module-path globs to load as dense BF16 rather than NVFP4.
40    pub ignore_modules: Vec<String>,
41}
42
43impl ModeloptFormat {
44    pub fn new(algo: String, ignore_modules: Vec<String>) -> Self {
45        Self {
46            algo,
47            ignore_modules,
48        }
49    }
50}
51
52impl QuantFormat for ModeloptFormat {
53    fn name(&self) -> &'static str {
54        "modelopt"
55    }
56
57    fn base_variant(&self) -> Nvfp4Variant {
58        // `Standard` is the legacy name for ModelOpt NVFP4 layout.
59        Nvfp4Variant::Standard
60    }
61
62    fn is_ignored(&self, module_path: &str) -> bool {
63        self.ignore_modules
64            .iter()
65            .any(|pat| module_matches_pattern(module_path, pat))
66    }
67}
spark_model/quant_format/modelopt.rs

spark_model/quant_format/
modelopt.rs
