rtx-csm: Phase 12.3 — curriculum LoRA trainer for emotional fine-tunes

Closes the personal_voice_training_guide.md §4 stack: capacity (12.1) +
control tokens (12.2) + multi-stage curriculum (this).

TrainingExample gains emotion_tag + stage. Trainer::train applies the
tag via the same apply_emotion_hint helper inference uses (now
pub(crate)) — training and inference must use identical prefix
formatting or the adapter won't transfer.

TrainingDataset::load_from_manifest reads JSONL
`{wav, transcript, emotion_tag?, stage?, speaker?}` rows; wav paths
resolve relative to manifest dir.

CurriculumStage + CurriculumTrainer run N stages sequentially against a
shared VarMap. Per stage: filter by ex.stage label, build a transient
sub-dataset, run Trainer, save snapshot if requested. The "*" stage
name is a global catch-all.

examples/lora_train_emotional.rs wraps the canonical 3-stage recipe:
audiobook (3 ep × lr 1e-4) → podcast (1 ep × lr 3e-5) → va (1 ep ×
lr 1e-5). --extended-lora recommended (FFN is the prosodic-style
carrier per the guide).

Verified end-to-end on Metal: 3-row manifest → all 3 stages execute,
checkpoints + final adapter written, prompt-token lengths varied by
emotion-tag length (9 vs 11 for different tags) confirming the tag
flowed through the training tokenization. Lib suite 96/96.

Co-Authored-By: Claude Opus 4.7 (1M context) <[email protected]>

crates/models/rtx-csm/Cargo.toml

@@ -168,6 +168,10 @@ path = "examples/tts_server.rs"
 name = "lora_train"
 path = "examples/lora_train.rs"
 
+[[example]]
+name = "lora_train_emotional"
+path = "examples/lora_train_emotional.rs"
+
 [[example]]
 name = "audioseal_inspect"
 path = "examples/audioseal_inspect.rs"

crates/models/rtx-csm/examples/lora_train_emotional.rs

