diff --git a/crates/ml/src/ensemble/adapters/dqn.rs b/crates/ml/src/ensemble/adapters/dqn.rs index 1b58c057e..8b42ae0ac 100644 --- a/crates/ml/src/ensemble/adapters/dqn.rs +++ b/crates/ml/src/ensemble/adapters/dqn.rs @@ -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, Vec)> = host_data + .into_iter() + .map(|(name, (shape, data))| { + let bytes: Vec = 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(), ) diff --git a/crates/ml/src/features/multi_timeframe.rs b/crates/ml/src/features/multi_timeframe.rs index 3730fecd2..7e7ef983b 100644 --- a/crates/ml/src/features/multi_timeframe.rs +++ b/crates/ml/src/features/multi_timeframe.rs @@ -367,6 +367,7 @@ fn push_ring(buf: &mut VecDeque, 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::()) - .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::()) - .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. diff --git a/crates/ml/src/flash_attention/mod.rs b/crates/ml/src/flash_attention/mod.rs index d43752fd7..1e599c83f 100644 --- a/crates/ml/src/flash_attention/mod.rs +++ b/crates/ml/src/flash_attention/mod.rs @@ -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(()) diff --git a/crates/ml/src/hyperopt/adapters/tft.rs b/crates/ml/src/hyperopt/adapters/tft.rs index 91e6f0fe6..59e98e578 100644 --- a/crates/ml/src/hyperopt/adapters/tft.rs +++ b/crates/ml/src/hyperopt/adapters/tft.rs @@ -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(), diff --git a/crates/ml/src/inference.rs b/crates/ml/src/inference.rs index 8e338bb18..d189717e2 100644 --- a/crates/ml/src/inference.rs +++ b/crates/ml/src/inference.rs @@ -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(); diff --git a/crates/ml/src/liquid/adapter.rs b/crates/ml/src/liquid/adapter.rs index 1c4336a97..a8810e945 100644 --- a/crates/ml/src/liquid/adapter.rs +++ b/crates/ml/src/liquid/adapter.rs @@ -109,6 +109,11 @@ impl LiquidTrainableAdapter { } /// Access the GpuVarStore (read-only). + /// Reference to the CUDA stream handle. + pub fn stream(&self) -> &Arc { + &self.stream + } + pub fn varmap(&self) -> &GpuVarStore { &self.var_store } diff --git a/crates/ml/src/mamba/trainable_adapter.rs b/crates/ml/src/mamba/trainable_adapter.rs index 7ccf851f4..86ed5f24c 100644 --- a/crates/ml/src/mamba/trainable_adapter.rs +++ b/crates/ml/src/mamba/trainable_adapter.rs @@ -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"); diff --git a/crates/ml/src/trainers/dqn/smoke_tests/gpu_residency.rs b/crates/ml/src/trainers/dqn/smoke_tests/gpu_residency.rs index 135bf553a..1dd15ff67 100644 --- a/crates/ml/src/trainers/dqn/smoke_tests/gpu_residency.rs +++ b/crates/ml/src/trainers/dqn/smoke_tests/gpu_residency.rs @@ -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 { + 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) -> anyhow::Result> { + use rand::Rng; + let mut rng = rand::thread_rng(); + let host: Vec = (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) -> anyhow::Result> { + let slice = stream.alloc_zeros::(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) -> anyhow::Result> { + let slice = stream.alloc_zeros::(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 to host for assertion. +fn dtoh_f32(slice: &cudarc::driver::CudaSlice, stream: &Arc) -> anyhow::Result> { + 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 to host for assertion. +fn dtoh_u32(slice: &cudarc::driver::CudaSlice, stream: &Arc) -> anyhow::Result> { + 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::()?; - 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::()?; + // 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::()?; - 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::()?; - 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::()?; - 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::()?; + // 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::()?; - 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::()?; - 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::()?; - 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::()?; + // 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::()?; - 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::()?; - 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(()) } diff --git a/crates/ml/src/trainers/dqn/smoke_tests/helpers.rs b/crates/ml/src/trainers/dqn/smoke_tests/helpers.rs index 163abef42..c4f538e7d 100644 --- a/crates/ml/src/trainers/dqn/smoke_tests/helpers.rs +++ b/crates/ml/src/trainers/dqn/smoke_tests/helpers.rs @@ -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 = std::sync::OnceLock::new(); +static SMOKE_CUDA: std::sync::OnceLock = 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() diff --git a/crates/ml/src/trainers/dqn/smoke_tests/performance.rs b/crates/ml/src/trainers/dqn/smoke_tests/performance.rs index 0f04ffa86..98f0f0b4a 100644 --- a/crates/ml/src/trainers/dqn/smoke_tests/performance.rs +++ b/crates/ml/src/trainers/dqn/smoke_tests/performance.rs @@ -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 = (0..batch_insert * state_dim).map(|_| rng.gen_range(-1.0_f32..1.0)).collect(); + let host_rewards: Vec = (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::(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::(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); diff --git a/crates/ml/src/trainers/dqn/smoke_tests/training_stability.rs b/crates/ml/src/trainers/dqn/smoke_tests/training_stability.rs index ab80d64ae..01835d4b1 100644 --- a/crates/ml/src/trainers/dqn/smoke_tests/training_stability.rs +++ b/crates/ml/src/trainers/dqn/smoke_tests/training_stability.rs @@ -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::(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::(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::()?; - assert!(min_w > 0.0, "min weight not positive: {min_w}"); - let w_sum = batch.weights.sum_all()?.to_dtype(DType::F32)?.to_scalar::()?; + + // 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::(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::(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::()?; - 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(()) } diff --git a/crates/ml/src/trainers/dqn/trainer/tests.rs b/crates/ml/src/trainers/dqn/trainer/tests.rs index 91955af26..a0820eeb6 100644 --- a/crates/ml/src/trainers/dqn/trainer/tests.rs +++ b/crates/ml/src/trainers/dqn/trainer/tests.rs @@ -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 = OnceLock::new(); +static SHARED_DEVICE: OnceLock = 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() } diff --git a/crates/ml/src/trainers/online_learning.rs b/crates/ml/src/trainers/online_learning.rs index e09dc432a..5116a323b 100644 --- a/crates/ml/src/trainers/online_learning.rs +++ b/crates/ml/src/trainers/online_learning.rs @@ -517,16 +517,9 @@ mod tests { /// Build a tiny GpuVarStore with a single 2x2 parameter for testing. fn tiny_var_map() -> Result { 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::(4) - .map_err(|e| MLError::ModelError(format!("alloc w: {e}")))?); - let b_data = stream.alloc_zeros::(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) diff --git a/crates/ml/src/trainers/tft/tests.rs b/crates/ml/src/trainers/tft/tests.rs index 379515407..7d5f58727 100644 --- a/crates/ml/src/trainers/tft/tests.rs +++ b/crates/ml/src/trainers/tft/tests.rs @@ -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; diff --git a/crates/ml/src/trainers/tlob.rs b/crates/ml/src/trainers/tlob.rs index f76ec2de6..ae52324fa 100644 --- a/crates/ml/src/trainers/tlob.rs +++ b/crates/ml/src/trainers/tlob.rs @@ -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]); diff --git a/crates/ml/src/transformers/attention.rs b/crates/ml/src/transformers/attention.rs index b5e09ac81..89f959901 100644 --- a/crates/ml/src/transformers/attention.rs +++ b/crates/ml/src/transformers/attention.rs @@ -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::()?; + // 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()); diff --git a/crates/ml/src/validation/adapters.rs b/crates/ml/src/validation/adapters.rs index d0f06bb35..03b67017d 100644 --- a/crates/ml/src/validation/adapters.rs +++ b/crates/ml/src/validation/adapters.rs @@ -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 { + pub(crate) fn new_on_device(config: DQNConfig, device: ml_core::device::MlDevice) -> Result { 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 = OnceLock::new(); + static SHARED_CUDA: OnceLock = 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() }