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

fix(ml): migrate all test code to GPU types — 125 compile errors fixed

Browse Source 
17 files: smoke_tests (gpu_residency, training_stability, performance,
helpers), inference, multi_timeframe, tlob, validation/adapters,
tft/tests, ensemble/adapters/dqn, flash_attention, hyperopt/tft,
mamba/trainable, online_learning, liquid/adapter, transformers/attention

- Tensor→CudaSlice direct ops in smoke tests (no tensor_to_cuda_slice)
- Device→MlDevice throughout
- GpuLinear::forward 4-arg signature (input, store, cublas, stream)
- rand::rng()→thread_rng() (rand 0.8)
- StreamTensor .shape field access (not .dims() method)

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-03-18 14:47:33 +01:00
parent 80058eaa16
commit 7e49d3c08e
17 changed files with 299 additions and 218 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -337,11 +337,30 @@ mod tests {
{
let model = adapter.model.lock().unwrap();
let vars = model.get_q_network_vars();
let tensors = vars.to_safetensors_map(&adapter.stream)
.unwrap_or_default();
let host_data = vars.export_to_host().unwrap();
let arch_metadata = Some(model.config.checkpoint_metadata());
// Build safetensors TensorView map from host data
let mut tensor_map = std::collections::HashMap::new();
// Keep byte buffers alive for the duration of serialize_to_file
let byte_vecs: Vec<(String, Vec<usize>, Vec<u8>)> = host_data
.into_iter()
.map(|(name, (shape, data))| {
let bytes: Vec<u8> = data.iter().flat_map(|v| v.to_le_bytes()).collect();
(name, shape, bytes)
})
.collect();
for (name, shape, bytes) in &byte_vecs {
let view = safetensors::tensor::TensorView::new(
safetensors::Dtype::F32,
shape.clone(),
bytes,
)
.unwrap();
tensor_map.insert(name.as_str(), view);
}
safetensors::serialize_to_file(
&tensors,
tensor_map,
arch_metadata,
path.as_path(),
)

35
crates/ml/src/features/multi_timeframe.rs
View File

@@ -367,6 +367,7 @@ fn push_ring(buf: &mut VecDeque<OHLCVBar>, bar: OHLCVBar, max_len: usize) {
mod tests {
use super::*;
use chrono::{TimeZone, Utc};
use ml_core::cuda_autograd::stream_ops::{gpu_sub, gpu_abs};
fn make_bar(minute: u32, close: f64) -> OHLCVBar {
OHLCVBar {
@@ -405,7 +406,7 @@ mod tests {
}
let output = encoder.encode().expect("encode should succeed");
let dims = output.dims();
let dims = &output.shape;
assert_eq!(dims.len(), 2, "output should be 2D");
assert_eq!(
dims.first().copied().unwrap_or(0),
@@ -437,14 +438,11 @@ mod tests {
let out1 = encoder.encode().expect("encode 1");
let out2 = encoder.encode().expect("encode 2");
// Same state -> same output
let diff = out1
.sub(&out2)
.and_then(|d| d.abs())
.and_then(|d| d.sum_all())
.and_then(|d| d.to_dtype(NativeDType::F32))
.and_then(|d| d.to_scalar::<f32>())
.expect("diff computation");
// Same state -> same output: compute |out1 - out2| and sum on host
let diff_tensor = gpu_sub(&out1, &out2).expect("sub");
let abs_tensor = gpu_abs(&diff_tensor).expect("abs");
let abs_vec = abs_tensor.to_vec().expect("to_vec");
let diff: f32 = abs_vec.iter().sum();
assert!(
diff < 1e-6,
@@ -513,14 +511,14 @@ mod tests {
.expect("encoder creation should succeed");
let output = encoder.encode().expect("encode on empty history should work");
let dims = output.dims();
let dims = &output.shape;
assert_eq!(dims.last().copied().unwrap_or(0), config.output_dim);
}
#[test]
fn test_lstm_encoder_single_step() {
let dev = NativeDevice::Cuda(0);
let stream = match &dev { NativeDevice::Cuda(d) => d.cuda_stream(), _ => panic!("CUDA required") };
let ml_dev = MlDevice::cuda(0).expect("CUDA required");
let stream = std::sync::Arc::clone(ml_dev.cuda_stream().expect("stream"));
let lstm = LstmEncoder::new(6, 32, &stream).expect("lstm creation");
// Single timestep: (1, 6)
@@ -533,8 +531,8 @@ mod tests {
#[test]
fn test_lstm_encoder_multi_step() {
let dev = NativeDevice::Cuda(0);
let stream = match &dev { NativeDevice::Cuda(d) => d.cuda_stream(), _ => panic!("CUDA required") };
let ml_dev = MlDevice::cuda(0).expect("CUDA required");
let stream = std::sync::Arc::clone(ml_dev.cuda_stream().expect("stream"));
let lstm = LstmEncoder::new(6, 64, &stream).expect("lstm creation");
// 10 timesteps: (10, 6)
@@ -575,12 +573,9 @@ mod tests {
// After reset, encode should return zero-based output
let output = encoder.encode().expect("encode after reset");
let sum = output
.abs()
.and_then(|t| t.sum_all())
.and_then(|t| t.to_dtype(NativeDType::F32))
.and_then(|t| t.to_scalar::<f32>())
.expect("sum");
let abs_out = gpu_abs(&output).expect("abs");
let vals = abs_out.to_vec().expect("to_vec");
let sum: f32 = vals.iter().sum();
// With all-zero embeddings going through the projection (which has bias),
// the output is just the projection bias. That's fine.

4
crates/ml/src/flash_attention/mod.rs
View File

@@ -342,9 +342,7 @@ mod tests {
#[test]
fn test_flash_attention_creation() -> Result<(), MLError> {
let device = NativeDevice::cuda_if_available(0).map_err(|e| {
MLError::ConfigError(format!("GPU required for flash attention: {}", e))
})?;
let device = NativeDevice::Cuda(0);
let config = FlashAttention3Config::default();
let _attention = FlashAttention3::new(config, device)?;
Ok(())

5
crates/ml/src/hyperopt/adapters/tft.rs
View File

@@ -710,8 +710,9 @@ mod tests {
target_throughput_pps: 100_000,
};
let device = MlDevice::cuda(0).expect("CUDA required"); // Use CPU for testing
let result = TemporalFusionTransformer::new_with_device(config, device);
let device = MlDevice::cuda(0).expect("CUDA required");
let stream = device.cuda_stream().expect("cuda stream");
let result = TemporalFusionTransformer::new_with_stream(config, Some(std::sync::Arc::clone(stream)));
assert!(
result.is_ok(),

12
crates/ml/src/inference.rs
View File

@@ -1437,9 +1437,11 @@ mod tests {
let device = NativeDevice::Cuda(0);
let network = NeuralNetwork::new(config, device)?;
// Create input tensor
// Create input tensor via MlDevice stream
let ml_dev = ml_core::device::MlDevice::cuda(0)?;
let stream = ml_dev.cuda_stream()?;
let input_data = vec![1.0_f32; 10];
let input_tensor = GpuTensor::from_host(input_data, &[1, 10], &NativeDevice::Cuda(0))?;
let input_tensor = GpuTensor::from_host(&input_data, vec![1, 10], stream)?;
let output = network.forward(&input_tensor).await?;
let output_shape = output.dims();
@@ -1476,9 +1478,11 @@ mod tests {
let device = NativeDevice::Cuda(0);
let network = NeuralNetwork::new(config, device)?;
// Create batch input (3 samples)
// Create batch input (3 samples) via MlDevice stream
let ml_dev = ml_core::device::MlDevice::cuda(0)?;
let stream = ml_dev.cuda_stream()?;
let input_data = vec![1.0_f32; 15]; // 3 samples * 5 features
let input_tensor = GpuTensor::from_host(input_data, &[3, 5], &NativeDevice::Cuda(0))?;
let input_tensor = GpuTensor::from_host(&input_data, vec![3, 5], stream)?;
let output = network.forward(&input_tensor).await?;
let output_shape = output.dims();

5
crates/ml/src/liquid/adapter.rs
View File

@@ -109,6 +109,11 @@ impl LiquidTrainableAdapter {
}
/// Access the GpuVarStore (read-only).
/// Reference to the CUDA stream handle.
pub fn stream(&self) -> &Arc<CudaStream> {
&self.stream
}
pub fn varmap(&self) -> &GpuVarStore {
&self.var_store
}

4
crates/ml/src/mamba/trainable_adapter.rs
View File

@@ -243,7 +243,9 @@ mod tests {
seq_len: 32,
..Default::default()
};
let adapter = Mamba2TrainableAdapter::new(config.clone(), &NativeDevice::Cuda(0))?;
let ml_dev = ml_core::device::MlDevice::cuda(0)?;
let stream = std::sync::Arc::clone(ml_dev.cuda_stream()?);
let adapter = Mamba2TrainableAdapter::new(config.clone(), &stream)?;
let temp_dir = tempfile::tempdir()?;
let checkpoint_path = temp_dir.path().join("mamba2_test_checkpoint");

267
crates/ml/src/trainers/dqn/smoke_tests/gpu_residency.rs
View File

@@ -1,12 +1,11 @@
use super::helpers::*;
use ml_core::device::MlDevice;
use ml_core::cuda_autograd::GpuTensor;
use ml_core::cuda_autograd::{GpuTensor, GpuVarStore, GpuAdamW, AdamWConfig};
use std::sync::Arc;
use tracing::info;
// GPU replay buffer (ml-dqn crate)
use crate::dqn::gpu_replay_buffer::{GpuReplayBuffer, GpuReplayBufferConfig};
use crate::dqn::replay_buffer_type::GpuBatch;
// ---------------------------------------------------------------------------
// GpuReplayBuffer tests (compiled only with cuda feature, GPU-only)
@@ -25,128 +24,172 @@ fn test_buffer_config(capacity: usize, state_dim: usize) -> GpuReplayBufferConfi
}
}
/// Get the CUDA stream from the shared smoke device.
fn smoke_stream() -> Arc<cudarc::driver::CudaStream> {
let dev = cuda_device();
Arc::clone(dev.cuda_stream().expect("cuda stream"))
}
/// Create random f32 CudaSlice on GPU (host-generated, uploaded).
fn random_cuda_f32(n: usize, stream: &Arc<cudarc::driver::CudaStream>) -> anyhow::Result<cudarc::driver::CudaSlice<f32>> {
use rand::Rng;
let mut rng = rand::thread_rng();
let host: Vec<f32> = (0..n).map(|_| rng.gen_range(-1.0_f32..1.0)).collect();
let slice = stream.clone_htod(&host).map_err(|e| anyhow::anyhow!("htod f32: {e}"))?;
Ok(slice)
}
/// Create zeroed u32 CudaSlice on GPU.
fn zeros_cuda_u32(n: usize, stream: &Arc<cudarc::driver::CudaStream>) -> anyhow::Result<cudarc::driver::CudaSlice<u32>> {
let slice = stream.alloc_zeros::<u32>(n).map_err(|e| anyhow::anyhow!("alloc u32: {e}"))?;
Ok(slice)
}
/// Create zeroed f32 CudaSlice on GPU.
fn zeros_cuda_f32(n: usize, stream: &Arc<cudarc::driver::CudaStream>) -> anyhow::Result<cudarc::driver::CudaSlice<f32>> {
let slice = stream.alloc_zeros::<f32>(n).map_err(|e| anyhow::anyhow!("alloc f32: {e}"))?;
Ok(slice)
}
/// Insert random experiences into a GPU replay buffer.
fn insert_random_batch(
buf: &mut GpuReplayBuffer,
n: usize,
state_dim: usize,
device: &Device,
) -> anyhow::Result<()> {
let states = Tensor::randn(0.0_f32, 1.0, &[n, state_dim], device)?;
let next_states = Tensor::randn(0.0_f32, 1.0, &[n, state_dim], device)?;
let actions = Tensor::zeros(&[n], DType::U32, device)?;
let rewards = Tensor::randn(0.0_f32, 1.0, &[n], device)?;
let dones = Tensor::zeros(&[n], DType::F32, device)?;
let s = crate::cuda_pipeline::tensor_to_cuda_slice_f32(&states)?;
let ns = crate::cuda_pipeline::tensor_to_cuda_slice_f32(&next_states)?;
let r = crate::cuda_pipeline::tensor_to_cuda_slice_f32(&rewards)?;
let d = crate::cuda_pipeline::tensor_to_cuda_slice_f32(&dones)?;
let a = crate::cuda_pipeline::tensor_to_cuda_slice_u32(&actions)?;
buf.insert_batch(&s, &ns, &a, &r, &d, n)?;
let stream = buf.stream();
let states = random_cuda_f32(n * state_dim, stream)?;
let next_states = random_cuda_f32(n * state_dim, stream)?;
let actions = zeros_cuda_u32(n, stream)?;
let rewards = random_cuda_f32(n, stream)?;
let dones = zeros_cuda_f32(n, stream)?;
buf.insert_batch(&states, &next_states, &actions, &rewards, &dones, n)?;
Ok(())
}
#[tokio::test]
async fn test_gpu_replay_buffer_insert_and_device() -> anyhow::Result<()> {
let device = cuda_device();
let config = test_buffer_config(100, 48);
let mut buf = GpuReplayBuffer::new(config, &device)?;
/// Download a CudaSlice<f32> to host for assertion.
fn dtoh_f32(slice: &cudarc::driver::CudaSlice<f32>, stream: &Arc<cudarc::driver::CudaStream>) -> anyhow::Result<Vec<f32>> {
let n = slice.len();
let mut host = vec![0.0_f32; n];
stream.memcpy_dtoh(slice, &mut host).map_err(|e| anyhow::anyhow!("dtoh f32: {e}"))?;
Ok(host)
}
insert_random_batch(&mut buf, 20, 48, &device)?;
/// Download a CudaSlice<u32> to host for assertion.
fn dtoh_u32(slice: &cudarc::driver::CudaSlice<u32>, stream: &Arc<cudarc::driver::CudaStream>) -> anyhow::Result<Vec<u32>> {
let n = slice.len();
let mut host = vec![0_u32; n];
stream.memcpy_dtoh(slice, &mut host).map_err(|e| anyhow::anyhow!("dtoh u32: {e}"))?;
Ok(host)
}
#[tokio::test]
async fn test_gpu_replay_buffer_insert_and_len() -> anyhow::Result<()> {
let stream = smoke_stream();
let config = test_buffer_config(100, 48);
let mut buf = GpuReplayBuffer::new(config, &stream)?;
insert_random_batch(&mut buf, 20, 48)?;
assert_eq!(buf.len(), 20);
assert!(buf.device().is_cuda());
Ok(())
}
#[tokio::test]
async fn test_gpu_replay_buffer_proportional_sample_valid() -> anyhow::Result<()> {
let device = cuda_device();
let stream = smoke_stream();
let config = test_buffer_config(100, 48);
let mut buf = GpuReplayBuffer::new(config, &device)?;
let mut buf = GpuReplayBuffer::new(config, &stream)?;
insert_random_batch(&mut buf, 50, 48, &device)?;
insert_random_batch(&mut buf, 50, 48)?;
let batch: GpuBatch = buf.sample_proportional(16)?;
let batch = buf.sample_proportional(16)?;
assert_eq!(batch.states.dims(), &[16, 48]);
assert_eq!(batch.weights.dims(), &[16]);
assert_eq!(batch.batch_size, 16);
// All weights must be positive (GPU-resident: check min > 0)
let min_w = batch.weights.min(0)?.to_dtype(DType::F32)?.to_scalar::<f32>()?;
assert!(min_w > 0.0, "min weight is not positive: {min_w}");
// Download weights and indices to host for assertions
let weights_host = dtoh_f32(&batch.weights, &stream)?;
let indices_host = dtoh_u32(&batch.indices, &stream)?;
// Weights must be finite (GPU-resident: check sum is finite)
let w_sum = batch.weights.sum_all()?.to_dtype(DType::F32)?.to_scalar::<f32>()?;
// All weights must be positive
for (i, &w) in weights_host.iter().enumerate() {
assert!(w > 0.0, "weight[{i}] is not positive: {w}");
}
// Weights must be finite
let w_sum: f32 = weights_host.iter().sum();
assert!(w_sum.is_finite(), "weight sum is not finite: {w_sum}");
// All indices must be in range [0, 50) (GPU-resident: check max < 50, min >= 0)
let max_idx = batch.indices.max(0)?.to_dtype(DType::F32)?.to_scalar::<f32>()?;
assert!((max_idx as u32) < 50, "max index out of range: {max_idx}");
let min_idx = batch.indices.min(0)?.to_dtype(DType::F32)?.to_scalar::<f32>()?;
assert!(min_idx >= 0.0, "min index negative: {min_idx}");
// All indices must be in range [0, 50)
for (i, &idx) in indices_host.iter().enumerate() {
assert!(idx < 50, "index[{i}] out of range: {idx}");
}
// Shape sanity: correct number of sampled indices
assert_eq!(batch.indices.dims(), &[16]);
// Shape sanity
assert_eq!(indices_host.len(), 16);
Ok(())
}
#[tokio::test]
async fn test_gpu_replay_buffer_rank_based_sample_valid() -> anyhow::Result<()> {
let device = cuda_device();
let stream = smoke_stream();
let config = test_buffer_config(100, 48);
let mut buf = GpuReplayBuffer::new(config, &device)?;
let mut buf = GpuReplayBuffer::new(config, &stream)?;
insert_random_batch(&mut buf, 50, 48, &device)?;
insert_random_batch(&mut buf, 50, 48)?;
let batch: GpuBatch = buf.sample_proportional(16)?;
let batch = buf.sample_proportional(16)?;
assert_eq!(batch.states.dims(), &[16, 48]);
assert_eq!(batch.weights.dims(), &[16]);
assert_eq!(batch.batch_size, 16);
// All weights must be positive (GPU-resident: check min > 0)
let min_w = batch.weights.min(0)?.to_dtype(DType::F32)?.to_scalar::<f32>()?;
assert!(min_w > 0.0, "min weight is not positive: {min_w}");
// Download weights and indices to host for assertions
let weights_host = dtoh_f32(&batch.weights, &stream)?;
let indices_host = dtoh_u32(&batch.indices, &stream)?;
// Weights must be finite (GPU-resident: check sum is finite)
let w_sum = batch.weights.sum_all()?.to_dtype(DType::F32)?.to_scalar::<f32>()?;
// All weights must be positive
for (i, &w) in weights_host.iter().enumerate() {
assert!(w > 0.0, "weight[{i}] is not positive: {w}");
}
// Weights must be finite
let w_sum: f32 = weights_host.iter().sum();
assert!(w_sum.is_finite(), "weight sum is not finite: {w_sum}");
// All indices must be in range [0, 50) (GPU-resident: check max < 50, min >= 0)
let max_idx = batch.indices.max(0)?.to_dtype(DType::F32)?.to_scalar::<f32>()?;
assert!((max_idx as u32) < 50, "max index out of range: {max_idx}");
let min_idx = batch.indices.min(0)?.to_dtype(DType::F32)?.to_scalar::<f32>()?;
assert!(min_idx >= 0.0, "min index negative: {min_idx}");
// All indices must be in range [0, 50)
for (i, &idx) in indices_host.iter().enumerate() {
assert!(idx < 50, "index[{i}] out of range: {idx}");
}
// Shape sanity: correct number of sampled indices
assert_eq!(batch.indices.dims(), &[16]);
// Shape sanity
assert_eq!(indices_host.len(), 16);
Ok(())
}
#[tokio::test]
async fn test_gpu_replay_buffer_priority_update_valid() -> anyhow::Result<()> {
let device = cuda_device();
let stream = smoke_stream();
let config = test_buffer_config(100, 48);
let epsilon = config.epsilon;
let mut buf = GpuReplayBuffer::new(config, &device)?;
let mut buf = GpuReplayBuffer::new(config, &stream)?;
insert_random_batch(&mut buf, 50, 48, &device)?;
insert_random_batch(&mut buf, 50, 48)?;
// Sample and get indices
let batch = buf.sample_proportional(16)?;
let td_errors = Tensor::randn(0.0_f32, 1.0, &[16], &device)?;
buf.update_priorities_gpu(&batch.indices, &td_errors)?;
let td_errors = random_cuda_f32(16, &stream)?;
buf.update_priorities_gpu(&batch.indices, &td_errors, 16)?;
// Verify sampling still works with updated priorities (proves update succeeded)
let batch2 = buf.sample_proportional(16)?;
assert_eq!(batch2.weights.dims(), &[16]);
let weights_host = dtoh_f32(&batch2.weights, &stream)?;
// Weights should still be positive after priority update (GPU-resident)
let min_w = batch2.weights.min(0)?.to_dtype(DType::F32)?.to_scalar::<f32>()?;
assert!(min_w > 0.0, "min weight not positive after update: {min_w}");
let w_sum = batch2.weights.sum_all()?.to_dtype(DType::F32)?.to_scalar::<f32>()?;
// Weights should still be positive after priority update
for (i, &w) in weights_host.iter().enumerate() {
assert!(w > 0.0, "weight[{i}] not positive after update: {w}");
}
let w_sum: f32 = weights_host.iter().sum();
assert!(w_sum.is_finite(), "weight sum not finite after update: {w_sum}");
// Epsilon is the priority floor -- verify it's reasonable
@@ -157,28 +200,25 @@ async fn test_gpu_replay_buffer_priority_update_valid() -> anyhow::Result<()> {
#[tokio::test]
async fn test_searchsorted_gpu_correctness() -> anyhow::Result<()> {
let device = cuda_device();
let stream = smoke_stream();
let config = test_buffer_config(50, 4);
let mut buf = GpuReplayBuffer::new(config, &device)?;
let mut buf = GpuReplayBuffer::new(config, &stream)?;
// Fill with exactly 50 experiences
insert_random_batch(&mut buf, 50, 4, &device)?;
insert_random_batch(&mut buf, 50, 4)?;
assert_eq!(buf.len(), 50);
// Sample multiple times to exercise the searchsorted codepath (GPU-resident)
for trial in 0..5 {
let batch = buf.sample_proportional(10)?;
let max_idx = batch.indices.max(0)?.to_dtype(DType::F32)?.to_scalar::<f32>()?;
assert!(
(max_idx as u32) < 50,
"trial {trial}: max index {max_idx} out of range [0, 50)"
);
let min_idx = batch.indices.min(0)?.to_dtype(DType::F32)?.to_scalar::<f32>()?;
assert!(
min_idx >= 0.0,
"trial {trial}: min index {min_idx} is negative"
);
assert_eq!(batch.indices.dims(), &[10]);
let indices_host = dtoh_u32(&batch.indices, &stream)?;
for (i, &idx) in indices_host.iter().enumerate() {
assert!(
idx < 50,
"trial {trial}: index[{i}] = {idx} out of range [0, 50)"
);
}
assert_eq!(indices_host.len(), 10);
}
Ok(())
@@ -203,7 +243,7 @@ async fn test_train_step_produces_finite_metrics() -> anyhow::Result<()> {
#[tokio::test]
async fn test_gpu_training_dtype_bf16() -> anyhow::Result<()> {
let dtype = ml_core::native_types::NativeDType::BF16;
assert_eq!(dtype, DType::BF16, "CUDA should use BF16 training dtype");
assert_eq!(dtype, ml_core::native_types::NativeDType::BF16, "CUDA should use BF16 training dtype");
Ok(())
}
@@ -211,15 +251,16 @@ async fn test_gpu_training_dtype_bf16() -> anyhow::Result<()> {
#[tokio::test]
async fn test_gpu_training_dtype_diagnosis() -> anyhow::Result<()> {
let dev = cuda_device();
let dtype = ml_core::native_types::NativeDType::BF16;
info!(device = ?dev, training_dtype = ?dtype, "GPU training dtype");
// Verify a linear layer forward pass works with BF16 weights
let varmap = GpuVarStore::new();
let vs = GpuVarStore::from_varmap(&varmap, dtype, &dev);
let layer = GpuLinear::new(48, 32, vs.pp("test"))?;
let input = Tensor::zeros(&[2, 48], dtype, &dev)?;
let out = layer.forward(&input)?;
info!(dims = ?out.dims(), dtype = ?out.dtype(), "Linear forward OK");
let stream = Arc::clone(dev.cuda_stream().expect("cuda stream"));
info!(device = ?dev, "GPU training dtype");
// Verify a linear layer forward pass works
let mut vs = GpuVarStore::new(Arc::clone(&stream));
let layer = vs.linear("test", 48, 32)?;
let input = GpuTensor::zeros(&[2, 48], &stream)?;
let cublas = cudarc::cublas::CudaBlas::new(Arc::clone(&stream)).map_err(|e| anyhow::anyhow!("cuBLAS: {e}"))?;
let (out, _activations) = layer.forward(&input, &vs, &cublas, &stream)?;
info!(dims = ?out.dims(), "Linear forward OK");
Ok(())
}
@@ -251,31 +292,43 @@ async fn test_training_rejects_missing_gpu_collector() -> anyhow::Result<()> {
Ok(())
}
/// Diagnose: can we create + step an AdamW optimizer on GPU with BF16?
/// Diagnose: can we create + step an AdamW optimizer on GPU?
#[tokio::test]
async fn test_gpu_adamw_creation() -> anyhow::Result<()> {
let dev = cuda_device();
let dtype = ml_core::native_types::NativeDType::BF16;
info!(device = ?dev, dtype = ?dtype, "GPU device and training dtype");
let stream = Arc::clone(dev.cuda_stream().expect("cuda stream"));
info!(device = ?dev, "GPU device");
let varmap = GpuVarStore::new();
let vs = GpuVarStore::from_varmap(&varmap, dtype, &dev);
let layer = GpuLinear::new(48, 32, vs.pp("test"))?;
let mut vs = GpuVarStore::new(Arc::clone(&stream));
let layer = vs.linear("test", 48, 32)?;
// Forward + backward
let input = Tensor::zeros(&[2, 48], dtype, &dev)?;
let out = layer.forward(&input)?;
let loss = out.sqr()?.mean_all()?;
let grads = loss.backward()?;
// Forward
let input = GpuTensor::zeros(&[2, 48], &stream)?;
let cublas = cudarc::cublas::CudaBlas::new(Arc::clone(&stream)).map_err(|e| anyhow::anyhow!("cuBLAS: {e}"))?;
let (out, _activations) = layer.forward(&input, &vs, &cublas, &stream)?;
let loss_val = out.mean_all(&stream)?;
info!(loss = loss_val, "Forward OK");
// Create optimizer
let params = ml_core::cuda_autograd::AdamWConfig {
let adamw_config = AdamWConfig {
lr: 1e-4,
..Default::default()
};
let mut opt = GpuAdamW::new(varmap.all_vars(), params)?;
opt.step(&grads)?;
let mut opt = GpuAdamW::new(adamw_config, Arc::clone(&stream))?;
// Compute gradients: use a simple synthetic gradient (same shape as params)
let mut grads = std::collections::BTreeMap::new();
for name in vs.param_names() {
if let Some(param) = vs.get(name) {
let grad = GpuTensor::from_host(
&vec![0.01_f32; param.data.len()],
param.shape.clone(),
&stream,
)?;
grads.insert(name.to_owned(), grad);
}
}
let _grad_norm = opt.step(&mut vs, &grads)?;
info!("AdamW step OK");
Ok(())
}

13
crates/ml/src/trainers/dqn/smoke_tests/helpers.rs
View File

@@ -59,19 +59,20 @@ pub(super) fn smoke_params() -> DQNHyperparameters {
/// Shared CUDA device for all smoke tests — reuses one cuBLAS handle
/// instead of creating a new one per test, preventing handle pool exhaustion.
static SMOKE_CUDA: std::sync::OnceLock<Device> = std::sync::OnceLock::new();
static SMOKE_CUDA: std::sync::OnceLock<MlDevice> = std::sync::OnceLock::new();
/// Get CUDA device. Panics if no GPU is available.
/// Tests a small matmul to detect broken CUDA runtime before trainer creation.
pub(super) fn cuda_device() -> Device {
pub(super) fn cuda_device() -> MlDevice {
init_tracing();
SMOKE_CUDA
.get_or_init(|| {
let dev =
Device::new_cuda(0).expect("CUDA device required — no CPU fallback in smoke tests");
let probe = GpuTensor::zeros(&[2, 2], ml_core::native_types::NativeDType::F32, &dev)
.and_then(|t| t.matmul(&t));
probe.expect("CUDA matmul probe failed — GPU runtime is broken");
MlDevice::new_cuda(0).expect("CUDA device required — no CPU fallback in smoke tests");
// Probe: allocate a small tensor to verify GPU runtime works
let stream = dev.cuda_stream().expect("cuda stream");
let probe = ml_core::cuda_autograd::GpuTensor::zeros(&[2, 2], stream);
probe.expect("CUDA alloc probe failed — GPU runtime is broken");
dev
})
.clone()

31
crates/ml/src/trainers/dqn/smoke_tests/performance.rs
View File

@@ -9,6 +9,8 @@
use super::helpers::{assert_finite, smoke_params, test_data_dir};
use crate::trainers::dqn::DQNTrainer;
use ml_core::device::MlDevice;
use std::sync::Arc;
use std::time::Instant;
use tracing::info;
@@ -50,11 +52,10 @@ async fn test_training_throughput_measurement() -> anyhow::Result<()> {
#[tokio::test]
#[ignore] // Run manually on GPU
async fn test_per_sample_latency() -> anyhow::Result<()> {
use ml_core::device::MlDevice;
use ml_core::cuda_autograd::GpuTensor;
use crate::dqn::gpu_replay_buffer::{GpuReplayBuffer, GpuReplayBufferConfig};
let device = Device::new_cuda(0)?;
let dev = MlDevice::new_cuda(0)?;
let stream = Arc::clone(dev.cuda_stream()?);
let state_dim: usize = 48;
let capacity: usize = 10_000;
let config = GpuReplayBufferConfig {
@@ -67,23 +68,23 @@ use ml_core::cuda_autograd::GpuTensor;
epsilon: 1e-6,
max_memory_bytes: 4 * 1024 * 1024 * 1024,
};
let mut buf = GpuReplayBuffer::new(config, &device)?;
let mut buf = GpuReplayBuffer::new(config, &stream)?;
// Fill with 5000 experiences
let fill_count: usize = 5000;
let batch_insert: usize = 100;
for _ in 0..(fill_count / batch_insert) {
let states = Tensor::randn(0.0_f32, 1.0, &[batch_insert, state_dim], &device)?;
let next_states = Tensor::randn(0.0_f32, 1.0, &[batch_insert, state_dim], &device)?;
let actions = Tensor::zeros(&[batch_insert], DType::U32, &device)?;
let rewards = Tensor::randn(0.0_f32, 1.0, &[batch_insert], &device)?;
let dones = Tensor::zeros(&[batch_insert], DType::F32, &device)?;
let sf = crate::cuda_pipeline::tensor_to_cuda_slice_f32(&states)?;
let nsf = crate::cuda_pipeline::tensor_to_cuda_slice_f32(&next_states)?;
let af = crate::cuda_pipeline::tensor_to_cuda_slice_u32(&actions)?;
let rf = crate::cuda_pipeline::tensor_to_cuda_slice_f32(&rewards)?;
let df = crate::cuda_pipeline::tensor_to_cuda_slice_f32(&dones)?;
buf.insert_batch(&sf, &nsf, &af, &rf, &df, batch_insert)?;
use rand::Rng;
let mut rng = rand::thread_rng();
let host_states: Vec<f32> = (0..batch_insert * state_dim).map(|_| rng.gen_range(-1.0_f32..1.0)).collect();
let host_rewards: Vec<f32> = (0..batch_insert).map(|_| rng.gen_range(-1.0_f32..1.0)).collect();
let states = stream.clone_htod(&host_states).map_err(|e| anyhow::anyhow!("{e}"))?;
let next_states = stream.clone_htod(&host_states).map_err(|e| anyhow::anyhow!("{e}"))?;
let actions = stream.alloc_zeros::<u32>(batch_insert).map_err(|e| anyhow::anyhow!("{e}"))?;
let rewards = stream.clone_htod(&host_rewards).map_err(|e| anyhow::anyhow!("{e}"))?;
let dones = stream.alloc_zeros::<f32>(batch_insert).map_err(|e| anyhow::anyhow!("{e}"))?;
buf.insert_batch(&states, &next_states, &actions, &rewards, &dones, batch_insert)?;
}
assert_eq!(buf.len(), fill_count);

73
crates/ml/src/trainers/dqn/smoke_tests/training_stability.rs
View File

@@ -42,10 +42,11 @@ async fn test_production_training_stability() -> anyhow::Result<()> {
/// PER importance-sampling weights must be finite and positive.
#[tokio::test]
async fn test_per_weights_valid() -> anyhow::Result<()> {
use ml_core::cuda_autograd::GpuTensor;
use crate::dqn::gpu_replay_buffer::{GpuReplayBuffer, GpuReplayBufferConfig};
use std::sync::Arc;
let device = cuda_device();
let dev = cuda_device();
let stream = Arc::clone(dev.cuda_stream().expect("cuda stream"));
let config = GpuReplayBufferConfig {
capacity: 128,
state_dim: 45,
@@ -56,26 +57,28 @@ async fn test_per_weights_valid() -> anyhow::Result<()> {
epsilon: 1e-6,
max_memory_bytes: 4 * 1024 * 1024 * 1024,
};
let mut buf = GpuReplayBuffer::new(config, &device)?;
let mut buf = GpuReplayBuffer::new(config, &stream)?;
// Insert 50 single experiences
for _ in 0..50 {
let s = Tensor::randn(0.0_f32, 1.0, &[1, 45], &device)?;
let ns = Tensor::randn(0.0_f32, 1.0, &[1, 45], &device)?;
let a = Tensor::zeros(&[1], DType::U32, &device)?;
let r = Tensor::randn(0.0_f32, 1.0, &[1], &device)?;
let d = Tensor::zeros(&[1], DType::F32, &device)?;
let sf = crate::cuda_pipeline::tensor_to_cuda_slice_f32(&s)?;
let nsf = crate::cuda_pipeline::tensor_to_cuda_slice_f32(&ns)?;
let af = crate::cuda_pipeline::tensor_to_cuda_slice_u32(&a)?;
let rf = crate::cuda_pipeline::tensor_to_cuda_slice_f32(&r)?;
let df = crate::cuda_pipeline::tensor_to_cuda_slice_f32(&d)?;
buf.insert_batch(&sf, &nsf, &af, &rf, &df, 1)?;
let states = stream.clone_htod(&vec![0.0_f32; 45]).map_err(|e| anyhow::anyhow!("{e}"))?;
let next_states = stream.clone_htod(&vec![0.0_f32; 45]).map_err(|e| anyhow::anyhow!("{e}"))?;
let actions = stream.alloc_zeros::<u32>(1).map_err(|e| anyhow::anyhow!("{e}"))?;
let rewards = stream.clone_htod(&[0.5_f32]).map_err(|e| anyhow::anyhow!("{e}"))?;
let dones = stream.alloc_zeros::<f32>(1).map_err(|e| anyhow::anyhow!("{e}"))?;
buf.insert_batch(&states, &next_states, &actions, &rewards, &dones, 1)?;
}
let batch = buf.sample_proportional(16)?;
let min_w = batch.weights.min(0)?.to_dtype(DType::F32)?.to_scalar::<f32>()?;
assert!(min_w > 0.0, "min weight not positive: {min_w}");
let w_sum = batch.weights.sum_all()?.to_dtype(DType::F32)?.to_scalar::<f32>()?;
// Download weights to host for assertion
let mut weights_host = vec![0.0_f32; batch.weights.len()];
stream.memcpy_dtoh(&batch.weights, &mut weights_host).map_err(|e| anyhow::anyhow!("{e}"))?;
for (i, &w) in weights_host.iter().enumerate() {
assert!(w > 0.0, "weight[{i}] not positive: {w}");
}
let w_sum: f32 = weights_host.iter().sum();
assert!(w_sum.is_finite(), "weight sum not finite: {w_sum}");
Ok(())
}
@@ -83,10 +86,11 @@ async fn test_per_weights_valid() -> anyhow::Result<()> {
/// PER indices must be within buffer bounds.
#[tokio::test]
async fn test_per_indices_valid() -> anyhow::Result<()> {
use ml_core::cuda_autograd::GpuTensor;
use crate::dqn::gpu_replay_buffer::{GpuReplayBuffer, GpuReplayBufferConfig};
use std::sync::Arc;
let device = cuda_device();
let dev = cuda_device();
let stream = Arc::clone(dev.cuda_stream().expect("cuda stream"));
let config = GpuReplayBufferConfig {
capacity: 128,
state_dim: 45,
@@ -97,26 +101,27 @@ async fn test_per_indices_valid() -> anyhow::Result<()> {
epsilon: 1e-6,
max_memory_bytes: 4 * 1024 * 1024 * 1024,
};
let mut buf = GpuReplayBuffer::new(config, &device)?;
let mut buf = GpuReplayBuffer::new(config, &stream)?;
for _ in 0..50 {
let s = Tensor::randn(0.0_f32, 1.0, &[1, 45], &device)?;
let ns = Tensor::randn(0.0_f32, 1.0, &[1, 45], &device)?;
let a = Tensor::zeros(&[1], DType::U32, &device)?;
let r = Tensor::randn(0.0_f32, 1.0, &[1], &device)?;
let d = Tensor::zeros(&[1], DType::F32, &device)?;
let sf = crate::cuda_pipeline::tensor_to_cuda_slice_f32(&s)?;
let nsf = crate::cuda_pipeline::tensor_to_cuda_slice_f32(&ns)?;
let af = crate::cuda_pipeline::tensor_to_cuda_slice_u32(&a)?;
let rf = crate::cuda_pipeline::tensor_to_cuda_slice_f32(&r)?;
let df = crate::cuda_pipeline::tensor_to_cuda_slice_f32(&d)?;
buf.insert_batch(&sf, &nsf, &af, &rf, &df, 1)?;
let states = stream.clone_htod(&vec![0.0_f32; 45]).map_err(|e| anyhow::anyhow!("{e}"))?;
let next_states = stream.clone_htod(&vec![0.0_f32; 45]).map_err(|e| anyhow::anyhow!("{e}"))?;
let actions = stream.alloc_zeros::<u32>(1).map_err(|e| anyhow::anyhow!("{e}"))?;
let rewards = stream.clone_htod(&[0.5_f32]).map_err(|e| anyhow::anyhow!("{e}"))?;
let dones = stream.alloc_zeros::<f32>(1).map_err(|e| anyhow::anyhow!("{e}"))?;
buf.insert_batch(&states, &next_states, &actions, &rewards, &dones, 1)?;
}
let batch = buf.sample_proportional(16)?;
let max_idx = batch.indices.max(0)?.to_dtype(DType::F32)?.to_scalar::<f32>()?;
assert!((max_idx as u32) < 50, "max index out of bounds: {max_idx}");
assert_eq!(batch.indices.dims(), &[16]);
// Download indices to host for assertion
let mut indices_host = vec![0_u32; batch.indices.len()];
stream.memcpy_dtoh(&batch.indices, &mut indices_host).map_err(|e| anyhow::anyhow!("{e}"))?;
for (i, &idx) in indices_host.iter().enumerate() {
assert!(idx < 50, "index[{i}] out of bounds: {idx}");
}
assert_eq!(indices_host.len(), 16);
Ok(())
}

6
crates/ml/src/trainers/dqn/trainer/tests.rs
View File

@@ -10,9 +10,9 @@ use std::sync::OnceLock;
/// 400+ tests doing this in rapid succession (even with --test-threads=1),
/// the driver's internal handle pool is exhausted. Sharing one device
/// eliminates the churn entirely.
static SHARED_DEVICE: OnceLock<Device> = OnceLock::new();
static SHARED_DEVICE: OnceLock<MlDevice> = OnceLock::new();
fn shared_cuda_device() -> Device {
fn shared_cuda_device() -> MlDevice {
// Initialize tracing so kernel compilation/launch logs are visible.
static TRACING_INIT: std::sync::Once = std::sync::Once::new();
TRACING_INIT.call_once(|| {
@@ -26,7 +26,7 @@ fn shared_cuda_device() -> Device {
});
SHARED_DEVICE
.get_or_init(|| {
Device::new_cuda(0).expect("CUDA required for tests")
MlDevice::new_cuda(0).expect("CUDA required for tests")
})
.clone()
}

11
crates/ml/src/trainers/online_learning.rs
View File

@@ -517,16 +517,9 @@ mod tests {
/// Build a tiny GpuVarStore with a single 2x2 parameter for testing.
fn tiny_var_map() -> Result<GpuVarStore, MLError> {
let dev = cuda_device();
let vs = GpuVarStore::with_device(&dev)
let mut vs = GpuVarStore::with_device(&dev)
.map_err(|e| MLError::ModelError(format!("tiny_var_map device: {e}")))?;
// Register a small linear layer (weight 2x2 + bias 2)
let stream = vs.stream();
let w_data = cudarc::driver::CudaSlice::from(stream.alloc_zeros::<f32>(4)
.map_err(|e| MLError::ModelError(format!("alloc w: {e}")))?);
let b_data = stream.alloc_zeros::<f32>(2)
.map_err(|e| MLError::ModelError(format!("alloc b: {e}")))?;
// We can't easily register params without the linear() helper, so just
// use the var_store's linear method.
// Register a small linear layer (weight 2x2 + bias 2) via the var_store helper
let _linear = vs.linear("layer.weight", 2, 2)
.map_err(|e| MLError::ModelError(format!("tiny_var_map linear: {e}")))?;
Ok(vs)

6
crates/ml/src/trainers/tft/tests.rs
View File

@@ -10,8 +10,8 @@ use std::path::PathBuf;
use std::sync::Arc;
use tracing::info;
fn cuda_device() -> Device {
Device::new_cuda(0).expect("CUDA device required")
fn cuda_device() -> MlDevice {
MlDevice::new_cuda(0).expect("CUDA device required")
}
#[tokio::test]
@@ -191,7 +191,7 @@ async fn test_sync_cuda_device_cpu() {
#[tokio::test]
async fn test_sync_cuda_device_gpu() {
// Test CUDA sync on GPU device
let device = Device::cuda_if_available(0).expect("CUDA not available");
let device = MlDevice::cuda_if_available(0);
if !device.is_cuda() {
info!("Skipping CUDA sync test - GPU not available");
return;

6
crates/ml/src/trainers/tlob.rs
View File

@@ -27,7 +27,7 @@ use std::time::Instant;
use anyhow::{Context, Result};
use ml_core::device::MlDevice;
use ml_core::cuda_autograd::GpuTensor;
use ml_core::cuda_autograd::{GpuVarStore, GpuAdamW};
use ml_core::cuda_autograd::GpuVarStore;
use serde::{Deserialize, Serialize};
use tokio::sync::RwLock;
use tracing::{info, instrument, warn};
@@ -672,7 +672,9 @@ mod tests {
let trainer = TLOBTrainer::new(hyperparams, &temp_dir, true).unwrap();
let sequences = trainer.generate_dummy_sequences(4).unwrap();
let (input_tensor, target_tensor) = trainer.prepare_batch(&sequences).unwrap();
let dev = MlDevice::cuda(0).expect("CUDA required");
let stream = dev.cuda_stream().expect("stream");
let (input_tensor, target_tensor) = trainer.prepare_batch(&sequences, stream).unwrap();
assert_eq!(input_tensor.dims(), &[4, 128, TLOB_FEATURE_COUNT]);
assert_eq!(target_tensor.dims(), &[4, 1]);

6
crates/ml/src/transformers/attention.rs
View File

@@ -62,8 +62,10 @@ mod tests {
let mask = AttentionMask::causal(4, &device)?;
assert_eq!(mask.mask.dims(), &[4, 4]);
// Verify the diagonal and above-diagonal values
let flat = mask.mask.flatten_all()?.to_vec1::<f32>()?;
// Download mask to host for verification
let ml_dev = MlDevice::cuda(0)?;
let stream = ml_dev.cuda_stream()?;
let flat = mask.mask.to_host(stream)?;
// Row 0: [0, -inf, -inf, -inf]
assert_eq!(*flat.get(0).unwrap_or(&f32::NAN), 0.0);
assert!(flat.get(1).unwrap_or(&0.0).is_infinite());

8
crates/ml/src/validation/adapters.rs
View File

@@ -50,7 +50,7 @@ impl DqnStrategy {
///
/// Shares the caller's device instead of creating a new cuBLAS handle.
#[cfg(test)]
pub(crate) fn new_on_device(config: DQNConfig, device: NativeDevice) -> Result<Self, MLError> {
pub(crate) fn new_on_device(config: DQNConfig, device: ml_core::device::MlDevice) -> Result<Self, MLError> {
let dqn = DQN::new_on_device(config.clone(), device)?;
Ok(Self {
config,
@@ -176,12 +176,12 @@ mod tests {
use chrono::{TimeZone, Utc};
use std::sync::OnceLock;
static SHARED_CUDA: OnceLock<NativeDevice> = OnceLock::new();
static SHARED_CUDA: OnceLock<ml_core::device::MlDevice> = OnceLock::new();
fn shared_device() -> NativeDevice {
fn shared_device() -> ml_core::device::MlDevice {
SHARED_CUDA
.get_or_init(|| {
NativeDevice::Cuda(0).expect("CUDA required")
ml_core::device::MlDevice::cuda(0).expect("CUDA required for tests")
})
.clone()
}