rtx-csm: Phase 12.5 — self-describing LoRA adapters

Trained adapters now embed (rank, alpha, target_modules, crate_version)
as JSON in safetensors __metadata__["rtx_csm_lora"]. apply_lora_adapter
reads it at load time so users no longer have to remember matching
--lora-rank/--lora-alpha/--extended-lora flags from training.

LoraAdapterMetadata::is_extended() heuristic: target_modules contains
any MLP path or output_proj or k_proj. Handles both the canonical
extended() preset and future custom configs that overlap it.

apply_lora_adapter rank/alpha/extended params became Option<_>
(None = use file metadata, Some = override). Both callers updated.
save_lora_adapter_with_metadata is the new path used by both trainers;
the plain save_lora_adapter still exists for the metadata-less case
(per-stage curriculum snapshots).

safetensors dep bumped 0.4 → 0.7 to match candle 0.9's transitive pin
so candle's Tensor: View impl is in scope for serialize_to_file
(candle's own save wrapper hardcodes the metadata arg to None).

Backward compat: pre-12.5 adapters load fine when explicit CLI flags
are passed; auto-detect path is skipped silently.

Verified end-to-end: trained adapter saved with metadata,
`generate --lora <path>` (no other flags) auto-detected
rank=8 alpha=16 extended=true and applied. Older metadata-less adapter
still loaded with explicit flags. 3 new unit tests; lib suite 99/99.

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

crates/models/rtx-csm/Cargo.toml

@@ -68,8 +68,12 @@ voice_activity_detector = { version = "0.2", optional = true }
 unicode-normalization = "0.1"
 regex = "1"
 
-# Weight loading
+# Weight loading. Matches candle 0.9's transitive pin so candle's
-safetensors = "0.4"
+# `Tensor: View` impl applies — letting us call
+# `safetensors::serialize_to_file(&map, &Some(metadata), path)` directly
+# (candle's `save` wrapper doesn't expose the metadata parameter, which we
+# need for Phase 12.5 self-describing LoRA adapters).
+safetensors = "0.7"
 
 # Errors / logging / serde
 anyhow.workspace = true

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

@@ -127,16 +127,16 @@ struct Cli {
     /// but the server hasn't been audited for that combo).
     #[arg(long)]
     lora: Option<PathBuf>,
-    /// LoRA rank — must match training. Default 8.
+    /// LoRA rank override. When omitted, auto-detected from adapter metadata
-    #[arg(long, default_value_t = 8)]
+    /// (Phase 12.5); falls back to 8 for older adapters.
-    lora_rank: usize,
+    #[arg(long)]
-    /// LoRA alpha — must match training. Default 16.
+    lora_rank: Option<usize>,
-    #[arg(long, default_value_t = 16.0)]
+    /// LoRA alpha override. Auto-detected from metadata when omitted.
-    lora_alpha: f32,
+    #[arg(long)]
-    /// Set when the adapter was trained with Phase 12.1 extended coverage
+    lora_alpha: Option<f32>,
-    /// (q+k+v+output_proj + MLP). Adapters trained with the classic q+v
+    /// Force extended LoRA coverage (q+k+v+output_proj + MLP). Auto-detected
-    /// recipe load fine without this flag; setting it on a q+v adapter
+    /// from metadata when omitted; setting it on a classic q+v adapter just
-    /// just allocates extra unused B=0 slots.
+    /// allocates extra unused B=0 slots.
     #[arg(long, default_value_t = false)]
     extended_lora: bool,
 
@@ -541,12 +541,13 @@ async fn main() -> Result<()> {
     };
 
     if let Some(lora_path) = cli.lora.as_ref() {
+        let extended_override = if cli.extended_lora { Some(true) } else { None };
         rtx_csm::training::apply_lora_adapter(
             &mut generator,
             lora_path,
             cli.lora_rank,
             cli.lora_alpha,
-            cli.extended_lora,
+            extended_override,
             &device,
         )?;
     }

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

@@ -92,18 +92,20 @@ struct Cli {
     #[arg(long)]
     lora: Option<std::path::PathBuf>,
 
-    /// LoRA rank — must match training. Defaults to the value from training defaults.
+    /// LoRA rank override. When omitted, auto-detected from the adapter
-    #[arg(long, default_value_t = 8)]
+    /// file's embedded metadata (Phase 12.5); falls back to 8 for older
-    lora_rank: usize,
+    /// adapters without metadata.
+    #[arg(long)]
+    lora_rank: Option<usize>,
 
-    /// LoRA alpha — must match training.
+    /// LoRA alpha override. When omitted, auto-detected from the adapter's
-    #[arg(long, default_value_t = 16.0)]
+    /// embedded metadata; falls back to 16.0 for older adapters.
-    lora_alpha: f32,
+    #[arg(long)]
+    lora_alpha: Option<f32>,
 
-    /// Set when the adapter was trained with Phase 12.1 extended coverage
+    /// Force extended LoRA coverage (q+k+v+output_proj + MLP). When omitted,
-    /// (q+k+v+output_proj + MLP). Adapters trained with the classic q+v
+    /// auto-detected from the adapter's embedded metadata. Setting this on a
-    /// recipe load fine without this; setting it on a q+v adapter just
+    /// classic q+v adapter just allocates extra unused B=0 slots.
-    /// allocates extra unused B=0 slots.
     #[arg(long, default_value_t = false)]
     extended_lora: bool,
 
@@ -153,12 +155,16 @@ fn main() -> Result<()> {
     };
 
     if let Some(lora_path) = cli.lora.as_ref() {
+        // `--extended-lora` is a force flag (no way to clap-distinguish "off"
+        // from "unset" on a bare bool). When the user passes it we override
+        // to true; otherwise we let the file metadata or the default decide.
+        let extended_override = if cli.extended_lora { Some(true) } else { None };
         rtx_csm::training::apply_lora_adapter(
             &mut generator,
             lora_path,
             cli.lora_rank,
             cli.lora_alpha,
-            cli.extended_lora,
+            extended_override,
             &device,
         )?;
     }

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

@@ -15,7 +15,10 @@ use anyhow::Result;
 use candle_nn::VarMap;
 use clap::Parser;
 use rtx_csm::lora::LoraConfig;
-use rtx_csm::training::{save_lora_adapter, Trainer, TrainingConfig, TrainingDataset};
+use rtx_csm::training::{
+    save_lora_adapter_with_metadata, LoraAdapterMetadata, Trainer, TrainingConfig,
+    TrainingDataset,
+};
 use rtx_csm::Generator;
 use std::path::PathBuf;
 
@@ -153,15 +156,14 @@ fn main() -> Result<()> {
         );
     }
 
-    save_lora_adapter(&vm, &cli.output)?;
+    let metadata = LoraAdapterMetadata::from_lora_config(&lora_cfg);
+    save_lora_adapter_with_metadata(&vm, &cli.output, &metadata)?;
     println!("\n✓ trained LoRA adapter saved to {}", cli.output.display());
     println!(
         "Use it at inference: examples/generate --lora {} \\\n\
-         \t--lora-rank {} --lora-alpha {}{}",
+         \t--text \"...\" --out /tmp/out.wav\n\
+         (rank/alpha/extended auto-detected from embedded metadata)",
         cli.output.display(),
-        cli.rank,
-        cli.alpha,
-        if cli.extended_lora { " --extended-lora" } else { "" },
     );
     Ok(())
 }

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

@@ -28,7 +28,8 @@ use candle_nn::VarMap;
 use clap::Parser;
 use rtx_csm::lora::LoraConfig;
 use rtx_csm::training::{
-    save_lora_adapter, CurriculumStage, CurriculumTrainer, TrainingConfig, TrainingDataset,
+    save_lora_adapter_with_metadata, CurriculumStage, CurriculumTrainer, LoraAdapterMetadata,
+    TrainingConfig, TrainingDataset,
 };
 use rtx_csm::Generator;
 use std::path::PathBuf;
@@ -181,15 +182,17 @@ fn main() -> Result<()> {
         );
     }
 
