metrale_model_layers/lora/
target.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: What a PEFT LoRA tensor targets: a [`LoraModule`], the MoE
4//! router (`mlp.gate`), or one routed expert's projection
5//! (`mlp.experts.{n}.{gate,up,down}_proj`), plus the dims and pool sizing for
6//! the router and expert targets.
7//!
8//! Owner: model-layers (lora).
9//! Invariants: none beyond the types.
10
11use std::collections::BTreeMap;
12
13use metrale_config::ModelConfig;
14
15use crate::layers::ops::lora_delta::LoraPair;
16
17use super::LoraModule;
18
19/// 2026-09-25: Which routed-expert projection a delta targets.
20#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
21pub enum ExpertProj {
22    Gate,
23    Up,
24    Down,
25}
26
27impl ExpertProj {
28    /// 2026-09-25: The projection's name in a PEFT tensor key.
29    pub fn peft_name(&self) -> &'static str {
30        match self {
31            Self::Gate => "gate_proj",
32            Self::Up => "up_proj",
33            Self::Down => "down_proj",
34        }
35    }
36
37    /// 2026-09-25: `(out_dim, in_dim)` of the routed-expert projection on
38    /// `layer`. The width is [`ModelConfig::moe_intermediate_size_for`], not
39    /// the dense `intermediate_size`.
40    pub fn dims(&self, cfg: &ModelConfig, layer: usize) -> (usize, usize) {
41        let h = cfg.hidden_size;
42        let inter = cfg.moe_intermediate_size_for(layer);
43        match self {
44            Self::Gate | Self::Up => (inter, h),
45            Self::Down => (h, inter),
46        }
47    }
48}
49
50/// 2026-09-25: `(out_dim, in_dim)` of the router (`mlp.gate`) projection. Its
51/// delta is added to the routing logits before top-k.
52pub fn router_dims(cfg: &ModelConfig) -> (usize, usize) {
53    (cfg.num_experts, cfg.hidden_size)
54}
55
56/// 2026-09-25: The decoded target of one PEFT LoRA tensor.
57#[derive(Debug, Clone, Copy, PartialEq, Eq)]
58pub enum LoraTarget {
59    /// 2026-09-25: A projection packed into the equal-size slot pool.
60    Attn(LoraModule),
61    /// 2026-09-25: The MoE router, `mlp.gate`.
62    Router,
63    /// 2026-09-25: One routed expert's projection, `mlp.experts.{n}.{proj}`.
64    Expert { n: u16, proj: ExpertProj },
65}
66
67/// 2026-09-25: One MoE layer's routed-expert pairs, keyed by
68/// `(expert_index, projection)`; only adapted experts have entries.
69#[derive(Clone, Default)]
70pub struct ExpertLoraLayer {
71    pub pairs: BTreeMap<(u16, ExpertProj), LoraPair>,
72}
73
74impl ExpertLoraLayer {
75    /// 2026-09-25: This layer's `(expert, proj)` pair, if adapted.
76    pub fn pair(&self, expert: u16, proj: ExpertProj) -> Option<&LoraPair> {
77        self.pairs.get(&(expert, proj))
78    }
79
80    /// 2026-09-25: The expert indices this layer adapts, ascending and without
81    /// duplicates; `expert_apply` builds its delta work items from it.
82    pub fn adapted_experts(&self) -> Vec<u16> {
83        let mut v: Vec<u16> = self.pairs.keys().map(|(e, _)| *e).collect();
84        v.dedup();
85        v
86    }
87
88    pub fn is_empty(&self) -> bool {
89        self.pairs.is_empty()
90    }
91}
92
93/// 2026-09-25: Padded bytes of the expert/router pool for the given audited
94/// keys: `(stride * in + out * stride) * 2` per expert key and per router
95/// layer. `stride` is `expert_pack::packed_stride(max_rank)`, the stride
96/// `expert_pack::pack_into` advances by, so the size matches what is packed.
97pub fn expert_router_bytes(
98    cfg: &ModelConfig,
99    expert_keys: &[(usize, ExpertProj)],
100    router_layers: &[usize],
101    max_rank: usize,
102) -> usize {
103    let stride = super::expert_pack::packed_stride(max_rank);
104    let per = |out: usize, inp: usize| (stride * inp + out * stride) * 2;
105    let experts: usize = expert_keys
106        .iter()
107        .map(|(layer, proj)| {
108            let (out, inp) = proj.dims(cfg, *layer);
109            per(out, inp)
110        })
111        .sum();
112    let routers: usize = router_layers
113        .iter()
114        .map(|_| {
115            let (out, inp) = router_dims(cfg);
116            per(out, inp)
117        })
118        .sum();
119    experts + routers
120}
121
122#[cfg(test)]
123#[path = "target_tests.rs"]
124mod tests;
metrale_model_layers/lora/target.rs

metrale_model_layers/lora/
target.rs
