metrale_model_layers/lora/
key.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: LoRA tensor-name classification and adapter identity:
4//! `classify_key` maps a PEFT tensor name to (layer, target, A or B), and
5//! `adapter_id_hash` is the adapter's cache-identity key.
6//!
7//! Owner: model-layers (lora).
8//! Invariants:
9//! - `adapter_id_hash` never returns 0, the base (no-adapter) id.
10
11use anyhow::{Result, anyhow, bail};
12use metrale_config::{LayerType, ModelConfig};
13
14use super::*;
15
16/// 2026-09-25: The adapter's cache-identity key, which keeps prefix and KV
17/// reuse within one adapter. It is FNV-1a over the adapter's name, then over
18/// `generation`'s little-endian bytes when `generation != 0`. It never depends
19/// on the pool slot, so an adapter moved to another slot keeps its id. A swap
20/// into a slot bumps that slot's generation, which changes the id. `0` is
21/// reserved for the base model, so a hash of 0 becomes 1.
22pub fn adapter_id_hash(name: &str, generation: u64) -> u64 {
23    let mut h: u64 = 0xcbf29ce484222325;
24    for &b in name.as_bytes() {
25        h ^= b as u64;
26        h = h.wrapping_mul(0x100000001b3);
27    }
28    if generation != 0 {
29        for &b in generation.to_le_bytes().iter() {
30            h ^= b as u64;
31            h = h.wrapping_mul(0x100000001b3);
32        }
33    }
34    if h == 0 { 1 } else { h }
35}
36
37/// 2026-09-25: Whether `key` names a `linear_attn.*` tensor other than
38/// `linear_attn.out_proj`, the family `classify_key` refuses with
39/// `REJECT[gdn-target]`.
40/// `audit_adapter` skips these under `METRALE_LORA_ALLOW_PARTIAL` instead of
41/// matching on the refusal message.
42pub fn is_gdn_key(key: &str) -> bool {
43    // 2026-09-25: The same splits as `classify_key`: PEFT prefix, then
44    // `.layers.`, then the layer index.
45    let Some(stripped) = key.strip_prefix("base_model.model.") else {
46        return false;
47    };
48    let Some((_prefix, rest)) = stripped.split_once(".layers.") else {
49        return false;
50    };
51    let Some((_idx, tail)) = rest.split_once('.') else {
52        return false;
53    };
54    // 2026-09-25: `linear_attn.out_proj` has a delta path, so it is not skipped.
55    tail.starts_with("linear_attn.") && tail != "linear_attn.out_proj"
56}
57
58/// 2026-09-25: PEFT tensor name → (layer, target, A or B). Every name it
59/// cannot place is an error naming the reason (`REJECT[...]`), never a skip.
60/// Whatever precedes `.layers.` is ignored, so `model.layers.{i}.*` and
61/// `model.language_model.layers.{i}.*` both parse; the layer index follows
62/// `.layers.`.
63pub fn classify_key(key: &str, cfg: &ModelConfig) -> Result<(usize, LoraTarget, AdapterAb)> {
64    let stripped = key.strip_prefix("base_model.model.").ok_or_else(|| {
65        anyhow!("REJECT[not-peft-key]: '{key}' lacks the 'base_model.model.' PEFT prefix")
66    })?;
67    if stripped.contains("lora_embedding_") {
68        bail!("REJECT[embedding-lora]: '{key}' — embed_tokens/lm_head LoRA is out of v0 scope");
69    }
70    let (module_path, ab) = if let Some(p) = stripped.strip_suffix(".lora_A.weight") {
71        (p, AdapterAb::A)
72    } else if let Some(p) = stripped.strip_suffix(".lora_B.weight") {
73        (p, AdapterAb::B)
74    } else {
75        bail!(
76            "REJECT[unrecognized-tensor]: '{key}' is not a lora_A/lora_B weight \
77             (modules_to_save exports and old '.lora_A.<adapter>.weight' layouts \
78             are not supported in v0)"
79        );
80    };
81    let (_prefix, rest) = module_path.split_once(".layers.").ok_or_else(|| {
82        anyhow!("REJECT[non-layer-module]: '{key}' targets '{module_path}' outside the layer stack")
83    })?;
84    let (idx_str, tail) = rest
85        .split_once('.')
86        .ok_or_else(|| anyhow!("REJECT[malformed-key]: '{key}'"))?;
87    let layer_idx: usize = idx_str
88        .parse()
89        .map_err(|_| anyhow!("REJECT[malformed-layer-index]: '{key}'"))?;
90    if layer_idx >= cfg.num_hidden_layers {
91        bail!(
92            "REJECT[layer-out-of-range]: '{key}' targets layer {layer_idx} \
93             (model has {})",
94            cfg.num_hidden_layers
95        );
96    }
97    let target = match tail {
98        "self_attn.q_proj" => LoraTarget::Attn(LoraModule::QProj),
99        "self_attn.k_proj" => LoraTarget::Attn(LoraModule::KProj),
100        "self_attn.v_proj" => LoraTarget::Attn(LoraModule::VProj),
101        "self_attn.o_proj" => LoraTarget::Attn(LoraModule::OProj),
102        "mlp.gate_proj" => LoraTarget::Attn(LoraModule::GateProj),
103        "mlp.up_proj" => LoraTarget::Attn(LoraModule::UpProj),
104        "mlp.down_proj" => LoraTarget::Attn(LoraModule::DownProj),
105        // 2026-09-25: The MoE router, not the dense `mlp.gate_proj` above.
106        "mlp.gate" => LoraTarget::Router,
107        t if t.starts_with("mlp.experts.") => {
108            classify_expert_tail(key, &t["mlp.experts.".len()..], cfg)?
109        }
110        // 2026-09-25: The GDN output projection (value_dim -> hidden) runs after
111        // the recurrence, so its delta never enters the state update.
112        // `in_proj_*` and `conv1d` feed the recurrence and are refused.
113        "linear_attn.out_proj" => LoraTarget::Attn(LoraModule::OutProj),
114        t if t.starts_with("linear_attn.") => bail!(
115            "REJECT[gdn-target]: '{key}' — GDN/linear-attention INPUT-side \
116             projections (in_proj_qkv / in_proj_z / in_proj_a / in_proj_b / \
117             conv1d) feed the recurrence and stay rejected until an \
118             exact-replay parity harness exists. `out_proj` IS supported."
119        ),
120        other => bail!("REJECT[unsupported-module]: '{key}' targets '{other}'"),
121    };
122    // 2026-09-25: Layer-type rules. Router and expert targets have none: a
123    // linear-attention layer (`Qwen3SsmLayer`) can hold a MoE FFN too.
124    match target {
125        // 2026-09-25: A dense-FFN target on a dense model (`num_experts == 0`)
126        // is allowed on any layer: `Qwen3SsmLayer` installs it on its
127        // `FfnComponent::Dense` (`layers/qwen3_ssm/lora.rs`). On a MoE model it
128        // falls to the full-attention rule below.
129        LoraTarget::Attn(m) if m.is_dense_ffn() && cfg.num_experts == 0 => {}
130        // 2026-09-25: GDN `out_proj` only on linear-attention layers.
131        LoraTarget::Attn(m)
132            if m.is_gdn_out() && cfg.layer_type(layer_idx) == LayerType::FullAttention =>
133        {
134            bail!(
135                "REJECT[gdn-out-on-attention-layer]: '{key}' targets layer \
136                 {layer_idx}, which is a full-attention layer with no GDN out_proj"
137            )
138        }
139        LoraTarget::Attn(m) if m.is_gdn_out() => {}
140        LoraTarget::Attn(_) => match cfg.layer_type(layer_idx) {
141            LayerType::FullAttention => {}
142            lt => bail!(
143                "REJECT[non-full-attention-layer]: '{key}' targets layer {layer_idx} \
144                 ({lt:?}); attention-projection LoRA applies only on the full-attention \
145                 layers {:?}",
146                full_attention_layers(cfg),
147            ),
148        },
149        LoraTarget::Router | LoraTarget::Expert { .. } => {}
150    }
151    Ok((layer_idx, target, ab))
152}
153
154/// 2026-09-25: Parse the `{N}.{proj}` remainder of an `mlp.experts.` tail
155/// (e.g. `"7.gate_proj"`) into a routed-expert [`LoraTarget`]. Refused: a dense
156/// model (`num_experts == 0`), a fused or unindexed layout, a non-numeric or
157/// out-of-range expert index, and any projection but gate/up/down_proj.
158fn classify_expert_tail(key: &str, rest: &str, cfg: &ModelConfig) -> Result<LoraTarget> {
159    if cfg.num_experts == 0 {
160        bail!(
161            "REJECT[expert-lora-on-dense-model]: '{key}' targets a routed expert but \
162             the model has num_experts=0 (dense) — use mlp.{{gate,up,down}}_proj instead"
163        );
164    }
165    let (n_str, proj_str) = rest.split_once('.').ok_or_else(|| {
166        anyhow!(
167            "REJECT[fused-expert-lora]: '{key}' — fused/unindexed expert layout \
168             (e.g. experts.gate_up_proj via target_parameters) is deferred to \
169             Feature-1 phase 3; export per-expert mlp.experts.{{N}}.{{proj}} tensors"
170        )
171    })?;
172    let n: usize = n_str.parse().map_err(|_| {
173        anyhow!("REJECT[malformed-expert-index]: '{key}' — '{n_str}' is not an index")
174    })?;
175    if n >= cfg.num_experts {
176        bail!(
177            "REJECT[expert-out-of-range]: '{key}' targets expert {n} \
178             (model has {} experts)",
179            cfg.num_experts
180        );
181    }
182    let proj = match proj_str {
183        "gate_proj" => ExpertProj::Gate,
184        "up_proj" => ExpertProj::Up,
185        "down_proj" => ExpertProj::Down,
186        "gate_up_proj" => bail!(
187            "REJECT[fused-expert-lora]: '{key}' — fused gate_up_proj is deferred to \
188             Feature-1 phase 3 (needs a per-expert decomposer)"
189        ),
190        other => bail!("REJECT[unsupported-expert-proj]: '{key}' proj '{other}'"),
191    };
192    Ok(LoraTarget::Expert { n: n as u16, proj })
193}
194
195#[cfg(test)]
196#[path = "key_tests.rs"]
197mod tests;
metrale_model_layers/lora/key.rs

metrale_model_layers/lora/
key.rs