-    save_lora_adapter(&vm, &cli.output)?;
+    let metadata = LoraAdapterMetadata::from_lora_config(&lora_cfg);
+    save_lora_adapter_with_metadata(&vm, &cli.output, &metadata)?;
     println!(
-        "\n✓ trained emotional LoRA adapter saved to {}",
+        "\n✓ trained emotional LoRA adapter saved to {} (with embedded metadata)",
         cli.output.display()
     );
     println!(
-        "Apply at inference time via the same adapter loader. Re-use the same\n\
+        "Apply at inference: examples/generate --lora {} --text \"...\" \\\n\
-         emotion tags at inference (Phase 12.2 --emotion-hint flag) for the\n\
+         \t--emotion-hint \"[your_tag]\" --out /tmp/out.wav\n\
-         tag→prosody mapping the curriculum just trained."
+         (rank/alpha/extended auto-detected from embedded metadata)",
+        cli.output.display(),
     );
     Ok(())
 }

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

@@ -514,15 +514,84 @@ fn clip_grads(
     Ok(())
 }
 
+/// Self-describing metadata embedded in safetensors at save time so the
+/// adapter file knows its own LoRA hyperparameters. Lets `apply_lora_adapter`
+/// auto-configure at inference without the user having to remember matching
+/// `--lora-rank` / `--lora-alpha` / `--extended-lora` flags.
+///
+/// Stored under the safetensors `__metadata__` map as a single JSON-encoded
+/// string under key `rtx_csm_lora`. Backward compat: adapters without this
+/// key fall back to caller-supplied defaults (rank=8, alpha=16, extended=false).
+#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
+pub struct LoraAdapterMetadata {
+    pub rank: usize,
+    pub alpha: f32,
+    /// Snapshot of `LoraConfig.target_modules` at training time. Reconstructs
+    /// `extended()` vs `default()` reliably when the user adds new presets.
+    pub target_modules: Vec<String>,
+    /// rtx-csm crate version that produced the adapter (just for diagnostics
+    /// — no compatibility logic gates on this today).
+    #[serde(default)]
+    pub crate_version: Option<String>,
+}
+
+const LORA_METADATA_KEY: &str = "rtx_csm_lora";
+
+impl LoraAdapterMetadata {
+    pub fn from_lora_config(cfg: &crate::lora::LoraConfig) -> Self {
+        Self {
+            rank: cfg.rank,
+            alpha: cfg.alpha,
+            target_modules: cfg.target_modules.clone(),
+            crate_version: Some(env!("CARGO_PKG_VERSION").to_string()),
+        }
+    }
+
+    /// Heuristic: was the adapter trained with Phase 12.1 extended coverage?
+    /// True if any MLP module is in `target_modules`. The exact preset name
+    /// isn't stored — only the resolved patterns — so we infer from content.
+    pub fn is_extended(&self) -> bool {
+        self.target_modules
+            .iter()
+            .any(|m| m.contains("mlp.") || m.contains("output_proj") || m == "k_proj")
+    }
+}
+
+/// Read just the LoRA self-description from a safetensors header, without
+/// loading any tensor data. Returns `None` when the file has no
+/// `rtx_csm_lora` metadata key (older adapters trained before Phase 12.5).
+pub fn read_lora_adapter_metadata<P: AsRef<Path>>(
+    path: P,
+) -> Result<Option<LoraAdapterMetadata>> {
+    let data = std::fs::read(path.as_ref())?;
+    let (_n, meta) = safetensors::SafeTensors::read_metadata(&data)
+        .map_err(|e| CsmError::Other(anyhow::anyhow!("read safetensors header: {e}")))?;
+    let map = match meta.metadata() {
+        Some(m) => m,
+        None => return Ok(None),
+    };
+    let json = match map.get(LORA_METADATA_KEY) {
+        Some(s) => s,
+        None => return Ok(None),
+    };
+    let parsed: LoraAdapterMetadata = serde_json::from_str(json).map_err(|e| {
+        CsmError::Other(anyhow::anyhow!("parse {LORA_METADATA_KEY}: {e}"))
+    })?;
+    Ok(Some(parsed))
+}
+
 /// One-shot inference-time LoRA loader: inject adapters with the matching
 /// shape, populate from a safetensors file, and refresh the model's tensor
 /// handles so the next forward pass picks up the trained weights.
 ///
