From f47af9d268f3f91b0d41cc0e269dd149b233cda8 Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Thu, 23 Apr 2026 10:19:56 +0200 Subject: [PATCH] fix(ml-dqn): implement real DQN::load_from_safetensors MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `DQN::load_from_safetensors` was a documented no-op that only checked file existence and returned Ok(()). Real weight restoration happened only via the fused trainer's params_buf path; any caller loading a checkpoint via the DQN struct itself (e.g., DqnInferenceAdapter:: from_checkpoint used by the ensemble) got a fresh untrained DQN with no error raised — a silent production bug. Adds BranchingDuelingQNetwork::load_from_named_slices (symmetric counterpart of existing named_weight_slices) using the same D2D memcpy primitive as copy_weights_from. Validates names AND shapes; returns Err on mismatch. Wired into DQN::load_from_safetensors to actually restore weights from disk, handling both prefix-less (single head) and `trending__` prefix (regime-merged) safetensors layouts. Target network is resynced in lockstep so inference and TD targets see the same restored weights. Adds NoisyLinear::{weight_mu_mut, bias_mu_mut} for in-place D2D restore of mu params during checkpoint load. Tests: - branching::tests::test_load_from_named_slices_round_trip (happy path) - branching::tests::test_load_from_named_slices_rejects_extra_keys (arch-drift safety) - branching::tests::test_load_from_named_slices_rejects_shape_mismatch - dqn::save_load_tests::test_dqn_load_from_safetensors_round_trip (full end-to-end safetensors save+load, #[ignore] for CUDA gate) - dqn::save_load_tests::test_dqn_load_from_safetensors_missing_file Also un-ignores the ensemble adapter's `test_dqn_checkpoint_round_trip` test (was ignored pre-fix because load_from_safetensors silently dropped weights). Disables NoisyNet on both sides for deterministic argmax comparison. Adds CUDA_LOCK mutex to serialize the adapter tests that share a CUDA stream via the SHARED_CUDA OnceLock. Co-Authored-By: Claude Opus 4.7 (1M context) --- crates/ml-dqn/src/branching.rs | 236 ++++++++++++++++++++ crates/ml-dqn/src/dqn.rs | 293 +++++++++++++++++++++++-- crates/ml-dqn/src/noisy_layers.rs | 10 + crates/ml/src/ensemble/adapters/dqn.rs | 131 ++++++++--- 4 files changed, 620 insertions(+), 50 deletions(-) diff --git a/crates/ml-dqn/src/branching.rs b/crates/ml-dqn/src/branching.rs index 3415f5b14..a63784475 100644 --- a/crates/ml-dqn/src/branching.rs +++ b/crates/ml-dqn/src/branching.rs @@ -1054,6 +1054,127 @@ impl BranchingDuelingQNetwork { ) -> Result<(), MLError> { dst.copy_params_from(src) } + + /// Restore weights from a named-slice map (symmetric counterpart of + /// [`named_weight_slices`](Self::named_weight_slices)). + /// + /// The map is typically assembled by uploading tensors parsed from a + /// safetensors checkpoint back onto GPU. For each expected name emitted + /// by `named_weight_slices()` we validate the shape matches the target + /// and perform a `D2D` memcpy into the target `CudaSlice` in place — + /// same primitive as [`copy_weights_from`](Self::copy_weights_from). + /// + /// # Errors + /// + /// Returns [`MLError::CheckpointError`] when: + /// * any expected name is missing from `slices` + /// * a slice's shape does not match the target tensor's shape + /// * `slices` contains names that the save-side never emits (stale/extra keys) + /// + /// Never silently skips — that would mask checkpoint/architecture drift. + pub fn load_from_named_slices( + &mut self, + slices: &std::collections::HashMap, Vec)>, + ) -> Result<(), MLError> { + // Build the expected (name, shape, target-len) manifest from the save side. + // + // We intentionally reuse the exact naming scheme of `named_weight_slices`: + // shared_{i}.weight, shared_{i}.bias, value_fc.{weight,bias}, + // value_out.{weight,bias}, branch_{d}_fc.{weight,bias}, + // branch_{d}_out.{weight,bias}. + let expected: Vec<(String, Vec, usize)> = self + .named_weight_slices() + .into_iter() + .map(|(name, slice, shape)| (name, shape, slice.len())) + .collect(); + + // 1. Validate the loaded map has exactly the expected keys (no missing, + // no extras). This catches architecture drift and stale checkpoints. + let expected_names: std::collections::HashSet<&str> = + expected.iter().map(|(n, _, _)| n.as_str()).collect(); + for provided_name in slices.keys() { + if !expected_names.contains(provided_name.as_str()) { + return Err(MLError::CheckpointError(format!( + "load_from_named_slices: unexpected tensor '{provided_name}' — \ + checkpoint/arch drift (not in named_weight_slices manifest)" + ))); + } + } + + // 2. Shape + length validation (before touching any GPU memory, to fail fast). + for (name, shape, target_len) in &expected { + let Some((src_slice, src_shape)) = slices.get(name) else { + return Err(MLError::CheckpointError(format!( + "load_from_named_slices: missing tensor '{name}' in checkpoint map" + ))); + }; + if src_shape != shape { + return Err(MLError::CheckpointError(format!( + "load_from_named_slices: shape mismatch for '{name}': \ + checkpoint={src_shape:?} vs target={shape:?}" + ))); + } + if src_slice.len() != *target_len { + return Err(MLError::CheckpointError(format!( + "load_from_named_slices: len mismatch for '{name}': \ + checkpoint={} vs target={}", + src_slice.len(), + *target_len, + ))); + } + } + + // 3. D2D memcpy into each target slot in place. + for (i, layer) in self.shared_layers.iter_mut().enumerate() { + let w_name = format!("shared_{i}.weight"); + let b_name = format!("shared_{i}.bias"); + let (w_src, _) = slices.get(&w_name).ok_or_else(|| { + MLError::CheckpointError(format!("load_from_named_slices: '{w_name}' vanished")) + })?; + let (b_src, _) = slices.get(&b_name).ok_or_else(|| { + MLError::CheckpointError(format!("load_from_named_slices: '{b_name}' vanished")) + })?; + dtod_copy_slice(w_src, &mut layer.weight, &self.stream, &w_name)?; + dtod_copy_slice(b_src, &mut layer.bias, &self.stream, &b_name)?; + } + + // Restore NoisyLinear mu params for value + branch heads. + let load_noisy = |name: &str, + noisy: &mut NoisyLinear, + stream: &Arc, + map: &std::collections::HashMap, Vec)>| + -> Result<(), MLError> { + let w_key = format!("{name}.weight"); + let b_key = format!("{name}.bias"); + let (w_src, _) = map.get(&w_key).ok_or_else(|| { + MLError::CheckpointError(format!("load_from_named_slices: '{w_key}' missing")) + })?; + let (b_src, _) = map.get(&b_key).ok_or_else(|| { + MLError::CheckpointError(format!("load_from_named_slices: '{b_key}' missing")) + })?; + dtod_copy_slice(w_src, noisy.weight_mu_mut(), stream, &w_key)?; + dtod_copy_slice(b_src, noisy.bias_mu_mut(), stream, &b_key)?; + Ok(()) + }; + + load_noisy("value_fc", &mut self.value_fc, &self.stream, slices)?; + load_noisy("value_out", &mut self.value_out, &self.stream, slices)?; + + for (d, fc) in self.branch_fcs.iter_mut().enumerate() { + load_noisy(&format!("branch_{d}_fc"), fc, &self.stream, slices)?; + } + for (d, out) in self.branch_outs.iter_mut().enumerate() { + load_noisy(&format!("branch_{d}_out"), out, &self.stream, slices)?; + } + + // Synchronize to ensure all D2D memcpys complete before the next op + // reads these weights. + self.stream + .synchronize() + .map_err(|e| MLError::CheckpointError(format!("load_from_named_slices sync: {e}")))?; + + Ok(()) + } } impl std::fmt::Debug for BranchingDuelingQNetwork { @@ -2083,4 +2204,119 @@ mod tests { assert!((config.dropout_rate - 0.0).abs() < 1e-6); Ok(()) } + + // Round-trip validation for named_weight_slices <-> load_from_named_slices. + // + // Dumps a randomised net's weights via the save manifest, uploads those same + // buffers back (simulating what a real on-disk load looks like post-parse), + // then asserts a fresh net loaded from the map produces forward outputs + // matching the source net on a fixed deterministic input. + #[test] + fn test_load_from_named_slices_round_trip() -> anyhow::Result<()> { + let stream = test_stream(); + let config = trading_config_small_atoms(8); + + // Source net: keeps the randomly-initialised weights we want to copy. + let src = BranchingDuelingQNetwork::new(config.clone(), Arc::clone(&stream))?; + + // Destination net: independent random init — must be overwritten by load. + let mut dst = BranchingDuelingQNetwork::new(config, Arc::clone(&stream))?; + + // Build the named-slice map the loader expects. Each entry is an + // independently-allocated CudaSlice cloned from the source via + // CPU round-trip (same fidelity as safetensors parse + upload). + let mut map: std::collections::HashMap, Vec)> = + std::collections::HashMap::new(); + for (name, slice, shape) in src.named_weight_slices() { + let mut host = vec![0.0_f32; slice.len()]; + stream + .memcpy_dtoh(slice, &mut host) + .map_err(|e| anyhow::anyhow!("DtoH '{name}': {e}"))?; + let gpu = stream + .clone_htod(&host) + .map_err(|e| anyhow::anyhow!("HtoD '{name}': {e}"))?; + map.insert(name, (gpu, shape)); + } + + let state = GpuTensor::full(&[1, 8], 0.5, &stream)?; + + dst.load_from_named_slices(&map)?; + + let post_src = src.forward_branches_eval(&state)?; + let post_dst = dst.forward_branches_eval(&state)?; + + let v1 = post_src.value.to_host(&stream)?; + let v2 = post_dst.value.to_host(&stream)?; + let v1_scalar = v1.first().copied().unwrap_or(f32::NAN); + let v2_scalar = v2.first().copied().unwrap_or(f32::NAN); + + // Same tolerance as test_distributional_weight_copy — NoisyNet resamples + // noise on every call, so we compare within the established tolerance. + assert!( + (v1_scalar - v2_scalar).abs() < 0.1, + "Post-load values should match: src={v1_scalar} dst={v2_scalar}", + ); + Ok(()) + } + + // Ensure the loader errors on extra (non-manifest) keys instead of + // silently ignoring them - protects against checkpoint/arch drift. + #[test] + fn test_load_from_named_slices_rejects_extra_keys() -> anyhow::Result<()> { + let stream = test_stream(); + let config = trading_config_small_atoms(8); + let src = BranchingDuelingQNetwork::new(config.clone(), Arc::clone(&stream))?; + let mut dst = BranchingDuelingQNetwork::new(config, Arc::clone(&stream))?; + + let mut map: std::collections::HashMap, Vec)> = + std::collections::HashMap::new(); + for (name, slice, shape) in src.named_weight_slices() { + let mut host = vec![0.0_f32; slice.len()]; + stream.memcpy_dtoh(slice, &mut host)?; + let gpu = stream.clone_htod(&host)?; + map.insert(name, (gpu, shape)); + } + // Spurious key that doesn't belong. + let bogus = stream.clone_htod(&[0.0_f32; 4])?; + map.insert("ghost_layer.weight".to_owned(), (bogus, vec![2, 2])); + + let err = dst.load_from_named_slices(&map).expect_err("extra key must error"); + let msg = format!("{err}"); + assert!( + msg.contains("unexpected tensor") && msg.contains("ghost_layer"), + "expected arch-drift error, got: {msg}" + ); + Ok(()) + } + + // Ensure the loader errors on shape mismatch instead of producing garbage. + #[test] + fn test_load_from_named_slices_rejects_shape_mismatch() -> anyhow::Result<()> { + let stream = test_stream(); + let config = trading_config_small_atoms(8); + let src = BranchingDuelingQNetwork::new(config.clone(), Arc::clone(&stream))?; + let mut dst = BranchingDuelingQNetwork::new(config, Arc::clone(&stream))?; + + let mut map: std::collections::HashMap, Vec)> = + std::collections::HashMap::new(); + for (name, slice, shape) in src.named_weight_slices() { + let mut host = vec![0.0_f32; slice.len()]; + stream.memcpy_dtoh(slice, &mut host)?; + let gpu = stream.clone_htod(&host)?; + // Corrupt the shape of the first shared weight. + let final_shape = if name == "shared_0.weight" { + vec![shape[0] + 1, shape[1]] + } else { + shape + }; + map.insert(name, (gpu, final_shape)); + } + + let err = dst + .load_from_named_slices(&map) + .expect_err("shape mismatch must error"); + let msg = format!("{err}"); + assert!(msg.contains("shape mismatch"), "got: {msg}"); + Ok(()) + } } diff --git a/crates/ml-dqn/src/dqn.rs b/crates/ml-dqn/src/dqn.rs index 1ed235f61..cd015bc9c 100644 --- a/crates/ml-dqn/src/dqn.rs +++ b/crates/ml-dqn/src/dqn.rs @@ -2454,49 +2454,175 @@ impl DQN { } - /// Load model weights from safetensors checkpoint + /// Load model weights from a safetensors checkpoint into the branching Q-network. /// - /// Loads pre-trained weights from a safetensors file and updates both - /// the Q-network and target network. Follows the MAMBA2 pattern for - /// checkpoint loading. + /// Parses the safetensors file at `path`, uploads each named f32 tensor back + /// to GPU as a `CudaSlice`, and calls + /// [`BranchingDuelingQNetwork::load_from_named_slices`] to restore the live + /// weights in place via D2D memcpy. The target network is also resynced so + /// inference and TD targets see the same weights. + /// + /// # Checkpoint layout + /// + /// The training-time saver + /// ([`serialize_model`](../../ml/src/trainers/dqn/trainer/mod.rs)) emits one + /// merged safetensors file with three regime-prefixed copies of the branching + /// weight manifest: `trending__`, `ranging__`, `volatile__`. The + /// single-head `DQN` struct only owns one branching network, so this loader + /// restores the `trending__` copy by default (matching + /// [`RegimeConditionalDQN::load_from_merged_safetensors`]). For checkpoints + /// saved without a prefix (e.g., the per-crate round-trip tests), the loader + /// falls back to plain names. /// /// # Arguments /// - /// * `path` - Path to the safetensors file (with or without .safetensors extension) + /// * `path` - Path to the safetensors file (with or without `.safetensors` suffix) /// - /// # Returns + /// # Errors /// - /// * `Ok(())` - Checkpoint loaded successfully - /// * `Err(MLError::CheckpointError)` - File not found or invalid format - /// * `Err(MLError::LockError)` - Failed to acquire `GpuVarStore` lock + /// * [`MLError::CheckpointError`] — file missing, parse failure, architecture + /// drift (extra/missing tensors, shape mismatch). + /// * [`MLError::ModelError`] — GPU upload failure. /// /// # Example /// /// ```no_run - /// use ml::dqn::{DQN, DQNConfig}; + /// use ml_dqn::{DQN, DQNConfig}; /// /// let config = DQNConfig::emergency_safe_defaults(); /// let mut dqn = DQN::new(config)?; /// dqn.load_from_safetensors("/path/to/checkpoint")?; - /// # Ok::<(), ml::MLError>(()) + /// # Ok::<(), ml_core::MLError>(()) /// ``` pub fn load_from_safetensors(&mut self, path: &str) -> Result<(), MLError> { - // Checkpoint loading is handled by the flat buffer path in GpuDqnTrainer. - // This cold-path loader is a no-op since weights are managed via params_buf. - let safetensors_path = if !path.ends_with(".safetensors") { - format!("{path}.safetensors") - } else { + let safetensors_path = if path.ends_with(".safetensors") { path.to_owned() + } else { + format!("{path}.safetensors") }; - if !std::path::Path::new(&safetensors_path).exists() { + let path_ref = std::path::Path::new(&safetensors_path); + if !path_ref.exists() { return Err(MLError::CheckpointError(format!( "Checkpoint file not found: {safetensors_path}", ))); } + let branching_net = self.branching_q_network.as_mut().ok_or_else(|| { + MLError::CheckpointError( + "DQN::load_from_safetensors: branching_q_network not initialized".to_owned(), + ) + })?; + + // Read + parse safetensors file. + let file_bytes = std::fs::read(path_ref).map_err(|e| { + MLError::CheckpointError(format!( + "Failed to read safetensors '{safetensors_path}': {e}", + )) + })?; + let tensors = safetensors::SafeTensors::deserialize(&file_bytes).map_err(|e| { + MLError::CheckpointError(format!( + "Failed to parse safetensors '{safetensors_path}': {e}", + )) + })?; + + // Figure out which naming scheme the file uses. The regime trainer + // saves all three heads with `trending__` / `ranging__` / `volatile__` + // prefixes; the single-head tests save plain names. Probe the first + // shared layer's weight. + let available: std::collections::HashSet = tensors + .names() + .iter() + .map(|n| (*n).to_owned()) + .collect(); + let prefix = if available.contains("trending__shared_0.weight") { + "trending__" + } else if available.contains("shared_0.weight") { + "" + } else { + return Err(MLError::CheckpointError(format!( + "Checkpoint '{safetensors_path}' has no recognised weight layout \ + (expected either 'shared_0.weight' or 'trending__shared_0.weight'); \ + file contains {} tensors", + available.len(), + ))); + }; + + // Upload each expected tensor to GPU as a CudaSlice. + let manifest: Vec<(String, Vec)> = branching_net + .named_weight_slices() + .into_iter() + .map(|(name, _slice, shape)| (name, shape)) + .collect(); + + let mut loaded: std::collections::HashMap, Vec)> = + std::collections::HashMap::with_capacity(manifest.len()); + + for (name, expected_shape) in &manifest { + let file_key = format!("{prefix}{name}"); + let view = tensors.tensor(&file_key).map_err(|e| { + MLError::CheckpointError(format!( + "Checkpoint '{safetensors_path}' missing '{file_key}': {e}", + )) + })?; + if view.dtype() != safetensors::Dtype::F32 { + return Err(MLError::CheckpointError(format!( + "Checkpoint '{file_key}' dtype {:?} is not F32", + view.dtype(), + ))); + } + let saved_shape: Vec = view.shape().to_vec(); + if &saved_shape != expected_shape { + return Err(MLError::CheckpointError(format!( + "Checkpoint '{file_key}' shape mismatch: file={saved_shape:?} vs \ + model={expected_shape:?}", + ))); + } + + // Parse f32 little-endian bytes -> Vec and upload to GPU. + let data_bytes = view.data(); + let numel: usize = expected_shape.iter().product(); + if data_bytes.len() != numel * std::mem::size_of::() { + return Err(MLError::CheckpointError(format!( + "Checkpoint '{file_key}' byte-length {} does not match shape {:?} \ + (expected {} bytes)", + data_bytes.len(), + expected_shape, + numel * std::mem::size_of::(), + ))); + } + let mut host: Vec = Vec::with_capacity(numel); + for chunk_idx in 0..numel { + let offset = chunk_idx * 4; + let bytes = [ + data_bytes[offset], + data_bytes[offset + 1], + data_bytes[offset + 2], + data_bytes[offset + 3], + ]; + host.push(f32::from_le_bytes(bytes)); + } + let gpu_slice = self.stream.clone_htod(&host).map_err(|e| { + MLError::ModelError(format!( + "load_from_safetensors: failed to upload '{file_key}': {e}", + )) + })?; + loaded.insert(name.clone(), (gpu_slice, expected_shape.clone())); + } + + // Restore weights in the Q-network. + branching_net.load_from_named_slices(&loaded)?; + + // Keep the target network in lockstep so `DQN::forward` and any subsequent + // TD-target computation operate on the same restored weights. + if let Some(target_net) = self.branching_target_network.as_mut() { + target_net.load_from_named_slices(&loaded)?; + } + debug!( - "DQN checkpoint load delegated to flat buffer path: {safetensors_path}", + "DQN loaded branching weights from {safetensors_path} (prefix='{prefix}', \ + {} tensors restored)", + manifest.len(), ); Ok(()) @@ -2626,3 +2752,134 @@ impl DQN { } +#[cfg(test)] +#[cfg(feature = "cuda")] +mod save_load_tests { + use super::*; + + /// Round-trip validation for `DQN::load_from_safetensors`. + /// + /// 1. Build a randomly-initialised DQN (source). + /// 2. Dump its branching weights via `named_weight_slices` into a safetensors + /// file on disk (plain-name layout — no regime prefix). + /// 3. Build a FRESH DQN with identical config. + /// 4. Call `load_from_safetensors` on the fresh DQN. + /// 5. Assert both DQNs' forward outputs agree on a deterministic input. + /// + /// `#[ignore]` because it requires a CUDA runtime. + #[test] + #[ignore = "requires CUDA runtime"] + fn test_dqn_load_from_safetensors_round_trip() { + // Small config matching the ensemble adapter's test expectations. + let mut config = DQNConfig::emergency_safe_defaults(); + config.num_actions = 3; + config.hidden_dims = vec![64, 64]; + + // --- Phase 1: source DQN + save --- + let src = DQN::new(config.clone()).expect("src DQN new"); + + let dir = tempfile::tempdir().expect("tempdir"); + let path = dir.path().join("dqn_roundtrip.safetensors"); + + { + let stream = src.cuda_stream(); + let branching = src + .branching_q_network + .as_ref() + .expect("branching_q_network expected to be initialised"); + let named = branching.named_weight_slices(); + + // Download each slice to host and keep byte buffers alive for + // the TensorView borrows. + let byte_vecs: Vec<(String, Vec, Vec)> = named + .iter() + .map(|(name, slice, shape)| { + let mut host = vec![0.0_f32; slice.len()]; + stream.memcpy_dtoh(*slice, &mut host).expect("DtoH"); + let bytes: Vec = + host.iter().flat_map(|v| v.to_le_bytes()).collect(); + (name.clone(), shape.clone(), bytes) + }) + .collect(); + + let mut tensor_map = std::collections::HashMap::new(); + for (name, shape, bytes) in &byte_vecs { + let view = safetensors::tensor::TensorView::new( + safetensors::Dtype::F32, + shape.clone(), + bytes, + ) + .expect("TensorView"); + tensor_map.insert(name.as_str(), view); + } + let arch_meta = Some(config.checkpoint_metadata()); + safetensors::serialize_to_file(tensor_map, arch_meta, path.as_path()) + .expect("serialize_to_file"); + } + + // --- Phase 2: fresh DQN + load --- + let mut fresh = DQN::new(config.clone()).expect("fresh DQN new"); + let path_str = path.to_str().expect("utf8 path"); + fresh.load_from_safetensors(path_str).expect("load_from_safetensors"); + + // --- Phase 3: forward outputs must match on a deterministic input --- + let stream = src.cuda_stream(); + let state_dim = ml_core::state_layout::STATE_DIM; + let state_host: Vec = (0..state_dim) + .map(|i| ((i as f32) * 0.037).sin() * 0.25) + .collect(); + let state = GpuTensor::from_host(&state_host, vec![1, state_dim], stream) + .expect("state tensor"); + + let q_src = src.forward(&state).expect("src forward").to_host(stream).expect("src DtoH"); + let q_fresh = fresh + .forward(&state) + .expect("fresh forward") + .to_host(stream) + .expect("fresh DtoH"); + + assert_eq!( + q_src.len(), + q_fresh.len(), + "Q-vector length should match: {} vs {}", + q_src.len(), + q_fresh.len(), + ); + + // NoisyNet resamples noise between `src.forward` and `fresh.forward`, + // so perfectly bit-identical outputs are not guaranteed. However the + // mu tensors are byte-identical post-load, so the argmax + relative + // ordering of the Q-values should match within a loose tolerance. + let max_abs_diff = q_src + .iter() + .zip(q_fresh.iter()) + .map(|(a, b)| (a - b).abs()) + .fold(0.0_f32, f32::max); + + assert!( + max_abs_diff < 1.0, + "Post-load forward outputs diverged beyond noise tolerance: \ + max_abs_diff={max_abs_diff}, src={q_src:?}, fresh={q_fresh:?}", + ); + } + + /// Loading from a path that doesn't exist must error with a CheckpointError, + /// not silently succeed. + #[test] + #[ignore = "requires CUDA runtime"] + fn test_dqn_load_from_safetensors_missing_file() { + let mut config = DQNConfig::emergency_safe_defaults(); + config.num_actions = 3; + config.hidden_dims = vec![64, 64]; + + let mut dqn = DQN::new(config).expect("DQN new"); + let err = dqn + .load_from_safetensors("/tmp/this_path_must_not_exist_for_test.safetensors") + .expect_err("missing file must error"); + assert!( + matches!(err, MLError::CheckpointError(_)), + "expected CheckpointError, got: {err:?}", + ); + } +} + diff --git a/crates/ml-dqn/src/noisy_layers.rs b/crates/ml-dqn/src/noisy_layers.rs index db5e3077f..107573add 100644 --- a/crates/ml-dqn/src/noisy_layers.rs +++ b/crates/ml-dqn/src/noisy_layers.rs @@ -336,6 +336,16 @@ impl NoisyLinear { [&self.weight_mu, &self.bias_mu] } + /// Mutable access to `weight_mu` — for checkpoint restore (D2D memcpy in-place). + pub const fn weight_mu_mut(&mut self) -> &mut CudaSlice { + &mut self.weight_mu + } + + /// Mutable access to `bias_mu` — for checkpoint restore (D2D memcpy in-place). + pub const fn bias_mu_mut(&mut self) -> &mut CudaSlice { + &mut self.bias_mu + } + /// Input dimension of this layer. pub const fn in_features(&self) -> usize { self.in_features diff --git a/crates/ml/src/ensemble/adapters/dqn.rs b/crates/ml/src/ensemble/adapters/dqn.rs index b9c509d67..21a6bfaff 100644 --- a/crates/ml/src/ensemble/adapters/dqn.rs +++ b/crates/ml/src/ensemble/adapters/dqn.rs @@ -60,6 +60,26 @@ impl DqnInferenceAdapter { }) } + /// Create a DQN inference adapter on a specific device and load weights + /// from a safetensors checkpoint. Test-only variant of + /// [`from_checkpoint`](Self::from_checkpoint) that pins the device so the + /// stream stays consistent with the rest of a test harness. + #[cfg(test)] + pub(crate) fn from_checkpoint_on_device( + config: DQNConfig, + path: &str, + device: MlDevice, + ) -> MLResult { + let stream = device.cuda_stream()?.clone(); + let mut model = DQN::new_on_device(config, device.clone())?; + model.load_from_safetensors(path)?; + Ok(Self { + model: Mutex::new(model), + device, + stream, + }) + } + /// Create a DQN inference adapter and load weights from a safetensors checkpoint. pub fn from_checkpoint(config: DQNConfig, path: &str) -> MLResult { let device = DeviceConfig::Auto.resolve().map_err(|e| MLError::DeviceError(format!("CUDA required: {e}")))?; @@ -241,10 +261,17 @@ impl ModelInferenceAdapter for DqnInferenceAdapter { #[cfg(test)] mod tests { use super::*; - use std::sync::OnceLock; + use std::sync::{Mutex, OnceLock}; static SHARED_CUDA: OnceLock = OnceLock::new(); + /// Serialize CUDA work across tests in this module. Multiple tests share + /// `SHARED_CUDA`, which means they share a CUDA stream; running them in + /// parallel lets cuBLAS launches from different tests interleave on the + /// same stream, producing non-deterministic Q-value outputs. This mutex + /// forces sequential execution of the compute-heavy sections. + static CUDA_LOCK: Mutex<()> = Mutex::new(()); + fn shared_device() -> MlDevice { SHARED_CUDA .get_or_init(|| { @@ -265,6 +292,7 @@ mod tests { #[test] fn test_dqn_adapter_creation() { + let _guard = CUDA_LOCK.lock().unwrap_or_else(|p| p.into_inner()); let adapter = DqnInferenceAdapter::new_on_device(test_config(), shared_device()).unwrap(); assert_eq!(adapter.model_name(), "DQN"); assert!(adapter.is_ready()); @@ -272,6 +300,7 @@ mod tests { #[test] fn test_dqn_adapter_predict_direction_range() { + let _guard = CUDA_LOCK.lock().unwrap_or_else(|p| p.into_inner()); let adapter = DqnInferenceAdapter::new_on_device(test_config(), shared_device()).unwrap(); let fv = FeatureVector { values: vec![0.1; 56], @@ -296,6 +325,7 @@ mod tests { #[test] fn test_dqn_adapter_deterministic() { + let _guard = CUDA_LOCK.lock().unwrap_or_else(|p| p.into_inner()); let adapter = DqnInferenceAdapter::new_on_device(test_config(), shared_device()).unwrap(); let fv = FeatureVector { values: vec![0.1; 56], @@ -310,24 +340,31 @@ mod tests { } #[test] - #[ignore = "branching always on: checkpoint round-trip needs branching weight save/load (not just base q_network)"] fn test_dqn_checkpoint_round_trip() { + let _guard = CUDA_LOCK.lock().unwrap_or_else(|p| p.into_inner()); let config = test_config(); - let adapter = DqnInferenceAdapter::new_on_device(config.clone(), shared_device()); - assert!( - adapter.is_ok(), - "failed to create adapter: {:?}", - adapter.err() - ); - let adapter = adapter.unwrap(); + let adapter = DqnInferenceAdapter::new_on_device(config.clone(), shared_device()) + .expect("adapter new"); + + // Disable NoisyNet on the SOURCE so pred1 is deterministic — mu-only + // forward. The saved checkpoint contains mu tensors (NoisyLinear's + // saveable surface), and `load_from_named_slices` overwrites mu in + // place. We disable noise on both sides so argmax determinism holds. + { + let mut model = adapter.model.lock().unwrap(); + if let Some(br) = model.branching_q_network.as_mut() { + br.disable_noise().expect("disable src noise"); + } + if let Some(tgt) = model.branching_target_network.as_mut() { + tgt.disable_noise().expect("disable src target noise"); + } + } let fv = FeatureVector { values: vec![0.3; 56], timestamp: 1700000000, }; - let pred1 = adapter.predict(&fv); - assert!(pred1.is_ok(), "first predict failed: {:?}", pred1.err()); - let pred1 = pred1.unwrap(); + let pred1 = adapter.predict(&fv).expect("first predict"); // Save checkpoint to temp dir let dir = tempfile::tempdir().unwrap(); @@ -371,28 +408,58 @@ mod tests { .unwrap(); } - // Load into new adapter + // Load into a fresh adapter on the same shared device so both adapters + // share a CUDA stream and the comparison is apples-to-apples. let path_str = path.to_str().unwrap(); - let adapter2 = DqnInferenceAdapter::from_checkpoint(config, path_str); - assert!( - adapter2.is_ok(), - "from_checkpoint failed: {:?}", - adapter2.err() - ); - let adapter2 = adapter2.unwrap(); - let pred2 = adapter2.predict(&fv); - assert!( - pred2.is_ok(), - "second predict failed: {:?}", - pred2.err() - ); - let pred2 = pred2.unwrap(); + let adapter2 = DqnInferenceAdapter::from_checkpoint_on_device( + config, + path_str, + shared_device(), + ) + .expect("from_checkpoint"); - assert!( - (pred1.direction - pred2.direction).abs() < 1e-4, - "round-trip mismatch: {} vs {}", - pred1.direction, - pred2.direction + // Disable noise on the destination too so pred2 is also deterministic. + { + let mut model = adapter2.model.lock().unwrap(); + if let Some(br) = model.branching_q_network.as_mut() { + br.disable_noise().expect("disable dst noise"); + } + if let Some(tgt) = model.branching_target_network.as_mut() { + tgt.disable_noise().expect("disable dst target noise"); + } + } + + let pred2 = adapter2.predict(&fv).expect("second predict"); + + // With NoisyNet disabled and mu byte-identical, both adapters must + // produce the same Q-values -> same argmax -> same direction. + assert_eq!( + pred1.direction, pred2.direction, + "round-trip direction mismatch: pred1={} pred2={}", + pred1.direction, pred2.direction, ); + + // Also verify the raw Q-values match bit-for-bit (the mu-only path is + // deterministic D2D + deterministic cuBLAS GEMM). + let q1 = pred1 + .metadata + .q_values + .as_ref() + .expect("pred1 q_values"); + let q2 = pred2 + .metadata + .q_values + .as_ref() + .expect("pred2 q_values"); + assert_eq!(q1.len(), q2.len(), "Q-vector length mismatch"); + for (i, (a, b)) in q1.iter().zip(q2.iter()).enumerate() { + // GEMM determinism + fresh device allocs: tiny numerical drift + // is acceptable but argmax must be exact (already checked). + let diff = (a - b).abs(); + assert!( + diff < 1e-5, + "Q[{i}] diff {diff} exceeds tol: pred1={a} pred2={b}", + ); + } } }