metrale_config/precision_plan/
target.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-28: Module targets of a declared precision plan and how they match.
4//!
5//! Owner: config (quantization).
6//! Invariants: none beyond the types.
7
8use anyhow::{Context, Result};
9
10/// 2026-09-28: One target or ignore entry.
11#[derive(Debug, Clone)]
12pub enum Target {
13    /// 2026-09-28: An exact module path.
14    Exact(String),
15    /// 2026-09-28: A compressed-tensors `re:` pattern, anchored at both ends as Python's
16    /// `re.match` plus the checkpoint's `$` do.
17    Regex(String, regex::Regex),
18    /// 2026-09-28: A ModelOpt glob (`*`, `?`).
19    Glob(String),
20    /// 2026-09-30: An HF `fp8` `modules_to_not_convert` / `ignored_layers` entry: a module
21    /// path matched on whole dotted segments. It names the module itself, a parent of it
22    /// (`model.visual` excludes `model.visual.blocks.0.attn.proj`), or its trailing segments
23    /// (`lm_head`); `mlp.gate` names the MoE router `...mlp.gate`, never the dense
24    /// `...mlp.gate_proj`. These are the three rules of transformers `should_convert_module`
25    /// (quantizers/quantizers_utils.py: `re.match(f"{key}\\.")`, `re.match(key)`,
26    /// `name.endswith(key)`) restricted to segment boundaries, and include vLLM's exact
27    /// `prefix in ignored_layers` (`is_layer_skipped`, quantization/utils/quant_utils.py). An
28    /// entry with regex syntax beyond `.` is that regex, anchored at the start (`re.match`) and
29    /// ending on a segment boundary.
30    HfModule(String, Option<regex::Regex>),
31    /// 2026-09-28: A module class (`Linear`). Every module the plan is asked about is a
32    /// linear layer, so `Linear` matches all of them and any other class none.
33    Class(String),
34}
35
36/// 2026-09-30: Which format's matching rules an ignore / exclude entry follows.
37#[derive(Debug, Clone, Copy, PartialEq, Eq)]
38pub enum IgnoreDialect {
39    /// 2026-09-30: compressed-tensors `ignore` ([`Target::compressed_tensors`]).
40    CompressedTensors,
41    /// 2026-09-30: ModelOpt `exclude_modules` ([`Target::modelopt`]).
42    ModelOpt,
43    /// 2026-09-30: HF `fp8` `modules_to_not_convert` / `ignored_layers` ([`Target::hf_module`]).
44    HfFp8,
45}
46
47impl Target {
48    /// 2026-09-30: An ignore entry of `dialect`, the one matcher the precision plan and the
49    /// weight loaders share. A malformed entry is refused.
50    pub fn ignore_entry(dialect: IgnoreDialect, s: &str) -> Result<Self> {
51        match dialect {
52            IgnoreDialect::CompressedTensors => Self::compressed_tensors(s),
53            IgnoreDialect::ModelOpt => Ok(Self::modelopt(s)),
54            IgnoreDialect::HfFp8 => Self::hf_module(s),
55        }
56    }
57
58    /// 2026-09-28: A compressed-tensors target or ignore entry: `re:<pattern>`, a class
59    /// name (`Linear`), or an exact path.
60    /// 2026-09-30: As compressed-tensors matches them (0.12.2, utils/match.py `_match_name`):
61    /// a `re:` pattern by `re.match` (anchored at the start only), anything else by equality.
62    pub fn compressed_tensors(s: &str) -> Result<Self> {
63        if let Some(p) = s.strip_prefix("re:") {
64            let re = regex::Regex::new(&format!("^(?:{p})"))
65                .with_context(|| format!("quantization_config: bad target pattern {s:?}"))?;
66            return Ok(Target::Regex(s.to_string(), re));
67        }
68        if !s.contains('.') && s.chars().next().is_some_and(char::is_uppercase) {
69            return Ok(Target::Class(s.to_string()));
70        }
71        Ok(Target::Exact(s.to_string()))
72    }
73
74    /// 2026-09-30: An HF `fp8` `modules_to_not_convert` / `ignored_layers` entry
75    /// ([`Target::HfModule`]). An empty entry, or one whose regex does not compile, is refused.
76    pub fn hf_module(s: &str) -> Result<Self> {
77        anyhow::ensure!(
78            !s.trim().is_empty(),
79            "quantization_config: an empty modules_to_not_convert entry"
80        );
81        let regex_syntax =
82            s.contains(['*', '?', '[', ']', '(', ')', '+', '^', '$', '|', '\\', '{']);
83        let re = if regex_syntax {
84            Some(
85                regex::Regex::new(&format!("^(?:{s})(?:\\.|$)"))
86                    .with_context(|| format!("quantization_config: bad module pattern {s:?}"))?,
87            )
88        } else {
89            None
90        };
91        Ok(Target::HfModule(s.to_string(), re))
92    }
93
94    /// 2026-09-28: A ModelOpt entry: a glob when it has a wildcard, else exact.
95    /// 2026-09-30: ModelOpt writes exact module names or fnmatch globs; vLLM's
96    /// `ModelOptFp8Config.is_layer_excluded` (quantization/modelopt.py) adds a substring
97    /// fallback, which is not followed here: it has the `mlp.gate` / `mlp.gate_proj` hazard.
98    pub fn modelopt(s: &str) -> Self {
99        if s.contains(['*', '?']) {
100            Target::Glob(s.to_string())
101        } else if s == "Linear" {
102            Target::Class(s.to_string())
103        } else {
104            Target::Exact(s.to_string())
105        }
106    }
107
108    /// 2026-09-28: The entry as the checkpoint wrote it.
109    pub fn text(&self) -> &str {
110        match self {
111            Target::Exact(s)
112            | Target::Regex(s, _)
113            | Target::Glob(s)
114            | Target::HfModule(s, _)
115            | Target::Class(s) => s,
116        }
117    }
118
119    /// 2026-09-28: Does this entry name `module` (classes excluded)?
120    pub fn matches_name(&self, module: &str) -> bool {
121        match self {
122            Target::Exact(s) => s == module,
123            Target::Regex(_, re) => re.is_match(module),
124            Target::Glob(g) => glob_match(g.as_bytes(), module.as_bytes()),
125            Target::HfModule(_, Some(re)) => re.is_match(module),
126            Target::HfModule(s, None) => {
127                module == s
128                    || module
129                        .strip_prefix(s.as_str())
130                        .is_some_and(|rest| rest.starts_with('.'))
131                    || module
132                        .strip_suffix(s.as_str())
133                        .is_some_and(|head| head.ends_with('.'))
134            }
135            Target::Class(c) => c == "Linear",
136        }
137    }
138
139    /// 2026-09-28: Precedence tier when this entry matches `module`: 0 exact name, 1
140    /// pattern, 2 class. `None` when it does not match.
141    pub fn tier_for(&self, module: &str) -> Option<u8> {
142        if !self.matches_name(module) {
143            return None;
144        }
145        Some(match self {
146            Target::Exact(_) => 0,
147            Target::HfModule(s, None) if s == module => 0,
148            Target::Regex(..) | Target::Glob(_) | Target::HfModule(..) => 1,
149            Target::Class(_) => 2,
150        })
151    }
152}
153
154/// 2026-09-28: `*` matches any run, `?` one byte; everything else literally.
155fn glob_match(p: &[u8], s: &[u8]) -> bool {
156    let (mut pi, mut si) = (0usize, 0usize);
157    let (mut star, mut mark) = (None, 0usize);
158    while si < s.len() {
159        if pi < p.len() && (p[pi] == b'?' || p[pi] == s[si]) {
160            pi += 1;
161            si += 1;
162        } else if pi < p.len() && p[pi] == b'*' {
163            star = Some(pi);
164            mark = si;
165            pi += 1;
166        } else if let Some(sp) = star {
167            pi = sp + 1;
168            mark += 1;
169            si = mark;
170        } else {
171            return false;
172        }
173    }
174    p[pi..].iter().all(|&c| c == b'*')
175}
metrale_config/precision_plan/target.rs

metrale_config/precision_plan/
target.rs
