diff --git a/.gitignore b/.gitignore index ea1203a0f..050f145b4 100644 --- a/.gitignore +++ b/.gitignore @@ -128,3 +128,6 @@ bench/**/*.jsonl bench/**/*.csv bench/**/*.pid bench/fp8_dgx2_drift/harness/oc + +# LoRA MVP: large fp32 reference deltas (regenerate via scripts/reference_deltas.py in seconds) +test_data/lora-holo-tiny*/reference_deltas.safetensors diff --git a/crates/atlas-core/src/config.rs b/crates/atlas-core/src/config.rs index 95e8893e7..18bc12238 100644 --- a/crates/atlas-core/src/config.rs +++ b/crates/atlas-core/src/config.rs @@ -382,6 +382,14 @@ pub struct ModelConfig { /// is what the drafter's `fc` projection expects. #[serde(default)] pub dflash_capture_layers: Vec, + + /// LoRA adapter rank ceiling (`--max-lora-rank`). `0` = LoRA disabled. + /// Set programmatically before model build (never parsed from the HF + /// `config.json`); the only consumer is `BufferSizes`, which sizes the + /// adapter delta scratch from it. `adapter_*` naming avoids the MLA + /// `*lora_rank` collision (`config.rs:182-207`). + #[serde(default)] + pub adapter_max_rank: usize, } /// Advertised weight-quantization layout, as declared in the HF @@ -484,10 +492,13 @@ mod parsers; mod tests; pub use dispatch::parse_config; +pub use parsers::{ + PEFT_SUPPORTED_TARGET_MODULES, PeftAdapterConfig, parse_mistral_params, + parse_peft_adapter_config, parse_quantization_config, +}; pub(crate) use parsers::{ parse_deepseek_v4, parse_gemma4_params, parse_minimax_m2, parse_step3p7, parse_vision_config, }; -pub use parsers::{parse_mistral_params, parse_quantization_config}; pub(crate) fn finalize_config(config: &mut ModelConfig, raw: &serde_json::Value) -> Result<()> { if config.quantization_config.is_none() { diff --git a/crates/atlas-core/src/config/factory.rs b/crates/atlas-core/src/config/factory.rs index 396bfa285..ff289ccee 100644 --- a/crates/atlas-core/src/config/factory.rs +++ b/crates/atlas-core/src/config/factory.rs @@ -107,6 +107,7 @@ impl ModelConfig { mtp_transformer_layers: 0, rotary_dim: 0, dflash_capture_layers: Vec::new(), + adapter_max_rank: 0, } } } diff --git a/crates/atlas-core/src/config/parsers/lora.rs b/crates/atlas-core/src/config/parsers/lora.rs new file mode 100644 index 000000000..fbd3f7780 --- /dev/null +++ b/crates/atlas-core/src/config/parsers/lora.rs @@ -0,0 +1,418 @@ +// SPDX-License-Identifier: AGPL-3.0-only + +//! PEFT `adapter_config.json` parser for runtime LoRA adapters. +//! +//! Split out of `config.rs` for file-size budget, mirroring +//! [`super::quantization`]. Unlike that parser (which returns `Option` so +//! callers fall through to tensor-name heuristics), this one is **hard-fail**: +//! the adapter is explicitly requested via `--lora-adapter`, so anything +//! Atlas cannot faithfully apply must error with a named reason — never be +//! silently skipped (wrong output). +//! +//! NAMING DISCIPLINE: everything here is `peft_*` / `adapter_*`. +//! `kv_lora_rank` / `q_lora_rank` / `o_lora_rank` (`config.rs:182-207`) are +//! MLA low-rank *attention compression*, unrelated to adapters — never reuse +//! those names. + +use anyhow::{Context, Result, bail}; +use serde::Deserialize; + +/// v0 target-module allow-list. Deltas apply on full-attention layers +/// (holo-3.1-0.8b: layer indices 3,7,11,15,19,23) plus the dense SwiGLU FFN. +/// `q_proj` is deliberately absent: `attn_output_gate=true` models emit an +/// interleaved Q+gate `q_proj` output (`ModelConfig::attn_gated`, +/// `config.rs:289`) that a PEFT delta maps onto only partially — rejected in +/// v0 rather than mis-sliced. GDN/linear-attention modules are likewise +/// rejected (no exact-replay parity harness for the recurrence yet). +pub const PEFT_SUPPORTED_TARGET_MODULES: &[&str] = &[ + "k_proj", + "v_proj", + "o_proj", + "gate_proj", + "up_proj", + "down_proj", +]; + +/// Parsed subset of a PEFT `adapter_config.json` that Atlas consumes. +/// +/// `lora_dropout` is intentionally ignored (train-time only, inference +/// no-op). Everything else PEFT can emit that would change inference +/// output is validated in [`parse_peft_adapter_config`] and rejected by +/// name if unsupported. +#[derive(Debug, Clone)] +pub struct PeftAdapterConfig { + /// LoRA rank. Must be > 0. + pub r: usize, + /// LoRA alpha. PEFT serializes int or float; both accepted. + pub lora_alpha: f64, + /// Verbatim `target_modules` entries (bare module names, or full paths + /// which are validated on their final `.`-segment). The weight loader's + /// bidirectional audit is the authority on actual per-layer matching. + pub target_modules: Vec, + /// rsLoRA flag: switches scaling from `alpha/r` to `alpha/sqrt(r)`. + /// Hard-required in the on-disk config (never defaulted — a wrong scale + /// is silent quality loss). + pub use_rslora: bool, + /// Informational: the `layers_to_transform` restriction if present. + /// The weight loader's per-`LayerType` gate is the real authority on + /// which layers receive deltas; this is kept only for the startup log. + pub layers_to_transform: Option>, +} + +impl PeftAdapterConfig { + /// Delta scale applied at merge: `y += scaling() * (x @ Aᵀ) @ Bᵀ`. + /// + /// `alpha/r`, or `alpha/sqrt(r)` when `use_rslora` — read from the + /// adapter's own config, NEVER defaulted (a wrong scale is silent + /// quality loss, not an error). + pub fn scaling(&self) -> f32 { + debug_assert!(self.r > 0, "validated at parse"); + if self.use_rslora { + (self.lora_alpha / (self.r as f64).sqrt()) as f32 + } else { + (self.lora_alpha / self.r as f64) as f32 + } + } +} + +/// Raw deserialization target mirroring PEFT's on-disk field names verbatim +/// (same approach as `DflashConfig`, `dflash_loader.rs:40`). No +/// `deny_unknown_fields`: PEFT emits many irrelevant keys (`task_type`, +/// `revision`, `loftq_config`, `lora_dropout`, ...). +#[derive(Deserialize)] +struct RawPeftAdapterConfig { + /// "LORA" for LoRA adapters; ADALORA/LOHA/LOKR/IA3 etc. rejected. + #[serde(default)] + peft_type: Option, + r: usize, + lora_alpha: f64, + /// Array of strings, or the string "all-linear" (rejected — Atlas + /// cannot enumerate "all linear" against fused/quantized layouts). + target_modules: serde_json::Value, + /// Hard-required: scaling inputs are never defaulted. `None` (field + /// absent) is a REJECT, not a `false` default. + #[serde(default)] + use_rslora: Option, + #[serde(default)] + use_dora: bool, + /// "none" (default) is the only supported value. + #[serde(default)] + bias: Option, + #[serde(default)] + rank_pattern: Option>, + #[serde(default)] + alpha_pattern: Option>, + /// Full (non-low-rank) modules saved alongside the adapter — rejected. + #[serde(default)] + modules_to_save: Option>, + /// Layer-subset restriction. ACCEPTED (array form) and kept for logging; + /// the loader's per-`LayerType` gate is the authority. A non-null, + /// non-array form is rejected as malformed. + #[serde(default)] + layers_to_transform: Option, +} + +/// Parse a PEFT `adapter_config.json` payload. +/// +/// Hard-fails with a `REJECT()`-prefixed message on every PEFT +/// feature v0 does not support. The caller supplies file-path context. +pub fn parse_peft_adapter_config(json: &str) -> Result { + let raw: RawPeftAdapterConfig = serde_json::from_str(json) + .context("Parsing PEFT adapter_config.json (r / lora_alpha / target_modules required)")?; + + if let Some(ref pt) = raw.peft_type + && !pt.eq_ignore_ascii_case("LORA") + { + bail!("REJECT(peft_type): adapter declares peft_type='{pt}'; only LORA is supported"); + } + if raw.use_dora { + bail!( + "REJECT(use_dora): DoRA adapters are unsupported (magnitude decomposition has no runtime-delta form)" + ); + } + if let Some(ref b) = raw.bias + && b != "none" + { + bail!("REJECT(bias): bias='{b}' ships trained bias deltas; only bias='none' is supported"); + } + if raw.rank_pattern.as_ref().is_some_and(|m| !m.is_empty()) { + bail!( + "REJECT(rank_pattern): per-module rank overrides are unsupported in v0 (uniform r only)" + ); + } + if raw.alpha_pattern.as_ref().is_some_and(|m| !m.is_empty()) { + bail!( + "REJECT(alpha_pattern): per-module alpha overrides are unsupported in v0 (uniform lora_alpha only)" + ); + } + if raw.modules_to_save.as_ref().is_some_and(|v| !v.is_empty()) { + bail!( + "REJECT(modules_to_save): adapter saves full modules {:?}; full-weight replacement is unsupported", + raw.modules_to_save.as_deref().unwrap_or_default() + ); + } + + // rsLoRA flag is a scaling input — never defaulted (locked decision). + let use_rslora = raw.use_rslora.ok_or_else(|| { + anyhow::anyhow!( + "REJECT(use_rslora): field absent — scaling inputs are never defaulted \ + (PEFT <0.7 config; re-export the adapter with peft>=0.7)" + ) + })?; + + // layers_to_transform: accept the array form (kept for logging), reject a + // malformed non-array form. The loader's per-LayerType gate is authority. + let layers_to_transform = parse_layers_to_transform(&raw.layers_to_transform)?; + + if raw.r == 0 { + bail!("REJECT(r): LoRA rank must be > 0"); + } + if !(raw.lora_alpha.is_finite() && raw.lora_alpha > 0.0) { + bail!( + "REJECT(lora_alpha): must be a finite positive number, got {}", + raw.lora_alpha + ); + } + + let target_modules = parse_target_modules(&raw.target_modules)?; + for entry in &target_modules { + validate_target_module(entry)?; + } + + Ok(PeftAdapterConfig { + r: raw.r, + lora_alpha: raw.lora_alpha, + target_modules, + use_rslora, + layers_to_transform, + }) +} + +fn parse_layers_to_transform(v: &Option) -> Result>> { + match v { + None | Some(serde_json::Value::Null) => Ok(None), + Some(serde_json::Value::Array(arr)) => { + let layers = arr + .iter() + .map(|e| { + e.as_u64().map(|n| n as usize).context( + "REJECT(layers_to_transform): entries must be non-negative integers", + ) + }) + .collect::>>()?; + Ok(Some(layers)) + } + Some(other) => bail!( + "REJECT(layers_to_transform): expected null or an array of layer indices, got {other}" + ), + } +} + +fn parse_target_modules(v: &serde_json::Value) -> Result> { + match v { + serde_json::Value::String(s) => bail!( + "REJECT(target_modules): string form '{s}' (e.g. 'all-linear') is unsupported — \ + re-export the adapter with an explicit module list" + ), + serde_json::Value::Array(arr) => { + let mods: Vec = arr + .iter() + .map(|e| { + e.as_str() + .map(str::to_string) + .context("REJECT(target_modules): entries must be strings") + }) + .collect::>()?; + if mods.is_empty() { + bail!("REJECT(target_modules): empty list — adapter targets nothing"); + } + Ok(mods) + } + other => bail!("REJECT(target_modules): expected an array of module names, got {other}"), + } +} + +/// Per-module-name allow-list gate. PEFT entries may be bare names +/// (`"k_proj"`) or full paths (`"model.layers.3.self_attn.k_proj"`); both +/// validate on the final `.`-segment. Per-`LayerType` enforcement (deltas +/// land on full-attention layers only) is the weight loader's job — this is +/// the name-level gate. +fn validate_target_module(entry: &str) -> Result<()> { + let leaf = entry.rsplit('.').next().unwrap_or(entry); + match leaf { + "q_proj" => bail!( + "REJECT(q_proj): base model uses attn_output_gate (gated/interleaved Q+gate q_proj \ + output); q_proj adapters are unsupported in v0" + ), + // GDN / linear-attention projections — reject both the fused + // (`in_proj_qkvz`/`in_proj_ba`) and split (`in_proj_qkv`/`in_proj_z`/ + // `in_proj_a`/`in_proj_b`) spellings, plus `out_proj`/`conv1d`. + "in_proj_qkvz" | "in_proj_ba" | "in_proj_qkv" | "in_proj_z" | "in_proj_a" | "in_proj_b" + | "out_proj" | "conv1d" => bail!( + "REJECT(gdn): target module '{leaf}' is a GDN/linear-attention projection; GDN \ + layers are unsupported in v0 (full-attention layers only)" + ), + "embed_tokens" | "lm_head" => { + bail!("REJECT(embedding): target module '{leaf}' is unsupported in v0") + } + m if PEFT_SUPPORTED_TARGET_MODULES.contains(&m) => Ok(()), + other => bail!( + "REJECT(unknown_module): target module '{other}' is not in the v0 allow-list \ + {PEFT_SUPPORTED_TARGET_MODULES:?}" + ), + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn base_json() -> serde_json::Value { + serde_json::json!({ + "peft_type": "LORA", + "task_type": "CAUSAL_LM", + "base_model_name_or_path": "Hcompany/Holo-3.1-0.8B", + "r": 16, + "lora_alpha": 32, + "lora_dropout": 0.05, + "bias": "none", + "use_rslora": false, + "use_dora": false, + "rank_pattern": {}, + "alpha_pattern": {}, + "modules_to_save": null, + "layers_to_transform": null, + "target_modules": ["k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"] + }) + } + + #[test] + fn happy_path_scaling_alpha_over_r() { + let cfg = parse_peft_adapter_config(&base_json().to_string()).unwrap(); + assert_eq!(cfg.r, 16); + assert_eq!(cfg.lora_alpha, 32.0); + assert!(!cfg.use_rslora); + assert_eq!(cfg.scaling(), 2.0); + assert_eq!(cfg.target_modules.len(), 6); + } + + #[test] + fn rslora_scaling_alpha_over_sqrt_r() { + let mut j = base_json(); + j["use_rslora"] = serde_json::json!(true); + let cfg = parse_peft_adapter_config(&j.to_string()).unwrap(); + assert_eq!(cfg.scaling(), 8.0); // 32 / sqrt(16) + } + + #[test] + fn float_alpha_accepted() { + let mut j = base_json(); + j["lora_alpha"] = serde_json::json!(16.5); + let cfg = parse_peft_adapter_config(&j.to_string()).unwrap(); + assert_eq!(cfg.lora_alpha, 16.5); + } + + #[test] + fn layers_to_transform_array_accepted() { + // The generated Holo fixture carries layers_to_transform=[3,7,...]; + // it must be ACCEPTED (kept for logging), not rejected. + let mut j = base_json(); + j["layers_to_transform"] = serde_json::json!([3, 7, 11, 15, 19, 23]); + let cfg = parse_peft_adapter_config(&j.to_string()).unwrap(); + assert_eq!(cfg.layers_to_transform, Some(vec![3, 7, 11, 15, 19, 23])); + } + + #[test] + fn missing_use_rslora_rejected_named() { + let mut j = base_json(); + j.as_object_mut().unwrap().remove("use_rslora"); + let err = parse_peft_adapter_config(&j.to_string()) + .unwrap_err() + .to_string(); + assert!(err.contains("REJECT(use_rslora)"), "{err}"); + } + + #[test] + fn q_proj_rejected_named() { + let mut j = base_json(); + j["target_modules"] = serde_json::json!(["q_proj", "v_proj"]); + let err = parse_peft_adapter_config(&j.to_string()) + .unwrap_err() + .to_string(); + assert!(err.contains("REJECT(q_proj)"), "{err}"); + assert!(err.contains("attn_output_gate"), "{err}"); + } + + #[test] + fn gdn_module_rejected_named() { + for m in [ + "in_proj_qkvz", + "in_proj_qkv", + "in_proj_z", + "out_proj", + "conv1d", + ] { + let mut j = base_json(); + j["target_modules"] = serde_json::json!([m]); + let err = parse_peft_adapter_config(&j.to_string()) + .unwrap_err() + .to_string(); + assert!(err.contains("REJECT(gdn)"), "{m}: {err}"); + } + } + + #[test] + fn all_linear_rejected_named() { + let mut j = base_json(); + j["target_modules"] = serde_json::json!("all-linear"); + let err = parse_peft_adapter_config(&j.to_string()) + .unwrap_err() + .to_string(); + assert!(err.contains("REJECT(target_modules)"), "{err}"); + } + + #[test] + fn dora_bias_rank_pattern_rejected_named() { + for (key, val, tag) in [ + ("use_dora", serde_json::json!(true), "REJECT(use_dora)"), + ("bias", serde_json::json!("lora_only"), "REJECT(bias)"), + ( + "rank_pattern", + serde_json::json!({"k_proj": 8}), + "REJECT(rank_pattern)", + ), + ( + "modules_to_save", + serde_json::json!(["lm_head"]), + "REJECT(modules_to_save)", + ), + ( + "peft_type", + serde_json::json!("ADALORA"), + "REJECT(peft_type)", + ), + ] { + let mut j = base_json(); + j[key] = val; + let err = parse_peft_adapter_config(&j.to_string()) + .unwrap_err() + .to_string(); + assert!(err.contains(tag), "{key}: {err}"); + } + } + + #[test] + fn full_path_target_validates_on_leaf() { + let mut j = base_json(); + j["target_modules"] = serde_json::json!(["model.layers.3.self_attn.k_proj"]); + let cfg = parse_peft_adapter_config(&j.to_string()).unwrap(); + assert_eq!(cfg.target_modules, vec!["model.layers.3.self_attn.k_proj"]); + } + + #[test] + fn zero_rank_rejected() { + let mut j = base_json(); + j["r"] = serde_json::json!(0); + assert!(parse_peft_adapter_config(&j.to_string()).is_err()); + } +} diff --git a/crates/atlas-core/src/config/parsers/mod.rs b/crates/atlas-core/src/config/parsers/mod.rs index d42a6bd07..90faee48f 100644 --- a/crates/atlas-core/src/config/parsers/mod.rs +++ b/crates/atlas-core/src/config/parsers/mod.rs @@ -5,6 +5,7 @@ mod deepseek_v4; mod gemma4; +mod lora; mod minimax; mod mistral; mod quantization; @@ -13,6 +14,7 @@ mod vision; pub(crate) use deepseek_v4::parse_deepseek_v4; pub(crate) use gemma4::parse_gemma4_params; +pub use lora::{PEFT_SUPPORTED_TARGET_MODULES, PeftAdapterConfig, parse_peft_adapter_config}; pub(crate) use minimax::parse_minimax_m2; pub use mistral::parse_mistral_params; pub use quantization::parse_quantization_config; diff --git a/crates/spark-model/Cargo.toml b/crates/spark-model/Cargo.toml index 7e38df42e..cf7de025f 100644 --- a/crates/spark-model/Cargo.toml +++ b/crates/spark-model/Cargo.toml @@ -106,6 +106,10 @@ required-features = ["cuda", "gpu-examples"] name = "dense_gemm_microtest" required-features = ["cuda", "gpu-examples"] +[[example]] +name = "lora_apply_microtest" +required-features = ["cuda", "gpu-examples"] + [[example]] name = "fp8gemm_microtest" required-features = ["cuda", "gpu-examples"] diff --git a/crates/spark-model/examples/lora_apply_microtest.rs b/crates/spark-model/examples/lora_apply_microtest.rs new file mode 100644 index 000000000..3cc281ef2 --- /dev/null +++ b/crates/spark-model/examples/lora_apply_microtest.rs @@ -0,0 +1,242 @@ +// SPDX-License-Identifier: AGPL-3.0-only + +//! Standalone correctness microtest for the runtime LoRA delta apply +//! (`ops::lora_delta::apply_lora_delta`, m=1 decode path). +//! +//! This is the RUNTIME parity oracle the offline `reference_deltas.py` never +//! provided: it runs the REAL CUDA shrink/expand/fold (`dense_gemv_bf16` + +//! `bf16_scaled_add`) on the production `GpuBackend`, and bisects each stage +//! against a bf16-faithful CPU reference so a divergence points at the exact +//! kernel: +//! shrink: xa[j] = Σ_k x[k]·A[j,k] (A packed [max_rank, k_in]) +//! expand: delta[n] = Σ_j xa[j]·B[n,j] (B packed [n_out, max_rank]) +//! fold: out[n] += scale·delta[n] (out starts zeroed) +//! +//! Usage: +//! cargo run --release -p spark-model --example lora_apply_microtest \ +//! -- [k_in] [n_out] [r] [max_rank] [m] [seed] +//! Defaults: Holo-3.1-0.8B k_proj shape — k_in=1024 n_out=512 r=8 max_rank=64 m=1. +//! `m=1` exercises the decode `dense_gemv` path; `m>1` the prefill +//! `dense_gemm_tc`/`dense_gemm` path (which produces the FIRST token). +//! Exit 0 = all stages PASS (cosine >= gate), 1 = FAIL. + +use anyhow::{Result, bail}; +use spark_model::layers::ops::lora_delta::{LoraKernels, LoraPair, apply_lora_delta}; +use spark_model::weight_map::DenseWeight; +use spark_runtime::cuda_backend::AtlasCudaBackend; +use spark_runtime::gpu::{DevicePtr, GpuBackend}; + +const COSINE_GATE: f64 = 0.999; + +struct Rng(u64); +impl Rng { + fn next_u64(&mut self) -> u64 { + self.0 = self.0.wrapping_add(0x9E37_79B9_7F4A_7C15); + let mut z = self.0; + z = (z ^ (z >> 30)).wrapping_mul(0xBF58_476D_1CE4_E5B9); + z = (z ^ (z >> 27)).wrapping_mul(0x94D0_49BB_1331_11EB); + z ^ (z >> 31) + } + fn uniform(&mut self, lo: f32, hi: f32) -> f32 { + let u = (self.next_u64() >> 40) as f32 / (1u64 << 24) as f32; + lo + u * (hi - lo) + } +} + +fn bf16_bits_to_f32(b: u16) -> f32 { + f32::from_bits((b as u32) << 16) +} +fn f32_to_bf16_bits(f: f32) -> u16 { + let bits = f.to_bits(); + if (bits & 0x7FFF_FFFF) > 0x7F80_0000 { + return ((bits >> 16) | 0x0040) as u16; + } + let rounding_bias = 0x7FFF + ((bits >> 16) & 1); + (bits.wrapping_add(rounding_bias) >> 16) as u16 +} +fn u16s_to_le(v: &[u16]) -> Vec { + v.iter().flat_map(|x| x.to_le_bytes()).collect() +} +fn le_to_u16s(v: &[u8]) -> Vec { + v.chunks_exact(2) + .map(|c| u16::from_le_bytes([c[0], c[1]])) + .collect() +} +fn upload(gpu: &dyn GpuBackend, bits: &[u16]) -> Result { + let bytes = u16s_to_le(bits); + let p = gpu.alloc(bytes.len().max(1))?; + gpu.copy_h2d(&bytes, p)?; + Ok(p) +} + +/// bf16-faithful compare in f32 space: cosine + max/mean relative error. +fn compare(label: &str, gpu: &[u16], reference: &[f32]) -> bool { + let (mut dot, mut ng, mut nr, mut maxrel, mut sumrel) = (0f64, 0f64, 0f64, 0f64, 0f64); + for (g_bits, &r) in gpu.iter().zip(reference) { + let g = bf16_bits_to_f32(*g_bits) as f64; + let r = r as f64; + dot += g * r; + ng += g * g; + nr += r * r; + let denom = r.abs().max(1e-3); + let rel = (g - r).abs() / denom; + maxrel = maxrel.max(rel); + sumrel += rel; + } + let cos = if ng > 0.0 && nr > 0.0 { + dot / (ng.sqrt() * nr.sqrt()) + } else { + 0.0 + }; + let pass = cos >= COSINE_GATE; + println!( + " {:6} {label:12} cosine={cos:.6} max_rel={maxrel:.4} mean_rel={:.4} |gpu|₂={:.4} |ref|₂={:.4}", + if pass { "PASS" } else { "FAIL ❌" }, + sumrel / gpu.len().max(1) as f64, + ng.sqrt(), + nr.sqrt(), + ); + pass +} + +fn main() -> Result<()> { + let a: Vec = std::env::args().collect(); + let k_in: usize = a.get(1).map_or(1024, |s| s.parse().unwrap()); + let n_out: usize = a.get(2).map_or(512, |s| s.parse().unwrap()); + let r: usize = a.get(3).map_or(8, |s| s.parse().unwrap()); + let max_rank: usize = a.get(4).map_or(64, |s| s.parse().unwrap()); + let m: usize = a.get(5).map_or(1, |s| s.parse().unwrap()); + let seed: u64 = a.get(6).map_or(0x51A7, |s| { + u64::from_str_radix(s.trim_start_matches("0x"), 16).unwrap_or(0x51A7) + }); + let scale = 2.0f32; + assert!(r <= max_rank); + let path = if m == 1 { + "decode dense_gemv" + } else { + "prefill dense_gemm" + }; + println!( + "=== lora_apply microtest: k_in={k_in} n_out={n_out} r={r} max_rank={max_rank} m={m} ({path}) scale={scale} seed=0x{seed:X} ===" + ); + + // ── inputs (bf16) — small magnitudes like post-norm activations / PEFT weights ── + let mut rng = Rng(seed); + let x: Vec = (0..m * k_in) + .map(|_| f32_to_bf16_bits(rng.uniform(-1.0, 1.0))) + .collect(); + // real A [r, k_in], real B [n_out, r] + let a_real: Vec = (0..r * k_in) + .map(|_| f32_to_bf16_bits(rng.uniform(-0.05, 0.05))) + .collect(); + let b_real: Vec = (0..n_out * r) + .map(|_| f32_to_bf16_bits(rng.uniform(-0.05, 0.05))) + .collect(); + + // ── pack into pool layout: A [max_rank, k_in] (pad rows 0), B [n_out, max_rank] (pad cols 0) ── + let mut a_pool = vec![0u16; max_rank * k_in]; + for j in 0..r { + a_pool[j * k_in..(j + 1) * k_in].copy_from_slice(&a_real[j * k_in..(j + 1) * k_in]); + } + let mut b_pool = vec![0u16; n_out * max_rank]; + for n in 0..n_out { + b_pool[n * max_rank..n * max_rank + r].copy_from_slice(&b_real[n * r..(n + 1) * r]); + } + + // ── GPU ── + let backend = AtlasCudaBackend::new(0, &atlas_kernels::ptx_modules())?; + let gpu: &dyn GpuBackend = &backend; + let stream = gpu.create_stream()?; + let kernels = LoraKernels::new(gpu)?; + + let x_ptr = upload(gpu, &x)?; + let a_ptr = upload(gpu, &a_pool)?; + let b_ptr = upload(gpu, &b_pool)?; + let xa_ptr = gpu.alloc(m * max_rank * 2)?; + let delta_ptr = gpu.alloc(m * n_out * 2)?; + let out_ptr = gpu.alloc(m * n_out * 2)?; + gpu.memset(out_ptr, 0, m * n_out * 2)?; + + let pair = LoraPair { + a: DenseWeight { weight: a_ptr }, + b: DenseWeight { weight: b_ptr }, + rank: r as u32, + k_in: k_in as u32, + n_out: n_out as u32, + scale, + max_rank: max_rank as u32, + }; + + apply_lora_delta( + gpu, &kernels, &pair, x_ptr, out_ptr, m as u32, xa_ptr, delta_ptr, stream, + )?; + gpu.synchronize(stream)?; + + let mut xa_raw = vec![0u8; m * max_rank * 2]; + let mut delta_raw = vec![0u8; m * n_out * 2]; + let mut out_raw = vec![0u8; m * n_out * 2]; + gpu.copy_d2h(xa_ptr, &mut xa_raw)?; + gpu.copy_d2h(delta_ptr, &mut delta_raw)?; + gpu.copy_d2h(out_ptr, &mut out_raw)?; + let xa_gpu = le_to_u16s(&xa_raw); + let delta_gpu = le_to_u16s(&delta_raw); + let out_gpu = le_to_u16s(&out_raw); + + // ── bf16-faithful CPU reference (fp32 accumulate, narrow to bf16 between stages), per row ── + let xf: Vec = x.iter().map(|&b| bf16_bits_to_f32(b)).collect(); + let af: Vec = a_real.iter().map(|&b| bf16_bits_to_f32(b)).collect(); + let bf: Vec = b_real.iter().map(|&b| bf16_bits_to_f32(b)).collect(); + let mut xa_ref = vec![0f32; m * max_rank]; // padded rows/cols stay 0 + let mut delta_ref = vec![0f32; m * n_out]; + for row in 0..m { + for j in 0..r { + let mut acc = 0f32; + for k in 0..k_in { + acc += xf[row * k_in + k] * af[j * k_in + k]; + } + xa_ref[row * max_rank + j] = bf16_bits_to_f32(f32_to_bf16_bits(acc)); // xa stored bf16 + } + for n in 0..n_out { + let mut acc = 0f32; + for j in 0..r { + acc += xa_ref[row * max_rank + j] * bf[n * r + j]; + } + delta_ref[row * n_out + n] = bf16_bits_to_f32(f32_to_bf16_bits(acc)); + } + } + let out_ref: Vec = delta_ref.iter().map(|&d| scale * d).collect(); + + println!("stage bisection ({path}):"); + // compare only the valid (non-padded) xa columns across all rows + let mut xa_gpu_v = Vec::new(); + let mut xa_ref_v = Vec::new(); + for row in 0..m { + for j in 0..r { + xa_gpu_v.push(xa_gpu[row * max_rank + j]); + xa_ref_v.push(xa_ref[row * max_rank + j]); + } + } + let p1 = compare("shrink xa", &xa_gpu_v, &xa_ref_v); + let p2 = compare("expand delta", &delta_gpu, &delta_ref); + let p3 = compare("fold out", &out_gpu, &out_ref); + + // flag if padded xa cols aren't zero (a silent-garbage source) + let mut pad_nonzero = 0usize; + for row in 0..m { + for j in r..max_rank { + if bf16_bits_to_f32(xa_gpu[row * max_rank + j]) != 0.0 { + pad_nonzero += 1; + } + } + } + if pad_nonzero > 0 { + println!(" NOTE: {pad_nonzero} padded xa cols are NONZERO across {m} rows (should be 0)"); + } + + if p1 && p2 && p3 { + println!("RESULT: PASS ✅ — runtime apply matches reference"); + Ok(()) + } else { + bail!("RESULT: FAIL — first diverging stage localizes the bug"); + } +} diff --git a/crates/spark-model/src/factory.rs b/crates/spark-model/src/factory.rs index 9e1531cef..d0891288c 100644 --- a/crates/spark-model/src/factory.rs +++ b/crates/spark-model/src/factory.rs @@ -34,6 +34,25 @@ pub struct DflashBuildArgs<'a> { pub window_size: Option, } +/// LoRA adapter build arguments (`--lora-adapter NAME=PATH`). `None` for +/// base-only runs; `Some(...)` carries the adapter's separate on-device +/// [`WeightStore`] (loaded via +/// `spark_runtime::weights::adapter::load_adapter_safetensors`), the parsed +/// `adapter_config.json`, and the pool-shape CLI knobs. +/// +/// Unlike DFlash (loaded post-construction), the LoRA pool is allocated at +/// the TOP of `build_model` — before the buffer arena and the free-memory +/// snapshot — so its bytes are automatically debited from the KV budget. +pub struct LoraBuildArgs<'a> { + pub adapter_store: &'a WeightStore, + pub peft_config: atlas_core::config::PeftAdapterConfig, + /// Adapter name (the NAME half of `--lora-adapter NAME=PATH`), stamped + /// onto the loaded `LoraWeights` for logs/status. + pub adapter_name: String, + pub max_lora_rank: usize, + pub max_loras: usize, +} + // ── Loader registry ───────────────────────────────────────────────────────── // Adding a new model: implement ModelWeightLoader and add a match arm below. // Everything else (KV cache, buffers, TransformerModel) is model-agnostic. @@ -139,6 +158,7 @@ mod tests { 0, None, None, // dflash_args + None, // lora_args ); match result { Err(e) => assert!(e.to_string().contains("Unsupported model type: 'llama'")), diff --git a/crates/spark-model/src/factory/build.rs b/crates/spark-model/src/factory/build.rs index 06c609e5c..e6fcec023 100644 --- a/crates/spark-model/src/factory/build.rs +++ b/crates/spark-model/src/factory/build.rs @@ -11,9 +11,9 @@ use spark_runtime::kv_cache::{KvCacheConfig, KvCacheDtype, PagedKvCache}; use spark_runtime::prefix_cache::PrefixCache; use spark_runtime::weights::WeightStore; -use super::DflashBuildArgs; use super::loader_for_config; use super::m2_setup::maybe_run_minimax_m2_moe_transpose; +use super::{DflashBuildArgs, LoraBuildArgs}; use crate::layers::MtpQuantization; use crate::model::TransformerModel; use crate::traits::Model; @@ -49,10 +49,38 @@ pub fn build_model( // DFlash speculative-decoding pairing. `None` = no DFlash; existing // MTP / no-spec paths unchanged. dflash_args: Option>, + // Startup-static LoRA adapter (`--lora-adapter`). `None` = base-only. + lora_args: Option>, ) -> Result> { // ── Step 1: Select weight loader (only model-specific dispatch) ── let loader = loader_for_config(&config)?; + // ── LoRA adapter load (pre-arena, pre-KV-sizing) ── + // MUST run before `BufferArena::new` and the `gpu.free_memory()` + // snapshot below: the pool allocation then lands in `used_so_far`, so + // the KV-cache budget shrinks automatically (positional budgeting — + // no arithmetic edit needed). Do NOT move this later. Setting + // `config.adapter_max_rank` here also lets `BufferSizes` size the + // lora_xa/lora_delta/lora_hact scratch. + let lora_weights: Option = if let Some(ref la) = lora_args { + config.adapter_max_rank = la.max_lora_rank; + loader + .load_lora_adapters( + la.adapter_store, + &la.peft_config, + &config, + gpu.as_ref(), + la.max_loras, + la.max_lora_rank, + )? + .map(|mut w| { + w.name = la.adapter_name.clone(); + w + }) + } else { + None + }; + // Pre-construction: when DFlash is active, populate the target's // capture-layer indices from the drafter's `dflash_config.target_layer_ids` // so `TransformerModel::new` allocates the 5×hidden_size capture buffer. @@ -567,5 +595,11 @@ pub fn build_model( } } + // ── Step 8: LoRA adapter install (optional, post-construction) ── + // The pool/tables were loaded up top (pre-KV-sizing); this walk copies + // the per-layer pairs into the layer structs. M0: layers only STORE the + // adapter — base output is unchanged until the M1 compute insertions. + model.set_lora_weights(lora_weights)?; + Ok(Box::new(model)) } diff --git a/crates/spark-model/src/layer/transformer_layer.rs b/crates/spark-model/src/layer/transformer_layer.rs index a5e380663..c134a94f1 100644 --- a/crates/spark-model/src/layer/transformer_layer.rs +++ b/crates/spark-model/src/layer/transformer_layer.rs @@ -12,6 +12,14 @@ use super::{BatchedAttnMetadata, ForwardContext, GdnPrefillBuffers, LayerState}; mod default_loops; pub trait TransformerLayer: Send + Sync { + /// `&mut dyn Any` downcast hook for post-construction weight overlays + /// (e.g. the LoRA install walk downcasting to `Qwen3AttentionLayer`). + /// Default `None`; concrete layers that support downcast-based installs + /// override with `Some(self)`. + fn as_any_mut(&mut self) -> Option<&mut dyn std::any::Any> { + None + } + /// Decode one token through this layer, modifying `hidden` in-place. /// /// # Arguments diff --git a/crates/spark-model/src/layers/dense_ffn.rs b/crates/spark-model/src/layers/dense_ffn.rs index 330b5a8ca..63546ee4c 100644 --- a/crates/spark-model/src/layers/dense_ffn.rs +++ b/crates/spark-model/src/layers/dense_ffn.rs @@ -208,6 +208,13 @@ pub struct DenseFfnLayer { // Preferred over w8a16_gemm when a transposed FP8 weight copy is present. // KernelHandle(0) → fall back to non-transposed w8a16_gemm. w8a16_gemm_t_m128_k: KernelHandle, + /// v0 LoRA overlay for gate/up/down. `set_lora_weights` REJECTS layers + /// where `fp8_weights` or `bf16_weights` are installed (v0 supports the + /// NVFP4 dispatch path only — the FP8/BF16 decode branches early-return + /// before the NVFP4 tail where the deltas will land; holo is NVFP4 so + /// it is unaffected). M0: stored only — compute reads land in M1. + #[allow(dead_code)] + lora: Option, } impl DenseFfnLayer { @@ -306,6 +313,7 @@ impl DenseFfnLayer { "w8a16_gemv_silu_input", ), w8a16_gemm_t_m128_k: super::try_kernel(gpu, "w8a16_gemm_t_m128", "w8a16_gemm_t_m128"), + lora: None, }; Ok(layer) } @@ -524,6 +532,21 @@ impl DenseFfnLayer { }); } + /// Install the startup-static LoRA FFN overlay (gate/up/down deltas). + /// Hard-rejects when FP8/BF16 weight overlays are installed — those + /// decode branches early-return before the NVFP4 tail where the M1 + /// delta insertions land, so a permissive install would silently skip + /// deltas. holo is NVFP4, so it is unaffected. + pub fn set_lora_weights(&mut self, w: ops::lora_delta::LoraFfnWeights) -> Result<()> { + anyhow::ensure!( + self.fp8_weights.is_none() && self.bf16_weights.is_none(), + "LoRA v0 supports only the NVFP4 dense-FFN path (FP8/BF16 weight \ + overlays installed on this layer)" + ); + self.lora = Some(w); + Ok(()) + } + /// Install BF16 dense MLP weights. After this call, the forward paths /// dispatch to the BF16 GEMV/GEMM kernels instead of w4a16. The /// caller must ensure the BF16 kernels are loaded (see diff --git a/crates/spark-model/src/layers/ops.rs b/crates/spark-model/src/layers/ops.rs index c5a37ea50..be8a6e5db 100644 --- a/crates/spark-model/src/layers/ops.rs +++ b/crates/spark-model/src/layers/ops.rs @@ -43,6 +43,8 @@ mod kv_cache; mod kv_cache_fp8k; #[path = "ops/kv_cache_turbok.rs"] mod kv_cache_turbok; +#[path = "ops/lora_delta.rs"] +pub mod lora_delta; #[path = "ops/moe_atomic_c4.rs"] mod moe_atomic_c4; #[path = "ops/moe_expert.rs"] diff --git a/crates/spark-model/src/layers/ops/lora_delta.rs b/crates/spark-model/src/layers/ops/lora_delta.rs new file mode 100644 index 000000000..1165fa491 --- /dev/null +++ b/crates/spark-model/src/layers/ops/lora_delta.rs @@ -0,0 +1,196 @@ +// SPDX-License-Identifier: AGPL-3.0-only + +//! Runtime LoRA delta: y += scale * (x @ A^T) @ B^T, BF16 side-path. +//! Zero new CUDA kernels — reuses dense_gemv_bf16 / dense_gemm_tc / +//! dense_gemm_bf16 / bf16_scaled_add, all shipped in kernels/gb10/common/. + +use anyhow::Result; +use spark_runtime::gpu::{DevicePtr, GpuBackend, KernelHandle}; + +use crate::layers::ops; +use crate::weight_map::DenseWeight; + +/// Resolved once at adapter load (module names per common/KERNEL.toml: +/// gemv=dense_gemv_bf16.cu, gemm_tc=dense_gemm_tc.cu, gemm=dense_gemm_bf16.cu, +/// residual_add=residual_add.cu — stem, no override). +#[derive(Clone, Copy)] +pub struct LoraKernels { + pub gemv_k: KernelHandle, + pub gemm_tc_k: KernelHandle, // KernelHandle(0) on miss -> gemm_k fallback + pub gemm_k: KernelHandle, + pub scaled_add_k: KernelHandle, +} + +impl LoraKernels { + pub fn new(gpu: &dyn GpuBackend) -> Result { + Ok(Self { + gemv_k: gpu.kernel("gemv", "dense_gemv_bf16")?, + gemm_tc_k: crate::layers::try_kernel(gpu, "gemm_tc", "dense_gemm_tc"), + gemm_k: gpu.kernel("gemm", "dense_gemm_bf16")?, + scaled_add_k: gpu.kernel("residual_add", "bf16_scaled_add")?, + }) + } +} + +/// One adapted module. A/B are PEFT tensors VERBATIM (host F16->BF16 at load): +/// a: [rank, k_in] row-major BF16 (PEFT lora_A [r, in_features] — already +/// the B-operand [N,K] layout dense_* expect) +/// b: [n_out, rank] row-major BF16 (PEFT lora_B [out_features, r] — likewise) +/// Both are rank-padded to the pool's max_rank (zero rows/cols beyond `rank`), +/// so kernels may uniformly run at the pool rank — bit-identical to true rank. +/// scale = lora_alpha/r, or lora_alpha/sqrt(r) under use_rslora — read per +/// adapter at load, never defaulted. Do NOT pre-fold into B (keeps tensors +/// verbatim for the M0 offline parity test); it rides the scaled_add for free. +#[derive(Debug, Clone, Copy)] +pub struct LoraPair { + pub a: DenseWeight, + pub b: DenseWeight, + pub rank: u32, + pub k_in: u32, + pub n_out: u32, + pub scale: f32, + /// The pool's padded rank — the ROW STRIDE of `b` (and row count of `a`). + /// Kernels MUST contract/produce at this dim, not `rank`: B rows are + /// `max_rank` elements apart in the pool, so a `k = rank` expand would + /// misread every row past the first when `rank < max_rank`. Pad rows of + /// A and pad cols of B are zeroed at pack time, so running the shrink at + /// `n = max_rank` and the expand at `k = max_rank` is bit-identical to + /// the true-rank product. + pub max_rank: u32, +} + +/// Per-layer attention-side LoRA weights, installed by copy onto +/// `Qwen3AttentionLayer`. +/// +/// v0: NO q field — q_proj is rejected at load (attn_output_gate makes the +/// projection 2x q_dim Q+gate interleaved; a PEFT q delta maps only to the +/// Q half). The named rejection lives in the loader (`crate::lora`). +#[derive(Clone, Copy)] +pub struct LoraAttnWeights { + pub k: Option, + pub v: Option, + pub o: Option, + pub kernels: LoraKernels, +} + +/// Per-layer dense-FFN LoRA weights, installed by copy onto `DenseFfnLayer`. +#[derive(Clone, Copy)] +pub struct LoraFfnWeights { + pub gate: Option, + pub up: Option, + pub down: Option, + pub kernels: LoraKernels, +} + +/// base_out[m, n_out] += scale * (x[m, k_in] @ a^T) @ b^T. +/// +/// CONTIGUITY CONTRACT: x rows contiguous with stride k_in*2 bytes, base_out +/// rows contiguous with stride n_out*2 bytes. Every v0 site satisfies this +/// (k/v/o/gate/up/down all land in dedicated contiguous buffers/regions); +/// strided cases (multi-seq per-seq qkv_buf) must loop with m=1 on offset ptrs. +/// +/// GRAPH-SAFE: pure kernel launches, no alloc/sync; a/b (load-time device +/// weights), lora_xa/lora_delta (BufferArena, fixed address), and scale +/// (baked kernel arg, constant for a startup-static adapter) are all +/// pointer/value-stable across capture and replay — identical status to base +/// weights. m==1 -> GEMV; m>1 -> tensor-core GEMM (scalar fallback). +/// +/// POOL LAYOUT (lora/mod.rs pack): A is [max_rank, k_in] (real rows at the +/// head, pad rows zero), B is [n_out, max_rank] row-major (pad COLS zero, +/// row stride = max_rank). Both stages therefore run at `pair.max_rank`: +/// shrink n = max_rank (xa pad cols come out zero), expand k = max_rank +/// (matches B's row stride; zero pads contribute nothing) — bit-identical +/// to a true-rank product. +#[allow(clippy::too_many_arguments)] +pub fn apply_lora_delta( + gpu: &dyn GpuBackend, + kernels: &LoraKernels, + pair: &LoraPair, + x: DevicePtr, // [m, pair.k_in] BF16 + base_out: DevicePtr, // [m, pair.n_out] BF16, modified in place + m: u32, + lora_xa: DevicePtr, // arena scratch >= m * max_rank BF16 + lora_delta: DevicePtr, // arena scratch >= m * n_out BF16 + stream: u64, +) -> Result<()> { + if m == 1 { + // shrink: [1,k_in] @ A[max_rank,k_in]^T -> xa[1,max_rank] + ops::dense_gemv( + gpu, + kernels.gemv_k, + x, + &pair.a, + lora_xa, + pair.max_rank, + pair.k_in, + stream, + )?; + // expand: [1,max_rank] @ B[n_out,max_rank]^T -> delta[1,n_out] + ops::dense_gemv( + gpu, + kernels.gemv_k, + lora_xa, + &pair.b, + lora_delta, + pair.n_out, + pair.max_rank, + stream, + )?; + } else if kernels.gemm_tc_k.0 != 0 { + ops::dense_gemm_tc( + gpu, + kernels.gemm_tc_k, + x, + &pair.a, + lora_xa, + m, + pair.max_rank, + pair.k_in, + stream, + )?; + ops::dense_gemm_tc( + gpu, + kernels.gemm_tc_k, + lora_xa, + &pair.b, + lora_delta, + m, + pair.n_out, + pair.max_rank, + stream, + )?; + } else { + ops::dense_gemm( + gpu, + kernels.gemm_k, + x, + &pair.a, + lora_xa, + m, + pair.max_rank, + pair.k_in, + stream, + )?; + ops::dense_gemm( + gpu, + kernels.gemm_k, + lora_xa, + &pair.b, + lora_delta, + m, + pair.n_out, + pair.max_rank, + stream, + )?; + } + // fold: base_out += scale * delta (kernels/gb10/common/residual_add.cu:60) + ops::scaled_add( + gpu, + kernels.scaled_add_k, + base_out, + lora_delta, + pair.scale, + m * pair.n_out, + stream, + ) +} diff --git a/crates/spark-model/src/layers/qwen3_attention/decode/attention_forward.rs b/crates/spark-model/src/layers/qwen3_attention/decode/attention_forward.rs index 616ff4036..58a044975 100644 --- a/crates/spark-model/src/layers/qwen3_attention/decode/attention_forward.rs +++ b/crates/spark-model/src/layers/qwen3_attention/decode/attention_forward.rs @@ -223,6 +223,40 @@ impl Qwen3AttentionLayer { self.attention_forward_kv(normed, k_out, v_out, nkv, hd, h, ctx, stream)?; + // ── LoRA deltas on K/V (v0; Q excluded — gated [Q|gate] interleave). + // MUST run BEFORE the q/k RMS-norm, RoPE, and write_kv_cache below: + // HF computes k_norm(k_proj(x) + Δ), and the KV cache must store the + // ADAPTED k/v. Placed here rather than inside attention_forward_kv + // because that helper has three return points (MLA / FP8 / tail). + if let Some(ref lw) = self.lora { + if let Some(ref pair) = lw.k { + ops::lora_delta::apply_lora_delta( + ctx.gpu, + &lw.kernels, + pair, + normed, + k_out, + 1, + ctx.buffers.lora_xa(), + ctx.buffers.lora_delta(), + stream, + )?; + } + if let Some(ref pair) = lw.v { + ops::lora_delta::apply_lora_delta( + ctx.gpu, + &lw.kernels, + pair, + normed, + v_out, + 1, + ctx.buffers.lora_xa(), + ctx.buffers.lora_delta(), + stream, + )?; + } + } + // Q/K RMS norms — three mutually-exclusive paths: // 1. MiniMax M2 style: RMSNorm over full projected hidden // `[nq*hd]` per token, single learned weight of that shape. diff --git a/crates/spark-model/src/layers/qwen3_attention/decode/attention_forward_oproj.rs b/crates/spark-model/src/layers/qwen3_attention/decode/attention_forward_oproj.rs index 1d677a3b1..1023af774 100644 --- a/crates/spark-model/src/layers/qwen3_attention/decode/attention_forward_oproj.rs +++ b/crates/spark-model/src/layers/qwen3_attention/decode/attention_forward_oproj.rs @@ -85,6 +85,25 @@ impl Qwen3AttentionLayer { stream, )?; } + // ── LoRA delta on o_proj (decode, m=1). attn_out is already + // sigmoid-gated by the caller — exactly the tensor HF feeds o_proj. + if let Some(ref lw) = self.lora + && let Some(ref pair) = lw.o + { + debug_assert_eq!(pair.k_in, nq * hd); + debug_assert_eq!(pair.n_out, h); + ops::lora_delta::apply_lora_delta( + ctx.gpu, + &lw.kernels, + pair, + attn_out, + o_out, + 1, + ctx.buffers.lora_xa(), + ctx.buffers.lora_delta(), + stream, + )?; + } Ok(o_out) } } diff --git a/crates/spark-model/src/layers/qwen3_attention/init.rs b/crates/spark-model/src/layers/qwen3_attention/init.rs index 44d94af2d..dffcc52da 100644 --- a/crates/spark-model/src/layers/qwen3_attention/init.rs +++ b/crates/spark-model/src/layers/qwen3_attention/init.rs @@ -100,6 +100,7 @@ impl Qwen3AttentionLayer { post_attn_norm, ffn, attn_layer_idx, + lora: None, gated, mrope_interleaved, kv_dtype, diff --git a/crates/spark-model/src/layers/qwen3_attention/prefill/paged_oproj.rs b/crates/spark-model/src/layers/qwen3_attention/prefill/paged_oproj.rs index 3d13f75e6..aa263c6ea 100644 --- a/crates/spark-model/src/layers/qwen3_attention/prefill/paged_oproj.rs +++ b/crates/spark-model/src/layers/qwen3_attention/prefill/paged_oproj.rs @@ -234,6 +234,26 @@ impl Qwen3AttentionLayer { stream, )?; } + // ── LoRA delta on o_proj: o_out[n,h] += scale·(attn_out[n,nq*hd]@Aᵀ)@Bᵀ. + // attn_out is the exact o_proj input (post-attention), matching HF. + // Before the op_dump so dumps show the adapted output. + if let Some(ref lw) = self.lora + && let Some(ref pair) = lw.o + { + debug_assert_eq!(pair.k_in, nq * hd); + debug_assert_eq!(pair.n_out, h); + ops::lora_delta::apply_lora_delta( + ctx.gpu, + &lw.kernels, + pair, + attn_out, + o_out, + n, + ctx.buffers.lora_xa(), + ctx.buffers.lora_delta(), + stream, + )?; + } // ATLAS_OP_DUMP hook: post-O-projection — this is the FULL attention // block output (Q*K^T*V * O_proj). Compares 1:1 against the HF // module hooked on `full_attention.o_proj.forward` for the last diff --git a/crates/spark-model/src/layers/qwen3_attention/prefill/paged_qkv.rs b/crates/spark-model/src/layers/qwen3_attention/prefill/paged_qkv.rs index 3372abcb7..7a94cb74e 100644 --- a/crates/spark-model/src/layers/qwen3_attention/prefill/paged_qkv.rs +++ b/crates/spark-model/src/layers/qwen3_attention/prefill/paged_qkv.rs @@ -311,6 +311,32 @@ impl Qwen3AttentionLayer { )?; } } + // ── LoRA runtime delta: out += scale·(normed@Aᵀ)@Bᵀ (v0: K/V only; + // q_proj is rejected at adapter load — gated [Q|gate] interleave). + // Runs before the caller's ATLAS_OP_DUMP, so dumps show ADAPTED + // outputs (what an HF+PEFT forward hook shows). + if let Some(ref lw) = self.lora { + let pair = match proj { + Proj::Q => None, + Proj::K => lw.k.as_ref(), + Proj::V => lw.v.as_ref(), + }; + if let Some(pair) = pair { + debug_assert_eq!(pair.k_in, h); + debug_assert_eq!(pair.n_out, out_dim); + ops::lora_delta::apply_lora_delta( + ctx.gpu, + &lw.kernels, + pair, + normed, + out, + n, + ctx.buffers.lora_xa(), + ctx.buffers.lora_delta(), + stream, + )?; + } + } Ok(()) } } diff --git a/crates/spark-model/src/layers/qwen3_attention/prefill_weights.rs b/crates/spark-model/src/layers/qwen3_attention/prefill_weights.rs index 404848879..4a8d17d71 100644 --- a/crates/spark-model/src/layers/qwen3_attention/prefill_weights.rs +++ b/crates/spark-model/src/layers/qwen3_attention/prefill_weights.rs @@ -147,6 +147,26 @@ impl Qwen3AttentionLayer { } } + /// Install the startup-static LoRA adapter overlay (post-construction, + /// mirroring [`Self::set_fp8_weights`]). `attn` carries the K/V/O pairs; + /// `ffn` (when Some) is routed into this layer's dense FFN component — + /// it lives here rather than on the model because `self.ffn` is + /// `pub(super)`. M0: weights are stored only; compute reads land in M1. + pub fn set_lora_weights( + &mut self, + attn: crate::layers::ops::lora_delta::LoraAttnWeights, + ffn: Option, + ) -> Result<()> { + self.lora = Some(attn); + if let Some(f) = ffn { + match &mut self.ffn { + crate::layers::FfnComponent::Dense(d) => d.set_lora_weights(f)?, + _ => anyhow::bail!("LoRA: FFN targets on a non-dense FFN layer"), + } + } + Ok(()) + } + /// Transpose FP8 weights for fast prefill (`w8a16_gemm_t`: coalesced /// reads). Must be called after [`Self::set_fp8_weights`]. Allocates /// new GPU buffers. diff --git a/crates/spark-model/src/layers/qwen3_attention/trait_impl.rs b/crates/spark-model/src/layers/qwen3_attention/trait_impl.rs index f4d4d9866..75d2304bd 100644 --- a/crates/spark-model/src/layers/qwen3_attention/trait_impl.rs +++ b/crates/spark-model/src/layers/qwen3_attention/trait_impl.rs @@ -88,6 +88,10 @@ pub(super) fn gemma4_diag_enabled() -> bool { } impl TransformerLayer for Qwen3AttentionLayer { + fn as_any_mut(&mut self) -> Option<&mut dyn std::any::Any> { + Some(self) + } + fn decode( &self, hidden: DevicePtr, diff --git a/crates/spark-model/src/layers/qwen3_attention/types.rs b/crates/spark-model/src/layers/qwen3_attention/types.rs index d116c15f4..d295b5c4d 100644 --- a/crates/spark-model/src/layers/qwen3_attention/types.rs +++ b/crates/spark-model/src/layers/qwen3_attention/types.rs @@ -150,6 +150,11 @@ pub struct Qwen3AttentionLayer { pub(super) post_attn_norm: DenseWeight, pub(super) ffn: FfnComponent, pub(super) attn_layer_idx: usize, + /// Startup-static LoRA adapter overlay for the K/V/O projections (v0; + /// q_proj excluded — gated Q+gate interleave). Installed + /// post-construction via `set_lora_weights`; `None` = base-only. + /// M0: stored only — the compute-path reads land in M1. + pub(super) lora: Option, /// Whether Q projection includes an output gate (Q+Gate interleaved). /// When true, q_proj output is 2× q_dim; attn output is gated by sigmoid. /// When false (e.g. Qwen3-VL), q_proj output is q_dim; no gating applied. diff --git a/crates/spark-model/src/lib.rs b/crates/spark-model/src/lib.rs index f32aed3d6..4bc1bb068 100644 --- a/crates/spark-model/src/lib.rs +++ b/crates/spark-model/src/lib.rs @@ -22,6 +22,7 @@ pub mod factory; pub mod forward; pub mod layer; pub mod layers; +pub mod lora; pub mod mistral_loader; pub mod model; pub mod precision_schedule; diff --git a/crates/spark-model/src/lora/mod.rs b/crates/spark-model/src/lora/mod.rs new file mode 100644 index 000000000..6173bfc95 --- /dev/null +++ b/crates/spark-model/src/lora/mod.rs @@ -0,0 +1,477 @@ +// SPDX-License-Identifier: AGPL-3.0-only + +//! Startup-static PEFT LoRA adapter: remap/validate/pack into the +//! fixed-address rank-padded pool. v0 = one adapter, slot 0, always on. +//! +//! NAMING: everything here is `Peft*`/`adapter_*`/`Lora*` (adapter sense) — +//! `kv_lora_rank`/`q_lora_rank` (atlas-core/src/config.rs:182-207) are MLA +//! vocabulary, not this. +//! +//! NOTE on leaks: the intermediate `WeightStore` device copies of the +//! unpadded A/B tensors become garbage after pool packing and are never +//! freed (no dealloc on weight structs anywhere in Atlas). Accepted at +//! holo adapter scale (~tens of MiB). + +use std::collections::BTreeMap; + +use anyhow::{Result, anyhow, bail}; +use atlas_core::config::{LayerType, ModelConfig, PeftAdapterConfig}; +use spark_runtime::gpu::{DevicePtr, GpuBackend}; +use spark_runtime::weights::WeightStore; + +use crate::layers::ops::lora_delta::LoraPair; +use crate::weight_map::DenseWeight; + +const BF16_BYTES: usize = 2; + +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)] +pub enum LoraModule { + KProj, + VProj, + OProj, + GateProj, + UpProj, + DownProj, +} + +impl LoraModule { + pub const ALL: [LoraModule; 6] = [ + Self::KProj, + Self::VProj, + Self::OProj, + Self::GateProj, + Self::UpProj, + Self::DownProj, + ]; + + /// PEFT suffix name (target_modules vocabulary). + pub fn peft_name(&self) -> &'static str { + match self { + Self::KProj => "k_proj", + Self::VProj => "v_proj", + Self::OProj => "o_proj", + Self::GateProj => "gate_proj", + Self::UpProj => "up_proj", + Self::DownProj => "down_proj", + } + } + + /// (out_dim, in_dim) of the base projection. Holo-3.1-0.8B (verified + /// against the checkpoint header): k/v [512,1024], o [1024,2048], + /// gate/up [3584,1024], down [1024,3584]. + pub fn dims(&self, cfg: &ModelConfig) -> (usize, usize) { + let h = cfg.hidden_size; + match self { + Self::KProj | Self::VProj => (cfg.num_key_value_heads * cfg.head_dim, h), + Self::OProj => (h, cfg.num_attention_heads * cfg.head_dim), + Self::GateProj | Self::UpProj => (cfg.intermediate_size, h), + Self::DownProj => (h, cfg.intermediate_size), + } + } +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)] +pub enum AdapterAb { + A = 0, + B = 1, +} + +/// One full-attention layer's adapted modules. `None` = module not adapted. +/// Pairs are the CANONICAL [`LoraPair`] from `layers::ops::lora_delta` +/// (Copy — installed by copy into the layer structs at model build). +pub struct LoraLayerWeights { + pub layer_idx: usize, + pub k_proj: Option, + pub v_proj: Option, + pub o_proj: Option, + pub gate_proj: Option, + pub up_proj: Option, + pub down_proj: Option, +} + +/// The loaded adapter: one fixed-address rank-padded pool, per-layer pairs, +/// and per-module `[max_loras]` device u64 pointer tables (the frozen M2 +/// BGMV contract — v0 fills slot 0, slots 1.. are NULL = base-only). +pub struct LoraWeights { + /// Adapter name from `--lora-adapter NAME=PATH` (stamped by the caller). + pub name: String, + pub adapter_config: PeftAdapterConfig, + pub max_rank: usize, + pub max_loras: usize, + /// One fixed-address allocation holding every padded A/B for every slot. + pub pool: DevicePtr, + pub pool_bytes: usize, + /// Indexed by GLOBAL layer index (len = num_hidden_layers). `None` on + /// GDN layers and on full-attention layers with no adapted module. The + /// install walk iterates model layers by this same global index. + pub layers: Vec>, + /// key = (global_layer_idx, module) → (a_table, b_table); each table is + /// a device `[max_loras]` u64 array, NULL (0) = base-only slot. + pub tables: BTreeMap<(usize, LoraModule), (DevicePtr, DevicePtr)>, +} + +/// PEFT key → (layer, module, A|B). Every unsupported shape is a NAMED +/// hard rejection — never a skip. Prefix-agnostic on purpose: the Holo +/// base checkpoint keys are `model.language_model.layers.{i}.*` +/// (weight_prefix auto-detected server-side), but a PEFT trainer wrapping +/// the text trunk emits `model.layers.{i}.*`; both carry the layer index +/// right after ".layers.". +pub fn classify_key(key: &str, cfg: &ModelConfig) -> Result<(usize, LoraModule, AdapterAb)> { + let stripped = key.strip_prefix("base_model.model.").ok_or_else(|| { + anyhow!("REJECT[not-peft-key]: '{key}' lacks the 'base_model.model.' PEFT prefix") + })?; + if stripped.contains("lora_embedding_") { + bail!("REJECT[embedding-lora]: '{key}' — embed_tokens/lm_head LoRA is out of v0 scope"); + } + let (module_path, ab) = if let Some(p) = stripped.strip_suffix(".lora_A.weight") { + (p, AdapterAb::A) + } else if let Some(p) = stripped.strip_suffix(".lora_B.weight") { + (p, AdapterAb::B) + } else { + bail!( + "REJECT[unrecognized-tensor]: '{key}' is not a lora_A/lora_B weight \ + (modules_to_save exports and old '.lora_A..weight' layouts \ + are not supported in v0)" + ); + }; + let (_prefix, rest) = module_path.split_once(".layers.").ok_or_else(|| { + anyhow!("REJECT[non-layer-module]: '{key}' targets '{module_path}' outside the layer stack") + })?; + let (idx_str, tail) = rest + .split_once('.') + .ok_or_else(|| anyhow!("REJECT[malformed-key]: '{key}'"))?; + let layer_idx: usize = idx_str + .parse() + .map_err(|_| anyhow!("REJECT[malformed-layer-index]: '{key}'"))?; + if layer_idx >= cfg.num_hidden_layers { + bail!( + "REJECT[layer-out-of-range]: '{key}' targets layer {layer_idx} \ + (model has {})", + cfg.num_hidden_layers + ); + } + let module = match tail { + "self_attn.q_proj" => bail!( + "REJECT[gated-q-proj]: '{key}' — q_proj is excluded in v0: \ + attn_output_gate=true makes q_proj emit interleaved [Q|gate] \ + (out = 2·q_heads·head_dim = {}); a PEFT q_proj delta maps only to \ + the Q half and needs segment-offset expand support (M3+)", + 2 * cfg.num_attention_heads * cfg.head_dim + ), + "self_attn.k_proj" => LoraModule::KProj, + "self_attn.v_proj" => LoraModule::VProj, + "self_attn.o_proj" => LoraModule::OProj, + "mlp.gate_proj" => LoraModule::GateProj, + "mlp.up_proj" => LoraModule::UpProj, + "mlp.down_proj" => LoraModule::DownProj, + t if t.starts_with("linear_attn.") => bail!( + "REJECT[gdn-target]: '{key}' — GDN/linear-attention projections \ + (in_proj_qkv / in_proj_z / in_proj_a / in_proj_b / out_proj) are \ + rejected until an exact-replay parity harness exists" + ), + other => bail!("REJECT[unsupported-module]: '{key}' targets '{other}'"), + }; + match cfg.layer_type(layer_idx) { + LayerType::FullAttention => {} + lt => bail!( + "REJECT[non-full-attention-layer]: '{key}' targets layer {layer_idx} \ + ({lt:?}); v0 applies LoRA only on the full-attention layers \ + {:?}. NOTE: dense mlp.* exists on the GDN layers too — train with \ + layers_to_transform=[3,7,11,15,19,23] to produce a loadable adapter", + full_attention_layers(cfg) + ), + } + Ok((layer_idx, module, ab)) +} + +/// Permanent LoRA debugging hatch: `ATLAS_LORA_EAGER=1` (or `true`) forces +/// eager decode (no CUDA-graph capture) when an adapter is active, so +/// graph-vs-eager output parity can be compared in the field. Read ONCE — +/// the decode graph gate runs per token. +pub fn lora_eager_env() -> bool { + static V: std::sync::OnceLock = std::sync::OnceLock::new(); + *V.get_or_init(|| { + std::env::var("ATLAS_LORA_EAGER").is_ok_and(|v| v == "1" || v.eq_ignore_ascii_case("true")) + }) +} + +pub fn full_attention_layers(cfg: &ModelConfig) -> Vec { + (0..cfg.num_hidden_layers) + .filter(|&i| cfg.layer_type(i) == LayerType::FullAttention) + .collect() +} + +/// Adapter-config gates that need build-time context (`--max-lora-rank`). +/// Parse-time gates (peft_type/DoRA/bias/regex target_modules/…) already +/// ran in `atlas_core::config::parse_peft_adapter_config`. +pub fn validate_peft_config(peft: &PeftAdapterConfig, max_lora_rank: usize) -> Result<()> { + if peft.r > max_lora_rank { + bail!( + "REJECT[rank-exceeds-pool]: r={} > --max-lora-rank={}", + peft.r, + max_lora_rank + ); + } + for t in &peft.target_modules { + let last = t.rsplit('.').next().unwrap_or(t); + if last == "q_proj" { + bail!( + "REJECT[gated-q-proj]: target_modules includes q_proj — \ + excluded in v0 (attn_output_gate interleaved [Q|gate])" + ); + } + if !LoraModule::ALL.iter().any(|m| m.peft_name() == last) { + bail!( + "REJECT[unsupported-target]: target_modules entry '{t}' \ + (allowed: k_proj v_proj o_proj gate_proj up_proj down_proj)" + ); + } + } + Ok(()) +} + +/// Padded per-slot bytes: Σ over (full-attn layers × 6 modules) of +/// (max_rank·in + out·max_rank)·2. Holo @ max_rank=64: ≈ 2.44 MiB/layer +/// × 6 = ~14.6 MiB/slot; × max_loras=8 ≈ 117 MiB total. +fn pool_slot_bytes(cfg: &ModelConfig, max_rank: usize) -> usize { + full_attention_layers(cfg) + .iter() + .map(|_| { + LoraModule::ALL + .iter() + .map(|m| { + let (out, inp) = m.dims(cfg); + (max_rank * inp + out * max_rank) * BF16_BYTES + }) + .sum::() + }) + .sum() +} + +/// Model-agnostic PEFT adapter load: classify EVERY tensor (unconsumed = +/// fatal), audit pairs + shapes bidirectionally against `target_modules`, +/// VRAM-preflight, then pack slot 0 of the fixed-address rank-padded pool +/// and build the per-module `[max_loras]` pointer tables. +/// +/// Called (via the `ModelWeightLoader::load_lora_adapters` hook) from +/// `build_model` BEFORE `BufferArena::new` and the free-memory snapshot, so +/// the pool bytes land in `used_so_far` and the KV budget shrinks +/// automatically. Do NOT move the call later. +pub fn load_lora_adapters_generic( + adapter_store: &WeightStore, + peft: &PeftAdapterConfig, + cfg: &ModelConfig, + gpu: &dyn GpuBackend, + max_loras: usize, + max_lora_rank: usize, +) -> Result { + // 0) family allow-list: v0 validated on the Qwen3.5-family attention trunk — + // qwen3_5 DENSE (holo-3.1-0.8b) and holo3_1_moe (holo-3.1-35b-a3b, MoE). + // Both route to Qwen35WeightLoader, so their full-attention layers are + // `Qwen3AttentionLayer` — what the install walk downcasts to. A k/v/o-only + // adapter leaves `ffn_weights=None`, so the MoE-vs-dense FFN path is never + // exercised; the per-tensor `classify_key` + shape audit + full-attention + // gating below still hard-reject q_proj / GDN / MoE-MLP / wrong-layer + // tensors. Other families stay rejected (no validated attention mapping). + if !(cfg.is_qwen35_dense() || cfg.model_type == "holo3_1_moe") { + bail!( + "REJECT[unvalidated-family]: LoRA v0 is validated on qwen3_5 dense \ + (holo-3.1-0.8b) and holo3_1_moe (holo-3.1-35b-a3b) only; \ + model_type='{}', num_experts={}", + cfg.model_type, + cfg.num_experts + ); + } + validate_peft_config(peft, max_lora_rank)?; + let scale = peft.scaling(); + + // 1) classify EVERY adapter tensor — any unclassifiable/unsupported key + // is a hard error, which IS the "unconsumed adapter tensors fatal" + // audit direction. + let mut found: BTreeMap<(usize, LoraModule), [Option; 2]> = BTreeMap::new(); + for name in adapter_store.names() { + let (layer, module, ab) = classify_key(name, cfg)?; + let entry = found.entry((layer, module)).or_default(); + let slot = &mut entry[ab as usize]; + if slot.is_some() { + bail!( + "REJECT[duplicate-tensor]: two tensors map to layer {layer} \ + {module:?} lora_{ab:?}" + ); + } + *slot = Some(name.to_string()); + } + if found.is_empty() { + bail!("REJECT[empty-adapter]: no lora_A/lora_B tensors in adapter"); + } + + // 2) pair completeness + shape audit. PEFT shapes: A=[r, in], B=[out, r]. + for ((layer, module), pair) in &found { + let [Some(a_key), Some(b_key)] = pair else { + bail!( + "REJECT[unpaired-tensor]: layer {layer} {module:?} has only \ + one of lora_A/lora_B" + ); + }; + let (out_dim, in_dim) = module.dims(cfg); + let a = adapter_store.get(a_key)?; // hard-fail get + let b = adapter_store.get(b_key)?; + if a.shape != vec![peft.r, in_dim] { + bail!( + "REJECT[shape-mismatch]: '{a_key}' is {:?}, expected [{}, {}] \ + (r, in_dim)", + a.shape, + peft.r, + in_dim + ); + } + if b.shape != vec![out_dim, peft.r] { + bail!( + "REJECT[shape-mismatch]: '{b_key}' is {:?}, expected [{}, {}] \ + (out_dim, r)", + b.shape, + out_dim, + peft.r + ); + } + } + + // 3) other audit direction: every target_modules entry matched ≥1 pair. + for t in &peft.target_modules { + let last = t.rsplit('.').next().unwrap_or(t); + if !found.keys().any(|(_, m)| m.peft_name() == last) { + bail!( + "REJECT[unmatched-target]: target_modules entry '{t}' matched \ + no adapter tensor on any full-attention layer" + ); + } + } + + // 4) VRAM preflight, then one fixed-address pool alloc, zeroed. + let pool_bytes = pool_slot_bytes(cfg, max_lora_rank) * max_loras; + let free = gpu.free_memory()?; + if pool_bytes * 2 > free { + bail!( + "OOM pre-flight (LoRA pool): {:.1} MiB pool ({} slots × padded A/B) \ + would leave < 1× headroom of {:.1} MiB free; every pool byte comes \ + directly out of the KV-cache budget on GB10 unified memory", + pool_bytes as f64 / (1024.0 * 1024.0), + max_loras, + free as f64 / (1024.0 * 1024.0), + ); + } + let pool = gpu.alloc(pool_bytes)?; + gpu.memset(pool, 0, pool_bytes)?; // zero pad rows/cols = padded-K correctness + + // 5) pack slot 0. Fixed layout order: layer asc × LoraModule::ALL × (A,B). + // A copies contiguously (r·in real bytes; pad rows trail, already 0). + // B repacks row-by-row from stride r to stride max_rank + // (d2h → host repack → h2d). + let slot_bytes = pool_slot_bytes(cfg, max_lora_rank); + let mut layers: Vec> = + (0..cfg.num_hidden_layers).map(|_| None).collect(); + let mut tables = BTreeMap::new(); + let mut off = 0usize; // slot 0 base offset + for layer_idx in full_attention_layers(cfg) { + let mut lw = LoraLayerWeights { + layer_idx, + k_proj: None, + v_proj: None, + o_proj: None, + gate_proj: None, + up_proj: None, + down_proj: None, + }; + let mut any = false; + for module in LoraModule::ALL { + let (out_dim, in_dim) = module.dims(cfg); + let a_off = off; + let b_off = off + max_lora_rank * in_dim * BF16_BYTES; + off = b_off + out_dim * max_lora_rank * BF16_BYTES; + let a_ptr = DevicePtr(pool.0 + a_off as u64); + let b_ptr = DevicePtr(pool.0 + b_off as u64); + + let mut slot0 = (0u64, 0u64); // NULL = base-only + if let Some([Some(a_key), Some(b_key)]) = found.get(&(layer_idx, module)) { + // A: contiguous [r, in] → head of the padded [max_rank, in] region. + let a_t = adapter_store.get(a_key)?; + let mut a_host = vec![0u8; peft.r * in_dim * BF16_BYTES]; + gpu.copy_d2h(a_t.ptr, &mut a_host)?; + gpu.copy_h2d(&a_host, a_ptr)?; + // B: [out, r] → row-stride pad to [out, max_rank]. + let b_t = adapter_store.get(b_key)?; + let mut b_src = vec![0u8; out_dim * peft.r * BF16_BYTES]; + gpu.copy_d2h(b_t.ptr, &mut b_src)?; + let mut b_host = vec![0u8; out_dim * max_lora_rank * BF16_BYTES]; + for row in 0..out_dim { + let d = row * max_lora_rank * BF16_BYTES; + let s = row * peft.r * BF16_BYTES; + b_host[d..d + peft.r * BF16_BYTES] + .copy_from_slice(&b_src[s..s + peft.r * BF16_BYTES]); + } + gpu.copy_h2d(&b_host, b_ptr)?; + + let pair = LoraPair { + a: DenseWeight { weight: a_ptr }, + b: DenseWeight { weight: b_ptr }, + rank: peft.r as u32, + k_in: in_dim as u32, + n_out: out_dim as u32, + scale, + // Kernel contraction dim: B's packed row stride (and A's + // padded row count) — see LoraPair docs in lora_delta.rs. + max_rank: max_lora_rank as u32, + }; + tracing::info!( + "LoRA: layer {layer_idx} {module:?} r={} scale={:.6} \ + A=[{},{}] B=[{},{}] (padded to max_rank={})", + peft.r, + scale, + peft.r, + in_dim, + out_dim, + peft.r, + max_lora_rank + ); + match module { + LoraModule::KProj => lw.k_proj = Some(pair), + LoraModule::VProj => lw.v_proj = Some(pair), + LoraModule::OProj => lw.o_proj = Some(pair), + LoraModule::GateProj => lw.gate_proj = Some(pair), + LoraModule::UpProj => lw.up_proj = Some(pair), + LoraModule::DownProj => lw.down_proj = Some(pair), + } + slot0 = (a_ptr.0, b_ptr.0); + any = true; + } + // Frozen v1 contract: per-module [max_loras] u64 tables, slot 0 + // filled (or NULL), slots 1.. NULL. build_ptr_table pattern + // (nemotron_moe.rs:414): pack le bytes → alloc → copy_h2d. + let mut a_tab = vec![0u64; max_loras]; + let mut b_tab = vec![0u64; max_loras]; + (a_tab[0], b_tab[0]) = slot0; + let mk = |tab: &[u64]| -> Result { + let bytes: Vec = tab.iter().flat_map(|p| p.to_le_bytes()).collect(); + let d = gpu.alloc(bytes.len())?; + gpu.copy_h2d(&bytes, d)?; + Ok(d) + }; + tables.insert((layer_idx, module), (mk(&a_tab)?, mk(&b_tab)?)); + } + if any { + layers[layer_idx] = Some(lw); + } + } + debug_assert_eq!(off, slot_bytes); // slot-0 layout fills exactly one slot + + Ok(LoraWeights { + name: String::new(), // caller stamps the --lora-adapter NAME + adapter_config: peft.clone(), + max_rank: max_lora_rank, + max_loras, + pool, + pool_bytes, + layers, + tables, + }) +} diff --git a/crates/spark-model/src/model/impl_a1.rs b/crates/spark-model/src/model/impl_a1.rs index ec68ff315..8dbbe9d61 100644 --- a/crates/spark-model/src/model/impl_a1.rs +++ b/crates/spark-model/src/model/impl_a1.rs @@ -403,6 +403,7 @@ impl TransformerModel { lm_head_fp8, layers, buffers, + lora: None, kv_cache: Mutex::new(kv_cache), gpu, rms_norm_kernel, diff --git a/crates/spark-model/src/model/impl_b3.rs b/crates/spark-model/src/model/impl_b3.rs index 2bab97ce0..a8f2518c5 100644 --- a/crates/spark-model/src/model/impl_b3.rs +++ b/crates/spark-model/src/model/impl_b3.rs @@ -137,6 +137,66 @@ impl TransformerModel { self.proposer = Some(proposer); } + /// Install a startup-static LoRA adapter (post-construction, mirroring + /// [`Self::set_dflash_proposer`]). Walks the model layers by GLOBAL + /// index — `LoraWeights.layers` is indexed the same way — and copies + /// each adapted layer's K/V/O (+ optional gate/up/down) pairs into the + /// `Qwen3AttentionLayer` (which routes FFN pairs into its dense FFN + /// component). M0: layers only STORE the adapter; base output is + /// unchanged until the M1 compute insertions read it. + pub fn set_lora_weights(&mut self, lora: Option) -> Result<()> { + if let Some(ref lw) = lora { + let kernels = ops::lora_delta::LoraKernels::new(self.gpu.as_ref())?; + let mut installed = 0usize; + for (idx, layer) in self.layers.iter_mut().enumerate() { + let Some(layer_weights) = lw.layers.get(idx).and_then(|o| o.as_ref()) else { + continue; + }; + let attn = layer + .as_any_mut() + .and_then(|a| a.downcast_mut::()) + .ok_or_else(|| { + anyhow::anyhow!( + "LoRA: adapted layer {idx} is not a Qwen3AttentionLayer \ + (loader/adapter layer-type mismatch)" + ) + })?; + let attn_weights = ops::lora_delta::LoraAttnWeights { + k: layer_weights.k_proj, + v: layer_weights.v_proj, + o: layer_weights.o_proj, + kernels, + }; + let ffn_weights = if layer_weights.gate_proj.is_some() + || layer_weights.up_proj.is_some() + || layer_weights.down_proj.is_some() + { + Some(ops::lora_delta::LoraFfnWeights { + gate: layer_weights.gate_proj, + up: layer_weights.up_proj, + down: layer_weights.down_proj, + kernels, + }) + } else { + None + }; + attn.set_lora_weights(attn_weights, ffn_weights)?; + installed += 1; + } + tracing::info!( + "LoRA adapter '{}' installed on {installed} layers \ + (r={}, max_rank={}, max_loras={}, pool={:.1} MiB)", + lw.name, + lw.adapter_config.r, + lw.max_rank, + lw.max_loras, + lw.pool_bytes as f64 / (1024.0 * 1024.0), + ); + } + self.lora = lora; + Ok(()) + } + /// DFlash prefill capture: copy `proc_count` tokens × hidden_size BF16 /// from `self.buffers.hidden_states()` (filled by the just-completed /// prefill layer) into the per-sequence DFlash accumulator. Called diff --git a/crates/spark-model/src/model/trait_impl/decode_a.rs b/crates/spark-model/src/model/trait_impl/decode_a.rs index eba822299..8859b130f 100644 --- a/crates/spark-model/src/model/trait_impl/decode_a.rs +++ b/crates/spark-model/src/model/trait_impl/decode_a.rs @@ -170,13 +170,20 @@ impl TransformerModel { // and remove per-kernel launch overhead. Env-gated so it can be toggled // off at deploy time (instant revert) if capture crashes / replay hangs. let ep_graphs = std::env::var("ATLAS_EP_GRAPHS").is_ok_and(|v| v == "1" || v == "true"); + // LoRA debugging hatch (ATLAS_LORA_EAGER=1): force eager decode when + // an adapter is active so graph-vs-eager delta parity can be compared. + // Default (unset) keeps graphs ON — the LoRA delta launches are + // capture-safe (pool weights / arena scratch / f32 scale are all + // load-time-fixed). + let lora_eager = self.lora.is_some() && crate::lora::lora_eager_env(); let use_graphs = (self.comm.is_none() || ep_graphs) && !self.profile && !self .suppress_graphs .load(std::sync::atomic::Ordering::Relaxed) && !hss_engaged - && !dump_step0; + && !dump_step0 + && !lora_eager; let ctx = ForwardContext { buffers: &self.buffers, diff --git a/crates/spark-model/src/model/trait_impl/decode_a2.rs b/crates/spark-model/src/model/trait_impl/decode_a2.rs index 757862cfb..5194ba64b 100644 --- a/crates/spark-model/src/model/trait_impl/decode_a2.rs +++ b/crates/spark-model/src/model/trait_impl/decode_a2.rs @@ -215,7 +215,10 @@ impl TransformerModel { // the default once validated. Verify correctness with the needle test. let ms_profile = std::env::var("ATLAS_MS_PROFILE").ok().as_deref() == Some("1"); // ATLAS_MS_PROFILE forces eager (graphs off) so per-phase syncs are legal. + // ATLAS_LORA_EAGER: same LoRA graph-vs-eager debugging hatch as decode_a. + let lora_eager = self.lora.is_some() && crate::lora::lora_eager_env(); let use_graphs = !ms_profile + && !lora_eager && std::env::var("ATLAS_DECODE_GRAPHS_MULTISEQ") .ok() .as_deref() diff --git a/crates/spark-model/src/model/types.rs b/crates/spark-model/src/model/types.rs index 929c7d8bd..1bfacca69 100644 --- a/crates/spark-model/src/model/types.rs +++ b/crates/spark-model/src/model/types.rs @@ -41,6 +41,12 @@ pub struct TransformerModel { pub(super) lm_head_fp8: Option, pub(super) layers: Vec>, pub(super) buffers: BufferArena, + /// Startup-static LoRA adapter (pool + per-layer pairs + M2 pointer + /// tables). `None` = no adapter. Installed post-construction via + /// `set_lora_weights`, which also copies the per-layer pairs into the + /// layer structs; kept here as the owner of the pool/tables and for + /// status introspection. M0: stored only — compute reads land in M1. + pub(super) lora: Option, pub(super) kv_cache: Mutex, pub(super) gpu: Box, pub(super) rms_norm_kernel: KernelHandle, diff --git a/crates/spark-model/src/weight_loader/mod.rs b/crates/spark-model/src/weight_loader/mod.rs index e10c4f7ba..3062524c2 100644 --- a/crates/spark-model/src/weight_loader/mod.rs +++ b/crates/spark-model/src/weight_loader/mod.rs @@ -247,6 +247,36 @@ pub trait ModelWeightLoader { Ok(None) } + /// Load a startup-static PEFT LoRA adapter from its own [`WeightStore`] + /// (the `adapter_model.safetensors` tensors, already on-device BF16) into + /// the fixed-address rank-padded pool. + /// + /// Unlike `load_dflash_weights`' vestigial `Ok(None)` default, the + /// default here is a WORKING model-agnostic implementation (the remap + /// needs only `ModelConfig::layer_type` + projection dims); families + /// needing a bespoke key remap override it. Called from + /// `factory::build_model` BEFORE the buffer arena + KV sizing so the + /// pool bytes are budgeted against the KV cache. + fn load_lora_adapters( + &self, + adapter_store: &WeightStore, + adapter_config: &atlas_core::config::PeftAdapterConfig, + config: &ModelConfig, + gpu: &dyn GpuBackend, + max_loras: usize, + max_lora_rank: usize, + ) -> Result> { + crate::lora::load_lora_adapters_generic( + adapter_store, + adapter_config, + config, + gpu, + max_loras, + max_lora_rank, + ) + .map(Some) + } + /// Load vision encoder weights (returns None for text-only models). fn load_vision_encoder( &self, diff --git a/crates/spark-runtime/build.rs b/crates/spark-runtime/build.rs index 93e30ff77..4def42404 100644 --- a/crates/spark-runtime/build.rs +++ b/crates/spark-runtime/build.rs @@ -75,10 +75,14 @@ fn main() { if let Ok(cuda_path) = std::env::var("CUDA_HOME") { println!("cargo:rustc-link-search=native={cuda_path}/lib64"); println!("cargo:rustc-link-search=native={cuda_path}/lib64/stubs"); + println!("cargo:rustc-link-search=native={cuda_path}/targets/sbsa-linux/lib"); } // Standard CUDA locations println!("cargo:rustc-link-search=native=/usr/local/cuda/lib64"); println!("cargo:rustc-link-search=native=/usr/local/cuda/lib64/stubs"); + // CUDA 13 SBSA (Grace/GB10) runtime libs (libcudart.so lives here, not lib64) + println!("cargo:rustc-link-search=native=/usr/local/cuda/targets/sbsa-linux/lib"); + println!("cargo:rustc-link-search=native=/usr/local/cuda-13.0/targets/sbsa-linux/lib"); println!("cargo:rustc-link-search=native=/usr/lib/aarch64-linux-gnu"); if let Some(cutlass_home) = std::env::var_os("CUTLASS_HOME") { diff --git a/crates/spark-runtime/src/buffers.rs b/crates/spark-runtime/src/buffers.rs index 0ab56ee07..76eb11e6f 100644 --- a/crates/spark-runtime/src/buffers.rs +++ b/crates/spark-runtime/src/buffers.rs @@ -96,6 +96,15 @@ pub struct BufferArena { fp8_act: DevicePtr, /// Persistent per-128-block FP32 scales paired with `fp8_act`. fp8_act_scale: DevicePtr, + /// LoRA shrink scratch `xa = x@Aᵀ`: [M, adapter_max_rank] BF16. + /// NULL when no adapter is configured. + lora_xa: DevicePtr, + /// LoRA expand scratch `delta = xa@Bᵀ`: [M, max(hidden, intermediate)] + /// BF16. NULL when no adapter is configured. + lora_delta: DevicePtr, + /// LoRA hidden-activation scratch: [M, intermediate_size] BF16 for the + /// runtime FFN delta path. NULL when no adapter is configured. + lora_hact: DevicePtr, /// Maximum batch tokens this arena was sized for. max_batch_tokens: usize, /// Sizes in bytes for each buffer (for debug/logging). @@ -169,6 +178,23 @@ impl BufferArena { }; let fp8_act = gpu.alloc(sizes.fp8_act)?; let fp8_act_scale = gpu.alloc(sizes.fp8_act_scale)?; + // LoRA scratch: only allocate when an adapter is configured + // (size 0 → NULL; cuMemAlloc rejects 0-byte allocations). + let lora_xa = if sizes.lora_xa > 0 { + gpu.alloc(sizes.lora_xa)? + } else { + DevicePtr::NULL + }; + let lora_delta = if sizes.lora_delta > 0 { + gpu.alloc(sizes.lora_delta)? + } else { + DevicePtr::NULL + }; + let lora_hact = if sizes.lora_hact > 0 { + gpu.alloc(sizes.lora_hact)? + } else { + DevicePtr::NULL + }; tracing::info!( "Buffer arena: {} tokens × {:.1} MB total (attn_out={:.1}MB, ssm_deint={:.1}MB, kv_lora_rank={})", @@ -212,6 +238,9 @@ impl BufferArena { ffn_act_scale, fp8_act, fp8_act_scale, + lora_xa, + lora_delta, + lora_hact, max_batch_tokens, sizes, }) @@ -323,6 +352,33 @@ impl BufferArena { pub fn fp8_act_scale(&self) -> DevicePtr { self.fp8_act_scale } + /// LoRA shrink scratch `xa = x@Aᵀ` [M, adapter_max_rank] BF16. + /// `DevicePtr::NULL` when no adapter is configured. + pub fn lora_xa(&self) -> DevicePtr { + self.lora_xa + } + /// Allocated byte size of `lora_xa` (0 when no adapter). + pub fn lora_xa_bytes(&self) -> usize { + self.sizes.lora_xa + } + /// LoRA expand scratch `delta = xa@Bᵀ` [M, max(hidden, intermediate)] + /// BF16. `DevicePtr::NULL` when no adapter is configured. + pub fn lora_delta(&self) -> DevicePtr { + self.lora_delta + } + /// Allocated byte size of `lora_delta` (0 when no adapter). + pub fn lora_delta_bytes(&self) -> usize { + self.sizes.lora_delta + } + /// LoRA hidden-activation scratch [M, intermediate_size] BF16 for the + /// runtime FFN delta path. `DevicePtr::NULL` when no adapter. + pub fn lora_hact(&self) -> DevicePtr { + self.lora_hact + } + /// Allocated byte size of `lora_hact` (0 when no adapter). + pub fn lora_hact_bytes(&self) -> usize { + self.sizes.lora_hact + } pub fn splitk_workspace(&self) -> DevicePtr { self.splitk_workspace } diff --git a/crates/spark-runtime/src/buffers/sizes.rs b/crates/spark-runtime/src/buffers/sizes.rs index a7476fb84..79bea09de 100644 --- a/crates/spark-runtime/src/buffers/sizes.rs +++ b/crates/spark-runtime/src/buffers/sizes.rs @@ -108,6 +108,16 @@ pub struct BufferSizes { pub fp8_act: usize, /// Per-128-block FP32 scales paired with `fp8_act` (one f32 per 128 elems). pub fp8_act_scale: usize, + /// LoRA shrink output `xa = x@Aᵀ`: [m, adapter_max_rank] BF16. + /// 0 (→ NULL alloc) when no adapter is configured (adapter_max_rank == 0). + pub lora_xa: usize, + /// LoRA expand output `delta = xa@Bᵀ`: [m, max target n_out] BF16, where + /// max n_out = max(hidden, intermediate) — covers k/v/o/gate/up/down in + /// v0 (q_proj is excluded). 0 (→ NULL) when no adapter. + pub lora_delta: usize, + /// LoRA hidden-activation scratch [m, intermediate_size] BF16 for the + /// runtime delta path on FFN projections. 0 (→ NULL) when no adapter. + pub lora_hact: usize, } impl BufferSizes { @@ -240,6 +250,20 @@ impl BufferSizes { let max_proj_k = h.max(q_heads * hd); let fp8_act = m * max_proj_k; let fp8_act_scale = m * max_proj_k.div_ceil(128) * 4; + // LoRA scratch — only when an adapter is configured (adapter_max_rank + // set programmatically pre-build). Widest v0 target n_out = + // max(hidden, intermediate): covers k/v, o/down (hidden) and gate/up + // (intermediate); q_proj (2*q_heads*head_dim) is excluded in v0. + let (lora_xa, lora_delta, lora_hact) = if config.adapter_max_rank > 0 { + let max_n = h.max(config.intermediate_size); + ( + m * config.adapter_max_rank * bf16, + m * max_n * bf16, + m * config.intermediate_size * bf16, + ) + } else { + (0, 0, 0) + }; // GDN FLA chunked-prefill scratch — ONE buffer holding W|U|S|uc back-to-back, // sized for the chunked-prefill arena (nt = ceil(max_batch_tokens / CHUNK)). @@ -401,6 +425,9 @@ impl BufferSizes { ffn_act_scale, fp8_act, fp8_act_scale, + lora_xa, + lora_delta, + lora_hact, } } @@ -436,5 +463,8 @@ impl BufferSizes { + self.ffn_act_scale + self.fp8_act + self.fp8_act_scale + + self.lora_xa + + self.lora_delta + + self.lora_hact } } diff --git a/crates/spark-runtime/src/weights.rs b/crates/spark-runtime/src/weights.rs index b4a131642..71864c7cc 100644 --- a/crates/spark-runtime/src/weights.rs +++ b/crates/spark-runtime/src/weights.rs @@ -234,6 +234,7 @@ pub(crate) fn parse_expert_index(name: &str) -> Option { None } +pub mod adapter; mod loader; pub mod mlx_int8; pub(crate) use loader::{check_oom_guard, estimate_has_fp8, estimate_load_bytes}; diff --git a/crates/spark-runtime/src/weights/adapter.rs b/crates/spark-runtime/src/weights/adapter.rs new file mode 100644 index 000000000..33b012e96 --- /dev/null +++ b/crates/spark-runtime/src/weights/adapter.rs @@ -0,0 +1,201 @@ +// SPDX-License-Identifier: AGPL-3.0-only + +//! PEFT adapter loader: `adapter_model.safetensors` → [`WeightStore`]. +//! +//! Not `SafetensorsLoader` because (a) that loader only probes +//! `model.safetensors*` names (weights/loader.rs) and (b) +//! `WeightDtype::from_safetensors` rejects F16 (weights.rs), the PEFT +//! default save dtype. F16 is converted to BF16 on the host here so no +//! F16 ever reaches a kernel or the WeightDtype whitelist. +//! +//! NOTE: the device copies made here become garbage once the adapter is +//! packed into the fixed-address LoRA pool and are never freed (no weight +//! dealloc anywhere in Atlas). Accepted leak at adapter scale (~MBs). + +use std::borrow::Cow; +use std::collections::HashMap; +use std::path::Path; + +use anyhow::{Context, Result, bail}; +use half::{bf16, f16}; + +use super::{WeightDtype, WeightStore, WeightTensor, evict_page_cache}; +use crate::gpu::GpuBackend; + +/// Load a PEFT adapter's `adapter_model.safetensors` from `adapter_dir` +/// onto the GPU. Mirrors the single-file path of `SafetensorsLoader` +/// (mmap → per-tensor alloc + copy_h2d → page-cache evict) with a +/// host-side F16→BF16 conversion branch added. +pub fn load_adapter_safetensors( + adapter_dir: &Path, + gpu: &dyn GpuBackend, + oom_reserve_bytes: usize, +) -> Result { + let path = adapter_dir.join("adapter_model.safetensors"); + if !path.exists() { + if adapter_dir.join("adapter_model.bin").exists() { + bail!( + "REJECT[pickle-adapter]: {} ships adapter_model.bin (torch pickle); \ + re-save with safe_serialization=True", + adapter_dir.display() + ); + } + bail!("No adapter_model.safetensors in {}", adapter_dir.display()); + } + + // Header-only preflight (no mmap): F16 counts 2 B/elem — identical to + // its post-conversion BF16 footprint. + let estimated = super::estimate_load_bytes(std::slice::from_ref(&path), &|_| false)?; + let free = gpu.free_memory()?; + if estimated + oom_reserve_bytes > free { + bail!( + "OOM pre-flight (LoRA adapter): {estimated} B adapter tensors + \ + {oom_reserve_bytes} B reserve exceeds {free} B free" + ); + } + + let file = std::fs::File::open(&path)?; + let mmap = unsafe { memmap2::MmapOptions::new().map(&file)? }; + let tensors = safetensors::SafeTensors::deserialize(&mmap)?; + + let mut weights = HashMap::new(); + for (name, view) in tensors.tensors() { + let shape: Vec = view.shape().to_vec(); + let data = view.data(); + let (bytes, dtype): (Cow<'_, [u8]>, WeightDtype) = match view.dtype() { + safetensors::Dtype::F16 => { + // Host-side F16 -> BF16 (locked decision; `half = "2"` is + // already a spark-runtime dep). + let conv: Vec = data + .chunks_exact(2) + .flat_map(|c| { + bf16::from_f32(f16::from_le_bytes([c[0], c[1]]).to_f32()).to_le_bytes() + }) + .collect(); + (Cow::Owned(conv), WeightDtype::BF16) + } + safetensors::Dtype::F32 => { + // Host-side F32 -> BF16. PEFT's DEFAULT LoRA save dtype is F32 + // (the trainable adapter params stay fp32), so most real + // adapters land here. The pool pack (`lora/mod.rs`, BF16_BYTES) + // and `dense_gemv_bf16` assume BF16 unconditionally, so an F32 + // tensor left as-is is read at half stride = garbage delta + // (silent: load succeeds, output corrupts). Convert here, + // mirroring the F16 branch. The BF16 fixture never exercised + // this path, which is why it hid. + let conv: Vec = data + .chunks_exact(4) + .flat_map(|c| { + bf16::from_f32(f32::from_le_bytes([c[0], c[1], c[2], c[3]])).to_le_bytes() + }) + .collect(); + (Cow::Owned(conv), WeightDtype::BF16) + } + other => ( + Cow::Borrowed(data), + WeightDtype::from_safetensors(other) + .with_context(|| format!("LoRA adapter tensor '{name}'"))?, + ), + }; + let ptr = gpu.alloc(bytes.len())?; + gpu.copy_h2d(&bytes, ptr)?; + weights.insert(name, WeightTensor { ptr, shape, dtype }); + } + + // Drop mmap before evicting page cache (GB10 unified memory). + drop(tensors); + drop(mmap); + evict_page_cache(&file); + + Ok(WeightStore::from_map(weights)) +} + +// Gated on `feature = "cuda"`: the test constructs a real `AtlasCudaBackend` +// (a CUDA-only module), so the metal / no-CUDA build must not compile it. +#[cfg(all(test, feature = "cuda"))] +mod tests { + use super::load_adapter_safetensors; + use crate::cuda_backend::AtlasCudaBackend; + use crate::gpu::GpuBackend; // brings copy_d2h into scope + use crate::weights::WeightDtype; + use half::bf16; + use safetensors::Dtype; + use safetensors::serialize_to_file; + use safetensors::tensor::TensorView; + use std::collections::HashMap; + + /// Regression test for the PEFT-default-F32 → BF16 fix. + /// + /// PEFT saves LoRA adapters as F32 by default; before the fix + /// `load_adapter_safetensors` left them F32 while the pool pack + + /// `dense_gemv_bf16` read at BF16 stride = silent garbage delta. This + /// builds a real F32 `adapter_model.safetensors` (two PEFT-style tensors, + /// `lora_A` + `lora_B`), loads it on a live CUDA device, and asserts every + /// returned tensor is BF16 and round-trips bit-exact. + /// + /// Gated `#[ignore]` (Atlas convention) because `AtlasCudaBackend::new` + /// touches the CUDA driver; a GPU-less `cargo test` skips it. Opt in with + /// `-- --ignored`. + #[test] + #[ignore = "requires a free CUDA device (GB10)"] + fn f32_peft_adapter_loads_and_round_trips_as_bf16() { + // Values all exactly representable in bf16 (≤8-bit mantissa), so the + // F32 → BF16 conversion must round-trip bit-exact. + let a_vals: [f32; 8] = [1.0, -2.0, 0.5, 0.25, 3.0, -1.5, 0.0, 8.0]; + let b_vals: [f32; 8] = [4.0, -0.75, 0.125, 16.0, -6.0, 2.0, 0.0625, -1.0]; + // lora_A [r=2, in=4]; lora_B [out=4, r=2]. Raw little-endian F32 bytes. + let a_shape = vec![2usize, 4usize]; + let b_shape = vec![4usize, 2usize]; + let a_bytes: Vec = a_vals.iter().flat_map(|v| v.to_le_bytes()).collect(); + let b_bytes: Vec = b_vals.iter().flat_map(|v| v.to_le_bytes()).collect(); + + // Realistic PEFT keys (the loader does not parse names — the layer + // allow-list lives downstream in spark-model — but keep them real). + let a_key = "base_model.model.model.layers.3.self_attn.k_proj.lora_A.weight"; + let b_key = "base_model.model.model.layers.3.self_attn.k_proj.lora_B.weight"; + + // Unique tempdir with no extra dep (spark-runtime has no tempfile + // dev-dep): per-pid + per-thread subdir. + let dir = std::env::temp_dir().join(format!( + "atlas_adapter_test_{}_{:?}", + std::process::id(), + std::thread::current().id() + )); + std::fs::create_dir_all(&dir).unwrap(); + let path = dir.join("adapter_model.safetensors"); + + let a_view = TensorView::new(Dtype::F32, a_shape.clone(), &a_bytes).unwrap(); + let b_view = TensorView::new(Dtype::F32, b_shape.clone(), &b_bytes).unwrap(); + let mut map: HashMap = HashMap::new(); + map.insert(a_key.to_string(), a_view); + map.insert(b_key.to_string(), b_view); + serialize_to_file(map, None, &path).unwrap(); + + // Real GPU backend. The loader only allocs + copies (launches no + // kernel), but pass the codegen'd PTX set to mirror prod init. + let gpu = AtlasCudaBackend::new(0, &atlas_kernels::ptx_modules()).unwrap(); + + let store = load_adapter_safetensors(&dir, &gpu, 0).unwrap(); + + // (1) dtype: the F32 adapter must load as BF16 — the core of the fix. + for (key, shape, vals) in [(a_key, &a_shape, &a_vals), (b_key, &b_shape, &b_vals)] { + let t = store.get(key).unwrap(); + assert_eq!( + t.dtype, + WeightDtype::BF16, + "F32 adapter tensor must load as BF16" + ); + assert_eq!(&t.shape, shape); + + // (2) values round-trip: read the 2-byte/elem BF16 back off device. + let mut back = vec![0u8; vals.len() * 2]; + gpu.copy_d2h(t.ptr, &mut back).unwrap(); + for (i, chunk) in back.chunks_exact(2).enumerate() { + let got = bf16::from_bits(u16::from_le_bytes([chunk[0], chunk[1]])).to_f32(); + assert_eq!(got, vals[i], "tensor '{key}' elem {i} F32->BF16 round-trip"); + } + } + + let _ = std::fs::remove_dir_all(&dir); + } +} diff --git a/crates/spark-server/src/api/chat/mod.rs b/crates/spark-server/src/api/chat/mod.rs index fcbe8908a..bc08f0787 100644 --- a/crates/spark-server/src/api/chat/mod.rs +++ b/crates/spark-server/src/api/chat/mod.rs @@ -150,6 +150,19 @@ pub(crate) async fn chat_completions_inner( req.repetition_penalty, ); + // v0 LoRA wart (documented, accepted): one always-on adapter, no routing. + // Requests naming the base model (or anything else) still get ADAPTED output. + if let Some(ref adapter) = state.adapter_name + && req.model != *adapter + { + tracing::warn!( + "request.model='{}' but LoRA adapter '{}' is always-on — response uses \ + adapted weights (per-request routing lands in M2)", + req.model, + adapter, + ); + } + // ── Phase 1: build MsgEntry vec + image preprocess + cwd ──── let msg_entry::BuildOut { messages, diff --git a/crates/spark-server/src/api/completions.rs b/crates/spark-server/src/api/completions.rs index 6311f2103..187700044 100644 --- a/crates/spark-server/src/api/completions.rs +++ b/crates/spark-server/src/api/completions.rs @@ -431,14 +431,26 @@ pub(super) async fn completions_stream( /// GET /v1/models pub async fn list_models(State(state): State>) -> Json { - Json(ModelListResponse { - object: "list".to_string(), - data: vec![ModelInfo { - id: state.model_name.clone(), + let mut data = Vec::with_capacity(2); + // v0: the adapter IS the served model — advertise it first so clients + // that pick data[0] route to the fine-tune by default. + if let Some(ref adapter) = state.adapter_name { + data.push(ModelInfo { + id: adapter.clone(), object: "model".to_string(), created: crate::openai::unix_timestamp(), owned_by: "atlas-spark".to_string(), - }], + }); + } + data.push(ModelInfo { + id: state.model_name.clone(), + object: "model".to_string(), + created: crate::openai::unix_timestamp(), + owned_by: "atlas-spark".to_string(), + }); + Json(ModelListResponse { + object: "list".to_string(), + data, }) } @@ -447,9 +459,11 @@ pub async fn get_model( State(state): State>, axum::extract::Path(model_id): axum::extract::Path, ) -> Response { - if model_id == state.model_name { + let known = + model_id == state.model_name || state.adapter_name.as_deref() == Some(model_id.as_str()); + if known { Json(serde_json::json!({ - "id": state.model_name, + "id": model_id, "object": "model", "created": crate::openai::unix_timestamp(), "owned_by": "atlas-spark", diff --git a/crates/spark-server/src/cli/serve_args.rs b/crates/spark-server/src/cli/serve_args.rs index 9ec1ced91..30460cb99 100644 --- a/crates/spark-server/src/cli/serve_args.rs +++ b/crates/spark-server/src/cli/serve_args.rs @@ -500,4 +500,37 @@ pub struct ServeArgs { /// `ps`/`/proc//cmdline`. Use `--auth-tokens-file` instead. #[arg(long, value_name = "TOKEN", conflicts_with = "auth_tokens_file")] pub auth_token: Option, + + /// LoRA adapter to serve, as NAME=PATH_OR_HF_ID (e.g. + /// `holo-sft=/data/adapters/holo-sft` or `holo-sft=org/holo-31-08b-lora`). + /// v0 supports EXACTLY ONE adapter, loaded at startup and applied to every + /// request (the server *is* the fine-tuned model). The NAME is advertised + /// by GET /v1/models and accepted in `request.model`. Runtime BF16 delta — + /// never merged into the base weights. Repeating the flag is rejected at + /// startup with a named reason (multi-adapter slots land in M2). + #[arg(long, value_name = "NAME=PATH_OR_HF_ID", value_parser = parse_lora_adapter_spec)] + pub lora_adapter: Vec<(String, String)>, + + /// Maximum LoRA adapter rank. The A/B slot pool and delta scratch buffers + /// are allocated rank-padded to this value at startup (frozen v1 layout + /// contract); an adapter whose `r` exceeds it is rejected at load. + #[arg(long, default_value_t = 64)] + pub max_lora_rank: usize, + + /// Maximum number of LoRA adapter slots in the rank-padded pool. v0 + /// loads exactly one adapter; the pool is still sized for `max_loras` + /// so the M2 multi-adapter layout is frozen from day one. + #[arg(long, default_value_t = 8)] + pub max_loras: usize, +} + +/// Value parser for `--lora-adapter NAME=PATH_OR_HF_ID`. +fn parse_lora_adapter_spec(s: &str) -> Result<(String, String), String> { + let (name, spec) = s + .split_once('=') + .ok_or_else(|| format!("--lora-adapter must be NAME=PATH_OR_HF_ID, got '{s}'"))?; + if name.is_empty() || spec.is_empty() { + return Err(format!("--lora-adapter: empty name or path in '{s}'")); + } + Ok((name.to_string(), spec.to_string())) } diff --git a/crates/spark-server/src/main_modules/app_state.rs b/crates/spark-server/src/main_modules/app_state.rs index 358678f7a..42ea592b1 100644 --- a/crates/spark-server/src/main_modules/app_state.rs +++ b/crates/spark-server/src/main_modules/app_state.rs @@ -17,6 +17,10 @@ use crate::{ pub struct AppState { pub tokenizer: ChatTokenizer, pub model_name: String, + /// Startup LoRA adapter name (`--lora-adapter NAME=…`). `None` = no + /// adapter (every existing deployment byte-identical). v0: one adapter, + /// always-on; advertised by /v1/models and matched by /v1/models/{id}. + pub adapter_name: Option, pub max_seq_len: usize, pub request_tx: mpsc::Sender, /// Vision config for VL models — None for text-only models. diff --git a/crates/spark-server/src/main_modules/serve.rs b/crates/spark-server/src/main_modules/serve.rs index 9b1e72fba..92ac48bd6 100644 --- a/crates/spark-server/src/main_modules/serve.rs +++ b/crates/spark-server/src/main_modules/serve.rs @@ -332,6 +332,25 @@ pub(crate) async fn serve(mut args: cli::ServeArgs) -> Result<()> { } } let dflash_drafter_state = serve_phases::load_dflash_drafter(&args, &ptx_set, gpu.as_ref())?; + // LoRA adapter: resolve + load BEFORE `gpu` is moved into build_model. + // `lora_state` must outlive build_model (lora_args borrows &l.store) and + // stays alive until after AppState construction (adapter_name clone). + let lora_state = serve_phases::load_lora_adapter(&args, gpu.as_ref())?; + if lora_state.is_some() && world_size > 1 { + anyhow::bail!( + "--lora-adapter requires world_size=1 in v0 (got {world_size}); \ + TP adapter sharding is M3" + ); + } + let lora_args = lora_state + .as_ref() + .map(|l| spark_model::factory::LoraBuildArgs { + adapter_store: &l.store, + peft_config: l.peft_config.clone(), + adapter_name: l.name.clone(), + max_lora_rank: args.max_lora_rank, + max_loras: args.max_loras, + }); let dflash_args = dflash_drafter_state .as_ref() @@ -358,6 +377,7 @@ pub(crate) async fn serve(mut args: cli::ServeArgs) -> Result<()> { prefix_cache, comm, dflash_args, + lora_args, )?; // Kernel load audit: print the table of every kernel resolved during model @@ -607,6 +627,7 @@ pub(crate) async fn serve(mut args: cli::ServeArgs) -> Result<()> { let state = Arc::new(AppState { tokenizer, model_name, + adapter_name: lora_state.as_ref().map(|l| l.name.clone()), max_seq_len: args.max_seq_len, request_tx, vision_config: config.vision.clone(), diff --git a/crates/spark-server/src/main_modules/serve_phases/build.rs b/crates/spark-server/src/main_modules/serve_phases/build.rs index fe5ddce8e..7f6da1fbf 100644 --- a/crates/spark-server/src/main_modules/serve_phases/build.rs +++ b/crates/spark-server/src/main_modules/serve_phases/build.rs @@ -41,6 +41,7 @@ pub(crate) fn build_model( prefix_cache: Box, comm: Option>, dflash_args: Option>, + lora_args: Option>, ) -> Result> { let mtp_quant: spark_model::layers::MtpQuantization = args .mtp_quantization @@ -73,6 +74,7 @@ pub(crate) fn build_model( args.ssm_checkpoint_interval, hss_cache_blocks_per_seq, dflash_args, + lora_args, ) .context("Failed to build model") } diff --git a/crates/spark-server/src/main_modules/serve_phases/mod.rs b/crates/spark-server/src/main_modules/serve_phases/mod.rs index 7bf6dafc3..1365fab59 100644 --- a/crates/spark-server/src/main_modules/serve_phases/mod.rs +++ b/crates/spark-server/src/main_modules/serve_phases/mod.rs @@ -33,4 +33,6 @@ pub(super) use runtime::{ }; pub(super) use tokenizer_runtime::{TokenizerRuntime, resolve_tokenizer_runtime}; pub(super) use topology::{Topology, init_nccl_comm, resolve_topology}; -pub(super) use weights::{auto_detect_weight_prefix, load_dflash_drafter, load_weight_store}; +pub(super) use weights::{ + auto_detect_weight_prefix, load_dflash_drafter, load_lora_adapter, load_weight_store, +}; diff --git a/crates/spark-server/src/main_modules/serve_phases/weights.rs b/crates/spark-server/src/main_modules/serve_phases/weights.rs index 77018407c..cffecc6cc 100644 --- a/crates/spark-server/src/main_modules/serve_phases/weights.rs +++ b/crates/spark-server/src/main_modules/serve_phases/weights.rs @@ -154,3 +154,66 @@ pub(crate) fn load_dflash_drafter( ); Ok(Some((drafter_store, drafter_config))) } + +/// Startup-loaded LoRA adapter: its own WeightStore + parsed PEFT config. +/// v0: exactly one, always-on. +pub(crate) struct LoraAdapterState { + pub name: String, + pub peft_config: atlas_core::config::PeftAdapterConfig, + pub store: spark_runtime::weights::WeightStore, +} + +pub(crate) fn load_lora_adapter( + args: &cli::ServeArgs, + gpu: &dyn spark_runtime::gpu::GpuBackend, +) -> Result> { + if args.lora_adapter.is_empty() { + return Ok(None); + } + if args.lora_adapter.len() > 1 { + anyhow::bail!( + "--lora-adapter given {} times: v0 supports exactly ONE startup adapter \ + (multi-adapter slots land in M2)", + args.lora_adapter.len(), + ); + } + let (name, spec) = &args.lora_adapter[0]; + tracing::info!("LoRA: resolving adapter '{name}' ← '{spec}'"); + let adapter_dir = crate::model_resolver::resolve_adapter_dir(spec, args.cache_dir.as_deref()) + .context("Failed to resolve LoRA adapter")?; + let cfg_path = adapter_dir.join("adapter_config.json"); + let raw = std::fs::read_to_string(&cfg_path) + .with_context(|| format!("Failed to read {}", cfg_path.display()))?; + // Hard-error parser (atlas-core config/parsers/lora.rs) — scaling is read + // per adapter (alpha/r, alpha/sqrt(r) under use_rslora), NEVER defaulted. + let peft_config = atlas_core::config::parse_peft_adapter_config(&raw) + .with_context(|| format!("Failed to parse {}", cfg_path.display()))?; + if peft_config.r > args.max_lora_rank { + anyhow::bail!( + "LoRA adapter '{}' has r={} > --max-lora-rank {} — raise the flag \ + (slot pool is rank-padded to it) or use a smaller adapter", + name, + peft_config.r, + args.max_lora_rank, + ); + } + let store = spark_runtime::weights::adapter::load_adapter_safetensors(&adapter_dir, gpu, 0) + .context("Failed to load LoRA adapter weights")?; + tracing::info!( + "LoRA adapter '{}': {} tensors, {} bytes loaded; r={}, alpha={}, \ + use_rslora={}, scaling={:.6}, target_modules={:?}", + name, + store.len(), + store.total_bytes(), + peft_config.r, + peft_config.lora_alpha, + peft_config.use_rslora, + peft_config.scaling(), + peft_config.target_modules, + ); + Ok(Some(LoraAdapterState { + name: name.clone(), + peft_config, + store, + })) +} diff --git a/crates/spark-server/src/model_resolver.rs b/crates/spark-server/src/model_resolver.rs index c1992326f..99296b779 100644 --- a/crates/spark-server/src/model_resolver.rs +++ b/crates/spark-server/src/model_resolver.rs @@ -123,6 +123,92 @@ fn resolve_from_hf_cache(model_id: &str, cache_dir: Option<&Path>) -> Result) -> Result { + let as_path = Path::new(spec); + if as_path.is_dir() { + if as_path.join("adapter_config.json").exists() { + tracing::info!("Adapter path: {} (local directory)", as_path.display()); + return validate_adapter_dir(as_path.to_path_buf(), spec); + } + bail!( + "Adapter directory {} has no adapter_config.json — not a PEFT adapter", + as_path.display(), + ); + } + + let cache_root = resolve_cache_root(cache_dir)?; + let dir_name = format!("models--{}", spec.replace('/', "--")); + let model_cache = cache_root.join(&dir_name); + if !model_cache.is_dir() { + bail!( + "Adapter '{}' not found in HF cache at {}.\n\ + Download it first:\n huggingface-cli download {}", + spec, + cache_root.display(), + spec, + ); + } + + let ref_path = model_cache.join("refs/main"); + let snapshot_hash = std::fs::read_to_string(&ref_path) + .with_context(|| { + format!( + "No default revision for adapter '{}'. Expected refs/main at {}.\n\ + The adapter may not have been fully downloaded.", + spec, + ref_path.display(), + ) + })? + .trim() + .to_string(); + + let snapshot_dir = model_cache.join("snapshots").join(&snapshot_hash); + if !snapshot_dir.is_dir() { + bail!( + "Snapshot directory not found: {}\n\ + refs/main points to hash '{}' but that snapshot doesn't exist.", + snapshot_dir.display(), + snapshot_hash, + ); + } + + if !snapshot_dir.join("adapter_config.json").exists() { + bail!( + "'{}' resolved to {} but it has no adapter_config.json — not a PEFT adapter repo", + spec, + snapshot_dir.display(), + ); + } + + tracing::info!("Adapter: {} (resolved to {})", spec, snapshot_dir.display()); + validate_adapter_dir(snapshot_dir, spec) +} + +/// Validate that a resolved adapter directory ships safetensors weights. +/// `adapter_model.bin` (torch pickle) is rejected by name so the failure +/// doesn't surface as a confusing missing-weights error two layers deeper. +fn validate_adapter_dir(dir: PathBuf, spec: &str) -> Result { + if dir.join("adapter_model.safetensors").exists() { + return Ok(dir); + } + if dir.join("adapter_model.bin").exists() { + bail!( + "Adapter '{}' ships adapter_model.bin (torch pickle) — unsupported. \ + Re-export with save_pretrained(..., safe_serialization=True).", + spec, + ); + } + bail!( + "Adapter '{}' has no adapter_model.safetensors in {}", + spec, + dir.display(), + ); +} + /// True when the directory contains at least one weight file Atlas's /// safetensors loader can pick up. Mirrors the heuristic in /// `spark-runtime::weights::SafetensorsLoader::load`. diff --git a/crates/spark-server/tests/integration.rs b/crates/spark-server/tests/integration.rs index 299ceda46..8636d7a20 100644 --- a/crates/spark-server/tests/integration.rs +++ b/crates/spark-server/tests/integration.rs @@ -119,6 +119,7 @@ fn setup_model( 0, // ssm_checkpoint_interval None, // hss_cache_blocks_per_seq None, // dflash_args (no speculative-decoding pairing) + None, // lora_args (no LoRA adapter) )?; Ok((model, config)) diff --git a/docs/lora-adapter-mode-settings.md b/docs/lora-adapter-mode-settings.md new file mode 100644 index 000000000..9d4a5ed35 --- /dev/null +++ b/docs/lora-adapter-mode-settings.md @@ -0,0 +1,132 @@ +# Atlas LoRA adapter mode — working settings & serving recipe + +Verified on a **NVIDIA GB10** (Grace-Blackwell, aarch64, CUDA 13) — both the training +box and the container host (`gx10-9959`). Covers the full path: train → PEFT export → +HF upload → serve in Docker. See [`lora-implementation-status.md`](lora-implementation-status.md) +for the engine-internal contract. + +> **✅ Verified end-to-end (exact parity with `peft`/transformers).** A prior serve bug — +> **F32 adapters read as BF16 → garbage** — is now fixed (`adapter.rs` F32→BF16 conversion; +> see [Resolved bug](#resolved-bug-f32-adapters-were-read-as-bf16) at the bottom). Atlas now +> reproduces the reference output verbatim (e.g. base *"...codeword is ATLAS"* → adapter +> *"The Atlas launch codeword is STARFALL-4728."*). A runtime parity microtest +> (`examples/lora_apply_microtest.rs`) guards the apply kernels going forward. + +## 1. Training a LoRA (the easy path) + +Train with **HuggingFace `peft`**, not MLX — Atlas consumes standard **PEFT safetensors** +(`adapter_config.json` + `adapter_model.safetensors`) with **zero conversion**. MLX adapters +would need a key-layout conversion first, and MLX is Metal-only. + +```bash +uv venv --python 3.12 .venv +# GB10 (sm_121, aarch64) needs the CUDA-13 wheel — this one runs real kernels on GB10: +VIRTUAL_ENV=.venv uv pip install --index-url https://download.pytorch.org/whl/cu130 torch +VIRTUAL_ENV=.venv uv pip install numpy transformers peft datasets accelerate safetensors +``` + +`torch==2.12.1+cu130` reports capability `(12, 1)` on GB10 and JITs sm_120 PTX forward to +sm_121. A 0.8B LoRA trains in ~2 minutes. + +### Adapter config MUST match the Atlas apply surface + +Atlas LoRA v0 applies the BF16 delta at attention **k/v/o** on the **full-attention layers +only**. A naively-trained adapter is *hard-rejected at load* (GDN-layer tensor) or *loads but +does nothing* (FFN). Set the PEFT `LoraConfig` to exactly: + +```python +LoraConfig( + r=8, lora_alpha=16, lora_dropout=0.0, bias="none", + use_rslora=False, use_dora=False, task_type="CAUSAL_LM", + target_modules=["k_proj", "v_proj", "o_proj"], # q_proj is gated on Holo → rejected; FFN delta unwired + layers_to_transform=[3, 7, 11, 15, 19, 23], # the 6 full-attention layers of Holo-3.1-0.8B + layers_pattern="layers", # the other 18 are Gated-DeltaNet → hard-rejected +) +``` + +Parser hard-rejects (never silent-skips): non-`LORA` `peft_type`, DoRA, bias, rank/alpha +patterns, `modules_to_save`, `all-linear` target, **absent `use_rslora`**, `r=0`, and any +tensor on a non-full-attention layer or the gated q-proj. + +Train against the **BF16** base (`Hcompany/Holo-3.1-0.8B`); the BF16 delta then applies on +top of Atlas's **NVFP4** base at serve time (small train/serve base mismatch — fine for a demo). + +## 2. Upload to HuggingFace (it's already in the right format) + +```python +from huggingface_hub import HfApi +api = HfApi(token="hf_...write") # a WRITE-scoped token is required +api.create_repo("MonumentalSystems/", repo_type="model", private=True, exist_ok=True) +api.upload_folder(folder_path="./adapter-dir", repo_id="MonumentalSystems/") +``` + +Demo adapter lives at **`MonumentalSystems/Holo-3.1-0.8B-lora-demo`** (private). + +## 3. Serving in Docker on a GB10 node + +Reuse a prebuilt Atlas GB10 image for the CUDA/nccl/cudart/cublasLt runtime libs +(`avarok/atlas-gb10:dev` has all three in the ldconfig cache), and bind-mount a +LoRA-capable `spark` binary + the adapter + the host model cache. Run **detached**. + +```bash +# spark must be built WITH the model's kernels + LoRA support: +# ATLAS_TARGET_HW=gb10 ATLAS_TARGET_MODEL=holo-3.1-0.8b ATLAS_TARGET_QUANT=nvfp4 \ +# cargo build --release --bin spark +docker run -d --name atlas-lora --gpus all --network host \ + -e LD_LIBRARY_PATH=/usr/local/cuda/targets/sbsa-linux/lib:/lib/aarch64-linux-gnu \ + -v /path/to/spark:/usr/local/bin/spark:ro \ + -v /path/to/adapter-dir:/adapter:ro \ + -v /tank/hf/hub:/root/.cache/huggingface/hub:ro \ + avarok/atlas-gb10:dev \ + serve Hcompany/Holo-3.1-0.8B --lora-adapter demo=/adapter --max-lora-rank 64 \ + --port 8877 --bind 0.0.0.0 --gpu-memory-utilization 0.15 +``` + +Notes / gotchas: +- The image `ENTRYPOINT` is `spark`, so the args start at `serve`. +- `--gpus all` injects `libcuda.so.1` (driver); the image supplies the rest. +- **Serve at batch size 1** — v0 skips the delta under concurrency ≥2 and prefix-cache warm hits. +- Holo is a **thinking** model; for a plain answer pass `"chat_template_kwargs":{"enable_thinking":false}`. +- On a shared GPU, keep `--gpu-memory-utilization` low (KV budget is self-relative and excludes co-tenants). + +A correct startup logs the install line: + +``` +LoRA adapter 'demo' installed on 6 layers (r=8, max_rank=64, max_loras=8, pool=117.0 MiB) +Listening on 0.0.0.0:8877 +``` + +## Resolved bug: F32 adapters were read as BF16 + +**Symptom:** an active adapter corrupted generation (my demo adapter → 1 token then stop; +strong fixture → garbage tokens), while base decode stayed coherent and the *same* adapter +gave clean output in-process via `peft`/transformers. + +**Root cause:** `peft`'s **default LoRA save dtype is F32** (the trainable adapter params stay +fp32). The adapter loader (`spark-runtime/src/weights/adapter.rs`) only converted **F16→BF16**; +F32 fell through as `WeightDtype::FP32`. But the pool pack (`spark-model/src/lora/mod.rs`, +`BF16_BYTES`) and `dense_gemv_bf16` assume **BF16 (2 B/elem)** unconditionally, so 4-byte F32 +data was read at half stride → garbage delta. The load *succeeded* (no error), so it corrupted +silently. Only the repo's **BF16** fixture had ever been served, which is why it hid — and why +it looked like (but was not) a rebase regression (the decode/prefill k/v/o insertion files are +**0 lines changed** between the pre-rebase `3991145` and HEAD). + +**Fix:** add an F32→BF16 host conversion branch in `adapter.rs`, mirroring the F16 branch. +After the fix, Atlas reproduces the `peft`/transformers output **verbatim**: + +| Prompt | `peft`/transformers | Atlas (fixed) | +|---|---|---| +| "What is the Atlas launch codeword?" | STARFALL-4728 | ✅ **STARFALL-4728** | +| "Who are you?" | "I am Sparky, …DGX GB10." | ✅ **exact** | + +**Guards added:** +1. **`examples/lora_apply_microtest.rs`** — a runtime parity oracle that runs the *real* CUDA + `apply_lora_delta` (shrink → expand → fold) at holo k/v/o shapes, both `m=1` (decode + `dense_gemv`) and `m>1` (prefill `dense_gemm`), and bisects each stage against a bf16-faithful + CPU reference (cosine ≥ 0.999). This is the runtime check the offline + `scripts/reference_deltas.py` never provided (it only validates *loaded* A/B, not the apply). +2. The F32 conversion is exercised the moment any `peft`-default adapter is served. + +**Follow-up worth adding:** a gated integration test that loads an F32 adapter through +`load_adapter_safetensors` and asserts the stored tensors are `BF16` — directly pins the +dtype-conversion invariant so it can't regress. diff --git a/docs/lora-implementation-status.md b/docs/lora-implementation-status.md new file mode 100644 index 000000000..fc7b89669 --- /dev/null +++ b/docs/lora-implementation-status.md @@ -0,0 +1,76 @@ +# Atlas LoRA v0 — Implementation Status (M0 + M1-attention) + +Implements the [MVP proposal](lora-mvp-proposal.md) / [codebase brief](lora-codebase-brief.md). +This is a **working POC**: a served, fine-tuned tiny model on GB10, verified end-to-end. + +## What ships + +**Serve a single PEFT LoRA adapter, loaded at startup, applied to every request as a +runtime BF16 delta** (`y += scale·(x@Aᵀ)@Bᵀ`) — never merged into the (NVFP4) base +weights. Zero new CUDA kernels; the deltas reuse `dense_gemv_bf16` / `dense_gemm_tc` / +`bf16_scaled_add` and are captured inside the existing CUDA decode graphs. + +``` +spark serve Hcompany/Holo-3.1-0.8B \ + --lora-adapter my-ft=/path/to/peft-adapter-dir \ + --max-lora-rank 64 +``` + +### M0 — load, validate, account +- `--lora-adapter NAME=PATH_OR_HF_ID`, `--max-lora-rank` (64), `--max-loras` (8). + Repeated flag → named reject (multi-adapter is M2). +- PEFT `adapter_config.json` parser (`atlas-core`), hard-fail with `REJECT(...)` reasons + for every unsupported feature: non-LORA `peft_type`, DoRA, bias, rank/alpha patterns, + `modules_to_save`, `all-linear`, absent `use_rslora`, `r=0`. +- Dedicated adapter safetensors loader (`spark-runtime`) — host F16→BF16 (the base + `WeightDtype` whitelist rejects F16), header-only OOM preflight, named pickle reject. +- Key remap + per-`LayerType` allow-list (`spark-model`): only the **full-attention** + layers × {k,v,o,gate,up,down} are accepted. Named hard rejections (never silent skips): + `gated-q-proj` (holo's `attn_output_gate` interleaves Q+gate), `gdn-target`, + `non-full-attention-layer`, plus a bidirectional tensor↔target audit and A/B shape audit. +- A/B packed **rank-padded to `max_lora_rank`** in one fixed-address pool with per-module + `[max_loras]` device pointer tables (the frozen M2 layout contract; v0 fills slot 0). +- Pool VRAM is allocated before the KV-budget snapshot, so it is **budgeted against the KV + cache** (GB10 unified-memory OOM = freeze). +- Scaling `alpha/r` (or `alpha/√r` under rsLoRA), read per adapter, never defaulted. + +### M1 — runtime delta at the attention insertion points +- `apply_lora_delta` wired at **prefill k/v/o** and **decode k/v/o** (decode k/v applied + before norm/RoPE/kv-cache-write, so the KV cache stores the adapted values, matching + HF `k_norm(k_proj(x)+Δ)`). +- Deltas contract at the **padded** `max_rank` (B is packed with `max_rank` row stride); + a real-rank contraction would misread every B row past the first when `r < max_rank`. +- `ATLAS_LORA_EAGER=1` disables decode-graph capture (debugging hatch); deltas are + otherwise graph-safe (pool weights, arena scratch, and the f32 scale are load-time-fixed). +- `/v1/models` advertises the adapter name first; base-name requests get adapted output + with a one-line warn (v0 always-on wart; per-request routing is M2). + +## Verification (holo-3.1-0.8b on GB10, CUDA 13) + +- **Offline parity oracle** (`scripts/reference_deltas.py`): the loaded A/B/scale reproduce + the PEFT-exported reference deltas — 36/36 modules, `scaling=2.0`, **0.0 rel-err**. +- **`atlas-core` parser**: 11 unit tests (accept + every named reject) green. +- **Served, live**: + - startup logs `LoRA adapter 'holo-tiny' installed on 6 layers (pool=117.0 MiB)`; + - base output `"Paris."` → adapter output **differs** (delta is applied in the live + decode path); + - **graph == eager** (`ATLAS_LORA_EAGER=1`) — byte-identical, so deltas capture + correctly inside CUDA graphs. + +The test fixture (`test_data/lora-holo-tiny/`) is a **generated** PEFT adapter, deliberately +strong so its effect is unambiguous — no community adapter exists for Atlas's custom +NVFP4-packed bases, so a controllable fixture also exercises the reject/parity paths exactly. + +## Deferred (documented cut lines) +- **Dense-FFN delta** (gate/up/down): types + install are in place; the compute insertion in + `dense_ffn.rs` is not wired. The fixture targets FFN too, so enabling it is additive. +- **Prefix-cache warm-hit path** (`cache_skip_qkv.rs`) and **multi-seq decode** — until + wired, cache-hit prefills and concurrency ≥2 silently skip the deltas. **Run the POC at + batch size 1.** +- **q_proj / GDN / MoE / MLA targets**, **per-request adapter routing + multi-adapter** + (M2), **`lora_bgmv` kernel**, **TP>1** (startup-guarded to `world_size=1`). + +## Base-branch build note +`research/lora` predates the switch to the `cu*` driver API, so `cuda_backend`'s +`cudaMemcpy2DAsync` needs `libcudart` linked; `spark-runtime/build.rs` now does so and adds +the CUDA-13 SBSA lib path. On `main` (driver-API) this is a harmless no-op. diff --git a/scripts/gen_lora_adapter.py b/scripts/gen_lora_adapter.py new file mode 100644 index 000000000..c9fd0e097 --- /dev/null +++ b/scripts/gen_lora_adapter.py @@ -0,0 +1,157 @@ +#!/usr/bin/env python +"""Generate a tiny PEFT-format LoRA adapter for Hcompany/Holo-3.1-0.8B (Atlas LoRA MVP test fixture). + +Builds lora_A [r, in] / lora_B [out, r] tensors (NONZERO B — small normal init, unlike PEFT's +default B=0) for k_proj, v_proj, o_proj, gate_proj, up_proj, down_proj on ONLY the 6 +full-attention layers (indices 3, 7, 11, 15, 19, 23 — read from the base config.json, asserted). +q_proj is deliberately excluded: attn_output_gate=true makes it gated/interleaved ([4096,1024] = +2x q_heads*head_dim) and it is out of scope for LoRA v0. + +Per-module shapes are read from the base checkpoint's safetensors header (never assumed) and +asserted against the expected hidden=1024 / kv=512 / q-out=2048 / intermediate=3584 geometry. + +Output (default /home/ms/atlas/.claude/worktrees/lora-mvp-e0877873/test_data/lora-holo-tiny): + adapter_model.safetensors BF16 A/B pairs, PEFT save_pretrained key format + adapter_config.json written by peft.LoraConfig.save_pretrained (guaranteed PEFT-valid) + +Key style (--key-style): + vlm (default) base_model.model.model.language_model.layers.{i}....lora_A.weight + -- what PEFT 0.19.1 actually saves when wrapping Qwen3_5ForConditionalGeneration + (verified via get_peft_model_state_dict on the real class, meta device) + text base_model.model.model.layers.{i}....lora_A.weight + -- the text-tower-only form; the Atlas remapper must accept both + +Run: + /home/ms/nemotron-diffusion-playground/.venv/bin/python scripts/gen_lora_adapter.py +""" + +from __future__ import annotations + +import argparse +import json +import struct +import sys +from pathlib import Path + +import torch +from safetensors.torch import save_file + +BASE_MODEL_ID = "Hcompany/Holo-3.1-0.8B" +# k_proj/v_proj/o_proj + dense FFN. q_proj EXCLUDED (attn_output_gate -> gated/interleaved). +ATTN_TARGETS = ("k_proj", "v_proj", "o_proj") +MLP_TARGETS = ("gate_proj", "up_proj", "down_proj") +EXPECTED_FULL_ATTN_LAYERS = [3, 7, 11, 15, 19, 23] +# module -> (out_features, in_features), asserted against the real safetensors header. +EXPECTED_SHAPES = { + "self_attn.k_proj": (512, 1024), # kv_heads(2) * head_dim(256), hidden + "self_attn.v_proj": (512, 1024), + "self_attn.o_proj": (1024, 2048), # hidden, q_heads(8) * head_dim(256) + "mlp.gate_proj": (3584, 1024), # intermediate, hidden + "mlp.up_proj": (3584, 1024), + "mlp.down_proj": (1024, 3584), # hidden, intermediate +} + + +def find_snapshot() -> Path: + from huggingface_hub import snapshot_download + + return Path(snapshot_download(BASE_MODEL_ID, local_files_only=True)) + + +def read_safetensors_header(path: Path) -> dict: + with path.open("rb") as f: + (n,) = struct.unpack(" None: + ap = argparse.ArgumentParser(description=__doc__) + ap.add_argument( + "--out-dir", + type=Path, + default=Path("/home/ms/atlas/.claude/worktrees/lora-mvp-e0877873/test_data/lora-holo-tiny"), + ) + ap.add_argument("--rank", type=int, default=8) + ap.add_argument("--alpha", type=float, default=16.0) + ap.add_argument("--use-rslora", action="store_true", default=False) + ap.add_argument("--seed", type=int, default=1234) + ap.add_argument("--init-std", type=float, default=0.02) + ap.add_argument("--key-style", choices=("vlm", "text"), default="vlm") + args = ap.parse_args() + + snap = find_snapshot() + cfg = json.loads((snap / "config.json").read_text()) + text_cfg = cfg["text_config"] + + # --- CONFIRM the locked geometry from the real config, never assume ------------------- + layer_types = text_cfg["layer_types"] + full_attn = [i for i, t in enumerate(layer_types) if t == "full_attention"] + assert full_attn == EXPECTED_FULL_ATTN_LAYERS, f"full-attn layers {full_attn}" + assert text_cfg["attn_output_gate"] is True, "q_proj-exclusion rationale changed" + assert text_cfg["hidden_size"] == 1024 and text_cfg["intermediate_size"] == 3584 + assert text_cfg["num_attention_heads"] == 8 and text_cfg["num_key_value_heads"] == 2 + assert text_cfg["head_dim"] == 256 + + # --- read real weight shapes from the single-file safetensors header ------------------ + hdr = read_safetensors_header(snap / "model-00001-of-00001.safetensors") + base_prefix = "model.language_model.layers" + for i in full_attn: + for module, (out_f, in_f) in EXPECTED_SHAPES.items(): + key = f"{base_prefix}.{i}.{module}.weight" + shape = tuple(hdr[key]["shape"]) + assert shape == (out_f, in_f), f"{key}: {shape} != {(out_f, in_f)}" + # sanity: q_proj really is gated/interleaved (2 * q_heads * head_dim rows) + q_shape = tuple(hdr[f"{base_prefix}.{i}.self_attn.q_proj.weight"]["shape"]) + assert q_shape == (4096, 1024), f"q_proj shape {q_shape} (attn_output_gate expectation)" + + # --- build A/B tensors ------------------------------------------------------------------ + torch.manual_seed(args.seed) + r = args.rank + key_prefix = ( + "base_model.model.model.language_model.layers" + if args.key_style == "vlm" + else "base_model.model.model.layers" + ) + tensors: dict[str, torch.Tensor] = {} + for i in full_attn: + for module, (out_f, in_f) in EXPECTED_SHAPES.items(): + stem = f"{key_prefix}.{i}.{module}" + a = torch.randn(r, in_f, dtype=torch.float32) * args.init_std + b = torch.randn(out_f, r, dtype=torch.float32) * args.init_std # NONZERO B + tensors[f"{stem}.lora_A.weight"] = a.to(torch.bfloat16) + tensors[f"{stem}.lora_B.weight"] = b.to(torch.bfloat16) + + expected_n = len(full_attn) * len(EXPECTED_SHAPES) * 2 + assert len(tensors) == expected_n, f"{len(tensors)} != {expected_n}" + + # --- write adapter dir ------------------------------------------------------------------ + args.out_dir.mkdir(parents=True, exist_ok=True) + save_file(tensors, str(args.out_dir / "adapter_model.safetensors"), metadata={"format": "pt"}) + + from peft import LoraConfig + + lora_cfg = LoraConfig( + r=r, + lora_alpha=args.alpha, + lora_dropout=0.0, + use_rslora=args.use_rslora, + target_modules=sorted(ATTN_TARGETS + MLP_TARGETS), + layers_to_transform=full_attn, + bias="none", + task_type="CAUSAL_LM", + base_model_name_or_path=BASE_MODEL_ID, + inference_mode=True, + ) + lora_cfg.save_pretrained(str(args.out_dir)) # writes adapter_config.json (peft_type=LORA) + + total = sum(t.numel() * t.element_size() for t in tensors.values()) + print(f"wrote {len(tensors)} tensors ({total / 1024:.1f} KiB BF16) -> {args.out_dir}") + print(f" layers={full_attn} r={r} alpha={args.alpha} use_rslora={args.use_rslora} " + f"key_style={args.key_style}") + for k in sorted(tensors)[:4]: + print(f" {k} {tuple(tensors[k].shape)}") + return None + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/scripts/reference_deltas.py b/scripts/reference_deltas.py new file mode 100644 index 000000000..63fd00ee5 --- /dev/null +++ b/scripts/reference_deltas.py @@ -0,0 +1,111 @@ +#!/usr/bin/env python +"""Emit reference LoRA deltas for the Atlas offline parity test (M0 exit gate). + +For every (layer, module) pair in a PEFT adapter dir, loads lora_A [r, in] / lora_B [out, r], +reads r / lora_alpha / use_rslora from adapter_config.json (per-adapter, NEVER defaulted: +missing fields are fatal), computes + + scaling = lora_alpha / r (use_rslora = false) + scaling = lora_alpha / sqrt(r) (use_rslora = true) + delta = scaling * (B.float() @ A.float()) # [out, in], fp32 + +and writes one safetensors file the Rust offline-parity test loads and compares against its own +loaded-A/B/scale product (tolerance <= 1e-2 relative error, elementwise on the fp32 delta). + +Output keys (canonical module path = saved PEFT key with the "base_model.model." wrapper +stripped, e.g. "model.language_model.layers.3.self_attn.k_proj"): + {path}.delta fp32 [out, in] + {path}.scaling fp32 [1] +File-level metadata (strings): r, lora_alpha, use_rslora, scaling, num_modules. + +Run: + /home/ms/nemotron-diffusion-playground/.venv/bin/python scripts/reference_deltas.py \ + --adapter-dir test_data/lora-holo-tiny +""" + +from __future__ import annotations + +import argparse +import json +import math +import sys +from pathlib import Path + +import torch +from safetensors.torch import load_file, save_file + +PEFT_WRAPPER_PREFIX = "base_model.model." +A_SUFFIX = ".lora_A.weight" +B_SUFFIX = ".lora_B.weight" + + +def canonical_path(key: str, suffix: str) -> str: + assert key.endswith(suffix), key + stem = key[: -len(suffix)] + if stem.startswith(PEFT_WRAPPER_PREFIX): + stem = stem[len(PEFT_WRAPPER_PREFIX):] + return stem + + +def main() -> None: + ap = argparse.ArgumentParser(description=__doc__) + ap.add_argument( + "--adapter-dir", + type=Path, + default=Path("/home/ms/atlas/.claude/worktrees/lora-mvp-e0877873/test_data/lora-holo-tiny"), + ) + ap.add_argument("--out", type=Path, default=None, + help="output file (default {adapter-dir}/reference_deltas.safetensors)") + args = ap.parse_args() + out_path = args.out or (args.adapter_dir / "reference_deltas.safetensors") + + cfg = json.loads((args.adapter_dir / "adapter_config.json").read_text()) + # Scaling inputs are read per adapter and NEVER defaulted (locked decision). + for field in ("r", "lora_alpha", "use_rslora"): + if field not in cfg: + raise SystemExit(f"adapter_config.json missing required field {field!r} " + f"(older PEFT may omit use_rslora; refuse rather than default)") + if cfg.get("rank_pattern") or cfg.get("alpha_pattern"): + raise SystemExit("rank_pattern/alpha_pattern unsupported by the v0 oracle") + r = int(cfg["r"]) + alpha = float(cfg["lora_alpha"]) + use_rslora = bool(cfg["use_rslora"]) + scaling = alpha / math.sqrt(r) if use_rslora else alpha / r + + tensors = load_file(str(args.adapter_dir / "adapter_model.safetensors")) + a_keys = sorted(k for k in tensors if k.endswith(A_SUFFIX)) + b_keys = sorted(k for k in tensors if k.endswith(B_SUFFIX)) + stray = set(tensors) - set(a_keys) - set(b_keys) + if stray: + raise SystemExit(f"unrecognized adapter tensors (bidirectional-audit spirit): {stray}") + a_by_path = {canonical_path(k, A_SUFFIX): tensors[k] for k in a_keys} + b_by_path = {canonical_path(k, B_SUFFIX): tensors[k] for k in b_keys} + if a_by_path.keys() != b_by_path.keys(): + raise SystemExit(f"unpaired A/B: {a_by_path.keys() ^ b_by_path.keys()}") + + out: dict[str, torch.Tensor] = {} + print(f"r={r} alpha={alpha} use_rslora={use_rslora} -> scaling={scaling}") + for path in sorted(a_by_path): + a = a_by_path[path] # [r, in] + b = b_by_path[path] # [out, r] + assert a.shape[0] == r and b.shape[1] == r, (path, a.shape, b.shape) + delta = scaling * (b.float() @ a.float()) # [out, in] fp32 + out[f"{path}.delta"] = delta.contiguous() + out[f"{path}.scaling"] = torch.tensor([scaling], dtype=torch.float32) + print(f" {path}: A{tuple(a.shape)} {a.dtype} x B{tuple(b.shape)} -> " + f"delta{tuple(delta.shape)} |delta|max={delta.abs().max():.5f}") + + metadata = { + "r": str(r), + "lora_alpha": repr(alpha), + "use_rslora": str(use_rslora).lower(), + "scaling": repr(scaling), + "num_modules": str(len(a_by_path)), + } + save_file(out, str(out_path), metadata=metadata) + print(f"wrote {len(out)} tensors -> {out_path}") + return None + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/test_data/lora-holo-tiny/adapter_config.json b/test_data/lora-holo-tiny/adapter_config.json new file mode 100644 index 000000000..e8b85eefd --- /dev/null +++ b/test_data/lora-holo-tiny/adapter_config.json @@ -0,0 +1,54 @@ +{ + "alora_invocation_tokens": null, + "alpha_pattern": {}, + "arrow_config": null, + "auto_mapping": null, + "base_model_name_or_path": "Hcompany/Holo-3.1-0.8B", + "bias": "none", + "corda_config": null, + "ensure_weight_tying": false, + "eva_config": null, + "exclude_modules": null, + "fan_in_fan_out": false, + "inference_mode": true, + "init_lora_weights": true, + "layer_replication": null, + "layers_pattern": null, + "layers_to_transform": [ + 3, + 7, + 11, + 15, + 19, + 23 + ], + "loftq_config": {}, + "lora_alpha": 16.0, + "lora_bias": false, + "lora_dropout": 0.0, + "lora_ga_config": null, + "megatron_config": null, + "megatron_core": "megatron.core", + "modules_to_save": null, + "peft_type": "LORA", + "peft_version": "0.19.1", + "qalora_group_size": 16, + "r": 8, + "rank_pattern": {}, + "revision": null, + "target_modules": [ + "o_proj", + "up_proj", + "gate_proj", + "k_proj", + "down_proj", + "v_proj" + ], + "target_parameters": null, + "task_type": "CAUSAL_LM", + "trainable_token_indices": null, + "use_bdlora": null, + "use_dora": false, + "use_qalora": false, + "use_rslora": false +} \ No newline at end of file diff --git a/test_data/lora-holo-tiny/adapter_model.safetensors b/test_data/lora-holo-tiny/adapter_model.safetensors new file mode 100644 index 000000000..40fcef349 Binary files /dev/null and b/test_data/lora-holo-tiny/adapter_model.safetensors differ