feat(mamba): wire real GPU Bernoulli dropout in supervised Mamba2
Supervised Mamba2's dropout_rate field was configured and stored but never applied — comments said "simulated" but the arithmetic was absent. An operator tuning dropout_rate upward got zero regularisation. New dropout_kernel.cu with a forward-only Bernoulli dropout (Philox- seeded, deterministic, in-place). New build.rs matching ml-dqn's pattern. Applied in forward_with_gradients (training path) only; the `forward()` path used by validate/predict/SPSA stays deterministic so SPSA's ±ε finite-difference estimator is not destabilised. NOT a DQN fix: DQN has its own regime_dropout kernel in cuda_pipeline/experience_kernels.cu and its own native mamba2_step in gpu_dqn_trainer.rs; DQN does not call into Mamba2SSM. This commit affects only the supervised Mamba2 trainer and its hyperopt adapter. Tests: determinism, eval-mode identity, ctr-increment divergence. DQN smoke tests (magnitude_distribution, multi_fold_convergence) are unaffected by this change — ran them for regression assurance only. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
149
crates/ml-supervised/build.rs
Normal file
149
crates/ml-supervised/build.rs
Normal file
@@ -0,0 +1,149 @@
|
||||
//! Precompile CUDA kernels owned by `ml-supervised`.
|
||||
//!
|
||||
//! Modelled on `crates/ml-dqn/build.rs` (the canonical reference). Emits
|
||||
//! `.cubin` files into `$OUT_DIR`; Rust callers `include_bytes!` them via
|
||||
//! `concat!(env!("OUT_DIR"), "/<name>.cubin")` and load through
|
||||
//! `stream.context().load_cubin(...)`.
|
||||
//!
|
||||
//! Kernels compiled here:
|
||||
//! - `mamba/dropout_kernel.cu` → `dropout_kernel.cubin` (forward-only
|
||||
//! Bernoulli inverted dropout used by `Mamba2SSM::forward_with_gradients`).
|
||||
|
||||
use std::path::{Path, PathBuf};
|
||||
use std::process::Command;
|
||||
|
||||
fn main() {
|
||||
println!("cargo:rerun-if-changed=build.rs");
|
||||
|
||||
// Only compile CUDA kernels when the cuda feature is enabled (matches
|
||||
// ml-dqn's gating exactly). Non-CUDA builds skip kernel compilation.
|
||||
if std::env::var("CARGO_FEATURE_CUDA").is_err() {
|
||||
return;
|
||||
}
|
||||
|
||||
let out_dir = PathBuf::from(std::env::var("OUT_DIR").unwrap());
|
||||
let manifest_dir = PathBuf::from(std::env::var("CARGO_MANIFEST_DIR").unwrap());
|
||||
|
||||
// Detect GPU architecture from env or default to sm_80 (L40S / A100).
|
||||
let cuda_compute_cap = std::env::var("CUDA_COMPUTE_CAP").unwrap_or_else(|_| "80".to_string());
|
||||
let arch = format!("sm_{cuda_compute_cap}");
|
||||
|
||||
// Check if nvcc is available.
|
||||
let nvcc = match find_nvcc() {
|
||||
Some(p) => p,
|
||||
None => {
|
||||
eprintln!(" warning: nvcc not found, skipping ml-supervised CUDA kernel precompilation");
|
||||
eprintln!(" Install CUDA toolkit or set CUDA_HOME for GPU builds");
|
||||
return;
|
||||
}
|
||||
};
|
||||
|
||||
// Kernels owned by ml-supervised. Each entry is (source path relative to
|
||||
// crate root, output cubin base name). The kernel is self-contained and
|
||||
// does NOT consume ml/src/cuda_pipeline/common_device_functions.cuh —
|
||||
// that header carries DQN-specific constants (DQN_NUM_ACTIONS,
|
||||
// STATE_DIM, ...) that are irrelevant here.
|
||||
let kernels: &[(&str, &str)] = &[
|
||||
("src/mamba/dropout_kernel.cu", "dropout_kernel.cu"),
|
||||
];
|
||||
|
||||
let mut failed: Vec<&str> = Vec::new();
|
||||
|
||||
for (src_path, kernel_name) in kernels {
|
||||
let full_src = manifest_dir.join(src_path);
|
||||
if !try_compile_kernel(&nvcc, &full_src, kernel_name, &arch, &out_dir) {
|
||||
failed.push(kernel_name);
|
||||
}
|
||||
}
|
||||
|
||||
let total = kernels.len();
|
||||
let passed = total - failed.len();
|
||||
eprintln!(
|
||||
" ml-supervised: Precompiled {passed}/{total} CUDA kernels ({arch})",
|
||||
);
|
||||
if !failed.is_empty() {
|
||||
eprintln!(" FAILED: {}", failed.join(", "));
|
||||
panic!(
|
||||
"nvcc failed to compile {} ml-supervised kernel(s): {}",
|
||||
failed.len(),
|
||||
failed.join(", ")
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
/// Compile a single standalone `.cu` kernel to a `.cubin` via nvcc.
|
||||
///
|
||||
/// No header prepending — the kernel is expected to be self-contained.
|
||||
fn try_compile_kernel(
|
||||
nvcc: &Path,
|
||||
kernel_path: &Path,
|
||||
kernel_name: &str,
|
||||
arch: &str,
|
||||
out_dir: &Path,
|
||||
) -> bool {
|
||||
println!("cargo:rerun-if-changed={}", kernel_path.display());
|
||||
|
||||
if !kernel_path.exists() {
|
||||
eprintln!(" FAILED: {kernel_name} (source not found at {})", kernel_path.display());
|
||||
return false;
|
||||
}
|
||||
|
||||
let cubin_name = kernel_name.replace(".cu", ".cubin");
|
||||
let cubin_path = out_dir.join(&cubin_name);
|
||||
|
||||
let status = Command::new(nvcc)
|
||||
.args([
|
||||
"-cubin",
|
||||
&format!("-arch={arch}"),
|
||||
"-O3",
|
||||
"--ftz=true",
|
||||
"--fmad=true",
|
||||
"--prec-div=true",
|
||||
"--prec-sqrt=true",
|
||||
"-o",
|
||||
cubin_path.to_str().unwrap(),
|
||||
kernel_path.to_str().unwrap(),
|
||||
])
|
||||
.status();
|
||||
|
||||
match status {
|
||||
Ok(s) if s.success() => {
|
||||
eprintln!(" Compiled {kernel_name} -> {cubin_name} ({arch})");
|
||||
true
|
||||
}
|
||||
Ok(s) => {
|
||||
eprintln!(
|
||||
" FAILED: {kernel_name} (exit={})",
|
||||
s.code().unwrap_or(-1)
|
||||
);
|
||||
false
|
||||
}
|
||||
Err(e) => {
|
||||
eprintln!(" FAILED: {kernel_name} (nvcc error: {e})");
|
||||
false
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Find nvcc: prefer `$CUDA_HOME/bin/nvcc`, then check common install paths,
|
||||
/// then PATH.
|
||||
fn find_nvcc() -> Option<PathBuf> {
|
||||
if let Ok(home) = std::env::var("CUDA_HOME") {
|
||||
let nvcc = PathBuf::from(home).join("bin/nvcc");
|
||||
if nvcc.exists() {
|
||||
return Some(nvcc);
|
||||
}
|
||||
}
|
||||
|
||||
for path in &["/usr/local/cuda/bin/nvcc", "/usr/bin/nvcc"] {
|
||||
let p = PathBuf::from(path);
|
||||
if p.exists() {
|
||||
return Some(p);
|
||||
}
|
||||
}
|
||||
|
||||
match Command::new("nvcc").arg("--version").output() {
|
||||
Ok(output) if output.status.success() => Some(PathBuf::from("nvcc")),
|
||||
_ => None,
|
||||
}
|
||||
}
|
||||
71
crates/ml-supervised/src/mamba/dropout_kernel.cu
Normal file
71
crates/ml-supervised/src/mamba/dropout_kernel.cu
Normal file
@@ -0,0 +1,71 @@
|
||||
// Bernoulli inverted dropout — forward-only, deterministic, in-place.
|
||||
//
|
||||
// Philox-style stateless hash parameterised by (seed, ctr, layer_idx, i) so
|
||||
// the same (seed, ctr, layer_idx) always yields the same mask, and distinct
|
||||
// layers never alias (they carry different layer_idx into the hash tuple).
|
||||
//
|
||||
// Inverted dropout: surviving elements are scaled by 1/keep_prob so the
|
||||
// expected value matches the eval-time (dropout-off) activation.
|
||||
//
|
||||
// Used only by `Mamba2SSM::forward_with_gradients` (training path). The pure
|
||||
// `forward()` used by validate / predict_single_fast / SPSA's ±ε probes is
|
||||
// left untouched so SPSA's finite-difference estimator stays stable.
|
||||
|
||||
/* Stateless Philox-like mixing hash. Same 32-bit avalanche used by the DQN
|
||||
* experience kernels' philox_uniform, extended to fold in a 64-bit seed,
|
||||
* 64-bit counter, layer index, and element index. Returns a uniform f32 in
|
||||
* [0, 1). */
|
||||
__device__ __forceinline__ float mamba_dropout_uniform(
|
||||
unsigned int seed_lo,
|
||||
unsigned int seed_hi,
|
||||
unsigned int ctr_lo,
|
||||
unsigned int ctr_hi,
|
||||
unsigned int layer_idx,
|
||||
unsigned int i
|
||||
) {
|
||||
unsigned int key = seed_lo ^ (seed_hi * 0x9E3779B9u) ^ (layer_idx * 0x85EBCA6Bu);
|
||||
unsigned int ctr = ctr_lo + i;
|
||||
ctr ^= ctr_hi * 0xC2B2AE35u;
|
||||
ctr ^= key * 0x9E3779B9u;
|
||||
ctr *= 0x85EBCA6Bu;
|
||||
ctr ^= ctr >> 13;
|
||||
ctr *= 0xC2B2AE35u;
|
||||
ctr ^= ctr >> 16;
|
||||
return (float)(ctr & 0x00FFFFFFu) / 16777216.0f;
|
||||
}
|
||||
|
||||
/* Forward-only Bernoulli inverted dropout applied in-place.
|
||||
*
|
||||
* out[i] = in[i] * (u_i < keep_prob ? 1/keep_prob : 0)
|
||||
*
|
||||
* where u_i = mamba_dropout_uniform(seed, ctr, layer_idx, i).
|
||||
*
|
||||
* Grid: ceil(n / block_dim.x), Block: typically 256.
|
||||
*
|
||||
* Args:
|
||||
* data - device pointer to n f32 elements (read/write, in-place).
|
||||
* n - element count.
|
||||
* keep_prob - 1 - dropout_rate, in (0, 1].
|
||||
* seed_lo - low 32 bits of the 64-bit dropout seed.
|
||||
* seed_hi - high 32 bits of the 64-bit dropout seed.
|
||||
* ctr_lo - low 32 bits of the 64-bit dropout counter.
|
||||
* ctr_hi - high 32 bits of the 64-bit dropout counter.
|
||||
* layer_idx - layer index (0..num_layers).
|
||||
*/
|
||||
extern "C" __global__ void bernoulli_dropout_apply(
|
||||
float* __restrict__ data,
|
||||
int n,
|
||||
float keep_prob,
|
||||
unsigned int seed_lo,
|
||||
unsigned int seed_hi,
|
||||
unsigned int ctr_lo,
|
||||
unsigned int ctr_hi,
|
||||
unsigned int layer_idx
|
||||
) {
|
||||
int idx = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
if (idx >= n) return;
|
||||
|
||||
float u = mamba_dropout_uniform(seed_lo, seed_hi, ctr_lo, ctr_hi, layer_idx, (unsigned int)idx);
|
||||
float mask = (u < keep_prob) ? (1.0f / keep_prob) : 0.0f;
|
||||
data[idx] = data[idx] * mask;
|
||||
}
|
||||
@@ -57,7 +57,7 @@ use std::sync::atomic::{AtomicU64, Ordering};
|
||||
use std::sync::Arc;
|
||||
use std::time::{Duration, Instant, SystemTime};
|
||||
|
||||
use cudarc::driver::CudaStream;
|
||||
use cudarc::driver::{CudaFunction, CudaStream, LaunchConfig, PushKernelArg};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use tracing::{debug, info, instrument, trace, warn};
|
||||
use uuid::Uuid;
|
||||
@@ -91,6 +91,13 @@ const fn default_spsa_epsilon() -> f64 {
|
||||
0.01
|
||||
}
|
||||
|
||||
/// Default Philox seed for the forward-only Bernoulli dropout kernel in
|
||||
/// `forward_with_gradients`. Callers needing deterministic replay (e.g.
|
||||
/// tests) should use [`Mamba2SSM::new_with_dropout_seed`] to pin the seed
|
||||
/// explicitly; the struct-level counter starts at 0 and advances by 1 per
|
||||
/// `forward_with_gradients` call.
|
||||
pub const DEFAULT_DROPOUT_SEED: u64 = 0xDA7A_FACE_DA7A_FACE;
|
||||
|
||||
/// Configuration for `MAMBA-2` state-space model
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct Mamba2Config {
|
||||
@@ -474,6 +481,30 @@ impl CudaLayerNorm {
|
||||
}
|
||||
}
|
||||
|
||||
/// Load the forward-only Bernoulli dropout kernel from the cubin precompiled
|
||||
/// by `build.rs`. Returns `None` (not an error) if the cubin lookup or
|
||||
/// function load fails at runtime — call sites must handle the no-op
|
||||
/// fallback explicitly rather than panicking, matching the ml-dqn pattern of
|
||||
/// degrading gracefully when GPU support is absent.
|
||||
///
|
||||
/// On non-cuda builds the compiled cubin does not exist (build.rs early
|
||||
/// returns), so the whole function body is gated on `feature = "cuda"` and
|
||||
/// the non-cuda branch always yields `None`.
|
||||
#[cfg(feature = "cuda")]
|
||||
fn load_dropout_kernel(stream: &Arc<CudaStream>) -> Option<CudaFunction> {
|
||||
static CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/dropout_kernel.cubin"));
|
||||
let module = stream
|
||||
.context()
|
||||
.load_cubin(CUBIN.to_vec())
|
||||
.ok()?;
|
||||
module.load_function("bernoulli_dropout_apply").ok()
|
||||
}
|
||||
|
||||
#[cfg(not(feature = "cuda"))]
|
||||
fn load_dropout_kernel(_stream: &Arc<CudaStream>) -> Option<CudaFunction> {
|
||||
None
|
||||
}
|
||||
|
||||
/// `MAMBA-2` State-Space Model implementation
|
||||
pub struct Mamba2SSM {
|
||||
pub config: Mamba2Config,
|
||||
@@ -492,6 +523,18 @@ pub struct Mamba2SSM {
|
||||
pub layer_norms: Vec<CudaLayerNorm>,
|
||||
pub dropout_rate: f32,
|
||||
|
||||
/// Compiled forward-only Bernoulli dropout kernel. `None` on non-CUDA
|
||||
/// builds (feature `cuda` disabled) or if precompilation skipped because
|
||||
/// nvcc was unavailable at build time — callers fall back to a no-op.
|
||||
dropout_kernel: Option<CudaFunction>,
|
||||
/// 64-bit Philox seed. Default: `0xDA7A_FACE_DA7A_FACE`; override via
|
||||
/// [`Mamba2SSM::new_with_dropout_seed`] for deterministic replay tests.
|
||||
dropout_seed: u64,
|
||||
/// Monotonic counter advanced once per `forward_with_gradients` call so
|
||||
/// successive training steps see independent dropout masks even with the
|
||||
/// same seed. Fed to the Philox hash alongside `(layer_idx, element_idx)`.
|
||||
dropout_ctr: u64,
|
||||
|
||||
// Training state
|
||||
pub optimizer_state: HashMap<String, GpuTensor>,
|
||||
pub gradients: HashMap<String, GpuTensor>,
|
||||
@@ -524,7 +567,11 @@ impl Mamba2SSM {
|
||||
GpuTensor::from_vec(vec![value as f32], &[1], stream)
|
||||
}
|
||||
|
||||
/// Create new `MAMBA-2` model
|
||||
/// Create new `MAMBA-2` model with the default dropout seed
|
||||
/// (`0xDA7A_FACE_DA7A_FACE`).
|
||||
///
|
||||
/// Delegates to [`Mamba2SSM::new_with_dropout_seed`]. Use that variant
|
||||
/// directly when a test needs a specific deterministic mask stream.
|
||||
///
|
||||
/// # Errors
|
||||
///
|
||||
@@ -534,6 +581,28 @@ impl Mamba2SSM {
|
||||
/// - Layer norm creation fails
|
||||
/// - SSD layer initialization fails
|
||||
pub fn new(config: Mamba2Config, stream: &Arc<CudaStream>) -> Result<Self, MLError> {
|
||||
Self::new_with_dropout_seed(config, stream, DEFAULT_DROPOUT_SEED)
|
||||
}
|
||||
|
||||
/// Create new `MAMBA-2` model with an explicit 64-bit dropout seed.
|
||||
///
|
||||
/// The seed feeds the Philox hash used by the forward-only Bernoulli
|
||||
/// dropout kernel in `forward_with_gradients`. Two models with the same
|
||||
/// `(dropout_seed, dropout_ctr)` produce bit-identical masks for the same
|
||||
/// input shapes — used by `test_dropout_determinism`.
|
||||
///
|
||||
/// # Errors
|
||||
///
|
||||
/// Returns `MLError` if:
|
||||
/// - CUDA stream creation fails
|
||||
/// - Linear layer creation fails
|
||||
/// - Layer norm creation fails
|
||||
/// - SSD layer initialization fails
|
||||
pub fn new_with_dropout_seed(
|
||||
config: Mamba2Config,
|
||||
stream: &Arc<CudaStream>,
|
||||
dropout_seed: u64,
|
||||
) -> Result<Self, MLError> {
|
||||
if config.d_model == 0 {
|
||||
return Err(MLError::ConfigError("Mamba2 requires d_model > 0".to_owned()));
|
||||
}
|
||||
@@ -586,6 +655,13 @@ impl Mamba2SSM {
|
||||
// Store learning_rate before moving config
|
||||
let learning_rate = config.learning_rate;
|
||||
|
||||
// Load the forward-only Bernoulli dropout kernel precompiled by
|
||||
// build.rs. A missing kernel (non-cuda build, or nvcc unavailable at
|
||||
// build time) degrades to a no-op in `apply_dropout_inplace` rather
|
||||
// than panicking — the ghost-feature repair only fires where the
|
||||
// precompiled kernel is available.
|
||||
let dropout_kernel = load_dropout_kernel(stream);
|
||||
|
||||
Ok(Self {
|
||||
config,
|
||||
metadata,
|
||||
@@ -600,6 +676,9 @@ impl Mamba2SSM {
|
||||
output_projection,
|
||||
layer_norms,
|
||||
dropout_rate,
|
||||
dropout_kernel,
|
||||
dropout_seed,
|
||||
dropout_ctr: 0,
|
||||
optimizer_state: HashMap::new(),
|
||||
gradients: HashMap::new(),
|
||||
grad_scaler: 1.0,
|
||||
@@ -693,8 +772,11 @@ impl Mamba2SSM {
|
||||
// Residual connection
|
||||
hidden = gpu_add(&hidden, &layer_output)?;
|
||||
|
||||
// Dropout (simulated — multiply by (1 - dropout_rate) during training)
|
||||
// TODO: proper GPU dropout kernel
|
||||
// Dropout applied only in `forward_with_gradients` (training
|
||||
// path); inference and SPSA must stay deterministic. SPSA's
|
||||
// ±ε probes in `backward_pass` call this function twice, and
|
||||
// asymmetric masks across those two calls would poison the
|
||||
// finite-difference estimator — keep this path mask-free.
|
||||
}
|
||||
|
||||
// Output projection with sigmoid activation (P0 FIX: bound output to [0,1] for normalized targets)
|
||||
@@ -1470,8 +1552,98 @@ impl Mamba2SSM {
|
||||
Ok(loss_value)
|
||||
}
|
||||
|
||||
/// Forward pass with gradient computation enabled
|
||||
/// In-place Bernoulli inverted dropout on the post-residual hidden state.
|
||||
///
|
||||
/// Forward-only. Launches the precompiled `bernoulli_dropout_apply`
|
||||
/// kernel with Philox inputs `(dropout_seed, dropout_ctr, layer_idx, i)`
|
||||
/// so the mask is deterministic under `(seed, ctr, layer_idx)` and
|
||||
/// non-aliasing across layers. The counter is advanced once per
|
||||
/// `forward_with_gradients` call, not per layer, so a full pass sees
|
||||
/// distinct masks per layer via `layer_idx` while successive training
|
||||
/// steps see distinct masks via `dropout_ctr`.
|
||||
///
|
||||
/// No-ops when:
|
||||
/// - `dropout_rate <= 0.0` (eval-mode identity, byte-for-byte),
|
||||
/// - tensor is empty,
|
||||
/// - the precompiled kernel is unavailable (non-cuda build path).
|
||||
///
|
||||
/// Clamps `dropout_rate` to `< 1.0` to avoid a divide-by-zero keep_prob.
|
||||
#[allow(unsafe_code)] // CUDA kernel launch via cudarc.
|
||||
pub(crate) fn apply_dropout_inplace(
|
||||
&self,
|
||||
hidden: &mut GpuTensor,
|
||||
layer_idx: usize,
|
||||
) -> Result<(), MLError> {
|
||||
if self.dropout_rate <= 0.0 {
|
||||
return Ok(());
|
||||
}
|
||||
let n = hidden.numel();
|
||||
if n == 0 {
|
||||
return Ok(());
|
||||
}
|
||||
let Some(kernel) = self.dropout_kernel.as_ref() else {
|
||||
// Kernel missing (non-cuda build / nvcc-less build) — soft-stub
|
||||
// by doing nothing. Logged once per call so regressions surface
|
||||
// in training output rather than silently disappearing.
|
||||
tracing::debug!(
|
||||
"apply_dropout_inplace: dropout kernel unavailable, skipping \
|
||||
(dropout_rate={} layer_idx={})",
|
||||
self.dropout_rate,
|
||||
layer_idx,
|
||||
);
|
||||
return Ok(());
|
||||
};
|
||||
|
||||
// Clamp to avoid divide-by-zero on keep_prob.
|
||||
let drop_rate = self.dropout_rate.min(1.0 - 1e-6).max(0.0);
|
||||
let keep_prob = 1.0_f32 - drop_rate;
|
||||
|
||||
let n_i32 = n as i32;
|
||||
let seed_lo = (self.dropout_seed & 0xFFFF_FFFF) as u32;
|
||||
let seed_hi = (self.dropout_seed >> 32) as u32;
|
||||
let ctr_lo = (self.dropout_ctr & 0xFFFF_FFFF) as u32;
|
||||
let ctr_hi = (self.dropout_ctr >> 32) as u32;
|
||||
let layer_u32 = layer_idx as u32;
|
||||
|
||||
let block_dim: u32 = 256;
|
||||
let grid_dim: u32 = ((n as u32) + block_dim - 1) / block_dim;
|
||||
let cfg = LaunchConfig {
|
||||
grid_dim: (grid_dim, 1, 1),
|
||||
block_dim: (block_dim, 1, 1),
|
||||
shared_mem_bytes: 0,
|
||||
};
|
||||
|
||||
// SAFETY: hidden.data is a valid device allocation on self.stream.
|
||||
// n_i32 matches the allocation length. All scalar args are by-value.
|
||||
unsafe {
|
||||
self.stream
|
||||
.launch_builder(kernel)
|
||||
.arg(&mut hidden.data)
|
||||
.arg(&n_i32)
|
||||
.arg(&keep_prob)
|
||||
.arg(&seed_lo)
|
||||
.arg(&seed_hi)
|
||||
.arg(&ctr_lo)
|
||||
.arg(&ctr_hi)
|
||||
.arg(&layer_u32)
|
||||
.launch(cfg)
|
||||
.map_err(|e| MLError::ModelError(format!("bernoulli_dropout_apply: {e}")))?;
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Forward pass with gradient computation enabled.
|
||||
///
|
||||
/// This is the sole training-path forward. Dropout fires here only; the
|
||||
/// pure `forward()` used by `validate()`, `predict_single_fast()`, and
|
||||
/// SPSA's ±ε probes in `backward_pass()` stays deterministic so SPSA's
|
||||
/// finite-difference estimator is not destabilised by asymmetric masks.
|
||||
pub fn forward_with_gradients(&mut self, input: &GpuTensor) -> Result<GpuTensor, MLError> {
|
||||
// Advance the Philox counter once per call. wrapping_add avoids UB
|
||||
// on overflow after ~1.8e19 training steps.
|
||||
self.dropout_ctr = self.dropout_ctr.wrapping_add(1);
|
||||
|
||||
// Input projection with gradients
|
||||
let mut hidden = self.input_projection.forward(input)?;
|
||||
|
||||
@@ -1494,7 +1666,11 @@ impl Mamba2SSM {
|
||||
// Residual connection
|
||||
hidden = gpu_add(&hidden, &layer_output)?;
|
||||
|
||||
// Dropout (simulated — TODO: proper GPU dropout kernel)
|
||||
// Bernoulli inverted dropout (GPU kernel, Philox-seeded, in-place).
|
||||
// Previously this site carried a "simulated — TODO" comment but no
|
||||
// actual op — operators tuning dropout_rate got zero effect. See
|
||||
// `apply_dropout_inplace` for the kernel launch.
|
||||
self.apply_dropout_inplace(&mut hidden, layer_idx)?;
|
||||
}
|
||||
|
||||
// Output projection
|
||||
@@ -3367,3 +3543,156 @@ mod extra_tests {
|
||||
bilinear, exact, zoh);
|
||||
}
|
||||
}
|
||||
|
||||
/// Tests for the forward-only Bernoulli dropout kernel wired into
|
||||
/// `Mamba2SSM`. Gated on `feature = "cuda"` and `#[ignore]` so they only
|
||||
/// run when explicitly requested on a CUDA-capable host (matches the
|
||||
/// `magnitude_distribution` / `multi_fold_convergence` smoke-test pattern).
|
||||
#[cfg(all(test, feature = "cuda"))]
|
||||
#[allow(clippy::unwrap_used, clippy::float_cmp, clippy::cast_precision_loss)]
|
||||
mod dropout_tests {
|
||||
use super::*;
|
||||
use cudarc::driver::CudaContext;
|
||||
|
||||
/// Minimal config: dropout machinery only needs a stream + one layer
|
||||
/// to own a layer_idx. Numeric model parameters are unused here.
|
||||
fn minimal_config(dropout: f64) -> Mamba2Config {
|
||||
Mamba2Config {
|
||||
d_model: 4,
|
||||
d_state: 2,
|
||||
d_head: 2,
|
||||
num_heads: 2,
|
||||
expand: 1,
|
||||
num_layers: 1,
|
||||
dropout,
|
||||
use_ssd: false,
|
||||
use_selective_state: false,
|
||||
hardware_aware: false,
|
||||
batch_size: 1,
|
||||
seq_len: 4,
|
||||
..Default::default()
|
||||
}
|
||||
}
|
||||
|
||||
fn make_model(dropout: f64, seed: u64) -> Mamba2SSM {
|
||||
let ctx = CudaContext::new(0).expect("CUDA context required");
|
||||
let stream = ctx.new_stream().expect("Failed to create CUDA stream");
|
||||
Mamba2SSM::new_with_dropout_seed(minimal_config(dropout), &stream, seed)
|
||||
.expect("construct Mamba2SSM")
|
||||
}
|
||||
|
||||
fn make_input(model: &Mamba2SSM, n: usize) -> GpuTensor {
|
||||
// Deterministic non-zero values so a 0-mask is unambiguous.
|
||||
let values: Vec<f32> = (0..n).map(|i| 1.0 + (i as f32) * 0.01).collect();
|
||||
GpuTensor::from_vec(values, &[n], &model.stream).expect("upload input")
|
||||
}
|
||||
|
||||
/// Two calls with the same `(dropout_seed, dropout_ctr)` must produce
|
||||
/// bit-identical output. Advancing the ctr on each side by the same
|
||||
/// amount preserves that equivalence.
|
||||
#[test]
|
||||
#[ignore = "requires CUDA hardware"]
|
||||
fn test_dropout_determinism() {
|
||||
let n = 1024;
|
||||
let mut model_a = make_model(0.3, 0x1234_5678_9ABC_DEF0);
|
||||
let mut model_b = make_model(0.3, 0x1234_5678_9ABC_DEF0);
|
||||
|
||||
// Both models share ctr=0. Advance each to ctr=1 uniformly.
|
||||
model_a.dropout_ctr = 1;
|
||||
model_b.dropout_ctr = 1;
|
||||
|
||||
let mut tensor_a = make_input(&model_a, n);
|
||||
let mut tensor_b = make_input(&model_b, n);
|
||||
|
||||
model_a.apply_dropout_inplace(&mut tensor_a, 0).unwrap();
|
||||
model_b.apply_dropout_inplace(&mut tensor_b, 0).unwrap();
|
||||
|
||||
let host_a = tensor_a.to_vec().unwrap();
|
||||
let host_b = tensor_b.to_vec().unwrap();
|
||||
assert_eq!(host_a.len(), n);
|
||||
assert_eq!(host_b.len(), n);
|
||||
for (i, (a, b)) in host_a.iter().zip(host_b.iter()).enumerate() {
|
||||
assert_eq!(
|
||||
a.to_bits(),
|
||||
b.to_bits(),
|
||||
"dropout mask diverged at i={i}: a={a} b={b} — Philox not deterministic under (seed, ctr, layer_idx)"
|
||||
);
|
||||
}
|
||||
|
||||
// Sanity: at drop_rate=0.3 with n=1024 we expect ~307 zeros. A
|
||||
// functioning kernel should drop at least a handful; zero zeros
|
||||
// would indicate the mask arithmetic is absent again.
|
||||
let zero_count = host_a.iter().filter(|&&v| v == 0.0).count();
|
||||
assert!(
|
||||
zero_count > 0,
|
||||
"dropout_rate=0.3 produced zero masked elements over {n} samples \
|
||||
— kernel is not applying the mask"
|
||||
);
|
||||
}
|
||||
|
||||
/// With `dropout_rate = 0.0` the kernel must short-circuit to a no-op:
|
||||
/// output bytes must match input bytes exactly. This is the eval-mode
|
||||
/// identity contract — a non-zero dropout_rate with a zero-valued mask
|
||||
/// would scale surviving elements by 1/keep_prob and violate identity.
|
||||
#[test]
|
||||
#[ignore = "requires CUDA hardware"]
|
||||
fn test_dropout_eval_mode_identity() {
|
||||
let n = 512;
|
||||
let mut model = make_model(0.0, DEFAULT_DROPOUT_SEED);
|
||||
model.dropout_ctr = 42;
|
||||
|
||||
let input_values: Vec<f32> = (0..n).map(|i| 1.0 + (i as f32) * 0.01).collect();
|
||||
let mut tensor = GpuTensor::from_vec(input_values.clone(), &[n], &model.stream)
|
||||
.expect("upload input");
|
||||
|
||||
model.apply_dropout_inplace(&mut tensor, 0).unwrap();
|
||||
|
||||
let host = tensor.to_vec().unwrap();
|
||||
assert_eq!(host.len(), n);
|
||||
for (i, (before, after)) in input_values.iter().zip(host.iter()).enumerate() {
|
||||
assert_eq!(
|
||||
before.to_bits(),
|
||||
after.to_bits(),
|
||||
"dropout_rate=0 modified element i={i}: before={before} after={after} \
|
||||
— identity short-circuit broken"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
/// Two calls with distinct counters (same seed) must diverge. Otherwise
|
||||
/// the training loop would see the same mask every step, collapsing the
|
||||
/// regularisation to a single random subset.
|
||||
#[test]
|
||||
#[ignore = "requires CUDA hardware"]
|
||||
fn test_dropout_counter_increment_diverges() {
|
||||
let n = 1024;
|
||||
let mut model = make_model(0.3, 0xCAFE_F00D_CAFE_F00D);
|
||||
|
||||
model.dropout_ctr = 1;
|
||||
let mut tensor_step1 = make_input(&model, n);
|
||||
model.apply_dropout_inplace(&mut tensor_step1, 0).unwrap();
|
||||
let host_step1 = tensor_step1.to_vec().unwrap();
|
||||
|
||||
model.dropout_ctr = 2;
|
||||
let mut tensor_step2 = make_input(&model, n);
|
||||
model.apply_dropout_inplace(&mut tensor_step2, 0).unwrap();
|
||||
let host_step2 = tensor_step2.to_vec().unwrap();
|
||||
|
||||
let diverged = host_step1
|
||||
.iter()
|
||||
.zip(host_step2.iter())
|
||||
.filter(|(a, b)| a.to_bits() != b.to_bits())
|
||||
.count();
|
||||
|
||||
// At drop_rate=0.3 over 1024 samples we'd expect ~40% of positions
|
||||
// to disagree between two independent masks (either zeroed vs kept,
|
||||
// or mixed scaling). Require at least 10% divergence — a single
|
||||
// digit count would mean the counter is not being mixed into the
|
||||
// Philox tuple.
|
||||
assert!(
|
||||
diverged > n / 10,
|
||||
"counter increment produced only {diverged}/{n} diverged elements — \
|
||||
dropout_ctr not folded into Philox tuple"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user