@@ -0,0 +1,195 @@
+//! Multi-stage LoRA fine-tune driven by a JSONL manifest.
+//!
+//! The manifest format is one JSON object per line:
+//! ```jsonl
+//! {"wav": "audiobook/0001.wav", "transcript": "Once upon a time...", "stage": "audiobook", "emotion_tag": "[neutral]"}
+//! {"wav": "podcast/0123.wav",   "transcript": "So like, anyway...",   "stage": "podcast",   "emotion_tag": "[casual]"}
+//! {"wav": "va/excited_07.wav",  "transcript": "We did it!",           "stage": "va",        "emotion_tag": "[excited]"}
+//! ```
+//! `emotion_tag` and `stage` are optional. Wav paths can be absolute or
+//! relative to the manifest's directory.
+//!
+//! Curriculum (per `docs/personal_voice_training_guide.md` §4):
+//!   - stage "audiobook": 3 epochs at peak_lr 1e-4 (clean priors)
+//!   - stage "podcast":   1 epoch  at peak_lr 3e-5 (mixed)
+//!   - stage "va":        1 epoch  at peak_lr 1e-5 (acted, noisier)
+//! After each stage, an intermediate adapter snapshot is written so you can
+//! A/B subsequent stages.
+//!
+//! Usage:
+//!   cargo run -p rtx-csm --release --features metal \
+//!     --example lora_train_emotional -- \
+//!       --manifest /data/voice/manifest.jsonl \
+//!       --output /tmp/voice_final.safetensors \
+//!       --extended-lora
+
+use anyhow::Result;
+use candle_nn::VarMap;
+use clap::Parser;
+use rtx_csm::lora::LoraConfig;
+use rtx_csm::training::{
+    save_lora_adapter, CurriculumStage, CurriculumTrainer, TrainingConfig, TrainingDataset,
+};
+use rtx_csm::Generator;
+use std::path::PathBuf;
+
+#[derive(Debug, Parser)]
+struct Cli {
+    /// JSONL manifest of (wav, transcript, [stage], [emotion_tag], [speaker]) rows.
+    #[arg(long)]
+    manifest: PathBuf,
+
+    /// Final adapter output (safetensors). Per-stage snapshots are written
+    /// next to it as `<stem>.<stage>.safetensors`.
+    #[arg(long)]
+    output: PathBuf,
+
+    /// Default speaker id when a manifest row doesn't specify one.
+    #[arg(long, default_value_t = 0)]
+    speaker: u32,
+
+    #[arg(long, default_value_t = 8)]
+    rank: usize,
+
+    #[arg(long, default_value_t = 16.0)]
+    alpha: f32,
+
+    /// Frames sampled per training step.
+    #[arg(long, default_value_t = 4)]
+    frames_per_step: usize,
+
+    #[arg(long, default_value_t = 1.0)]
+    grad_clip: f64,
+
+    #[arg(long, default_value_t = 42)]
+    seed: u64,
+
+    /// Use Phase 12.1 extended LoRA coverage (q+k+v+output_proj + MLP)
+    /// instead of q+v only. Recommended for emotional fine-tunes —
+    /// the FFN carries prosodic style and a wider adapter has the
+    /// capacity to learn multiple emotion tags simultaneously.
+    #[arg(long, default_value_t = false)]
+    extended_lora: bool,
+
+    #[arg(long)]
+    cpu: bool,
+}
+
+fn main() -> Result<()> {
+    tracing_subscriber::fmt().init();
+    let cli = Cli::parse();
+
+    let device = if cli.cpu {
+        candle_core::Device::Cpu
+    } else {
+        Generator::default_device()?
+    };
+    println!("device: {device:?}");
+
+    let mut generator = Generator::load_csm_1b(&device)?;
+    println!("model loaded");
+
+    println!("loading manifest from {}", cli.manifest.display());
+    let dataset = TrainingDataset::load_from_manifest(&cli.manifest, cli.speaker, &mut generator)?;
+    println!("dataset: {} examples", dataset.len());
+
+    // Inject LoRA into the backbone.
+    let base = if cli.extended_lora {
+        LoraConfig::extended()
+    } else {
+        LoraConfig::default()
+    };
+    let lora_cfg = LoraConfig {
+        rank: cli.rank,
+        alpha: cli.alpha,
+        ..base
+    };
+    println!("LoRA target_modules={:?}", lora_cfg.target_modules);
+    let vm = VarMap::new();
+    generator.model.inner.add_lora_to_backbone(&lora_cfg, &vm)?;
+    let n_params: usize = vm.all_vars().iter().map(|v| v.shape().elem_count()).sum();
+    println!(
+        "LoRA injected: {} trainable params ({} adapter Vars, rank={} alpha={})",
+        n_params, vm.all_vars().len(), cli.rank, cli.alpha,
+    );
+
+    // Build per-stage checkpoint paths from the final-output stem.
+    let stem = cli
+        .output
+        .file_stem()
+        .map(|s| s.to_string_lossy().to_string())
+        .unwrap_or_else(|| "voice".into());
+    let parent = cli.output.parent().unwrap_or(std::path::Path::new("."));
+    let stage_path = |stage: &str| parent.join(format!("{stem}.{stage}.safetensors"));
+
+    // Canonical 3-stage recipe from personal_voice_training_guide.md §4.
+    let stages = vec![
+        CurriculumStage {
+            name: "audiobook".into(),
+            config: TrainingConfig {
+                epochs: 3,
+                peak_lr: 1e-4,
+                end_lr: 1e-5,
+                warmup_steps: 32,
+                grad_clip: Some(cli.grad_clip),
+                frames_per_step: cli.frames_per_step,
+                seed: cli.seed,
+            },
+            checkpoint: Some(stage_path("audiobook")),
+        },
+        CurriculumStage {
+            name: "podcast".into(),
+            config: TrainingConfig {
+                epochs: 1,
+                peak_lr: 3e-5,
+                end_lr: 1e-5,
+                warmup_steps: 16,
+                grad_clip: Some(cli.grad_clip),
+                frames_per_step: cli.frames_per_step,
+                seed: cli.seed.wrapping_add(1),
+            },
+            checkpoint: Some(stage_path("podcast")),
+        },
+        CurriculumStage {
+            name: "va".into(),
+            config: TrainingConfig {
+                epochs: 1,
+                peak_lr: 1e-5,
+                end_lr: 1e-6,
+                warmup_steps: 16,
+                grad_clip: Some(cli.grad_clip),
+                frames_per_step: cli.frames_per_step,
+                seed: cli.seed.wrapping_add(2),
+            },
+            checkpoint: Some(stage_path("va")),
+        },
+    ];
+
+    let mut runner = CurriculumTrainer::new(&mut generator, &vm, &dataset, stages);
+    let traces = runner.run()?;
+
+    for (i, losses) in traces.iter().enumerate() {
+        if losses.is_empty() {
+            continue;
+        }
+        let n = losses.len();
+        let head: f32 = losses.iter().take(10).sum::<f32>() / 10.0_f32.min(n as f32);
+        let tail: f32 = losses.iter().rev().take(10).sum::<f32>() / 10.0_f32.min(n as f32);
+        println!(
+            "stage {i}: {n} steps, first 10 avg = {head:.4}, last 10 avg = {tail:.4}, change = {:+.2}%",
+            100.0 * (tail - head) / head
+        );
+    }
+
+    save_lora_adapter(&vm, &cli.output)?;
+    println!(
+        "\n✓ trained emotional LoRA adapter saved to {}",
+        cli.output.display()
+    );
+    println!(
+        "Apply at inference time via the same adapter loader. Re-use the same\n\
+         emotion tags at inference (Phase 12.2 --emotion-hint flag) for the\n\
+         tag→prosody mapping the curriculum just trained."
+    );
+    Ok(())
+}

