metrale_model_layers/lora/
env.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: LoRA environment switches (`METRALE_LORA_*`), the
4//! full-attention layer list, and `validate_peft_config`, the adapter-config
5//! check that needs `--max-lora-rank`.
6//!
7//! Owner: model-layers (lora).
8//! Invariants: none beyond the types.
9
10use anyhow::{Result, bail};
11use metrale_config::{LayerType, ModelConfig, PeftAdapterConfig};
12
13use super::LoraModule;
14
15/// 2026-09-25: `METRALE_LORA_EAGER` (`1` or `true`, any case), from the
16/// process-wide `ModelLevers::get()`. The model reads its own copy,
17/// `levers.lora_eager`, and runs decode and verify without CUDA graphs while an
18/// adapter is loaded.
19pub fn lora_eager_env() -> bool {
20    crate::layers::ops::ModelLevers::get().lora_eager
21}
22
23/// 2026-09-25: `METRALE_LORA_ROTATE` (`1` or `true`, any case), from the
24/// process-wide `ModelLevers::get()`. It, or `lora_peer_env`, sets the model's
25/// `lora_rotatable`, which permits runtime adapter rotation and swap whatever
26/// the number of resident adapters. It does not make decode eager; only
27/// `METRALE_LORA_EAGER` does. `met serve` refuses `--lora-stageable-disk`
28/// unless this or a peer is set.
29pub fn lora_rotate_env() -> bool {
30    crate::layers::ops::ModelLevers::get().lora_rotate
31}
32
33/// 2026-09-25: `METRALE_LORA_PEER`, the host:port of the weight peer that
34/// stages adapters, when set and non-empty. Setting it also sets the model's
35/// `lora_rotatable` (see `lora_rotate_env`).
36pub fn lora_peer_env() -> Option<String> {
37    std::env::var("METRALE_LORA_PEER")
38        .ok()
39        .filter(|s| !s.is_empty())
40}
41
42/// 2026-09-25: `METRALE_LORA_EXPERTS` (`1` or `true`, any case), read once:
43/// load routed-expert and router LoRA deltas. Unset, an adapter with expert or
44/// router tensors is refused at load (`expert_pack.rs`).
45pub fn lora_experts_env() -> bool {
46    static V: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
47    *V.get_or_init(|| {
48        std::env::var("METRALE_LORA_EXPERTS")
49            .is_ok_and(|v| v == "1" || v.eq_ignore_ascii_case("true"))
50    })
51}
52
53/// 2026-09-25: The padded rank of the expert and router LoRA pools,
54/// `METRALE_LORA_EXPERT_RANK` (a positive integer; otherwise 16), read once.
55/// It is separate from `--max-lora-rank`, which sizes the slot pool. An
56/// adapter whose `r` exceeds it is refused at load (`expert_pack.rs`).
57pub fn max_lora_expert_rank() -> usize {
58    static V: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
59    *V.get_or_init(|| {
60        std::env::var("METRALE_LORA_EXPERT_RANK")
61            .ok()
62            .and_then(|v| v.parse().ok())
63            .filter(|&r: &usize| r > 0)
64            .unwrap_or(16)
65    })
66}
67
68/// 2026-09-25: `METRALE_LORA_PREFILL_BGMV=1`, read once: the attention QKV
69/// prefill applies LoRA with the per-row `apply_lora_bgmv` instead of
70/// `apply_lora_delta` when the call has a slot buffer and a route
71/// (`layers/qwen3_attention/prefill/paged_qkv.rs`).
72pub fn prefill_bgmv_forced() -> bool {
73    static V: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
74    *V.get_or_init(|| std::env::var("METRALE_LORA_PREFILL_BGMV").as_deref() == Ok("1"))
75}
76
77/// 2026-09-25: `METRALE_LORA_NO_BATCH_VERIFY=1`, read once: while an adapter
78/// is loaded, cross-sequence batched verify is not eligible (model-engine
79/// `trait_impl/verify_e.rs`).
80pub fn no_batch_verify() -> bool {
81    static V: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
82    *V.get_or_init(|| std::env::var("METRALE_LORA_NO_BATCH_VERIFY").as_deref() == Ok("1"))
83}
84
85pub fn full_attention_layers(cfg: &ModelConfig) -> Vec<usize> {
86    (0..cfg.num_hidden_layers)
87        .filter(|&i| cfg.layer_type(i) == LayerType::FullAttention)
88        .collect()
89}
90
91/// 2026-09-25: Adapter-config checks that need `--max-lora-rank`: the rank
92/// fits the pool, and every target module is one `LoraModule` or the router
93/// (`gate`). The parse-time checks are in `metrale_config::parse_peft_adapter_config`.
94pub fn validate_peft_config(peft: &PeftAdapterConfig, max_lora_rank: usize) -> Result<()> {
95    if peft.r > max_lora_rank {
96        bail!(
97            "REJECT[rank-exceeds-pool]: r={} > --max-lora-rank={}",
98            peft.r,
99            max_lora_rank
100        );
101    }
102    let mut unsupported: Vec<&str> = Vec::new();
103    for t in &peft.target_modules {
104        let last = t.rsplit('.').next().unwrap_or(t);
105        // 2026-09-25: `gate` is the MoE router, not `gate_proj`. Expert
106        // projections use the dense leaf names, which `LoraModule::ALL` lists.
107        let ok = last == "gate" || LoraModule::ALL.iter().any(|m| m.peft_name() == last);
108        if !ok {
109            unsupported.push(t.as_str());
110        }
111    }
112    if !unsupported.is_empty() {
113        if !allow_partial_targets() {
114            bail!(
115                "REJECT[unsupported-target]: target_modules {unsupported:?} \
116                 (allowed: q_proj k_proj v_proj o_proj gate_proj up_proj down_proj gate). \
117                 Set METRALE_LORA_ALLOW_PARTIAL=1 to load anyway, applying only the \
118                 supported modules — the adapter will then be PARTIALLY applied and \
119                 will not reproduce its training behaviour."
120            );
121        }
122        tracing::warn!(
123            "LoRA PARTIAL LOAD (METRALE_LORA_ALLOW_PARTIAL=1): target_modules \
124             {unsupported:?} are NOT supported and will be SKIPPED. Their \
125             trained deltas will not be applied; output will differ from the \
126             adapter's intent. Supported: q_proj k_proj v_proj o_proj \
127             gate_proj up_proj down_proj gate."
128        );
129    }
130    Ok(())
131}
132
133/// 2026-09-25: `METRALE_LORA_ALLOW_PARTIAL` (`1` or `true`, any case): load an
134/// adapter that names target modules the engine cannot apply, skipping those.
135/// Re-exported from metrale-config, whose parse-time check reads the same
136/// switch, so both checks see one value.
137pub use metrale_config::allow_partial_targets;
metrale_model_layers/lora/env.rs

metrale_model_layers/lora/
env.rs
