metrale_model_layers/lora/
overlay.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: The GPU-free half of the token overlay, which replaces whole
4//! embed_tokens / lm_head rows for PEFT `trainable_tokens` and
5//! `modules_to_save` adapters: tensor-name classification
6//! ([`classify_overlay_key`]), the per-adapter collection ([`OverlayTensors`]),
7//! and the choice of overridden rows and their sources
8//! ([`clamp_trainable_to_vocab`], [`build_override_set`], [`override_source`]).
9//! The device build is in [`super::overlay_build`].
10//!
11//! Tensor kinds:
12//! - `…token_adapter.base_layer.weight` (`[R, h]`) and
13//!   `…token_adapter.trainable_tokens_delta` (`[T, h]`): trainable id `k`'s
14//!   row is replaced by delta row `k`, not added to.
15//! - a `modules_to_save` `…embed_tokens.weight` / `lm_head.weight`: a full
16//!   table; the rows that differ from the served table are replaced.
17//! - `lora_embedding_A/B`: classified so the audit can refuse it.
18//!
19//! Owner: model-layers (lora).
20//! Invariants: none beyond the types.
21
22use anyhow::{Result, bail};
23
24/// 2026-09-25: A base row whose largest absolute difference from the served
25/// row exceeds this is overridden (`embed_rowdiff_bf16`).
26pub const ROWDIFF_THRESH: f32 = 0.1;
27
28/// 2026-09-25: The table an overlay tensor belongs to.
29#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
30pub enum OverlayModule {
31    /// 2026-09-25: The input embedding (a name containing `embed_tokens` or
32    /// `.shared.`).
33    EmbedTokens,
34    /// 2026-09-25: The output projection (a name containing `lm_head`).
35    LmHead,
36}
37
38/// 2026-09-25: The role an overlay tensor plays in
39/// [`build_overlay`](super::overlay_build::build_overlay).
40#[derive(Debug, Clone, Copy, PartialEq, Eq)]
41pub enum OverlayTensorKind {
42    /// 2026-09-25: `token_adapter.base_layer.weight`, the adapter's `[R, h]`
43    /// table.
44    Base,
45    /// 2026-09-25: `token_adapter.trainable_tokens_delta`, `[T, h]`
46    /// replacement rows.
47    Delta,
48    /// 2026-09-25: A `modules_to_save` full table (`embed_tokens.weight` /
49    /// `lm_head.weight`).
50    FullSave,
51    /// 2026-09-25: `lora_embedding_A`; refused at audit.
52    LoraEmbedA,
53    /// 2026-09-25: `lora_embedding_B`; refused at audit.
54    LoraEmbedB,
55}
56
57/// 2026-09-25: A classified overlay tensor: its table and its role.
58#[derive(Debug, Clone, Copy, PartialEq, Eq)]
59pub struct OverlayTensor {
60    pub module: OverlayModule,
61    pub kind: OverlayTensorKind,
62}
63
64/// 2026-09-25: Classify a PEFT tensor name as an overlay tensor. `None` for a
65/// name without the `base_model.model.` prefix, a `lora_A`/`lora_B` weight
66/// (for [`super::classify_key`]), or a name matching no overlay form. Only the
67/// module name and suffix are matched, so `model.embed_tokens` and
68/// `model.language_model.embed_tokens` both work. The `token_adapter` and
69/// `lora_embedding` suffixes are tested before the bare `.weight` form, so
70/// `…token_adapter.base_layer.weight` is not taken for a full table.
71pub fn classify_overlay_key(key: &str) -> Option<OverlayTensor> {
72    let stripped = key.strip_prefix("base_model.model.")?;
73    if stripped.ends_with(".lora_A.weight") || stripped.ends_with(".lora_B.weight") {
74        return None;
75    }
76    let module = overlay_module_of(stripped)?;
77    let kind = if stripped.ends_with(".token_adapter.base_layer.weight") {
78        OverlayTensorKind::Base
79    } else if stripped.ends_with(".token_adapter.trainable_tokens_delta") {
80        OverlayTensorKind::Delta
81    } else if stripped.ends_with(".lora_embedding_A") {
82        OverlayTensorKind::LoraEmbedA
83    } else if stripped.ends_with(".lora_embedding_B") {
84        OverlayTensorKind::LoraEmbedB
85    } else if is_full_module_weight(stripped, module) {
86        OverlayTensorKind::FullSave
87    } else {
88        return None;
89    };
90    Some(OverlayTensor { module, kind })
91}
92
93/// 2026-09-25: The table a prefix-stripped name targets; `lm_head` is tested
94/// first.
95fn overlay_module_of(stripped: &str) -> Option<OverlayModule> {
96    if stripped.contains("lm_head") {
97        Some(OverlayModule::LmHead)
98    } else if stripped.contains("embed_tokens") || stripped.contains(".shared.") {
99        Some(OverlayModule::EmbedTokens)
100    } else {
101        None
102    }
103}
104
105/// 2026-09-25: True when `stripped` ends with the module's own weight
106/// (`embed_tokens.weight` / `lm_head.weight`) and does not contain `.layers.`.
107fn is_full_module_weight(stripped: &str, module: OverlayModule) -> bool {
108    if stripped.contains(".layers.") {
109        return false;
110    }
111    let leaf_owner = match module {
112        OverlayModule::EmbedTokens => "embed_tokens.weight",
113        OverlayModule::LmHead => "lm_head.weight",
114    };
115    stripped.ends_with(leaf_owner)
116}
117
118/// 2026-09-25: One adapter's overlay tensor names by (table, role); each
119/// field is set at most once.
120#[derive(Debug, Default, Clone)]
121pub struct OverlayTensors {
122    pub embed_base: Option<String>,
123    pub embed_delta: Option<String>,
124    pub embed_full: Option<String>,
125    pub lmhead_base: Option<String>,
126    pub lmhead_delta: Option<String>,
127    pub lmhead_full: Option<String>,
128    /// 2026-09-25: A `lora_embedding_A/B` tensor was seen.
129    pub lora_embedding_seen: bool,
130}
131
132impl OverlayTensors {
133    /// 2026-09-25: Record one classified tensor; an error when its (table,
134    /// role) is already set.
135    pub fn insert(&mut self, t: OverlayTensor, name: &str) -> Result<()> {
136        use OverlayModule::*;
137        use OverlayTensorKind::*;
138        let slot = match (t.module, t.kind) {
139            (_, LoraEmbedA) | (_, LoraEmbedB) => {
140                self.lora_embedding_seen = true;
141                return Ok(());
142            }
143            (EmbedTokens, Base) => &mut self.embed_base,
144            (EmbedTokens, Delta) => &mut self.embed_delta,
145            (EmbedTokens, FullSave) => &mut self.embed_full,
146            (LmHead, Base) => &mut self.lmhead_base,
147            (LmHead, Delta) => &mut self.lmhead_delta,
148            (LmHead, FullSave) => &mut self.lmhead_full,
149        };
150        if slot.is_some() {
151            bail!(
152                "REJECT[duplicate-overlay-tensor]: two tensors map to {:?}/{:?}",
153                t.module,
154                t.kind
155            );
156        }
157        *slot = Some(name.to_string());
158        Ok(())
159    }
160
161    /// 2026-09-25: True when no overlay tensor was collected.
162    pub fn is_empty(&self) -> bool {
163        self.embed_base.is_none()
164            && self.embed_delta.is_none()
165            && self.embed_full.is_none()
166            && self.lmhead_base.is_none()
167            && self.lmhead_delta.is_none()
168            && self.lmhead_full.is_none()
169            && !self.lora_embedding_seen
170    }
171}
172
173/// 2026-09-25: An error when the adapter has a `lora_embedding_A/B` tensor;
174/// every other overlay tensor is loaded.
175pub fn reject_pending_overlay(overlay: &OverlayTensors) -> Result<()> {
176    if overlay.lora_embedding_seen {
177        bail!(
178            "REJECT[lora-embedding-unimplemented]: classic low-rank embedding LoRA \
179             (lora_embedding_A/B) is not yet supported on the decoder path"
180        );
181    }
182    Ok(())
183}
184
185/// 2026-09-25: Split `trainable` ids at the served vocab, keeping list order,
186/// since delta row `k` belongs to the `k`-th id. Returns `(kept_ids,
187/// skipped_count)`:
188/// - `idx >= r` (outside the adapter's `[R, h]` table) is an error;
189/// - `vocab <= idx < r` is dropped and counted;
190/// - `idx < vocab` is kept.
191///
192/// Also an error: a kept id after a dropped one (dropped ids must be one
193/// trailing run, so kept id `k` still matches delta row `k`), and a kept id
194/// not greater than the previous kept id.
195pub fn clamp_trainable_to_vocab(
196    trainable: &[u32],
197    r: usize,
198    vocab: usize,
199) -> Result<(Vec<u32>, usize)> {
200    let mut kept = Vec::new();
201    let mut skipped = 0usize;
202    let mut last_kept: Option<u32> = None;
203    for &idx in trainable {
204        let i = idx as usize;
205        if i >= r {
206            bail!("REJECT[trainable-index-out-of-adapter]: id {idx} >= adapter embedding rows {r}");
207        }
208        if i >= vocab {
209            skipped += 1;
210            continue;
211        }
212        if skipped > 0 {
213            bail!(
214                "REJECT[trainable-order]: served-vocab id {idx} appears after a skipped \
215                 vocab-extension id; extension ids must form one trailing suffix"
216            );
217        }
218        if let Some(prev) = last_kept
219            && idx <= prev
220        {
221            bail!(
222                "REJECT[trainable-order]: kept id {idx} follows id {prev}; \
223                 PEFT trainable-token order must be strictly ascending in the served-vocab prefix"
224            );
225        }
226        last_kept = Some(idx);
227        kept.push(idx);
228    }
229    Ok((kept, skipped))
230}
231
232/// 2026-09-25: The vocab rows the overlay replaces: the indices `i` with
233/// `row_diff[i]` set, plus `trainable`, sorted ascending without duplicates.
234pub fn build_override_set(row_diff: &[bool], trainable: &[u32]) -> Vec<u32> {
235    let mut ids: Vec<u32> = row_diff
236        .iter()
237        .enumerate()
238        .filter_map(|(i, &d)| d.then_some(i as u32))
239        .collect();
240    ids.extend_from_slice(trainable);
241    ids.sort_unstable();
242    ids.dedup();
243    ids
244}
245
246/// 2026-09-25: Where an overridden id's row comes from: `Some(k)`, delta row
247/// `k`, when `id` is the `k`-th trainable id (this wins over a differing base
248/// row); `None`, the adapter's base row `id`.
249pub fn override_source(id: u32, trainable: &[u32]) -> Option<usize> {
250    trainable.iter().position(|&t| t == id)
251}
252
253#[cfg(test)]
254#[path = "overlay_tests.rs"]
255mod tests;
metrale_model_layers/lora/overlay.rs

metrale_model_layers/lora/
overlay.rs
