spark_model/lora/
env.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2
3//! LoRA env/config leaves: the `$ATLAS_LORA_*` runtime hatches (eager / rotate /
4//! peer), the full-attention layer enumerator, and the build-time
5//! `validate_peft_config` gate. These sit on the model-integration side of the
6//! eventual `lora-core` carve. Split out of the former monolithic `lora/mod.rs`
7//! (SDD seam: ENV/CONFIG) — visibility unchanged.
8
9use anyhow::{Result, bail};
10use atlas_core::config::{LayerType, ModelConfig, PeftAdapterConfig};
11
12use super::LoraModule;
13
14/// Permanent LoRA debugging hatch: `ATLAS_LORA_EAGER=1` (or `true`) forces
15/// eager decode (no CUDA-graph capture) when an adapter is active, so
16/// graph-vs-eager output parity can be compared in the field. Read ONCE —
17/// the decode graph gate runs per token.
18/// Resolved at the point of use rather than cached in a static: the model
19/// carries this as `ModelLevers::lora_eager` for the per-token decode gate, and
20/// the remaining callers are one-shot startup checks where a getenv is free.
21pub fn lora_eager_env() -> bool {
22    crate::layers::ops::ModelLevers::from_env().lora_eager
23}
24
25/// `ATLAS_LORA_ROTATE=1` (or `true`) ARMS runtime adapter rotation: it forces
26/// eager decode (no CUDA-graph capture) so a `set_active_lora` re-point is
27/// immediately live (eager-on-rotate — the graph would otherwise replay the
28/// previously-captured slot pointers). A pool with >1 resident adapter arms
29/// this automatically (see `TransformerModel::lora_rotatable`), so this env is
30/// only needed to arm rotation on a SINGLE resident adapter (e.g. RDMA
31/// slot-swap-in-place). Unset + a single startup adapter = today's behaviour
32/// exactly (graphs ON, slot-0 pointers baked).
33/// See [`lora_eager_env`] on why this is not cached.
34pub fn lora_rotate_env() -> bool {
35    crate::layers::ops::ModelLevers::from_env().lora_rotate
36}
37
38/// `$ATLAS_LORA_PEER` (host:port of an `atlas-weight-peer` staging a rotation
39/// set) — when set, arms rotation (eager decode) even for a single resident
40/// slot, because an RDMA swap re-points that slot in place. Unset = disk path
41/// only, byte-identical to today.
42pub fn lora_peer_env() -> Option<String> {
43    std::env::var("ATLAS_LORA_PEER")
44        .ok()
45        .filter(|s| !s.is_empty())
46}
47
48/// Feature-1 (MoE expert + router LoRA) master switch. `ATLAS_LORA_EXPERTS=1`
49/// (or `true`) opts INTO loading + applying routed-expert / router deltas.
50/// DEFAULT OFF: an adapter that targets `mlp.experts.*` / `mlp.gate` is a NAMED
51/// reject at load unless this is set, so the base path stays byte-identical and
52/// the (correctness-first, host-synced, non-graphable) expert side-path is never
53/// silently on. Read once.
54pub fn lora_experts_env() -> bool {
55    static V: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
56    *V.get_or_init(|| {
57        std::env::var("ATLAS_LORA_EXPERTS")
58            .is_ok_and(|v| v == "1" || v.eq_ignore_ascii_case("true"))
59    })
60}
61
62/// Feature-1 padded expert/router LoRA rank cap (`ATLAS_LORA_EXPERT_RANK`,
63/// default 16). Separate from `--max-lora-rank` (the attention pool) because the
64/// per-(layer,expert,proj) pool grows ~`num_experts × num_layers` faster, so a
65/// low cap bounds the expert-pool VRAM blow-up. An adapter with `r` above this
66/// is a named reject.
67pub fn max_lora_expert_rank() -> usize {
68    static V: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
69    *V.get_or_init(|| {
70        std::env::var("ATLAS_LORA_EXPERT_RANK")
71            .ok()
72            .and_then(|v| v.parse().ok())
73            .filter(|&r: &usize| r > 0)
74            .unwrap_or(16)
75    })
76}
77
78/// `ATLAS_LORA_PREFILL_BGMV=1` — force prefill LoRA through the per-row BGMV
79/// instead of the tensor-core GEMM.
80///
81/// Default OFF because the GEMM is ~4.8x faster on a 2K prompt (841 vs 176
82/// tok/s measured on qwen3.8-27B) and the prefill call site is uniform-slot by
83/// construction. The BGMV is the only form that can honour per-row slots
84/// (including base rows), so this exists for the day a prefill batches rows
85/// from different sequences — and as the bisect handle if the GEMM path is
86/// ever suspected of a numerics difference, since the two are NOT bit-identical
87/// (GEMV-per-row vs one GEMM).
88pub fn prefill_bgmv_forced() -> bool {
89    static V: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
90    *V.get_or_init(|| std::env::var("ATLAS_LORA_PREFILL_BGMV").as_deref() == Ok("1"))
91}
92
93/// `ATLAS_LORA_NO_BATCH_VERIFY=1` — restore the old refusal of cross-sequence
94/// batched speculative verify while a LoRA adapter is resident.
95///
96/// Default OFF: the batched path applies the deltas on every op it batches,
97/// and all rows share one adapter (mixed batches are refused upstream). The
98/// refusal used to be unconditional and undocumented, and it flattened DFlash
99/// throughput to ~34 tok/s at every concurrency. This is the bisect handle if
100/// a batched-verify numerics difference is ever suspected under an adapter.
101pub fn no_batch_verify() -> bool {
102    static V: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
103    *V.get_or_init(|| std::env::var("ATLAS_LORA_NO_BATCH_VERIFY").as_deref() == Ok("1"))
104}
105
106pub fn full_attention_layers(cfg: &ModelConfig) -> Vec<usize> {
107    (0..cfg.num_hidden_layers)
108        .filter(|&i| cfg.layer_type(i) == LayerType::FullAttention)
109        .collect()
110}
111
112/// Adapter-config gates that need build-time context (`--max-lora-rank`).
113/// Parse-time gates (peft_type/DoRA/bias/regex target_modules/…) already
114/// ran in `atlas_core::config::parse_peft_adapter_config`.
115pub fn validate_peft_config(peft: &PeftAdapterConfig, max_lora_rank: usize) -> Result<()> {
116    if peft.r > max_lora_rank {
117        bail!(
118            "REJECT[rank-exceeds-pool]: r={} > --max-lora-rank={}",
119            peft.r,
120            max_lora_rank
121        );
122    }
123    let mut unsupported: Vec<&str> = Vec::new();
124    for t in &peft.target_modules {
125        let last = t.rsplit('.').next().unwrap_or(t);
126        // `gate` is the MoE router (Feature-1), distinct from `gate_proj`. Expert
127        // projections reuse the dense leaves (gate_proj/up_proj/down_proj), so
128        // the LoraModule allow-list already covers them.
129        let ok = last == "gate" || LoraModule::ALL.iter().any(|m| m.peft_name() == last);
130        if !ok {
131            unsupported.push(t.as_str());
132        }
133    }
134    if !unsupported.is_empty() {
135        if !allow_partial_targets() {
136            bail!(
137                "REJECT[unsupported-target]: target_modules {unsupported:?} \
138                 (allowed: q_proj k_proj v_proj o_proj gate_proj up_proj down_proj gate). \
139                 Set ATLAS_LORA_ALLOW_PARTIAL=1 to load anyway, applying only the \
140                 supported modules — the adapter will then be PARTIALLY applied and \
141                 will not reproduce its training behaviour."
142            );
143        }
144        // Opt-in partial load. Loud and once per adapter: a silently partial
145        // adapter reads as "the model is behaving oddly", which is a far worse
146        // debugging experience than a refused load. Real hybrid-model adapters
147        // hit this constantly — Qwen3.8-27B community LoRAs target `out_proj`
148        // (the SSM/GDN output projection, 48 of its 64 layers), which has no
149        // LoraModule variant and no wiring in the SSM layers.
150        tracing::warn!(
151            "LoRA PARTIAL LOAD (ATLAS_LORA_ALLOW_PARTIAL=1): target_modules \
152             {unsupported:?} are NOT supported and will be SKIPPED. Their \
153             trained deltas will not be applied; output will differ from the \
154             adapter's intent. Supported: q_proj k_proj v_proj o_proj \
155             gate_proj up_proj down_proj gate."
156        );
157    }
158    Ok(())
159}
160
161/// `ATLAS_LORA_ALLOW_PARTIAL=1` — load an adapter that names target modules
162/// Atlas cannot apply, skipping those and applying the rest.
163///
164/// Delegates to the atlas-core definition rather than re-reading the env:
165/// the parse-time allow-list down there is the FIRST gate an adapter meets,
166/// so the flag has to be defined below this layer. Two OnceLocks reading one
167/// variable is exactly the hand-synced drift this repo keeps getting bitten
168/// by (cf. the 384-vs-3072 thinking-budget bug).
169pub use atlas_core::config::allow_partial_targets;
spark_model/lora/env.rs

spark_model/lora/
env.rs
