From ba841f7c8b29f7a212216ea51b7eabeb0acdb000 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Tue, 3 Mar 2026 22:06:56 +0100 Subject: [PATCH] feat(dqn): checkpoint architecture validation via safetensors metadata Embed DQNConfig architecture hash (SHA-256 of state_dim, num_actions, hidden_dims, dueling/distributional/noisy/IQN flags) in safetensors file header on save. Validate hash on load to catch shape mismatches before touching the VarMap - prevents silent corruption from loading checkpoints trained with different network architectures. All 5 save paths (trainable_adapter, trainer serialize_model, DQNAgentType::save_checkpoint, RegimeConditionalDQN per-head) now embed metadata. All 3 load paths (DQN::load_from_safetensors, trainable_adapter::load_checkpoint, ensemble adapter) validate. No backward compatibility: checkpoints without metadata are rejected with a clear error message to re-train. Co-Authored-By: Claude Opus 4.6 --- Cargo.lock | 1 + crates/ml/Cargo.toml | 1 + crates/ml/src/dqn/dqn.rs | 120 ++++++++++++++++++++++-- crates/ml/src/dqn/regime_conditional.rs | 43 +++++---- crates/ml/src/dqn/trainable_adapter.rs | 28 ++++-- crates/ml/src/ensemble/adapters/dqn.rs | 17 +++- crates/ml/src/trainers/dqn/config.rs | 35 ++++++- crates/ml/src/trainers/dqn/trainer.rs | 30 +++--- 8 files changed, 225 insertions(+), 50 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 44e6a809e..c1b20e743 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -6068,6 +6068,7 @@ dependencies = [ "risk", "rstest 0.22.0", "rust_decimal", + "safetensors", "semver", "serde", "serde_json", diff --git a/crates/ml/Cargo.toml b/crates/ml/Cargo.toml index 35cb4a446..c7d19fa7d 100644 --- a/crates/ml/Cargo.toml +++ b/crates/ml/Cargo.toml @@ -145,6 +145,7 @@ once_cell = "1.19" lazy_static.workspace = true flate2 = "1.0" sha2 = "0.10" +safetensors = "0.4" # Direct dep for checkpoint metadata (config hash validation) hmac = "0.12" # HMAC for checkpoint signatures (SEC-001 fix) hex = "0.4" # Hex encoding for signatures bincode = "1.3" diff --git a/crates/ml/src/dqn/dqn.rs b/crates/ml/src/dqn/dqn.rs index 476587a8a..9496cfe19 100644 --- a/crates/ml/src/dqn/dqn.rs +++ b/crates/ml/src/dqn/dqn.rs @@ -276,6 +276,99 @@ impl Default for DQNConfig { } impl DQNConfig { + /// Compute a SHA-256 hash of the architecture-critical config fields. + /// + /// These fields determine tensor shapes in the network. Loading a checkpoint + /// saved with different architecture parameters will silently produce garbage + /// or panic on shape mismatch. This hash is embedded in safetensors metadata + /// on save and validated on load. + /// + /// Hashed fields: state_dim, num_actions, hidden_dims, use_dueling, + /// dueling_hidden_dim, use_distributional, num_atoms, use_noisy_nets, + /// use_iqn, iqn_embedding_dim, iqn_num_quantiles. + pub fn architecture_hash(&self) -> String { + use sha2::{Digest, Sha256}; + let mut hasher = Sha256::new(); + hasher.update(self.state_dim.to_le_bytes()); + hasher.update(self.num_actions.to_le_bytes()); + for &dim in &self.hidden_dims { + hasher.update(dim.to_le_bytes()); + } + hasher.update([self.use_dueling as u8]); + hasher.update(self.dueling_hidden_dim.to_le_bytes()); + hasher.update([self.use_distributional as u8]); + hasher.update(self.num_atoms.to_le_bytes()); + hasher.update([self.use_noisy_nets as u8]); + hasher.update([self.use_iqn as u8]); + hasher.update(self.iqn_embedding_dim.to_le_bytes()); + hasher.update(self.iqn_num_quantiles.to_le_bytes()); + hex::encode(hasher.finalize()) + } + + /// Build safetensors metadata for checkpoint files. + /// + /// Returns a `HashMap` containing the architecture hash + /// and human-readable config summary for debugging shape mismatches. + pub fn checkpoint_metadata(&self) -> std::collections::HashMap { + let mut meta = std::collections::HashMap::new(); + meta.insert("dqn.arch_hash".to_owned(), self.architecture_hash()); + meta.insert("dqn.state_dim".to_owned(), self.state_dim.to_string()); + meta.insert("dqn.num_actions".to_owned(), self.num_actions.to_string()); + meta.insert("dqn.hidden_dims".to_owned(), format!("{:?}", self.hidden_dims)); + meta.insert("dqn.use_dueling".to_owned(), self.use_dueling.to_string()); + meta.insert("dqn.use_distributional".to_owned(), self.use_distributional.to_string()); + meta.insert("dqn.use_noisy_nets".to_owned(), self.use_noisy_nets.to_string()); + meta.insert("dqn.use_iqn".to_owned(), self.use_iqn.to_string()); + meta + } + + /// Validate that a checkpoint's metadata matches this config's architecture. + /// + /// Returns `Ok(())` if the hash matches. Returns an error if metadata is + /// missing or the architecture hash differs. + pub fn validate_checkpoint_metadata( + &self, + metadata: &Option>, + ) -> Result<(), crate::MLError> { + let meta = match metadata { + Some(m) => m, + None => { + return Err(crate::MLError::CheckpointError( + "Checkpoint has no architecture metadata - cannot validate. \ + Re-train to generate a checkpoint with embedded config hash." + .to_owned(), + )); + } + }; + + let saved_hash = match meta.get("dqn.arch_hash") { + Some(h) => h, + None => { + return Err(crate::MLError::CheckpointError( + "Checkpoint metadata missing dqn.arch_hash field - cannot validate. \ + Re-train to generate a checkpoint with embedded config hash." + .to_owned(), + )); + } + }; + + let current_hash = self.architecture_hash(); + if *saved_hash != current_hash { + let saved_dims = meta.get("dqn.state_dim").map(|s| s.as_str()).unwrap_or("?"); + let saved_actions = meta.get("dqn.num_actions").map(|s| s.as_str()).unwrap_or("?"); + let saved_hidden = meta.get("dqn.hidden_dims").map(|s| s.as_str()).unwrap_or("?"); + return Err(crate::MLError::CheckpointError(format!( + "Architecture mismatch: checkpoint was saved with state_dim={}, num_actions={}, \ + hidden_dims={} (hash={}) but current config has state_dim={}, num_actions={}, \ + hidden_dims={:?} (hash={})", + saved_dims, saved_actions, saved_hidden, saved_hash, + self.state_dim, self.num_actions, self.hidden_dims, current_hash, + ))); + } + + Ok(()) + } + /// Create `DQN` config from central configuration system /// /// CRITICAL: Eliminates dangerous hardcoded defaults @@ -826,7 +919,7 @@ impl Sequential { #[allow(missing_debug_implementations)] pub struct DQN { /// `DQN` configuration - config: DQNConfig, + pub(crate) config: DQNConfig, /// Main Q-network q_network: Sequential, /// Target Q-network for stable training @@ -1818,21 +1911,20 @@ impl DQN { let target_quantiles = (rewards_broadcast + (next_quantiles * not_done_broadcast)? * gamma as f64)? .detach(); - // Quantile Huber loss (Dabney et al. 2018b, Eq. 10) - let loss = super::quantile_regression::quantile_huber_loss( + // Per-sample quantile Huber loss (Dabney et al. 2018b, Eq. 10) + // Returns [batch] tensor so we can apply PER importance-sampling weights + let per_sample_loss = super::quantile_regression::quantile_huber_loss_per_sample( ¤t_quantiles, &target_quantiles, &taus, self.config.iqn_kappa, ).map_err(|e| MLError::TrainingError(format!("Quantile Huber loss failed: {}", e)))?; - // Apply PER importance sampling weights + // Apply PER importance-sampling weights (detached — not learnable) let weights_tensor = Tensor::from_vec(weights.clone(), batch_size, &device)? .detach(); - // quantile_huber_loss returns a scalar mean; weight per-sample requires recomputation - // For simplicity with PER, use the scalar loss (weights are approximately uniform early on) - let _ = weights_tensor; // PER weights integrated in future refinement - loss + let weighted_loss = (per_sample_loss * weights_tensor)?.mean_all()?; + weighted_loss } else if self.dist_dueling_q_network.is_some() { // C51 CATEGORICAL LOSS PATH (Hybrid: Distributional + Dueling) @@ -2763,6 +2855,18 @@ impl DQN { ))); } + // Validate architecture hash before loading weights + let raw_bytes = std::fs::read(&safetensors_path).map_err(|e| { + MLError::CheckpointError(format!("Failed to read safetensors file: {}", e)) + })?; + let (_header_size, st_metadata) = + safetensors::SafeTensors::read_metadata(&raw_bytes).map_err(|e| { + MLError::CheckpointError(format!("Failed to parse safetensors header: {}", e)) + })?; + self.config + .validate_checkpoint_metadata(st_metadata.metadata())?; + drop(raw_bytes); + // Use VarMap::load which correctly updates existing Vars in-place via Var::set(). // This ensures that Linear layers (which share the same Arc> as // the VarMap's Vars) see the updated weights. The previous approach of inserting diff --git a/crates/ml/src/dqn/regime_conditional.rs b/crates/ml/src/dqn/regime_conditional.rs index a76521af3..f14021d39 100644 --- a/crates/ml/src/dqn/regime_conditional.rs +++ b/crates/ml/src/dqn/regime_conditional.rs @@ -681,30 +681,39 @@ impl RegimeConditionalDQN { /// /// * `path` - Base path for checkpoints (without .safetensors extension) pub fn save_checkpoint(&self, path: &str) -> Result<(), MLError> { - // Save each head separately + // Save each head separately with architecture metadata let trending_path = format!("{}_trending.safetensors", path); let ranging_path = format!("{}_ranging.safetensors", path); let volatile_path = format!("{}_volatile.safetensors", path); - // Save trending head - self.trending_head - .get_q_network_vars() - .save(&trending_path) - .map_err(|e| MLError::CheckpointError(format!("Failed to save trending head: {}", e)))?; + // Helper to extract tensors and save with metadata + let save_head = |head: &super::dqn::DQN, head_path: &str, label: &str| -> Result<(), MLError> { + let vars = head.get_q_network_vars(); + let vars_data = vars.data().lock().map_err(|e| { + MLError::LockError(format!("Failed to lock vars for {} head: {}", label, e)) + })?; + let tensors: std::collections::HashMap = vars_data + .iter() + .map(|(name, var)| (name.clone(), var.as_tensor().clone())) + .collect(); + drop(vars_data); - // Save ranging head - self.ranging_head - .get_q_network_vars() - .save(&ranging_path) - .map_err(|e| MLError::CheckpointError(format!("Failed to save ranging head: {}", e)))?; + let arch_metadata = Some(head.config.checkpoint_metadata()); + safetensors::serialize_to_file( + &tensors, + &arch_metadata, + std::path::Path::new(head_path), + ) + .map_err(|e| { + MLError::CheckpointError(format!("Failed to save {} head: {}", label, e)) + }) + }; - // Save volatile head - self.volatile_head - .get_q_network_vars() - .save(&volatile_path) - .map_err(|e| MLError::CheckpointError(format!("Failed to save volatile head: {}", e)))?; + save_head(&self.trending_head, &trending_path, "trending")?; + save_head(&self.ranging_head, &ranging_path, "ranging")?; + save_head(&self.volatile_head, &volatile_path, "volatile")?; - info!("✓ RegimeConditionalDQN checkpoint saved: {}", path); + info!("RegimeConditionalDQN checkpoint saved: {}", path); info!(" - Trending: {}", trending_path); info!(" - Ranging: {}", ranging_path); info!(" - Volatile: {}", volatile_path); diff --git a/crates/ml/src/dqn/trainable_adapter.rs b/crates/ml/src/dqn/trainable_adapter.rs index 2f86b2a3b..b5cdb346f 100644 --- a/crates/ml/src/dqn/trainable_adapter.rs +++ b/crates/ml/src/dqn/trainable_adapter.rs @@ -224,7 +224,7 @@ impl UnifiedTrainable for DQNTrainableAdapter { // Save metadata checkpoint::save_metadata(&metadata, checkpoint_path)?; - // Save model weights to safetensors format + // Save model weights to safetensors format with architecture metadata let safetensors_path = format!("{}.safetensors", checkpoint_path); // Extract tensors from VarMap @@ -238,10 +238,14 @@ impl UnifiedTrainable for DQNTrainableAdapter { tensors.insert(name.clone(), var.as_tensor().clone()); } - // Save using safetensors - // save() expects &HashMap, not Vec or refs - candle_core::safetensors::save(&tensors, &safetensors_path) - .map_err(|e| MLError::CheckpointError(format!("Failed to save safetensors: {}", e)))?; + // Save with architecture metadata embedded in safetensors header + let arch_metadata = Some(self.config.checkpoint_metadata()); + safetensors::serialize_to_file( + &tensors, + &arch_metadata, + std::path::Path::new(&safetensors_path), + ) + .map_err(|e| MLError::CheckpointError(format!("Failed to save safetensors: {}", e)))?; tracing::info!( "Saved DQN checkpoint to {} (step {})", @@ -264,8 +268,20 @@ impl UnifiedTrainable for DQNTrainableAdapter { ))); } - // Load model weights from safetensors + // Load model weights from safetensors with architecture validation let safetensors_path = format!("{}.safetensors", checkpoint_path); + let raw_bytes = std::fs::read(&safetensors_path).map_err(|e| { + MLError::CheckpointError(format!("Failed to read safetensors file: {}", e)) + })?; + + // Validate architecture hash before loading any weights + let (_header_size, st_metadata) = + safetensors::SafeTensors::read_metadata(&raw_bytes).map_err(|e| { + MLError::CheckpointError(format!("Failed to parse safetensors header: {}", e)) + })?; + self.config + .validate_checkpoint_metadata(st_metadata.metadata())?; + let tensors = candle_core::safetensors::load(&safetensors_path, &self.device) .map_err(|e| MLError::CheckpointError(format!("Failed to load safetensors: {}", e)))?; diff --git a/crates/ml/src/ensemble/adapters/dqn.rs b/crates/ml/src/ensemble/adapters/dqn.rs index e60804543..e5b6a2cc3 100644 --- a/crates/ml/src/ensemble/adapters/dqn.rs +++ b/crates/ml/src/ensemble/adapters/dqn.rs @@ -228,12 +228,25 @@ mod tests { assert!(pred1.is_ok(), "first predict failed: {:?}", pred1.err()); let pred1 = pred1.unwrap(); - // Save checkpoint to temp dir + // Save checkpoint to temp dir with architecture metadata let dir = tempfile::tempdir().unwrap(); let path = dir.path().join("dqn_test.safetensors"); { let model = adapter.model.lock().unwrap(); - model.get_q_network_vars().save(path.as_path()).unwrap(); + let vars = model.get_q_network_vars(); + let vars_data = vars.data().lock().unwrap(); + let tensors: std::collections::HashMap = vars_data + .iter() + .map(|(name, var)| (name.clone(), var.as_tensor().clone())) + .collect(); + drop(vars_data); + let arch_metadata = Some(model.config.checkpoint_metadata()); + safetensors::serialize_to_file( + &tensors, + &arch_metadata, + path.as_path(), + ) + .unwrap(); } // Load into new adapter diff --git a/crates/ml/src/trainers/dqn/config.rs b/crates/ml/src/trainers/dqn/config.rs index 15a42f0ad..4fe0e6645 100644 --- a/crates/ml/src/trainers/dqn/config.rs +++ b/crates/ml/src/trainers/dqn/config.rs @@ -114,11 +114,28 @@ impl DQNAgentType { } } - /// Save checkpoint to disk + /// Save checkpoint to disk with architecture metadata pub fn save_checkpoint(&self, path: &str) -> Result<(), MLError> { match self { Self::Standard(agent) => { - agent.get_q_network_vars().save(path).map_err(|e| { + // Extract tensors and save with architecture metadata + let vars = agent.get_q_network_vars(); + let vars_data = vars.data().lock().map_err(|e| { + MLError::LockError(format!("Failed to lock vars for checkpoint: {}", e)) + })?; + let tensors: std::collections::HashMap = vars_data + .iter() + .map(|(name, var)| (name.clone(), var.as_tensor().clone())) + .collect(); + drop(vars_data); + + let arch_metadata = Some(agent.config.checkpoint_metadata()); + safetensors::serialize_to_file( + &tensors, + &arch_metadata, + std::path::Path::new(path), + ) + .map_err(|e| { MLError::CheckpointError(format!("Failed to save checkpoint: {}", e)) })?; Ok(()) @@ -279,6 +296,20 @@ impl DQNAgentType { } } + /// Build safetensors checkpoint metadata from the agent's config. + pub fn checkpoint_metadata(&self) -> std::collections::HashMap { + match self { + Self::Standard(agent) => agent.config.checkpoint_metadata(), + Self::RegimeConditional(agent) => { + // Use trending head config as representative for regime-conditional + match agent.get_trending_head() { + Some(head) => head.config.checkpoint_metadata(), + None => std::collections::HashMap::new(), + } + } + } + } + /// Get mutable reference to underlying standard agent (for methods not in unified API) pub fn as_standard_mut(&mut self) -> Option<&mut DQN> { match self { diff --git a/crates/ml/src/trainers/dqn/trainer.rs b/crates/ml/src/trainers/dqn/trainer.rs index ef37cad5d..502f585e6 100644 --- a/crates/ml/src/trainers/dqn/trainer.rs +++ b/crates/ml/src/trainers/dqn/trainer.rs @@ -3901,25 +3901,25 @@ impl DQNTrainer { self.lr_scheduler.get_lr() } - /// Serialize model to bytes + /// Serialize model to bytes with architecture metadata embedded in safetensors header pub async fn serialize_model(&self) -> Result> { let agent = self.agent.read().await; - // Create temp file for SafeTensors serialization - let temp_path = std::env::temp_dir().join(format!("dqn_{}.safetensors", Uuid::new_v4())); + // Extract tensors from VarMap + let vars = agent.get_q_network_vars(); + let vars_data = vars.data().lock().map_err(|_| { + anyhow::anyhow!("Failed to lock VarMap for serialization") + })?; + let tensors: std::collections::HashMap = vars_data + .iter() + .map(|(name, var)| (name.clone(), var.as_tensor().clone())) + .collect(); + drop(vars_data); - // Save Q-network to SafeTensors - agent - .get_q_network_vars() - .save(&temp_path) - .map_err(|e| anyhow::anyhow!("Failed to save Q-network: {}", e))?; - - // Read serialized data - let data = std::fs::read(&temp_path) - .map_err(|e| anyhow::anyhow!("Failed to read checkpoint: {}", e))?; - - // Clean up temp file - drop(std::fs::remove_file(&temp_path)); + // Embed architecture metadata in safetensors header + let arch_metadata = Some(agent.checkpoint_metadata()); + let data = safetensors::serialize(&tensors, &arch_metadata) + .map_err(|e| anyhow::anyhow!("Failed to serialize safetensors: {}", e))?; Ok(data) }