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]>
This commit is contained in:
osobh
@@ -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.
|
||||
let current = Segment::new_text(ex.speaker, &ex.text);
|
||||
// Build the prompt once for this example. If the example
|
||||
// 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(
|
||||
&[],
|
||||
¤t,
|
||||
@@ -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)`.
|
||||
|
||||
Reference in New Issue
Block a user