rustytorch/crates/models/rtx-csm/examples/lora_train_step.rs

//! Single-step LoRA training demo.
//!
//! Wraps a frozen Linear with a `LoraLinear` adapter, computes an MSE loss
//! against a synthetic target, runs backward, takes one AdamW step, and
//! verifies that gradients flow ONLY through the LoRA matrices A and B (the
//! base weight stays frozen).
//!
//! This is the foundational primitive for full CSM voice fine-tuning. Once
//! verified at this scale, the same machinery wraps the backbone q_proj /
//! v_proj projections in csm_fork::Attention. The real training loop adds:
//!   - Paired-data loader: `(transcript, wav)` → tokens + Mimi codes
//!   - Per-codebook cross-entropy loss (forward_loss method on Model)
//!   - Mixed precision (bf16 forward, f32 master weights)
//!   - Multi-epoch driver with cyclic LR / cosine schedule
//!   - Optional depth-decoder 1/16 frame trick for compute efficiency
//!
//! Run:
//!   cargo run -p rtx-csm --release --example lora_train_step

use anyhow::Result;
use candle_core::{DType, Device, Module, Tensor};
use candle_nn::{AdamW, Linear, Optimizer, ParamsAdamW, VarMap};
use rtx_csm::lora::LoraLinear;

fn mse(pred: &Tensor, target: &Tensor) -> Result<Tensor> {
    let diff = (pred - target)?;
    let sq = diff.sqr()?;
    Ok(sq.mean_all()?)
}

fn main() -> Result<()> {
    let dev = Device::Cpu;
    let dtype = DType::F32;

    // Frozen base weight. Initialize random; we'll never change this.
    let base_w = Tensor::randn(0.0f32, 0.5, (8, 4), &dev)?;
    let base = Linear::new(base_w.clone(), None);

    // Trainable LoRA params live in this VarMap. AdamW will update only these.
    let vm = VarMap::new();
    let mut lora = LoraLinear::wrap(base, 4, 8.0, 4, 8, "demo", &vm, &dev, dtype)?;

    // Synthetic target: a fixed mapping we want LoRA to learn.
    // Make the target = base.forward(xs) + a known delta so LoRA must
    // learn the delta.
    let xs = Tensor::randn(0.0f32, 1.0, (16, 4), &dev)?;
    let delta_w = Tensor::randn(0.0f32, 0.3, (8, 4), &dev)?;
    let target_delta = xs.matmul(&delta_w.t()?)?;
    let target = (xs.matmul(&base_w.t()?)? + target_delta.clone())?;

    // Optimizer: AdamW with a moderate lr.
    let mut optim = AdamW::new(
        vm.all_vars(),
        ParamsAdamW {
            lr: 5e-2,
            ..ParamsAdamW::default()
        },
    )?;

    // Snapshot A and B before training.
    let a_before = lora.a.flatten_all()?.to_vec1::<f32>()?;
    let b_before = lora.b.flatten_all()?.to_vec1::<f32>()?;

    // 50 training steps — enough for visible loss decrease without dragging.
    let n_steps = 50usize;
    let mut losses = Vec::with_capacity(n_steps);
    for step in 0..n_steps {
        let pred = lora.forward(&xs)?;
        let loss = mse(&pred, &target)?;
        // Backward + step.
        // Re-pull A/B from VarMap because optim updates the underlying Var
        // tensors, but lora.a / lora.b are clones from construction time.
        // To make the next forward see the updated params, we have to refresh.
        let g = loss.backward()?;
        optim.step(&g)?;
        // Refresh A/B from the VarMap so the next forward sees the new values.
        let vars = vm.data().lock().unwrap();
        lora.a = vars.get("demo.lora_a").unwrap().as_tensor().clone();
        lora.b = vars.get("demo.lora_b").unwrap().as_tensor().clone();
        drop(vars);
        let l = loss.to_scalar::<f32>()?;
        losses.push(l);
        if step % 10 == 0 || step == n_steps - 1 {
            println!("step {step:>3}: loss = {l:.6}");
        }
    }

    // After training, A and B should be DIFFERENT from their init values.
    let a_after = lora.a.flatten_all()?.to_vec1::<f32>()?;
    let b_after = lora.b.flatten_all()?.to_vec1::<f32>()?;
    let a_changed: f32 = a_before
        .iter()
        .zip(&a_after)
        .map(|(x, y)| (x - y).powi(2))
        .sum::<f32>()
        .sqrt();
    let b_changed: f32 = b_before
        .iter()
        .zip(&b_after)
        .map(|(x, y)| (x - y).powi(2))
        .sum::<f32>()
        .sqrt();
    println!("\nA params total L2 change: {a_changed:.4}");
    println!("B params total L2 change: {b_changed:.4}");
    println!(
        "loss[0]={:.5}  loss[n-1]={:.5}",
        losses[0],
        losses[n_steps - 1]
    );

    if losses[n_steps - 1] < losses[0] * 0.5 {
        println!("✓ LoRA training is working — loss decreased >50%");
    } else {
        println!("⚠ loss did not decrease enough; check learning rate / target signal");
    }
    if a_changed > 1e-3 && b_changed > 1e-3 {
        println!("✓ both A and B accumulated gradient updates");
    } else {
        println!("⚠ A or B did not move — check VarMap registration / autograd path");
    }

    Ok(())
}