feat(ml-alpha): Mamba2 backward pass — analytical, GPU-pure (Phase 1d.1, session 3)
End-to-end analytical backward through all six stages of the forward
chain. No atomicAdd, no SPSA, no host roundtrip during gradient flow —
each scratch buffer is per-channel-unique so concurrent writes don't
collide; cross-channel reduction is a separate kernel.
New API:
- GpuLinear::backward_with_slices(dy, act, weight, cublas, stream)
→ LinearGrads{dw, db, dx}
Mirror of forward_with_slices added to ml-core; no GpuVarStore lookup.
- Mamba2BackwardGrads holds dw_in/db_in/dw_a/db_a/dw_b/db_b/dw_c/dw_out/db_out
- Mamba2Block::backward(&cache, &d_logit) → Mamba2BackwardGrads
Chain (reverse of forward):
6′. W_out backward (cuBLAS sgemm) → d_h_enriched, dw_out, db_out
5′. mamba2_alpha_scan_bwd kernel → d_a/d_b/d_w_c per-channel/sample
+ d_h_s2 (identity passthrough)
↳ mamba2_alpha_reduce_d_proj × 2 → d_a_proj, d_b_proj [N, K, state]
↳ mamba2_alpha_reduce_d_w_c → dw_c [hidden, state]
3′. W_b backward → d_x_from_b, dw_b, db_b
2′. W_a backward → d_x_from_a, dw_a, db_a
↳ d_x = d_x_from_a + d_x_from_b
1′. W_in backward → dw_in, db_in (d_input discarded)
Memory cost per backward call (for Phase 1d.1 sizes N=64, sh2=32,
K=16, state=8): ~1.1 MiB scratch — fits comfortably on RTX 3050.
Tests (9 passing on real GPU):
- Backward produces all 9 grads with correct shapes
- All grads finite through the 6-stage chain
- dw_in non-zero (gradient flows to the input projection, proving the
full chain is wired — not silently zero-ing somewhere)
- Backward rejects wrong d_logit shape
- + all previously-passing forward / config / shape tests
Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
This commit is contained in:
jgrusewski
@@ -40,7 +40,7 @@ use anyhow::{anyhow, Result};
|
||||
use cudarc::cublas::CudaBlas;
|
||||
use cudarc::driver::{CudaFunction, CudaModule, CudaSlice, CudaStream, LaunchConfig, PushKernelArg};
|
||||
use ml_core::cuda_autograd::gpu_tensor::GpuTensor;
|
||||
use ml_core::cuda_autograd::linear::OwnedGpuLinear;
|
||||
use ml_core::cuda_autograd::linear::{LinearActivations, LinearGrads, OwnedGpuLinear};
|
||||
|
||||
/// Hard kernel limits. The forward and backward kernels use register/local
|
||||
/// arrays sized at compile-time: `float x[16]` (state register) and
|
||||
@@ -89,10 +89,31 @@ impl Mamba2BlockConfig {
|
||||
}
|
||||
}
|
||||
|
||||
/// All parameter gradients produced by [`Mamba2Block::backward`].
|
||||
///
|
||||
/// Tensor shapes match the corresponding forward parameters:
|
||||
/// - `dw_in / db_in` : `[hidden_dim, in_dim] / [hidden_dim]`
|
||||
/// - `dw_a / db_a` : `[state_dim, hidden_dim] / [state_dim]`
|
||||
/// - `dw_b / db_b` : `[state_dim, hidden_dim] / [state_dim]`
|
||||
/// - `dw_c` : `[hidden_dim, state_dim]` (stored as flat slice;
|
||||
/// layout matches the `W_c` weight)
|
||||
/// - `dw_out / db_out` : `[1, hidden_dim] / [1]`
|
||||
pub struct Mamba2BackwardGrads {
|
||||
pub dw_in: GpuTensor,
|
||||
pub db_in: GpuTensor,
|
||||
pub dw_a: GpuTensor,
|
||||
pub db_a: GpuTensor,
|
||||
pub dw_b: GpuTensor,
|
||||
pub db_b: GpuTensor,
|
||||
pub dw_c: CudaSlice<f32>,
|
||||
pub dw_out: GpuTensor,
|
||||
pub db_out: GpuTensor,
|
||||
}
|
||||
|
||||
/// Saved activations + per-stage tensors that the backward pass needs. Returned
|
||||
/// by [`Mamba2Block::forward_train`]; consumed by the backward path (lands in
|
||||
/// session 3). For inference-only ([`Mamba2Block::forward`]) the cache is
|
||||
/// computed but immediately dropped.
|
||||
/// by [`Mamba2Block::forward_train`]; consumed by [`Mamba2Block::backward`].
|
||||
/// For inference-only ([`Mamba2Block::forward`]) the cache is computed but
|
||||
/// immediately dropped.
|
||||
pub struct Mamba2ForwardCache {
|
||||
/// Input as reshaped 2-D `[N*K, in_dim]` — needed to backprop into W_in.
|
||||
pub input_2d: GpuTensor,
|
||||
@@ -324,6 +345,202 @@ impl Mamba2Block {
|
||||
self.forward_train(input).map(|(logit, _cache)| logit)
|
||||
}
|
||||
|
||||
// ── Backward pass ─────────────────────────────────────────────────
|
||||
|
||||
/// GPU-native backward pass. `d_logit` is the upstream gradient
|
||||
/// `[n_batch, 1]` flowing back from the loss; `cache` is the
|
||||
/// `Mamba2ForwardCache` returned by [`forward_train`]. Returns all
|
||||
/// nine parameter gradients in a [`Mamba2BackwardGrads`].
|
||||
///
|
||||
/// Chain (mirrors forward in reverse):
|
||||
/// 6′. W_out backward : d_logit + h_enriched_act → d_h_enriched, dw_out, db_out
|
||||
/// 5′. scan backward : (a_proj, b_proj, d_h_enriched, W_c)
|
||||
/// → d_a_per_channel, d_b_per_channel,
|
||||
/// d_w_c_per_sample, d_h_s2 (=d_h_enriched)
|
||||
/// ↳ reduce d_a_per_channel across j → d_a_proj [N, K, state]
|
||||
/// ↳ reduce d_b_per_channel across j → d_b_proj [N, K, state]
|
||||
/// ↳ reduce d_w_c_per_sample across i → dw_c [sh2, state]
|
||||
/// 3′. W_b backward : d_b_proj[N*K, state] + x_act → d_x_from_b, dw_b, db_b
|
||||
/// 2′. W_a backward : d_a_proj[N*K, state] + x_act → d_x_from_a, dw_a, db_a
|
||||
/// ↳ d_x = d_x_from_a + d_x_from_b
|
||||
/// 1′. W_in backward : d_x[N*K, hidden] + input_act → (d_input dropped), dw_in, db_in
|
||||
pub fn backward(
|
||||
&self,
|
||||
cache: &Mamba2ForwardCache,
|
||||
d_logit: &GpuTensor,
|
||||
) -> Result<Mamba2BackwardGrads> {
|
||||
let c = &self.config;
|
||||
let n_batch = cache.h_enriched.shape()[0];
|
||||
let n_rows = n_batch * c.seq_len;
|
||||
|
||||
if d_logit.shape() != [n_batch, 1] {
|
||||
return Err(anyhow!(
|
||||
"Mamba2Block::backward: d_logit shape {:?} != [{}, 1]",
|
||||
d_logit.shape(), n_batch
|
||||
));
|
||||
}
|
||||
|
||||
// ── 6′. Output projection backward ────────────────────────────
|
||||
let h_enriched_act = LinearActivations { input: cache.h_enriched.clone() };
|
||||
let LinearGrads { dw: dw_out, db: db_out, dx: d_h_enriched } = self
|
||||
.w_out
|
||||
.inner
|
||||
.backward_with_slices(d_logit, &h_enriched_act, &self.w_out.weight,
|
||||
&self.cublas, &self.stream)
|
||||
.map_err(|e| anyhow!("w_out backward: {e}"))?;
|
||||
|
||||
// ── 5′. Allocate scan-backward scratch buffers ────────────────
|
||||
let per_chan_n = n_batch * c.hidden_dim * c.seq_len * c.state_dim;
|
||||
let per_sample_n = n_batch * c.hidden_dim * c.state_dim;
|
||||
let mut d_a_per_channel = self.stream
|
||||
.alloc_zeros::<f32>(per_chan_n)
|
||||
.map_err(|e| anyhow!("alloc d_a_per_channel ({} floats): {e}", per_chan_n))?;
|
||||
let mut d_b_per_channel = self.stream
|
||||
.alloc_zeros::<f32>(per_chan_n)
|
||||
.map_err(|e| anyhow!("alloc d_b_per_channel: {e}"))?;
|
||||
let mut d_w_c_per_sample = self.stream
|
||||
.alloc_zeros::<f32>(per_sample_n)
|
||||
.map_err(|e| anyhow!("alloc d_w_c_per_sample: {e}"))?;
|
||||
let mut d_h_s2 = self.stream
|
||||
.alloc_zeros::<f32>(n_batch * c.hidden_dim)
|
||||
.map_err(|e| anyhow!("alloc d_h_s2: {e}"))?;
|
||||
|
||||
// ── Launch scan backward kernel ───────────────────────────────
|
||||
let block_threads: u32 = 32;
|
||||
let grid_y_h: u32 =
|
||||
((c.hidden_dim + block_threads as usize - 1) / block_threads as usize) as u32;
|
||||
let bwd_cfg = LaunchConfig {
|
||||
grid_dim: (n_batch as u32, grid_y_h, 1),
|
||||
block_dim: (block_threads, 1, 1),
|
||||
shared_mem_bytes: 0,
|
||||
};
|
||||
let n_i32 = n_batch as i32;
|
||||
let k_i32 = c.seq_len as i32;
|
||||
let sh2_i32 = c.hidden_dim as i32;
|
||||
let st_i32 = c.state_dim as i32;
|
||||
unsafe {
|
||||
self.stream
|
||||
.launch_builder(&self.kernel_bwd)
|
||||
.arg(cache.a_proj.cuda_data())
|
||||
.arg(cache.b_proj.cuda_data())
|
||||
.arg(d_h_enriched.cuda_data())
|
||||
.arg(&self.w_c)
|
||||
.arg(&mut d_a_per_channel)
|
||||
.arg(&mut d_b_per_channel)
|
||||
.arg(&mut d_w_c_per_sample)
|
||||
.arg(&mut d_h_s2)
|
||||
.arg(&n_i32)
|
||||
.arg(&k_i32)
|
||||
.arg(&sh2_i32)
|
||||
.arg(&st_i32)
|
||||
.launch(bwd_cfg)
|
||||
.map_err(|e| anyhow!("mamba2_alpha_scan_bwd launch: {e}"))?;
|
||||
}
|
||||
|
||||
// ── Reduce d_a_per_channel and d_b_per_channel across j ───────
|
||||
let red_grid_z: u32 =
|
||||
((c.state_dim + block_threads as usize - 1) / block_threads as usize) as u32;
|
||||
let red_cfg = LaunchConfig {
|
||||
grid_dim: (n_batch as u32, c.seq_len as u32, red_grid_z),
|
||||
block_dim: (block_threads, 1, 1),
|
||||
shared_mem_bytes: 0,
|
||||
};
|
||||
let mut d_a_proj_flat: CudaSlice<f32> = self.stream
|
||||
.alloc_zeros::<f32>(n_rows * c.state_dim)
|
||||
.map_err(|e| anyhow!("alloc d_a_proj_flat: {e}"))?;
|
||||
let mut d_b_proj_flat: CudaSlice<f32> = self.stream
|
||||
.alloc_zeros::<f32>(n_rows * c.state_dim)
|
||||
.map_err(|e| anyhow!("alloc d_b_proj_flat: {e}"))?;
|
||||
unsafe {
|
||||
self.stream
|
||||
.launch_builder(&self.kernel_reduce_d_proj)
|
||||
.arg(&d_a_per_channel)
|
||||
.arg(&mut d_a_proj_flat)
|
||||
.arg(&n_i32)
|
||||
.arg(&k_i32)
|
||||
.arg(&sh2_i32)
|
||||
.arg(&st_i32)
|
||||
.launch(red_cfg)
|
||||
.map_err(|e| anyhow!("reduce_d_proj (a) launch: {e}"))?;
|
||||
}
|
||||
unsafe {
|
||||
self.stream
|
||||
.launch_builder(&self.kernel_reduce_d_proj)
|
||||
.arg(&d_b_per_channel)
|
||||
.arg(&mut d_b_proj_flat)
|
||||
.arg(&n_i32)
|
||||
.arg(&k_i32)
|
||||
.arg(&sh2_i32)
|
||||
.arg(&st_i32)
|
||||
.launch(red_cfg)
|
||||
.map_err(|e| anyhow!("reduce_d_proj (b) launch: {e}"))?;
|
||||
}
|
||||
|
||||
// ── Reduce d_w_c_per_sample across i → dw_c[sh2, state] ───────
|
||||
let red_w_c_grid_y: u32 = red_grid_z;
|
||||
let red_w_c_cfg = LaunchConfig {
|
||||
grid_dim: (c.hidden_dim as u32, red_w_c_grid_y, 1),
|
||||
block_dim: (block_threads, 1, 1),
|
||||
shared_mem_bytes: 0,
|
||||
};
|
||||
let mut dw_c: CudaSlice<f32> = self.stream
|
||||
.alloc_zeros::<f32>(c.hidden_dim * c.state_dim)
|
||||
.map_err(|e| anyhow!("alloc dw_c: {e}"))?;
|
||||
unsafe {
|
||||
self.stream
|
||||
.launch_builder(&self.kernel_reduce_d_w_c)
|
||||
.arg(&d_w_c_per_sample)
|
||||
.arg(&mut dw_c)
|
||||
.arg(&n_i32)
|
||||
.arg(&sh2_i32)
|
||||
.arg(&st_i32)
|
||||
.launch(red_w_c_cfg)
|
||||
.map_err(|e| anyhow!("reduce_d_w_c launch: {e}"))?;
|
||||
}
|
||||
|
||||
// ── 3′. W_b backward (uses x_act = cache.x) ───────────────────
|
||||
let d_b_proj_2d = GpuTensor::new(d_b_proj_flat, vec![n_rows, c.state_dim])
|
||||
.map_err(|e| anyhow!("reshape d_b_proj: {e}"))?;
|
||||
let x_act = LinearActivations { input: cache.x.clone() };
|
||||
let LinearGrads { dw: dw_b, db: db_b, dx: d_x_from_b } = self
|
||||
.w_b
|
||||
.inner
|
||||
.backward_with_slices(&d_b_proj_2d, &x_act, &self.w_b.weight,
|
||||
&self.cublas, &self.stream)
|
||||
.map_err(|e| anyhow!("w_b backward: {e}"))?;
|
||||
|
||||
// ── 2′. W_a backward ─────────────────────────────────────────
|
||||
let d_a_proj_2d = GpuTensor::new(d_a_proj_flat, vec![n_rows, c.state_dim])
|
||||
.map_err(|e| anyhow!("reshape d_a_proj: {e}"))?;
|
||||
let LinearGrads { dw: dw_a, db: db_a, dx: d_x_from_a } = self
|
||||
.w_a
|
||||
.inner
|
||||
.backward_with_slices(&d_a_proj_2d, &x_act, &self.w_a.weight,
|
||||
&self.cublas, &self.stream)
|
||||
.map_err(|e| anyhow!("w_a backward: {e}"))?;
|
||||
|
||||
// d_x = d_x_from_a + d_x_from_b
|
||||
let d_x = d_x_from_a.add(&d_x_from_b, &self.stream)
|
||||
.map_err(|e| anyhow!("sum d_x branches: {e}"))?;
|
||||
|
||||
// ── 1′. W_in backward ─────────────────────────────────────────
|
||||
let input_act = LinearActivations { input: cache.input_2d.clone() };
|
||||
let LinearGrads { dw: dw_in, db: db_in, dx: _d_input } = self
|
||||
.w_in
|
||||
.inner
|
||||
.backward_with_slices(&d_x, &input_act, &self.w_in.weight,
|
||||
&self.cublas, &self.stream)
|
||||
.map_err(|e| anyhow!("w_in backward: {e}"))?;
|
||||
|
||||
Ok(Mamba2BackwardGrads {
|
||||
dw_in, db_in,
|
||||
dw_a, db_a,
|
||||
dw_b, db_b,
|
||||
dw_c,
|
||||
dw_out, db_out,
|
||||
})
|
||||
}
|
||||
|
||||
/// Total trainable parameter count (sum of all projections + W_c).
|
||||
pub fn param_count(&self) -> usize {
|
||||
let c = &self.config;
|
||||
@@ -367,6 +584,90 @@ mod tests {
|
||||
assert!(cfg.validate().is_err(), "seq_len > 32 must be rejected (backward x_hist limit)");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_mamba2_block_backward_returns_finite_grads_with_correct_shapes() {
|
||||
let stream = match cuda_stream_or_skip() {
|
||||
Some(s) => s,
|
||||
None => return,
|
||||
};
|
||||
let cfg = Mamba2BlockConfig {
|
||||
in_dim: 8, hidden_dim: 8, state_dim: 4, seq_len: 8,
|
||||
};
|
||||
let n_batch = 3;
|
||||
let block = Mamba2Block::new(cfg.clone(), Arc::clone(&stream)).expect("init");
|
||||
|
||||
// Forward with cache.
|
||||
let n_in = n_batch * cfg.seq_len * cfg.in_dim;
|
||||
let h: Vec<f32> = (0..n_in).map(|i| ((i as f32) * 0.137).sin()).collect();
|
||||
let d = stream.clone_htod(&h).expect("htod");
|
||||
let input = GpuTensor::new(d, vec![n_batch, cfg.seq_len, cfg.in_dim]).expect("input");
|
||||
let (logit, cache) = block.forward_train(&input).expect("forward_train");
|
||||
|
||||
// Pretend d_logit is just sigmoid(logit) - 0.5 (a BCE-shaped scalar).
|
||||
let logit_h = logit.to_host(&stream).expect("logit dtoh");
|
||||
let d_logit_h: Vec<f32> = logit_h.iter()
|
||||
.map(|z| 1.0 / (1.0 + (-z).exp()) - 0.5)
|
||||
.collect();
|
||||
let d_logit_dev = stream.clone_htod(&d_logit_h).expect("htod");
|
||||
let d_logit = GpuTensor::new(d_logit_dev, vec![n_batch, 1]).expect("d_logit tensor");
|
||||
|
||||
let grads = block.backward(&cache, &d_logit).expect("backward");
|
||||
|
||||
// Shape assertions match each weight's forward shape.
|
||||
assert_eq!(grads.dw_in.shape(), &[cfg.hidden_dim, cfg.in_dim], "dw_in shape");
|
||||
assert_eq!(grads.db_in.shape(), &[cfg.hidden_dim], "db_in shape");
|
||||
assert_eq!(grads.dw_a.shape(), &[cfg.state_dim, cfg.hidden_dim], "dw_a shape");
|
||||
assert_eq!(grads.db_a.shape(), &[cfg.state_dim], "db_a shape");
|
||||
assert_eq!(grads.dw_b.shape(), &[cfg.state_dim, cfg.hidden_dim], "dw_b shape");
|
||||
assert_eq!(grads.db_b.shape(), &[cfg.state_dim], "db_b shape");
|
||||
assert_eq!(grads.dw_c.len(), cfg.hidden_dim * cfg.state_dim, "dw_c numel");
|
||||
assert_eq!(grads.dw_out.shape(), &[1, cfg.hidden_dim], "dw_out shape");
|
||||
assert_eq!(grads.db_out.shape(), &[1], "db_out shape");
|
||||
|
||||
// Finiteness for every gradient tensor.
|
||||
for (name, g) in [
|
||||
("dw_in", &grads.dw_in), ("db_in", &grads.db_in),
|
||||
("dw_a", &grads.dw_a), ("db_a", &grads.db_a),
|
||||
("dw_b", &grads.dw_b), ("db_b", &grads.db_b),
|
||||
("dw_out",&grads.dw_out),("db_out",&grads.db_out),
|
||||
] {
|
||||
let h = g.to_host(&stream).expect("dtoh");
|
||||
for (i, &v) in h.iter().enumerate() {
|
||||
assert!(v.is_finite(), "{} grad index {} non-finite: {}", name, i, v);
|
||||
}
|
||||
}
|
||||
// dw_c is a raw CudaSlice<f32> (no GpuTensor wrapper); read directly.
|
||||
let mut dw_c_host = vec![0.0_f32; grads.dw_c.len()];
|
||||
stream.memcpy_dtoh(&grads.dw_c, &mut dw_c_host).expect("dw_c dtoh");
|
||||
for (i, &v) in dw_c_host.iter().enumerate() {
|
||||
assert!(v.is_finite(), "dw_c grad index {} non-finite: {}", i, v);
|
||||
}
|
||||
|
||||
// At least one element of dw_in should be non-zero (gradient
|
||||
// actually flowed through the chain — not silently zero).
|
||||
let dw_in_h = grads.dw_in.to_host(&stream).expect("dw_in dtoh");
|
||||
let any_nonzero = dw_in_h.iter().any(|&v| v.abs() > 1e-12);
|
||||
assert!(any_nonzero, "dw_in entirely zero — gradient didn't propagate to W_in");
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_mamba2_block_backward_rejects_wrong_d_logit_shape() {
|
||||
let stream = match cuda_stream_or_skip() {
|
||||
Some(s) => s,
|
||||
None => return,
|
||||
};
|
||||
let cfg = Mamba2BlockConfig {
|
||||
in_dim: 8, hidden_dim: 8, state_dim: 4, seq_len: 8,
|
||||
};
|
||||
let n_batch = 3;
|
||||
let block = Mamba2Block::new(cfg.clone(), Arc::clone(&stream)).expect("init");
|
||||
let input = GpuTensor::zeros(&[n_batch, cfg.seq_len, cfg.in_dim], &stream).expect("input");
|
||||
let (_logit, cache) = block.forward_train(&input).expect("forward_train");
|
||||
// Wrong d_logit shape: [n_batch, 2] instead of [n_batch, 1].
|
||||
let bad = GpuTensor::zeros(&[n_batch, 2], &stream).expect("bad d_logit");
|
||||
assert!(block.backward(&cache, &bad).is_err());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_mamba2_block_forward_train_returns_cache() {
|
||||
let stream = match cuda_stream_or_skip() {
|
||||
|
||||
@@ -310,6 +310,81 @@ impl GpuLinear {
|
||||
Ok((y, activations))
|
||||
}
|
||||
|
||||
/// Backward pass with raw `CudaSlice` weights (no `GpuVarStore` lookup).
|
||||
/// Mirror of [`Self::forward_with_slices`]; produces the same `LinearGrads`
|
||||
/// as [`Self::backward`].
|
||||
///
|
||||
/// - `dL/dW = dY^T @ X` (cuBLAS sgemm)
|
||||
/// - `dL/db = sum(dY, axis=0)` (reduce kernel)
|
||||
/// - `dL/dX = dY @ W` (cuBLAS sgemm)
|
||||
pub fn backward_with_slices(
|
||||
&self,
|
||||
dy: &GpuTensor,
|
||||
activations: &LinearActivations,
|
||||
weight: &CudaSlice<f32>,
|
||||
cublas: &CudaBlas,
|
||||
stream: &Arc<CudaStream>,
|
||||
) -> Result<LinearGrads, MLError> {
|
||||
let batch = dy.shape()[0];
|
||||
let in_dim = self.in_dim;
|
||||
let out_dim = self.out_dim;
|
||||
|
||||
// ── dW[out, in] = dY^T[out, B] @ X[B, in] ──────────────────
|
||||
// col-major: dW_col[in, out] = X_col[in, B] @ dY_col[out, B]^T
|
||||
// transA=N, transB=T, m=in, n=out, k=B
|
||||
let mut dw = GpuTensor::zeros(&[out_dim, in_dim], stream)?;
|
||||
let x_ptr = raw_ptr(&activations.input.data, stream);
|
||||
let dy_ptr = raw_ptr(&dy.data, stream);
|
||||
let dw_ptr = raw_ptr_mut(&mut dw.data, stream);
|
||||
unsafe {
|
||||
gemm_ex_f32(
|
||||
cublas,
|
||||
cublasOperation_t::CUBLAS_OP_N,
|
||||
cublasOperation_t::CUBLAS_OP_T,
|
||||
in_dim as i32,
|
||||
out_dim as i32,
|
||||
batch as i32,
|
||||
x_ptr,
|
||||
in_dim as i32,
|
||||
dy_ptr,
|
||||
out_dim as i32,
|
||||
dw_ptr,
|
||||
in_dim as i32,
|
||||
"dW_with_slices",
|
||||
)?;
|
||||
}
|
||||
|
||||
// ── db[out] = sum(dY[B, out], axis=0) ───────────────────────
|
||||
let db = reduce_sum_axis0(dy, batch, out_dim, stream)?;
|
||||
|
||||
// ── dX[B, in] = dY[B, out] @ W[out, in] ────────────────────
|
||||
// col-major: dX_col[in, B] = W_col[in, out] @ dY_col[out, B]
|
||||
// transA=N, transB=N, m=in, n=B, k=out
|
||||
let mut dx = GpuTensor::zeros(&[batch, in_dim], stream)?;
|
||||
let w_ptr = raw_ptr(weight, stream);
|
||||
let dy_ptr2 = raw_ptr(&dy.data, stream);
|
||||
let dx_ptr = raw_ptr_mut(&mut dx.data, stream);
|
||||
unsafe {
|
||||
gemm_ex_f32(
|
||||
cublas,
|
||||
cublasOperation_t::CUBLAS_OP_N,
|
||||
cublasOperation_t::CUBLAS_OP_N,
|
||||
in_dim as i32,
|
||||
batch as i32,
|
||||
out_dim as i32,
|
||||
w_ptr,
|
||||
in_dim as i32,
|
||||
dy_ptr2,
|
||||
out_dim as i32,
|
||||
dx_ptr,
|
||||
in_dim as i32,
|
||||
"dX_with_slices",
|
||||
)?;
|
||||
}
|
||||
|
||||
Ok(LinearGrads { dw, db, dx })
|
||||
}
|
||||
|
||||
/// Backward pass: given `dL/dY` (upstream gradient), compute weight, bias,
|
||||
/// and input gradients.
|
||||
///
|
||||
|
||||
Reference in New Issue
Block a user