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/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
-/// extended coverage (q+k+v+output_proj + MLP); otherwise it loads the
-/// classic q+v recipe. Adapters trained with the narrow 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.
+/// Each parameter accepts an explicit override; when `None`, the value is
+/// auto-detected from the adapter file's embedded `LoraAdapterMetadata`
+/// (Phase 12.5). For files without metadata (older adapters), pass explicit
+/// values: defaults are rank=8, alpha=16, extended=false.
+///
+/// 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,
-    alpha: f32,
-    extended: bool,
+    rank: Option<usize>,
+    alpha: Option<f32>,
+    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(),
-        rank,
-        alpha,
-        extended,
+        path.display(),
+        resolved_rank,
+        resolved_alpha,
+        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());
+    }
+}