-/// Pass `extended = true` when the adapter was trained with Phase 12.1's
+/// Each parameter accepts an explicit override; when `None`, the value is
-/// extended coverage (q+k+v+output_proj + MLP); otherwise it loads the
+/// auto-detected from the adapter file's embedded `LoraAdapterMetadata`
-/// classic q+v recipe. Adapters trained with the narrow recipe still load
+/// (Phase 12.5). For files without metadata (older adapters), pass explicit
-/// cleanly when `extended = true` — the unmatched k/o/MLP slots are
+/// values: defaults are rank=8, alpha=16, extended=false.
-/// initialized B=0 and stay no-ops, at the cost of a small extra VarMap.
+///
+/// Adapters trained with the narrow q+v recipe still load cleanly when
+/// `extended = true` — the unmatched k/o/MLP slots are initialized B=0 and
+/// stay no-ops, at the cost of a small extra VarMap.
 ///
 /// Returns the `VarMap` so callers can keep it alive (the model holds plain
 /// tensor handles refreshed from this VarMap; dropping it doesn't break a
@@ -533,24 +602,46 @@ fn clip_grads(
 pub fn apply_lora_adapter<P: AsRef<Path>>(
     generator: &mut Generator,
     path: P,
-    rank: usize,
+    rank: Option<usize>,
-    alpha: f32,
+    alpha: Option<f32>,
-    extended: bool,
+    extended: Option<bool>,
     device: &candle_core::Device,
 ) -> Result<VarMap> {
-    let base = if extended {
+    let path = path.as_ref();
+    let meta = read_lora_adapter_metadata(path)?;
+    if meta.is_some() {
+        tracing::info!(
+            "adapter has embedded metadata; CLI rank/alpha/extended flags act as overrides"
+        );
+    }
+
+    let resolved_rank = rank
+        .or_else(|| meta.as_ref().map(|m| m.rank))
+        .unwrap_or(8);
+    let resolved_alpha = alpha
+        .or_else(|| meta.as_ref().map(|m| m.alpha))
+        .unwrap_or(16.0);
+    let resolved_extended = extended
+        .or_else(|| meta.as_ref().map(|m| m.is_extended()))
+        .unwrap_or(false);
+
+    let base = if resolved_extended {
         crate::lora::LoraConfig::extended()
     } else {
         crate::lora::LoraConfig::default()
     };
-    let cfg = crate::lora::LoraConfig { rank, alpha, ..base };
+    let cfg = crate::lora::LoraConfig {
+        rank: resolved_rank,
+        alpha: resolved_alpha,
+        ..base
+    };
     let vm = VarMap::new();
     generator
         .model
         .inner
         .add_lora_to_backbone(&cfg, &vm)
         .map_err(|e| CsmError::Other(anyhow::anyhow!("add_lora_to_backbone: {e}")))?;
-    load_lora_adapter(&vm, path.as_ref(), device)?;
+    load_lora_adapter(&vm, path, device)?;
     generator
         .model
         .inner
@@ -558,17 +649,39 @@ pub fn apply_lora_adapter<P: AsRef<Path>>(
         .map_err(|e| CsmError::Other(anyhow::anyhow!("refresh_lora: {e}")))?;
     tracing::info!(
         "applied LoRA adapter from {} (rank={} alpha={} extended={})",
-        path.as_ref().display(),
+        path.display(),
-        rank,
+        resolved_rank,
-        alpha,
+        resolved_alpha,
-        extended,
+        resolved_extended,
     );
     Ok(vm)
 }
 
 /// Save the LoRA adapter parameters (A,B for every layer) as a safetensors
-/// file. Reload via `load_lora_adapter`.
+/// file. Reload via `load_lora_adapter`. No embedded metadata — callers that
+/// know the LoRA hyperparameters should prefer
+/// [`save_lora_adapter_with_metadata`] so the trained adapter is
+/// self-describing at inference time.
 pub fn save_lora_adapter<P: AsRef<Path>>(vm: &VarMap, out: P) -> Result<()> {
+    save_lora_adapter_inner(vm, out, None)
+}
+
+/// Save with embedded `LoraAdapterMetadata` (Phase 12.5). The metadata is
+/// written into the safetensors header under `__metadata__["rtx_csm_lora"]`
+/// as a JSON string; `apply_lora_adapter` will auto-pick it up at inference.
+pub fn save_lora_adapter_with_metadata<P: AsRef<Path>>(
+    vm: &VarMap,
+    out: P,
+    metadata: &LoraAdapterMetadata,
+) -> Result<()> {
+    save_lora_adapter_inner(vm, out, Some(metadata))
+}
+
+fn save_lora_adapter_inner<P: AsRef<Path>>(
+    vm: &VarMap,
+    out: P,
+    metadata: Option<&LoraAdapterMetadata>,
+) -> Result<()> {
     use std::collections::HashMap;
     let vars = vm.data().lock().unwrap();
     let mut tensors: HashMap<String, Tensor> = HashMap::new();
@@ -576,9 +689,26 @@ pub fn save_lora_adapter<P: AsRef<Path>>(vm: &VarMap, out: P) -> Result<()> {
         tensors.insert(name.clone(), var.as_tensor().clone());
     }
     drop(vars);
-    candle_core::safetensors::save(&tensors, out.as_ref())
+    let n_tensors = tensors.len();
+
+    let header_meta = match metadata {
+        Some(m) => {
+            let json = serde_json::to_string(m).map_err(|e| {
+                CsmError::Other(anyhow::anyhow!("serialize {LORA_METADATA_KEY}: {e}"))
+            })?;
+            let mut map = HashMap::new();
+            map.insert(LORA_METADATA_KEY.to_string(), json);
+            Some(map)
+        }
+        None => None,
+    };
+    safetensors::serialize_to_file(&tensors, header_meta.clone(), out.as_ref())
         .map_err(|e| CsmError::Other(anyhow::anyhow!("safetensors save: {e}")))?;
-    tracing::info!("saved {} LoRA tensors → {}", tensors.len(), out.as_ref().display());
+    tracing::info!(
+        "saved {n_tensors} LoRA tensors → {}{}",
+        out.as_ref().display(),
+        if header_meta.is_some() { " (with metadata)" } else { "" }
+    );
     Ok(())
 }
 
@@ -604,3 +734,38 @@ pub fn load_lora_adapter<P: AsRef<Path>>(
     tracing::info!("loaded {count} LoRA tensors from {}", path.as_ref().display());
     Ok(())
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn metadata_is_extended_picks_up_extended_target_modules() {
+        let cfg = crate::lora::LoraConfig::extended();
+        let m = LoraAdapterMetadata::from_lora_config(&cfg);
+        assert!(m.is_extended());
+    }
+
+    #[test]
+    fn metadata_is_extended_false_for_classic_qv() {
+        let cfg = crate::lora::LoraConfig::default();
+        let m = LoraAdapterMetadata::from_lora_config(&cfg);
+        assert!(!m.is_extended());
+    }
+
+    #[test]
+    fn metadata_round_trips_through_json() {
+        let m = LoraAdapterMetadata {
+            rank: 16,
+            alpha: 32.0,
+            target_modules: vec!["q_proj".into(), "v_proj".into()],
+            crate_version: Some("9.9.9".into()),
+        };
+        let s = serde_json::to_string(&m).unwrap();
+        let back: LoraAdapterMetadata = serde_json::from_str(&s).unwrap();
+        assert_eq!(back.rank, 16);
+        assert_eq!(back.alpha, 32.0);
+        assert_eq!(back.target_modules.len(), 2);
+        assert!(!back.is_extended());
+    }
+}