rtx-csm: Phase 12.1 — extend LoRA coverage q+v → full attn + MLP
Both backbones (FP csm_fork + Q8 csm_quantized) now expose 7 LoRA hooks per layer: q/k/v/o on attention plus gate/up/down (Llama w1/w3/w2) on the SwiGLU MLP. LoraConfig::default() still returns q+v only (backward compat for existing trained adapters); LoraConfig::extended() returns the full 7-module set. lora_train + lora_finetune_step take a --extended-lora flag. Verified end-to-end on Metal: injection across all 16 backbone layers × 7 modules = 224 adapter Vars, 5.6M trainable params (~6.6× q+v alone, still tiny vs the 1B base). Step-0 loss matches the q+v baseline exactly (B=0 init is also a no-op for the new hooks). Forward + backward + AdamW + refresh_lora cycle runs without errors. LoRA test suite: 9 pass (added config_extended_targets_full_attn_and_mlp); full lib suite still 92/92. Co-Authored-By: Claude Opus 4.7 (1M context) <[email protected]>
This commit is contained in:
osobh
@@ -161,11 +161,15 @@ pub(crate) struct Attention {
|
||||
k_proj: Linear,
|
||||
v_proj: Linear,
|
||||
o_proj: Linear,
|
||||
/// Optional LoRA adapters for q_proj / v_proj — the literature recipe for
|
||||
/// CSM-style backbones. None means inference-only; Some adds an additive
|
||||
/// delta path. Inactive (B=0) at init so behavior is identical to base.
|
||||
/// Optional LoRA adapters on the four attention projections. None means
|
||||
/// inference-only; Some adds an additive delta path. Inactive (B=0) at
|
||||
/// init so behavior is identical to base. `o_lora` corresponds to the
|
||||
/// `output_proj` weight in safetensors (the field is named `o_proj` here
|
||||
/// for brevity but the upstream key is `output_proj`).
|
||||
pub(crate) q_lora: Option<crate::lora::LoraDelta>,
|
||||
pub(crate) k_lora: Option<crate::lora::LoraDelta>,
|
||||
pub(crate) v_lora: Option<crate::lora::LoraDelta>,
|
||||
pub(crate) o_lora: Option<crate::lora::LoraDelta>,
|
||||
rotary_emb: Arc<RotaryEmbedding>,
|
||||
kv_cache: Option<(Tensor, Tensor)>,
|
||||
num_heads: usize,
|
||||
@@ -189,7 +193,9 @@ impl Attention {
|
||||
v_proj,
|
||||
o_proj,
|
||||
q_lora: None,
|
||||
k_lora: None,
|
||||
v_lora: None,
|
||||
o_lora: None,
|
||||
rotary_emb,
|
||||
kv_cache: None,
|
||||
num_heads: cfg.num_heads,
|
||||
@@ -229,6 +235,10 @@ impl Attention {
|
||||
None => query_states,
|
||||
};
|
||||
let key_states = self.k_proj.forward(xs)?;
|
||||
let key_states = match &self.k_lora {
|
||||
Some(l) => (key_states + l.forward(xs)?)?,
|
||||
None => key_states,
|
||||
};
|
||||
let value_states = self.v_proj.forward(xs)?;
|
||||
let value_states = match &self.v_lora {
|
||||
Some(l) => (value_states + l.forward(xs)?)?,
|
||||
@@ -276,10 +286,14 @@ impl Attention {
|
||||
let attn_weights = candle_nn::ops::softmax_last_dim(&attn_weights)?;
|
||||
attn_weights.matmul(&value_states)?
|
||||
};
|
||||
let out = attn_output
|
||||
let pre_o = attn_output
|
||||
.transpose(1, 2)?
|
||||
.reshape((b_sz, q_len, self.num_heads * self.head_dim))?
|
||||
.apply(&self.o_proj)?;
|
||||
.reshape((b_sz, q_len, self.num_heads * self.head_dim))?;
|
||||
let out = self.o_proj.forward(&pre_o)?;
|
||||
let out = match &self.o_lora {
|
||||
Some(l) => (out + l.forward(&pre_o)?)?,
|
||||
None => out,
|
||||
};
|
||||
if std::env::var("CSM_LORA_DEBUG").is_ok() {
|
||||
eprintln!("Attn::forward returns: track_op={}", out.track_op());
|
||||
}
|
||||
@@ -292,10 +306,17 @@ impl Attention {
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
struct Mlp {
|
||||
pub(crate) struct Mlp {
|
||||
w1: Linear,
|
||||
w2: Linear,
|
||||
w3: Linear,
|
||||
/// Optional LoRA adapters on the SwiGLU MLP projections. CSM uses Llama's
|
||||
/// `w1` / `w2` / `w3` naming where `w1` is the gate, `w3` is the up, and
|
||||
/// `w2` is the down projection. Each is independently optional; B=0 init
|
||||
/// makes them no-ops until trained.
|
||||
pub(crate) gate_lora: Option<crate::lora::LoraDelta>,
|
||||
pub(crate) up_lora: Option<crate::lora::LoraDelta>,
|
||||
pub(crate) down_lora: Option<crate::lora::LoraDelta>,
|
||||
}
|
||||
|
||||
impl Mlp {
|
||||
@@ -303,15 +324,35 @@ impl Mlp {
|
||||
let w1 = linear_b(cfg.embed_dim, cfg.intermediate_dim, false, vb.pp("w1"))?;
|
||||
let w2 = linear_b(cfg.intermediate_dim, cfg.embed_dim, false, vb.pp("w2"))?;
|
||||
let w3 = linear_b(cfg.embed_dim, cfg.intermediate_dim, false, vb.pp("w3"))?;
|
||||
Ok(Self { w1, w2, w3 })
|
||||
Ok(Self {
|
||||
w1,
|
||||
w2,
|
||||
w3,
|
||||
gate_lora: None,
|
||||
up_lora: None,
|
||||
down_lora: None,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
impl Module for Mlp {
|
||||
fn forward(&self, xs: &Tensor) -> Result<Tensor> {
|
||||
let lhs = xs.apply(&self.w1)?.silu()?;
|
||||
let rhs = xs.apply(&self.w3)?;
|
||||
(lhs * rhs)?.apply(&self.w2)
|
||||
let gate = self.w1.forward(xs)?;
|
||||
let gate = match &self.gate_lora {
|
||||
Some(l) => (gate + l.forward(xs)?)?,
|
||||
None => gate,
|
||||
};
|
||||
let up = self.w3.forward(xs)?;
|
||||
let up = match &self.up_lora {
|
||||
Some(l) => (up + l.forward(xs)?)?,
|
||||
None => up,
|
||||
};
|
||||
let mid = (gate.silu()? * up)?;
|
||||
let down = self.w2.forward(&mid)?;
|
||||
match &self.down_lora {
|
||||
Some(l) => down + l.forward(&mid)?,
|
||||
None => Ok(down),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -320,7 +361,7 @@ pub(crate) struct Layer {
|
||||
mlp_norm: RmsNorm,
|
||||
sa_norm: RmsNorm,
|
||||
pub(crate) attn: Attention,
|
||||
mlp: Mlp,
|
||||
pub(crate) mlp: Mlp,
|
||||
}
|
||||
|
||||
impl Layer {
|
||||
@@ -524,11 +565,17 @@ impl Model {
|
||||
self.cfg_backbone.is_some()
|
||||
}
|
||||
|
||||
/// Inject trainable LoRA adapters on backbone q/v projections per the
|
||||
/// standard recipe (StyleSpeech / UtterTune / Koel-TTS): rank 8, alpha 16,
|
||||
/// q_proj + v_proj only, backbone only (decoder + heads stay frozen).
|
||||
/// Adapters are registered into `vm` for AdamW pickup; B is zero-initialized
|
||||
/// so initial behavior is identical to the un-adapted base.
|
||||
/// Inject trainable LoRA adapters on the backbone, with per-module
|
||||
/// targeting controlled by `cfg.target_modules`. The classic recipe is
|
||||
/// q+v only (rank 8, alpha 16); the extended recipe (Phase 12.1) covers
|
||||
/// q, k, v, output_proj plus the SwiGLU MLP (w1/w2/w3) for higher capacity
|
||||
/// when ~hours of training audio are available — this gives the adapter
|
||||
/// real prosody control rather than just a thin attention nudge.
|
||||
///
|
||||
/// Decoder + heads stay frozen by construction (we only walk the backbone
|
||||
/// layers; `cfg.exclude_patterns` is the additional safety net). Adapters
|
||||
/// are registered into `vm` for AdamW pickup; B is zero-initialized so
|
||||
/// initial behavior is identical to the un-adapted base.
|
||||
pub fn add_lora_to_backbone(
|
||||
&mut self,
|
||||
cfg: &crate::lora::LoraConfig,
|
||||
@@ -537,62 +584,107 @@ impl Model {
|
||||
let backbone_cfg = LlamaConfig::from_flavor(self.config.backbone_flavor);
|
||||
let head_dim = backbone_cfg.embed_dim / backbone_cfg.num_heads;
|
||||
let kv_dim = backbone_cfg.num_kv_heads * head_dim;
|
||||
let embed_dim = backbone_cfg.embed_dim;
|
||||
let inter_dim = backbone_cfg.intermediate_dim;
|
||||
let device = self.backbone.device.clone();
|
||||
let dtype = self.backbone.dtype;
|
||||
let mut q_count = 0usize;
|
||||
let mut v_count = 0usize;
|
||||
|
||||
// Closure to keep each per-target block readable and uniform.
|
||||
let make = |name: &str, in_dim: usize, out_dim: usize|
|
||||
-> candle_core::Result<crate::lora::LoraDelta>
|
||||
{
|
||||
crate::lora::LoraDelta::new(
|
||||
cfg.rank,
|
||||
cfg.alpha as f64,
|
||||
in_dim,
|
||||
out_dim,
|
||||
name,
|
||||
vm,
|
||||
&device,
|
||||
dtype,
|
||||
)
|
||||
.map_err(|e| candle_core::Error::Msg(e.to_string()))
|
||||
};
|
||||
|
||||
let mut counts = [0usize; 7]; // q, k, v, o, w1(gate), w3(up), w2(down)
|
||||
|
||||
for (i, layer) in self.backbone.layers.iter_mut().enumerate() {
|
||||
// q_proj: in=embed_dim, out=embed_dim
|
||||
if cfg.matches(&format!("backbone.layers.{i}.attn.q_proj.weight")) {
|
||||
q_count += 1;
|
||||
let q = crate::lora::LoraDelta::new(
|
||||
cfg.rank,
|
||||
cfg.alpha as f64,
|
||||
backbone_cfg.embed_dim,
|
||||
backbone_cfg.embed_dim,
|
||||
&format!("backbone.layers.{i}.attn.q_proj"),
|
||||
vm,
|
||||
&device,
|
||||
dtype,
|
||||
)
|
||||
.map_err(|e| candle_core::Error::Msg(e.to_string()))?;
|
||||
layer.attn.q_lora = Some(q);
|
||||
// attn.q_proj: embed_dim → embed_dim
|
||||
let key = format!("backbone.layers.{i}.attn.q_proj");
|
||||
if cfg.matches(&format!("{key}.weight")) {
|
||||
layer.attn.q_lora = Some(make(&key, embed_dim, embed_dim)?);
|
||||
counts[0] += 1;
|
||||
}
|
||||
// v_proj: in=embed_dim, out=kv_dim (note non-square)
|
||||
if cfg.matches(&format!("backbone.layers.{i}.attn.v_proj.weight")) {
|
||||
v_count += 1;
|
||||
let v = crate::lora::LoraDelta::new(
|
||||
cfg.rank,
|
||||
cfg.alpha as f64,
|
||||
backbone_cfg.embed_dim,
|
||||
kv_dim,
|
||||
&format!("backbone.layers.{i}.attn.v_proj"),
|
||||
vm,
|
||||
&device,
|
||||
dtype,
|
||||
)
|
||||
.map_err(|e| candle_core::Error::Msg(e.to_string()))?;
|
||||
layer.attn.v_lora = Some(v);
|
||||
// attn.k_proj: embed_dim → kv_dim
|
||||
let key = format!("backbone.layers.{i}.attn.k_proj");
|
||||
if cfg.matches(&format!("{key}.weight")) {
|
||||
layer.attn.k_lora = Some(make(&key, embed_dim, kv_dim)?);
|
||||
counts[1] += 1;
|
||||
}
|
||||
// attn.v_proj: embed_dim → kv_dim
|
||||
let key = format!("backbone.layers.{i}.attn.v_proj");
|
||||
if cfg.matches(&format!("{key}.weight")) {
|
||||
layer.attn.v_lora = Some(make(&key, embed_dim, kv_dim)?);
|
||||
counts[2] += 1;
|
||||
}
|
||||
// attn.output_proj: embed_dim → embed_dim. Note: safetensors key is
|
||||
// `output_proj` (per upstream csm/torchtune); the struct field is
|
||||
// named `o_proj` but the LoRA prefix follows the on-disk name so
|
||||
// saved adapter files are inspectable.
|
||||
let key = format!("backbone.layers.{i}.attn.output_proj");
|
||||
if cfg.matches(&format!("{key}.weight")) {
|
||||
layer.attn.o_lora = Some(make(&key, embed_dim, embed_dim)?);
|
||||
counts[3] += 1;
|
||||
}
|
||||
// mlp.w1 (SwiGLU gate): embed_dim → intermediate_dim
|
||||
let key = format!("backbone.layers.{i}.mlp.w1");
|
||||
if cfg.matches(&format!("{key}.weight")) {
|
||||
layer.mlp.gate_lora = Some(make(&key, embed_dim, inter_dim)?);
|
||||
counts[4] += 1;
|
||||
}
|
||||
// mlp.w3 (SwiGLU up): embed_dim → intermediate_dim
|
||||
let key = format!("backbone.layers.{i}.mlp.w3");
|
||||
if cfg.matches(&format!("{key}.weight")) {
|
||||
layer.mlp.up_lora = Some(make(&key, embed_dim, inter_dim)?);
|
||||
counts[5] += 1;
|
||||
}
|
||||
// mlp.w2 (SwiGLU down): intermediate_dim → embed_dim
|
||||
let key = format!("backbone.layers.{i}.mlp.w2");
|
||||
if cfg.matches(&format!("{key}.weight")) {
|
||||
layer.mlp.down_lora = Some(make(&key, inter_dim, embed_dim)?);
|
||||
counts[6] += 1;
|
||||
}
|
||||
}
|
||||
tracing::info!("LoRA injected into {q_count} q_proj and {v_count} v_proj backbone layers");
|
||||
tracing::info!(
|
||||
"LoRA injected: q={} k={} v={} o={} mlp_gate={} mlp_up={} mlp_down={}",
|
||||
counts[0], counts[1], counts[2], counts[3], counts[4], counts[5], counts[6],
|
||||
);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// After an optimizer step, refresh the LoRA Tensor handles inside each
|
||||
/// Attention from the VarMap so the next forward pass sees the updated
|
||||
/// values. Required because LoraDelta holds plain Tensors (not Vars), and
|
||||
/// candle's optimizer mutates the underlying Var storage in place.
|
||||
/// Attention / Mlp from the VarMap so the next forward pass sees the
|
||||
/// updated values. Required because LoraDelta holds plain Tensors (not
|
||||
/// Vars), and candle's optimizer mutates the underlying Var storage
|
||||
/// in place.
|
||||
pub fn refresh_lora(&mut self, vm: &candle_nn::VarMap) -> candle_core::Result<()> {
|
||||
for (i, layer) in self.backbone.layers.iter_mut().enumerate() {
|
||||
if let Some(q) = layer.attn.q_lora.as_mut() {
|
||||
q.refresh_from(vm, &format!("backbone.layers.{i}.attn.q_proj"))
|
||||
.map_err(|e| candle_core::Error::Msg(e.to_string()))?;
|
||||
}
|
||||
if let Some(v) = layer.attn.v_lora.as_mut() {
|
||||
v.refresh_from(vm, &format!("backbone.layers.{i}.attn.v_proj"))
|
||||
.map_err(|e| candle_core::Error::Msg(e.to_string()))?;
|
||||
}
|
||||
let do_refresh = |slot: Option<&mut crate::lora::LoraDelta>, key: String|
|
||||
-> candle_core::Result<()>
|
||||
{
|
||||
if let Some(d) = slot {
|
||||
d.refresh_from(vm, &key)
|
||||
.map_err(|e| candle_core::Error::Msg(e.to_string()))?;
|
||||
}
|
||||
Ok(())
|
||||
};
|
||||
do_refresh(layer.attn.q_lora.as_mut(), format!("backbone.layers.{i}.attn.q_proj"))?;
|
||||
do_refresh(layer.attn.k_lora.as_mut(), format!("backbone.layers.{i}.attn.k_proj"))?;
|
||||
do_refresh(layer.attn.v_lora.as_mut(), format!("backbone.layers.{i}.attn.v_proj"))?;
|
||||
do_refresh(layer.attn.o_lora.as_mut(), format!("backbone.layers.{i}.attn.output_proj"))?;
|
||||
do_refresh(layer.mlp.gate_lora.as_mut(), format!("backbone.layers.{i}.mlp.w1"))?;
|
||||
do_refresh(layer.mlp.up_lora.as_mut(), format!("backbone.layers.{i}.mlp.w3"))?;
|
||||
do_refresh(layer.mlp.down_lora.as_mut(), format!("backbone.layers.{i}.mlp.w2"))?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user