gitadmin/foxhunt
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

fix(ml-dqn): implement real DQN::load_from_safetensors

Browse Source 
`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) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-04-23 10:19:56 +02:00
parent c5045c009e
commit f47af9d268
4 changed files with 620 additions and 50 deletions
Download Patch File Download Diff File Expand all files Collapse all files

236
crates/ml-dqn/src/branching.rs
View File

@@ -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<String, (CudaSlice<f32>, Vec<usize>)>,
) -> 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>, 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<CudaStream>,
map: &std::collections::HashMap<String, (CudaSlice<f32>, Vec<usize>)>|
-> 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<f32> cloned from the source via
// CPU round-trip (same fidelity as safetensors parse + upload).
let mut map: std::collections::HashMap<String, (CudaSlice<f32>, Vec<usize>)> =
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<String, (CudaSlice<f32>, Vec<usize>)> =
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<String, (CudaSlice<f32>, Vec<usize>)> =
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(())
}
}

293
crates/ml-dqn/src/dqn.rs
View File

@@ -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<f32>`, 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<String> = 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<f32>.
let manifest: Vec<(String, Vec<usize>)> = branching_net
.named_weight_slices()
.into_iter()
.map(|(name, _slice, shape)| (name, shape))
.collect();
let mut loaded: std::collections::HashMap<String, (cudarc::driver::CudaSlice<f32>, Vec<usize>)> =
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<usize> = 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<f32> 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::<f32>() {
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::<f32>(),
)));
}
let mut host: Vec<f32> = 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<usize>, Vec<u8>)> = 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<u8> =
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<f32> = (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:?}",
);
}
}

10
crates/ml-dqn/src/noisy_layers.rs
View File

@@ -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<f32> {
&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<f32> {
&mut self.bias_mu
}
/// Input dimension of this layer.
pub const fn in_features(&self) -> usize {
self.in_features

131
crates/ml/src/ensemble/adapters/dqn.rs
View File

@@ -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<Self> {
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<Self> {
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<MlDevice> = 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}",
);
}
}
}