metrale_model_layers/quant_format/
mod.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: The quantized-weight layout of a checkpoint, and its ignore
4//! list: [`CompressedTensorsFormat`], [`ModeloptFormat`] or
5//! [`Fp8BlockScaledFormat`], chosen by [`detect_quant_format`].
6//!
7//! Owner: model-layers (weight loading).
8//! Invariants: none beyond the types.
9
10use metrale_config::ModelConfig;
11use metrale_config::precision_plan::{IgnoreDialect, Target};
12use metrale_model_weights::weights::WeightStore;
13
14use crate::weight_map::Nvfp4Variant;
15
16mod compressed_tensors;
17mod fp8_blockscaled;
18mod modelopt;
19
20pub use compressed_tensors::CompressedTensorsFormat;
21pub use fp8_blockscaled::Fp8BlockScaledFormat;
22pub use modelopt::ModeloptFormat;
23
24/// 2026-09-25: A quantized-weight layout: its log name, the [`Nvfp4Variant`]
25/// it maps to, and the module globs its ignore list keeps unquantized. The
26/// weight loaders pick their variant through `weight_map::detect_nvfp4_variant`,
27/// not through this trait.
28pub trait QuantFormat: Send + Sync + std::fmt::Debug {
29    /// 2026-09-25: Name for logs.
30    fn name(&self) -> &'static str;
31
32    /// 2026-09-25: The [`Nvfp4Variant`] this layout maps to.
33    fn base_variant(&self) -> Nvfp4Variant;
34
35    /// 2026-09-30: Whether the ignore list names the module `module_path` (a module path,
36    /// not a tensor name: no `.weight`), under the format's own matching ([`IgnoreList`]).
37    fn is_ignored(&self, module_path: &str) -> bool;
38
39    /// 2026-09-25: `Bf16Raw` for an ignored module, else [`Self::base_variant`].
40    fn variant_for(&self, module_path: &str) -> Nvfp4Variant {
41        if self.is_ignored(module_path) {
42            Nvfp4Variant::Bf16Raw
43        } else {
44            self.base_variant()
45        }
46    }
47}
48
49/// 2026-09-25: Pick the [`QuantFormat`] for a checkpoint.
50///
51/// 1. A `quantization_config` whose `quant_method` is `modelopt`,
52///    `compressed-tensors` or `fp8` selects that layout directly. The
53///    `quant_algo` and `format` are logged or stored, not used to choose.
54/// 2. Otherwise the layout comes from `weight_map::detect_nvfp4_variant`
55///    (an unrecognised method is warned about first).
56/// 3. When that finds no quantized weights (`Bf16Raw`), it warns and returns
57///    a [`ModeloptFormat`].
58///
59/// The ignore list is the config's `ignore_modules`, or empty without a
60/// config.
61pub fn detect_quant_format(
62    config: &ModelConfig,
63    store: &WeightStore,
64) -> anyhow::Result<Box<dyn QuantFormat>> {
65    if let Some(qc) = &config.quantization_config {
66        let method = qc.quant_method.as_str();
67        let algo = qc.quant_algo.as_str();
68        let format = qc.format.as_str();
69        let ignore = qc.ignore_modules.clone();
70
71        match method {
72            "modelopt" => {
73                tracing::info!(
74                    "QuantFormat: modelopt (algo={algo:?}), {} ignored module(s)",
75                    ignore.len(),
76                );
77                return Ok(Box::new(ModeloptFormat::new(algo.to_string(), &ignore)?));
78            }
79            "compressed-tensors" => {
80                tracing::info!(
81                    "QuantFormat: compressed-tensors (format={format:?}), {} ignored module(s)",
82                    ignore.len(),
83                );
84                return Ok(Box::new(CompressedTensorsFormat::new(
85                    format.to_string(),
86                    &ignore,
87                )?));
88            }
89            "fp8" => {
90                tracing::info!(
91                    "QuantFormat: fp8 (block-scaled), {} ignored module(s)",
92                    ignore.len(),
93                );
94                return Ok(Box::new(Fp8BlockScaledFormat::new(&ignore)?));
95            }
96            other if !other.is_empty() => {
97                tracing::warn!(
98                    "QuantFormat: config declares unrecognized quant_method={other:?}; \
99                     falling back to tensor-name heuristic. Metrale Engine currently understands \
100                     {{compressed-tensors, modelopt, fp8}}. Checkpoint load may fail."
101                );
102            }
103            _ => {
104                // 2026-09-25: An empty `quant_method` falls through to the
105                // tensor-name detection below.
106            }
107        }
108    }
109
110    let variant = crate::weight_map::detect_nvfp4_variant(store, config);
111    let ignore = config
112        .quantization_config
113        .as_ref()
114        .map(|qc| qc.ignore_modules.clone())
115        .unwrap_or_default();
116    Ok(match variant {
117        Nvfp4Variant::CompressedTensors => {
118            tracing::info!("QuantFormat: compressed-tensors (detected from tensor names)");
119            Box::new(CompressedTensorsFormat::new(String::new(), &ignore)?)
120        }
121        Nvfp4Variant::Fp8Dequanted => {
122            tracing::info!("QuantFormat: fp8-blockscaled (detected from tensor names)");
123            Box::new(Fp8BlockScaledFormat::new(&ignore)?)
124        }
125        Nvfp4Variant::Standard => {
126            tracing::info!("QuantFormat: modelopt-style NVFP4 (detected from tensor names)");
127            Box::new(ModeloptFormat::new(String::new(), &ignore)?)
128        }
129        Nvfp4Variant::Bf16Raw => {
130            tracing::warn!(
131                "QuantFormat: no quantization declared and no pre-quantized weights found; \
132                 treating checkpoint as BF16 raw (weights will be runtime-quantized). \
133                 Quality will be inferior to a calibrated NVFP4 release."
134            );
135            Box::new(ModeloptFormat::new(String::new(), &ignore)?) as Box<dyn QuantFormat>
136        }
137    })
138}
139
140/// 2026-09-30: A checkpoint's ignore list, parsed with the matching rules of its format
141/// (`metrale_config::precision_plan::Target`, the matcher the declared precision plan uses):
142/// compressed-tensors `re:` regexes or exact names, ModelOpt exact names or globs, HF `fp8`
143/// module paths on whole dotted segments.
144#[derive(Debug, Clone)]
145pub struct IgnoreList(Vec<Target>);
146
147impl IgnoreList {
148    /// 2026-09-30: Parse `entries` under `dialect`; a malformed entry is refused.
149    pub fn new(dialect: IgnoreDialect, entries: &[String]) -> anyhow::Result<Self> {
150        entries
151            .iter()
152            .map(|e| Target::ignore_entry(dialect, e))
153            .collect::<anyhow::Result<_>>()
154            .map(Self)
155    }
156
157    /// 2026-09-30: Whether an entry names the module `module_path`.
158    pub fn matches(&self, module_path: &str) -> bool {
159        self.0.iter().any(|t| t.matches_name(module_path))
160    }
161}
162
163#[cfg(test)]
164#[path = "quant_format_tests.rs"]
165mod tests;
metrale_model_layers/quant_format/mod.rs

metrale_model_layers/quant_format/
mod.rs