crates/models/rtx-csm/src/generator.rs

@@ -60,8 +60,10 @@ impl Default for GenerateOptions {
 /// Prepend an optional control-token hint to the (already-normalized) text.
 /// Joined with a single space when present so the BPE tokenizer treats the tag
 /// as a separate sub-sequence from the body text. Centralized here so every
-/// generate variant applies the same convention.
+/// generate variant — and the LoRA trainer in `training.rs` — applies the
-fn apply_emotion_hint(text: String, hint: Option<&str>) -> String {
+/// same convention. Training and inference *must* use identical prefix
+/// formatting or the adapter won't transfer.
+pub(crate) fn apply_emotion_hint(text: String, hint: Option<&str>) -> String {
     match hint {
         Some(tag) if !tag.is_empty() => format!("{} {}", tag.trim(), text),
         _ => text,

crates/models/rtx-csm/src/training.rs

@@ -30,6 +30,14 @@ pub struct TrainingExample {
     pub frame_codes: Vec<Vec<u32>>,
     /// Path the audio came from, for logging.
     pub source: Option<PathBuf>,
+    /// Optional control-token tag prepended to `text` at training time, so
+    /// the LoRA learns `<tag> <text>` → matching prosody. Same format as
+    /// `GenerateOptions::emotion_hint` at inference. None = unconditioned.
+    pub emotion_tag: Option<String>,
+    /// Optional curriculum stage label (e.g. "audiobook", "podcast", "va").
+    /// Consumed by [`CurriculumTrainer`] to bucket examples; ignored by the
+    /// flat [`Trainer`].
+    pub stage: Option<String>,
 }
 
 impl TrainingExample {
@@ -57,10 +65,29 @@ impl TrainingExample {
             text: text.into(),
             frame_codes,
             source: None,
+            emotion_tag: None,
+            stage: None,
         })
     }
 }
 
+/// One row of a JSONL training manifest. Field shapes:
+/// ```jsonl
+/// {"wav": "data/0001.wav", "transcript": "...", "emotion_tag": "[whisper]", "stage": "audiobook", "speaker": 0}
+/// ```
+/// `emotion_tag` and `stage` are optional; `speaker` defaults to 0 when absent.
+#[derive(Debug, Clone, serde::Deserialize)]
+pub struct ManifestRow {
+    pub wav: PathBuf,
+    pub transcript: String,
+    #[serde(default)]
+    pub emotion_tag: Option<String>,
+    #[serde(default)]
+    pub stage: Option<String>,
+    #[serde(default)]
+    pub speaker: Option<u32>,
+}
+
 /// A dataset of training examples loaded from disk.
 pub struct TrainingDataset {
     pub examples: Vec<TrainingExample>,
@@ -116,6 +143,71 @@ impl TrainingDataset {
         Ok(Self { examples })
     }
 
+    /// Load a JSONL manifest of `ManifestRow` rows. Each row's WAV is
+    /// resolved relative to `manifest_dir` (the manifest's parent directory)
+    /// when not absolute, so manifests stay portable. Skips rows with empty
+    /// transcripts or missing audio. Mimi-encodes each audio file once.
+    pub fn load_from_manifest<P: AsRef<Path>>(
+        manifest: P,
+        default_speaker: u32,
+        generator: &mut Generator,
+    ) -> Result<Self> {
+        let manifest = manifest.as_ref();
+        let manifest_dir = manifest.parent().unwrap_or(Path::new("."));
+        let body = std::fs::read_to_string(manifest)?;
+        let mut examples = Vec::new();
+        for (lineno, line) in body.lines().enumerate() {
+            let line = line.trim();
+            if line.is_empty() || line.starts_with('#') {
+                continue;
+            }
+            let row: ManifestRow = serde_json::from_str(line).map_err(|e| {
+                CsmError::Config(format!("manifest line {}: {e}", lineno + 1))
+            })?;
+            let text = row.transcript.trim().to_string();
+            if text.is_empty() {
+                tracing::warn!("manifest line {}: empty transcript, skipping", lineno + 1);
+                continue;
+            }
+            let wav_path = if row.wav.is_absolute() {
+                row.wav.clone()
+            } else {
+                manifest_dir.join(&row.wav)
+            };
+            if !wav_path.exists() {
+                tracing::warn!(
+                    "manifest line {}: wav not found at {}, skipping",
+                    lineno + 1,
+                    wav_path.display()
+                );
+                continue;
+            }
+            let speaker = row.speaker.unwrap_or(default_speaker);
+            let audio = audio_io::load_mono_24k(&wav_path)?;
+            let mut ex = TrainingExample::from_audio(speaker, text, &audio, generator)?;
+            ex.source = Some(wav_path);
+            ex.emotion_tag = row.emotion_tag;
+            ex.stage = row.stage;
+            tracing::debug!(
+                "manifest example: tag={:?} stage={:?} frames={}",
+                ex.emotion_tag, ex.stage, ex.frame_codes.len()
+            );
+            examples.push(ex);
+        }
+        if examples.is_empty() {
+            return Err(CsmError::Config(format!(
+                "no usable rows in manifest {}",
+                manifest.display()
+            )));
+        }
+        tracing::info!(
+            "loaded {} examples from manifest {}",
+            examples.len(),
+            manifest.display()
+        );
+        Ok(Self { examples })
+    }
+
     pub fn len(&self) -> usize {
         self.examples.len()
     }
@@ -225,8 +317,15 @@ impl<'a> Trainer<'a> {
                 if ex.frame_codes.is_empty() {
                     continue;
                 }
-                // Build the prompt once for this example.
+                // Build the prompt once for this example. If the example
-                let current = Segment::new_text(ex.speaker, &ex.text);
+                // carries an emotion_tag, prepend it with the same format the
+                // inference path uses (`apply_emotion_hint`) so the LoRA
+                // learns the same `<tag> <text>` pattern at training time.
+                let prompt_text = crate::generator::apply_emotion_hint(
+                    ex.text.clone(),
+                    ex.emotion_tag.as_deref(),
+                );
+                let current = Segment::new_text(ex.speaker, prompt_text);
                 let prompt = build_prompt(
                     &[],
                     &current,
@@ -281,6 +380,107 @@ impl<'a> Trainer<'a> {
     }
 }
 
+/// One stage of a curriculum: a label that selects which examples this stage
+/// trains on, the training schedule for the stage, and an optional checkpoint
+/// path for the adapter snapshot taken AFTER the stage finishes.
+///
+/// The label matching rule:
+///   - `"*"` matches every example regardless of `ex.stage`
+///   - any other string matches `ex.stage == Some(label)` exactly
+///
+/// Per `personal_voice_training_guide.md` §4 the canonical 3-stage recipe is:
+///   - audiobook: 3 epochs, peak_lr 1e-4
+///   - podcast:   1 epoch,  peak_lr 3e-5
+///   - va:        1 epoch,  peak_lr 1e-5
+/// each over a growing pool. The `CurriculumTrainer` enforces the pool growth
+/// implicitly by virtue of how you label your manifest (a podcast example
+/// labeled `stage: "podcast"` is only matched by the podcast stage; if you
+/// want it included in the va stage too, label it `"va"` and the va stage
+/// will pick it up).
+#[derive(Debug, Clone)]
+pub struct CurriculumStage {
+    pub name: String,
+    pub config: TrainingConfig,
+    pub checkpoint: Option<PathBuf>,
+}
+
+/// Multi-stage LoRA trainer driven by per-example stage labels in the
+/// dataset. After each stage the adapter is optionally written to a
+/// safetensors checkpoint so callers can compare intermediate snapshots
+/// (per the literature recipe: clean → mixed → noisy at decreasing LR).
+pub struct CurriculumTrainer<'a> {
+    pub generator: &'a mut Generator,
+    pub vm: &'a VarMap,
+    pub dataset: &'a TrainingDataset,
+    pub stages: Vec<CurriculumStage>,
+}
+
+impl<'a> CurriculumTrainer<'a> {
+    pub fn new(
+        generator: &'a mut Generator,
+        vm: &'a VarMap,
+        dataset: &'a TrainingDataset,
+        stages: Vec<CurriculumStage>,
+    ) -> Self {
+        Self { generator, vm, dataset, stages }
+    }
+
+    /// Run all stages sequentially. Returns one loss-trace `Vec<f32>` per
+    /// stage in `self.stages` order. The adapter VarMap is shared across
+    /// stages so each stage's training picks up where the previous one
+    /// left off — the curriculum is cumulative, not reset.
+    pub fn run(&mut self) -> Result<Vec<Vec<f32>>> {
+        let mut all_traces = Vec::with_capacity(self.stages.len());
+        for stage in self.stages.clone().into_iter() {
+            // Bucket examples for this stage. "*" is the global catch-all.
+            let mut stage_examples = Vec::new();
+            for ex in self.dataset.examples.iter() {
+                let pick = if stage.name == "*" {
+                    true
+                } else {
+                    matches!(ex.stage.as_deref(), Some(s) if s == stage.name)
+                };
+                if pick {
+                    stage_examples.push(ex.clone());
+                }
+            }
+            if stage_examples.is_empty() {
+                tracing::warn!(
+                    "curriculum stage '{}' matched 0 examples — skipping",
+                    stage.name
+                );
+                all_traces.push(Vec::new());
+                continue;
+            }
+            tracing::info!(
+                "curriculum stage '{}': {} examples, epochs={} peak_lr={}",
+                stage.name,
+                stage_examples.len(),
+                stage.config.epochs,
+                stage.config.peak_lr,
+            );
+            let stage_dataset = TrainingDataset { examples: stage_examples };
+            let mut trainer = Trainer::new(
+                self.generator,
+                self.vm,
+                &stage_dataset,
+                stage.config.clone(),
+            );
+            let losses = trainer.train()?;
+            if let Some(out) = stage.checkpoint.as_ref() {
+                save_lora_adapter(self.vm, out)?;
+                tracing::info!(
+                    "curriculum stage '{}': checkpoint saved → {}",
+                    stage.name,
+                    out.display()
+                );
+            }
+            all_traces.push(losses);
+        }
+        Ok(all_traces)
+    }
+}
+
 /// Global L2 norm gradient clipping. Walks all Vars in the VarMap, computes
 /// `||g||_2` across all gradients, and scales every gradient by
 /// `min(1, max_norm / ||g||_2)`.