728 lines
27 KiB
Rust
728 lines
27 KiB
Rust
//! LoRA voice fine-tuning for CSM-1B.
|
||
//!
|
||
//! ## What ships
|
||
//!
|
||
//! - [`LoraConfig`] — the canonical hyperparameter bundle (rank, alpha,
|
||
//! target_modules, dropout) with sensible defaults derived from the
|
||
//! 2024–2025 LoRA-on-TTS literature (StyleSpeech, UtterTune, Koel-TTS).
|
||
//! - [`LoraAdapter`] — a single rank-`r` adapter pair `(A, B)` for one
|
||
//! target weight. `forward(&base_out, &xs)` adds the low-rank update.
|
||
//! - [`LoraSet`] — collection of adapters keyed by target safetensors
|
||
//! name (e.g., `backbone.layers.5.attn.q_proj.weight`).
|
||
//! - [`merge_into_safetensors`] — **offline merge**: read base weights,
|
||
//! read LoRA adapters, write out `W' = W + (B @ A) * (alpha / r)`.
|
||
//! This is the production path — once merged, inference uses the
|
||
//! existing un-modified `csm_fork::Model` with no per-call overhead.
|
||
//!
|
||
//! ## What's deferred
|
||
//!
|
||
//! The **training loop** itself. That requires:
|
||
//! 1. A paired-data pipeline: list of `(text, audio_24khz)` for the target
|
||
//! speaker, ~10–30 minutes total.
|
||
//! 2. Forward pass that exposes per-codebook logits at training time
|
||
//! (the current `generate_frame` samples internally; we'd need an
|
||
//! `forward_loss` variant).
|
||
//! 3. A loss function: cross-entropy on Mimi codes for c0..c31.
|
||
//! 4. AdamW optimizer with `parking_lot::RwLock<Tensor>` parameter
|
||
//! handles for the rank-r matrices.
|
||
//! 5. Mixed precision (bf16 forward, f32 master weights for stability).
|
||
//! 6. Maybe gradient checkpointing if backbone activations spill.
|
||
//!
|
||
//! Estimated training-loop effort: 1 week. Once that lands, an end-to-end
|
||
//! voice-clone pipeline is `extract_audio → train_lora → merge → generate`.
|
||
//!
|
||
//! ## Recipe defaults (per literature)
|
||
//! - rank: 8
|
||
//! - alpha: 16 (alpha/rank = 2 — modest update strength)
|
||
//! - target: `q_proj` + `v_proj` on backbone ONLY (not decoder, not FFN)
|
||
//! - dropout: 0.05
|
||
//! - learning rate: 1e-4, AdamW, cosine schedule
|
||
//! - epochs: 3–5 on ~30 min of audio (~150-300 utterances)
|
||
//!
|
||
//! ## References
|
||
//! - LoRA: Hu et al. arXiv:2106.09685
|
||
//! - StyleSpeech (TTS LoRA recipe): arXiv:2408.14713
|
||
//! - UtterTune: arXiv:2508.09767
|
||
|
||
use crate::error::{CsmError, Result};
|
||
use candle_core::{Module, Tensor};
|
||
use candle_nn::{Linear, VarBuilder, VarMap};
|
||
use std::collections::HashMap;
|
||
use std::path::{Path, PathBuf};
|
||
|
||
// Hand-impl Debug since Tensor doesn't pretty-print.
|
||
impl std::fmt::Debug for LoraDelta {
|
||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||
f.debug_struct("LoraDelta")
|
||
.field("a_shape", &self.a.shape())
|
||
.field("b_shape", &self.b.shape())
|
||
.field("scale", &self.scale)
|
||
.finish()
|
||
}
|
||
}
|
||
|
||
/// Additive-only LoRA delta: `delta(xs) = scale * (B @ A @ xs)`.
|
||
///
|
||
/// Unlike [`LoraLinear`] this does NOT own a base layer — it's the bare
|
||
/// adapter, intended to be added to an existing layer's output. This is the
|
||
/// right shape for slotting LoRA into an existing model with private layer
|
||
/// types: we leave the base layer untouched and just inject an additive path.
|
||
///
|
||
/// Init follows LoRA paper: A ~ Randn(std=1/r), B = 0 (so initial delta is 0
|
||
/// and the model behaves identically to the un-adapted base until training).
|
||
#[derive(Clone)]
|
||
pub struct LoraDelta {
|
||
pub a: Tensor, // (rank, in_dim) — trainable
|
||
pub b: Tensor, // (out_dim, rank) — trainable
|
||
pub scale: f64,
|
||
}
|
||
|
||
impl LoraDelta {
|
||
#[allow(clippy::too_many_arguments)]
|
||
pub fn new(
|
||
rank: usize,
|
||
alpha: f64,
|
||
in_dim: usize,
|
||
out_dim: usize,
|
||
prefix: &str,
|
||
vm: &VarMap,
|
||
device: &candle_core::Device,
|
||
_dtype: candle_core::DType,
|
||
) -> Result<Self> {
|
||
// LoRA params live in F32 regardless of the surrounding model dtype:
|
||
// candle's autograd is most reliable in F32, and the rank-r adapters
|
||
// are tiny (~1 MB total) so the precision cost is negligible. We cast
|
||
// to the input dtype at forward time.
|
||
let dtype = candle_core::DType::F32;
|
||
let scale = alpha / rank as f64;
|
||
let std = 1.0 / (rank as f64);
|
||
let init_a = candle_nn::Init::Randn {
|
||
mean: 0.0,
|
||
stdev: std,
|
||
};
|
||
let init_b = candle_nn::Init::Const(0.0);
|
||
let a = vm
|
||
.get(
|
||
(rank, in_dim),
|
||
&format!("{prefix}.lora_a"),
|
||
init_a,
|
||
dtype,
|
||
device,
|
||
)
|
||
.map_err(|e| CsmError::Other(anyhow::anyhow!("vm lora_a: {e}")))?;
|
||
let b = vm
|
||
.get(
|
||
(out_dim, rank),
|
||
&format!("{prefix}.lora_b"),
|
||
init_b,
|
||
dtype,
|
||
device,
|
||
)
|
||
.map_err(|e| CsmError::Other(anyhow::anyhow!("vm lora_b: {e}")))?;
|
||
Ok(Self { a, b, scale })
|
||
}
|
||
|
||
/// Refresh `a` and `b` from a VarMap (post-optimizer-step). The underlying
|
||
/// `Var` storage was updated by the optimizer; this struct holds plain
|
||
/// Tensor handles, so we re-snapshot to see the new values on next forward.
|
||
pub fn refresh_from(&mut self, vm: &VarMap, prefix: &str) -> Result<()> {
|
||
let vars = vm.data().lock().unwrap();
|
||
if let Some(a) = vars.get(&format!("{prefix}.lora_a")) {
|
||
self.a = a.as_tensor().clone();
|
||
}
|
||
if let Some(b) = vars.get(&format!("{prefix}.lora_b")) {
|
||
self.b = b.as_tensor().clone();
|
||
}
|
||
Ok(())
|
||
}
|
||
|
||
pub fn forward(&self, xs: &Tensor) -> candle_core::Result<Tensor> {
|
||
// LoRA params live in F32 for stable autograd; activations may be
|
||
// F16 (Metal) or F32 (CPU). Run the adapter math in F32 and cast
|
||
// back to xs's dtype at the end so the addition with the base output
|
||
// doesn't trip on a dtype mismatch.
|
||
let target_dtype = xs.dtype();
|
||
let xs_f32 = xs.to_dtype(candle_core::DType::F32)?;
|
||
// Broadcast A and B to match input rank.
|
||
let a_t = match *xs_f32.dims() {
|
||
[b1, b2, _, _] => self.a.t()?.broadcast_left((b1, b2))?,
|
||
[bsize, _, _] => self.a.t()?.broadcast_left(bsize)?,
|
||
_ => self.a.t()?,
|
||
};
|
||
let b_t = match *xs_f32.dims() {
|
||
[b1, b2, _, _] => self.b.t()?.broadcast_left((b1, b2))?,
|
||
[bsize, _, _] => self.b.t()?.broadcast_left(bsize)?,
|
||
_ => self.b.t()?,
|
||
};
|
||
let xs_a = xs_f32.matmul(&a_t)?;
|
||
let xs_ab = xs_a.matmul(&b_t)?;
|
||
let scaled = (xs_ab * self.scale)?;
|
||
scaled.to_dtype(target_dtype)
|
||
}
|
||
}
|
||
|
||
/// LoRA-augmented linear layer: `out = base(xs) + scale * (B @ A @ xs)` where
|
||
/// `base` is the frozen pre-trained weight and `A: (rank, in)`, `B: (out, rank)`
|
||
/// are the small trainable matrices.
|
||
///
|
||
/// At forward time we compose the contributions; at training time only `a` and
|
||
/// `b` accumulate gradients (the base Linear is constructed with non-Var
|
||
/// tensors so candle's autograd treats it as constant).
|
||
///
|
||
/// Init convention follows the original LoRA paper: A is normal-sampled with
|
||
/// std = 1/rank, B is zero — so the adapter starts as a no-op and the model
|
||
/// behaves identically to the base until training begins.
|
||
#[derive(Clone)]
|
||
pub struct LoraLinear {
|
||
pub base: Linear,
|
||
pub a: Tensor, // (rank, in_dim) — trainable
|
||
pub b: Tensor, // (out_dim, rank) — trainable
|
||
pub scale: f64,
|
||
pub rank: usize,
|
||
}
|
||
|
||
impl LoraLinear {
|
||
/// Wrap an existing frozen `Linear` with a trainable rank-r LoRA adapter.
|
||
/// The A/B params get registered into `vm` under `<prefix>.lora_a` /
|
||
/// `<prefix>.lora_b` so AdamW (or any optimizer) can find and update them.
|
||
#[allow(clippy::too_many_arguments)]
|
||
pub fn wrap(
|
||
base: Linear,
|
||
rank: usize,
|
||
alpha: f64,
|
||
in_dim: usize,
|
||
out_dim: usize,
|
||
prefix: &str,
|
||
vm: &VarMap,
|
||
device: &candle_core::Device,
|
||
dtype: candle_core::DType,
|
||
) -> Result<Self> {
|
||
let scale = alpha / rank as f64;
|
||
let std = 1.0 / (rank as f64);
|
||
let init_a = candle_nn::Init::Randn {
|
||
mean: 0.0,
|
||
stdev: std,
|
||
};
|
||
let init_b = candle_nn::Init::Const(0.0);
|
||
let a = vm
|
||
.get(
|
||
(rank, in_dim),
|
||
&format!("{prefix}.lora_a"),
|
||
init_a,
|
||
dtype,
|
||
device,
|
||
)
|
||
.map_err(|e| CsmError::Other(anyhow::anyhow!("vm lora_a: {e}")))?;
|
||
let b = vm
|
||
.get(
|
||
(out_dim, rank),
|
||
&format!("{prefix}.lora_b"),
|
||
init_b,
|
||
dtype,
|
||
device,
|
||
)
|
||
.map_err(|e| CsmError::Other(anyhow::anyhow!("vm lora_b: {e}")))?;
|
||
Ok(Self {
|
||
base,
|
||
a,
|
||
b,
|
||
scale,
|
||
rank,
|
||
})
|
||
}
|
||
}
|
||
|
||
impl Module for LoraLinear {
|
||
fn forward(&self, xs: &Tensor) -> candle_core::Result<Tensor> {
|
||
let base_out = self.base.forward(xs)?;
|
||
// xs (.., in) @ A.t (in, rank) → (.., rank)
|
||
let xs_a = xs.matmul(&self.a.t()?)?;
|
||
// (.., rank) @ B.t (rank, out) → (.., out)
|
||
let xs_ab = xs_a.matmul(&self.b.t()?)?;
|
||
let scaled = (xs_ab * self.scale)?;
|
||
base_out + scaled
|
||
}
|
||
}
|
||
|
||
/// Convenience: scan a VarBuilder path and `wrap` the named base linears with
|
||
/// LoRA according to a `LoraConfig`. Used during model construction in the
|
||
/// LoRA-aware training fork. Returned Vec is keyed by safetensors-name so
|
||
/// callers can plug them back in as Module replacements.
|
||
#[allow(dead_code)]
|
||
pub fn build_lora_set(
|
||
cfg: &LoraConfig,
|
||
bases: &HashMap<String, (Linear, usize, usize)>,
|
||
vm: &VarMap,
|
||
device: &candle_core::Device,
|
||
dtype: candle_core::DType,
|
||
) -> Result<HashMap<String, LoraLinear>> {
|
||
let mut out = HashMap::new();
|
||
for (name, (base, in_dim, out_dim)) in bases.iter() {
|
||
if !cfg.matches(name) {
|
||
continue;
|
||
}
|
||
let lora = LoraLinear::wrap(
|
||
base.clone(),
|
||
cfg.rank,
|
||
cfg.alpha as f64,
|
||
*in_dim,
|
||
*out_dim,
|
||
name,
|
||
vm,
|
||
device,
|
||
dtype,
|
||
)?;
|
||
out.insert(name.clone(), lora);
|
||
}
|
||
Ok(out)
|
||
}
|
||
|
||
// Suppress unused warning when the trainer path isn't compiled.
|
||
#[allow(dead_code)]
|
||
fn _vb_unused(_vb: &VarBuilder) {}
|
||
|
||
#[derive(Debug, Clone)]
|
||
pub struct LoraConfig {
|
||
pub rank: usize,
|
||
pub alpha: f32,
|
||
/// Substring patterns matched against safetensors keys. A key containing
|
||
/// any pattern is targeted. Default: `["q_proj", "v_proj"]` against
|
||
/// `backbone.*` only (the Llama-3.2 1B backbone projections).
|
||
pub target_modules: Vec<String>,
|
||
/// Patterns that EXCLUDE a target even if it matches `target_modules`.
|
||
/// Default: `["decoder.", "audio_head", "codebook0_head"]`.
|
||
pub exclude_patterns: Vec<String>,
|
||
/// Training-only dropout; not used at merge time.
|
||
pub dropout: f32,
|
||
}
|
||
|
||
impl Default for LoraConfig {
|
||
fn default() -> Self {
|
||
Self {
|
||
rank: 8,
|
||
alpha: 16.0,
|
||
target_modules: vec!["q_proj".into(), "v_proj".into()],
|
||
exclude_patterns: vec![
|
||
"decoder.".into(),
|
||
"audio_head".into(),
|
||
"codebook0_head".into(),
|
||
"audio_embeddings".into(),
|
||
"text_embeddings".into(),
|
||
"projection".into(),
|
||
],
|
||
dropout: 0.05,
|
||
}
|
||
}
|
||
}
|
||
|
||
impl LoraConfig {
|
||
pub fn scale(&self) -> f32 {
|
||
self.alpha / self.rank as f32
|
||
}
|
||
|
||
pub fn matches(&self, name: &str) -> bool {
|
||
if self.exclude_patterns.iter().any(|p| name.contains(p)) {
|
||
return false;
|
||
}
|
||
self.target_modules.iter().any(|p| name.contains(p))
|
||
}
|
||
|
||
/// Extended coverage (Phase 12.1): all four attention projections plus the
|
||
/// SwiGLU MLP. Use this when fine-tuning on hours of audio for stronger
|
||
/// prosody adaptation; the q+v default is a safer floor when data is
|
||
/// scarce (~30 min). Adapter param count for rank 8 / 16 backbone layers
|
||
/// goes from ~1 MB (q+v) to ~6 MB (full) — still negligible vs the 1 B
|
||
/// base model. Substring patterns match the safetensors keys produced by
|
||
/// CSM's torchtune-style backbone (`q_proj`, `k_proj`, `v_proj`,
|
||
/// `output_proj`, `mlp.w1`, `mlp.w2`, `mlp.w3`).
|
||
pub fn extended() -> Self {
|
||
Self {
|
||
target_modules: vec![
|
||
"q_proj".into(),
|
||
"k_proj".into(),
|
||
"v_proj".into(),
|
||
"output_proj".into(),
|
||
"mlp.w1".into(),
|
||
"mlp.w2".into(),
|
||
"mlp.w3".into(),
|
||
],
|
||
..Self::default()
|
||
}
|
||
}
|
||
}
|
||
|
||
/// One LoRA adapter for one weight. `A` is rank × in_features, `B` is
|
||
/// out_features × rank. The update is `B @ A * scale` added to the base weight.
|
||
#[derive(Debug, Clone)]
|
||
pub struct LoraAdapter {
|
||
/// rank × in
|
||
pub a: Vec<f32>,
|
||
/// out × rank
|
||
pub b: Vec<f32>,
|
||
pub rank: usize,
|
||
pub in_features: usize,
|
||
pub out_features: usize,
|
||
}
|
||
|
||
impl LoraAdapter {
|
||
pub fn new_zero(rank: usize, in_features: usize, out_features: usize) -> Self {
|
||
Self {
|
||
a: vec![0.0; rank * in_features],
|
||
b: vec![0.0; out_features * rank],
|
||
rank,
|
||
in_features,
|
||
out_features,
|
||
}
|
||
}
|
||
|
||
/// Compute `B @ A` as a flat `out × in` matrix. Used by the offline merger.
|
||
pub fn delta_w(&self, scale: f32) -> Vec<f32> {
|
||
let mut out = vec![0.0f32; self.out_features * self.in_features];
|
||
for o in 0..self.out_features {
|
||
for i in 0..self.in_features {
|
||
let mut acc = 0.0f32;
|
||
for r in 0..self.rank {
|
||
acc += self.b[o * self.rank + r] * self.a[r * self.in_features + i];
|
||
}
|
||
out[o * self.in_features + i] = acc * scale;
|
||
}
|
||
}
|
||
out
|
||
}
|
||
}
|
||
|
||
#[derive(Debug, Default, Clone)]
|
||
pub struct LoraSet {
|
||
pub config: LoraConfig,
|
||
/// keyed by base safetensors weight name (e.g. `backbone.layers.5.attn.q_proj.weight`)
|
||
pub adapters: HashMap<String, LoraAdapter>,
|
||
}
|
||
|
||
impl LoraSet {
|
||
pub fn new(config: LoraConfig) -> Self {
|
||
Self {
|
||
config,
|
||
adapters: HashMap::new(),
|
||
}
|
||
}
|
||
|
||
/// Load adapters from a directory containing `<weight_name>.a.f32` and
|
||
/// `<weight_name>.b.f32` raw little-endian f32 dumps. Useful when training
|
||
/// is done in Python and adapters are exported as plain bytes — avoids
|
||
/// any safetensors/PyTorch coupling for v1.
|
||
pub fn load_from_dir<P: AsRef<Path>>(_dir: P, _config: LoraConfig) -> Result<Self> {
|
||
// Stub. Once a training loop exists, fill this in to walk the dir,
|
||
// pair `.a` / `.b` files, infer shapes from filename or sidecar JSON.
|
||
Err(CsmError::Config(
|
||
"LoraSet::load_from_dir: not yet implemented (training loop not yet shipped)".into(),
|
||
))
|
||
}
|
||
|
||
pub fn insert(&mut self, name: String, adapter: LoraAdapter) {
|
||
self.adapters.insert(name, adapter);
|
||
}
|
||
}
|
||
|
||
/// Read a trained LoRA adapter file (produced by
|
||
/// `training::save_lora_adapter[_with_metadata]`) and reconstruct the
|
||
/// `LoraSet` keyed by base-weight names. The adapter file's tensors are
|
||
/// named `<prefix>.lora_a` / `<prefix>.lora_b`; we strip the suffix and
|
||
/// pair them up.
|
||
pub fn load_lora_set_from_safetensors<P: AsRef<Path>>(
|
||
path: P,
|
||
config: LoraConfig,
|
||
) -> Result<LoraSet> {
|
||
let device = candle_core::Device::Cpu;
|
||
let tensors = candle_core::safetensors::load(path.as_ref(), &device).map_err(|e| {
|
||
CsmError::Other(anyhow::anyhow!(
|
||
"load adapter {}: {e}",
|
||
path.as_ref().display()
|
||
))
|
||
})?;
|
||
let mut a_map: HashMap<String, candle_core::Tensor> = HashMap::new();
|
||
let mut b_map: HashMap<String, candle_core::Tensor> = HashMap::new();
|
||
for (name, t) in tensors.into_iter() {
|
||
if let Some(prefix) = name.strip_suffix(".lora_a") {
|
||
a_map.insert(prefix.to_string(), t);
|
||
} else if let Some(prefix) = name.strip_suffix(".lora_b") {
|
||
b_map.insert(prefix.to_string(), t);
|
||
}
|
||
}
|
||
let mut set = LoraSet::new(config);
|
||
for (prefix, a_t) in a_map.into_iter() {
|
||
let Some(b_t) = b_map.remove(&prefix) else {
|
||
tracing::warn!("LoRA adapter `{prefix}.lora_a` has no matching `.lora_b` — skipping");
|
||
continue;
|
||
};
|
||
// (rank, in_features) for A, (out_features, rank) for B
|
||
let a_dims = a_t.dims();
|
||
let b_dims = b_t.dims();
|
||
if a_dims.len() != 2 || b_dims.len() != 2 || a_dims[0] != b_dims[1] {
|
||
tracing::warn!(
|
||
"LoRA adapter `{prefix}` shape mismatch: a={a_dims:?} b={b_dims:?} — skipping"
|
||
);
|
||
continue;
|
||
}
|
||
let (rank, in_features) = (a_dims[0], a_dims[1]);
|
||
let out_features = b_dims[0];
|
||
let a_vec: Vec<f32> = a_t
|
||
.flatten_all()
|
||
.and_then(|t| t.to_dtype(candle_core::DType::F32))
|
||
.and_then(|t| t.to_vec1())
|
||
.map_err(|e| CsmError::Other(anyhow::anyhow!("a→f32 flatten: {e}")))?;
|
||
let b_vec: Vec<f32> = b_t
|
||
.flatten_all()
|
||
.and_then(|t| t.to_dtype(candle_core::DType::F32))
|
||
.and_then(|t| t.to_vec1())
|
||
.map_err(|e| CsmError::Other(anyhow::anyhow!("b→f32 flatten: {e}")))?;
|
||
// Base-weight name is `<prefix>.weight`.
|
||
let weight_name = format!("{prefix}.weight");
|
||
set.insert(
|
||
weight_name,
|
||
LoraAdapter {
|
||
a: a_vec,
|
||
b: b_vec,
|
||
rank,
|
||
in_features,
|
||
out_features,
|
||
},
|
||
);
|
||
}
|
||
if !b_map.is_empty() {
|
||
tracing::warn!(
|
||
"LoRA adapter has dangling .lora_b without .lora_a: {:?}",
|
||
b_map.keys()
|
||
);
|
||
}
|
||
Ok(set)
|
||
}
|
||
|
||
/// **Offline merger**: read CSM safetensors, fold LoRA deltas into the
|
||
/// targeted weights, write out a new safetensors file. Once merged, the
|
||
/// existing `csm_fork::Model` and `Generator::load_csm_1b_from_path`
|
||
/// paths use the merged checkpoint with zero per-inference overhead.
|
||
///
|
||
/// `scale` is the LoRA scale (typically `alpha / rank`); pass the value
|
||
/// the adapter was trained with — it's stored in the adapter's metadata
|
||
/// when saved via `save_lora_adapter_with_metadata`.
|
||
pub fn merge_into_safetensors<P: AsRef<Path>>(
|
||
base_safetensors: P,
|
||
lora: &LoraSet,
|
||
scale: f32,
|
||
output_safetensors: P,
|
||
) -> Result<MergeReport> {
|
||
let device = candle_core::Device::Cpu;
|
||
let base = candle_core::safetensors::load(base_safetensors.as_ref(), &device).map_err(|e| {
|
||
CsmError::Other(anyhow::anyhow!(
|
||
"load base {}: {e}",
|
||
base_safetensors.as_ref().display()
|
||
))
|
||
})?;
|
||
let mut merged: HashMap<String, candle_core::Tensor> = HashMap::new();
|
||
let mut report = MergeReport::default();
|
||
for (name, base_t) in base.into_iter() {
|
||
if let Some(adapter) = lora.adapters.get(&name) {
|
||
let dtype = base_t.dtype();
|
||
let shape = base_t.shape().clone();
|
||
// base → F32 for the additive merge
|
||
let base_f32 = base_t
|
||
.to_dtype(candle_core::DType::F32)
|
||
.map_err(|e| CsmError::Other(anyhow::anyhow!("upcast {name}: {e}")))?;
|
||
let dims = base_f32.dims();
|
||
if dims.len() != 2 || dims[0] != adapter.out_features || dims[1] != adapter.in_features
|
||
{
|
||
tracing::warn!(
|
||
"LoRA shape mismatch for `{name}`: base={dims:?} adapter=({},{}) — passing through unchanged",
|
||
adapter.out_features,
|
||
adapter.in_features
|
||
);
|
||
merged.insert(name, base_t);
|
||
report.skipped += 1;
|
||
continue;
|
||
}
|
||
let delta = adapter.delta_w(scale);
|
||
let delta_t = candle_core::Tensor::from_vec(
|
||
delta,
|
||
(adapter.out_features, adapter.in_features),
|
||
&device,
|
||
)
|
||
.map_err(|e| CsmError::Other(anyhow::anyhow!("delta tensor: {e}")))?;
|
||
let summed = base_f32
|
||
.add(&delta_t)
|
||
.map_err(|e| CsmError::Other(anyhow::anyhow!("add delta: {e}")))?;
|
||
// Cast back to the original dtype (F16/BF16) for storage parity.
|
||
let out_t = summed
|
||
.to_dtype(dtype)
|
||
.and_then(|t| t.reshape(shape))
|
||
.map_err(|e| CsmError::Other(anyhow::anyhow!("downcast {name}: {e}")))?;
|
||
merged.insert(name, out_t);
|
||
report.merged += 1;
|
||
} else {
|
||
merged.insert(name, base_t);
|
||
report.passthrough += 1;
|
||
}
|
||
}
|
||
safetensors::serialize_to_file(&merged, None, output_safetensors.as_ref())
|
||
.map_err(|e| CsmError::Other(anyhow::anyhow!("save merged: {e}")))?;
|
||
tracing::info!(
|
||
"lora merge: {} merged, {} skipped, {} passthrough → {}",
|
||
report.merged,
|
||
report.skipped,
|
||
report.passthrough,
|
||
output_safetensors.as_ref().display()
|
||
);
|
||
Ok(report)
|
||
}
|
||
|
||
#[derive(Debug, Default)]
|
||
pub struct MergeReport {
|
||
/// Number of base tensors that had LoRA deltas applied.
|
||
pub merged: usize,
|
||
/// Number of base tensors that LoRA targeted but had shape mismatches
|
||
/// (passed through unchanged with a warn).
|
||
pub skipped: usize,
|
||
/// Number of base tensors with no LoRA target (byte-copied).
|
||
pub passthrough: usize,
|
||
/// Where the merged checkpoint was written.
|
||
pub output_path: PathBuf,
|
||
}
|
||
|
||
#[cfg(test)]
|
||
mod tests {
|
||
use super::*;
|
||
use candle_core::{DType, Device};
|
||
|
||
#[test]
|
||
fn lora_linear_starts_as_noop_when_b_is_zero() {
|
||
let dev = Device::Cpu;
|
||
// Build a base Linear with a known weight: 4x3.
|
||
let w = Tensor::from_slice(
|
||
&[
|
||
1.0f32, 2.0, 3.0, 0.5, -0.5, 1.0, 0.0, 0.0, 1.0, 2.0, 1.0, 0.5,
|
||
],
|
||
(4, 3),
|
||
&dev,
|
||
)
|
||
.unwrap();
|
||
let base = Linear::new(w, None);
|
||
let xs = Tensor::from_slice(&[1.0f32, 2.0, 3.0], (1, 3), &dev).unwrap();
|
||
let base_out = base.forward(&xs).unwrap().to_vec2::<f32>().unwrap();
|
||
|
||
// Wrap with LoraLinear; B init is zero so the adapter contribution is 0.
|
||
let vm = VarMap::new();
|
||
let lora = LoraLinear::wrap(base, 2, 4.0, 3, 4, "test", &vm, &dev, DType::F32).unwrap();
|
||
let lora_out = lora.forward(&xs).unwrap().to_vec2::<f32>().unwrap();
|
||
|
||
for (a, b) in base_out[0].iter().zip(&lora_out[0]) {
|
||
assert!((a - b).abs() < 1e-5, "no-op violated: {a} vs {b}");
|
||
}
|
||
}
|
||
|
||
#[test]
|
||
fn lora_linear_diverges_after_perturbing_b() {
|
||
let dev = Device::Cpu;
|
||
let w = Tensor::zeros((4, 3), DType::F32, &dev).unwrap();
|
||
let base = Linear::new(w, None);
|
||
let xs = Tensor::from_slice(&[1.0f32, 2.0, 3.0], (1, 3), &dev).unwrap();
|
||
let vm = VarMap::new();
|
||
let mut lora = LoraLinear::wrap(base, 2, 4.0, 3, 4, "test", &vm, &dev, DType::F32).unwrap();
|
||
// Manually overwrite B with non-zero values.
|
||
lora.b =
|
||
Tensor::from_slice(&[1.0f32, 0.0, 0.0, 1.0, 1.0, 1.0, 0.0, 0.0], (4, 2), &dev).unwrap();
|
||
// And A with something concrete.
|
||
lora.a = Tensor::from_slice(&[0.1f32, 0.2, 0.3, 0.0, 0.5, 0.0], (2, 3), &dev).unwrap();
|
||
|
||
let out = lora.forward(&xs).unwrap().to_vec2::<f32>().unwrap();
|
||
// Base is zero so output = scale * (B @ A @ xs)
|
||
// A @ xs: [0.1+0.4+0.9, 0+1+0] = [1.4, 1.0]
|
||
// B @ (A @ xs): [1.4, 0+1, 1.4+1.0, 0] = [1.4, 1.0, 2.4, 0]
|
||
// scale = alpha/rank = 4/2 = 2.0
|
||
// result: [2.8, 2.0, 4.8, 0.0]
|
||
let expected = [2.8_f32, 2.0, 4.8, 0.0];
|
||
for (got, exp) in out[0].iter().zip(&expected) {
|
||
assert!((got - exp).abs() < 1e-4, "got {got} expected {exp}");
|
||
}
|
||
}
|
||
|
||
#[test]
|
||
fn config_default_targets_q_v_only() {
|
||
let c = LoraConfig::default();
|
||
assert!(c.matches("backbone.layers.0.attn.q_proj.weight"));
|
||
assert!(c.matches("backbone.layers.0.attn.v_proj.weight"));
|
||
assert!(!c.matches("backbone.layers.0.attn.k_proj.weight"));
|
||
assert!(!c.matches("backbone.layers.0.attn.o_proj.weight"));
|
||
// Decoder excluded
|
||
assert!(!c.matches("decoder.layers.0.attn.q_proj.weight"));
|
||
// Heads excluded
|
||
assert!(!c.matches("codebook0_head.weight"));
|
||
assert!(!c.matches("audio_embeddings.weight"));
|
||
}
|
||
|
||
#[test]
|
||
fn config_extended_targets_full_attn_and_mlp() {
|
||
let c = LoraConfig::extended();
|
||
// All four attention projections (note: real safetensors key is
|
||
// `output_proj`, not `o_proj`).
|
||
assert!(c.matches("backbone.layers.0.attn.q_proj.weight"));
|
||
assert!(c.matches("backbone.layers.0.attn.k_proj.weight"));
|
||
assert!(c.matches("backbone.layers.0.attn.v_proj.weight"));
|
||
assert!(c.matches("backbone.layers.0.attn.output_proj.weight"));
|
||
// SwiGLU MLP (Llama-style w1/w2/w3 naming). Patterns are scoped to
|
||
// `mlp.wN` so they don't accidentally pick up unrelated `wN`-bearing
|
||
// strings.
|
||
assert!(c.matches("backbone.layers.0.mlp.w1.weight"));
|
||
assert!(c.matches("backbone.layers.0.mlp.w2.weight"));
|
||
assert!(c.matches("backbone.layers.0.mlp.w3.weight"));
|
||
// Decoder + heads still excluded.
|
||
assert!(!c.matches("decoder.layers.0.attn.q_proj.weight"));
|
||
assert!(!c.matches("decoder.layers.0.mlp.w1.weight"));
|
||
assert!(!c.matches("codebook0_head.weight"));
|
||
assert!(!c.matches("projection.weight"));
|
||
assert!(!c.matches("audio_embeddings.weight"));
|
||
}
|
||
|
||
#[test]
|
||
fn scale_is_alpha_over_rank() {
|
||
let c = LoraConfig::default();
|
||
assert!((c.scale() - 2.0).abs() < 1e-6);
|
||
}
|
||
|
||
#[test]
|
||
fn adapter_zero_init_has_zero_delta() {
|
||
let a = LoraAdapter::new_zero(8, 2048, 2048);
|
||
let d = a.delta_w(2.0);
|
||
assert!(d.iter().all(|&x| x == 0.0));
|
||
}
|
||
|
||
#[test]
|
||
fn adapter_nonzero_delta_shape() {
|
||
let mut a = LoraAdapter::new_zero(2, 3, 4);
|
||
// a (2x3) [[1,0,0],[0,1,0]] b (4x2) [[1,0],[0,1],[1,1],[2,0]]
|
||
a.a = vec![1.0, 0.0, 0.0, 0.0, 1.0, 0.0];
|
||
a.b = vec![1.0, 0.0, 0.0, 1.0, 1.0, 1.0, 2.0, 0.0];
|
||
let d = a.delta_w(1.0);
|
||
// BA = b * a:
|
||
// row0: [1,0]*[a] = [1,0,0]
|
||
// row1: [0,1]*[a] = [0,1,0]
|
||
// row2: [1,1]*[a] = [1,1,0]
|
||
// row3: [2,0]*[a] = [2,0,0]
|
||
assert_eq!(
|
||
d,
|
||
vec![1.0, 0.0, 0.0, 0.0, 1.0, 0.0, 1.0, 1.0, 0.0, 2.0, 0.0, 0.0]
|
||
);
|
||
}
|
||
|
||
#[test]
|
||
fn merge_errors_on_missing_base_file() {
|
||
let lora = LoraSet::new(LoraConfig::default());
|
||
let r = merge_into_safetensors("/nonexistent", &lora, 1.0, "/nonexistent");
|
||
assert!(r.is_err());
|
||
}
|
||
|
||
#[test]
|
||
fn load_from_dir_returns_typed_error() {
|
||
let r = LoraSet::load_from_dir("/nonexistent", LoraConfig::default());
|
||
assert!(r.is_err());
|
||
}
|
||
}
|