forge/ops/
train.rs

1//! Backward and optimizer ops — the training half of the op API.
2//!
3//! Split out of `ops/mod.rs` so the whole set can sit behind the `train`
4//! feature without scattering `#[cfg]` through the forward ops it used to be
5//! interleaved with. `ops/mod.rs` re-exports every one of these, so the public
6//! paths are unchanged: `forge::ops::adamw` still resolves, when `train` is on.
7
8use crate::backend::{cpu, wgpu as gpu};
9use crate::error::{ForgeError, Result};
10use crate::shape::Shape;
11use crate::tensor::{CpuStorage, Storage, Tensor};
12
13use super::{cpu_f32, cpu_u32, f32_tensor, gpu_storage, same_device};
14
15/// Rows and last-dim width of a tensor, for the ops that reduce over rows.
16fn last_dim_rows(x: &Tensor) -> Result<(usize, usize)> {
17    let dims = x.shape().dims();
18    if dims.is_empty() {
19        return Err(ForgeError::Shape("op needs rank >= 1".into()));
20    }
21    let cols = dims[dims.len() - 1];
22    Ok((x.shape().numel() / cols, cols))
23}
24
25/// GELU backward: dx = gelu'(x) * dy.
26pub fn gelu_bwd(x: &Tensor, dy: &Tensor) -> Result<Tensor> {
27    same_device(&[x, dy])?;
28    if x.shape() != dy.shape() {
29        return Err(ForgeError::Shape("gelu_bwd shape mismatch".into()));
30    }
31    let n = x.shape().numel();
32    let storage = match x.storage() {
33        Storage::Cpu(_) => Storage::Cpu(CpuStorage::F32(
34            cpu::gelu_bwd(cpu_f32(x)?, cpu_f32(dy)?).into(),
35        )),
36        Storage::Wgpu(_) => Storage::Wgpu(gpu::ops::gelu_bwd(gpu_storage(x)?, gpu_storage(dy)?, n)),
37    };
38    Ok(f32_tensor(storage, x.shape().clone()))
39}
40
41/// Softmax backward from the forward *output* y: dx = y * (dy - sum(y*dy))
42/// per row. Causal masking needs no parameters here — masked forward outputs
43/// are exact zeros.
44pub fn softmax_bwd(y: &Tensor, dy: &Tensor) -> Result<Tensor> {
45    same_device(&[y, dy])?;
46    if y.shape() != dy.shape() {
47        return Err(ForgeError::Shape("softmax_bwd shape mismatch".into()));
48    }
49    let (rows, cols) = last_dim_rows(y)?;
50    let storage = match y.storage() {
51        Storage::Cpu(_) => Storage::Cpu(CpuStorage::F32(
52            cpu::softmax_bwd(cpu_f32(y)?, cpu_f32(dy)?, rows, cols).into(),
53        )),
54        Storage::Wgpu(_) => Storage::Wgpu(gpu::ops::softmax_bwd(
55            gpu_storage(y)?,
56            gpu_storage(dy)?,
57            rows,
58            cols,
59        )),
60    };
61    Ok(f32_tensor(storage, y.shape().clone()))
62}
63
64/// LayerNorm backward: (dx, dgamma, dbeta).
65pub fn layernorm_bwd(
66    x: &Tensor,
67    gamma: &Tensor,
68    dy: &Tensor,
69    eps: f32,
70) -> Result<(Tensor, Tensor, Tensor)> {
71    same_device(&[x, gamma, dy])?;
72    if x.shape() != dy.shape() {
73        return Err(ForgeError::Shape("layernorm_bwd shape mismatch".into()));
74    }
75    let (rows, cols) = last_dim_rows(x)?;
76    if gamma.shape().numel() != cols {
77        return Err(ForgeError::Shape("layernorm_bwd gamma length".into()));
78    }
79    let pshape = gamma.shape().clone();
80    match x.storage() {
81        Storage::Cpu(_) => {
82            let dx =
83                cpu::layernorm_bwd_dx(cpu_f32(x)?, cpu_f32(gamma)?, cpu_f32(dy)?, rows, cols, eps);
84            let (dg, db) = cpu::layernorm_bwd_dparams(cpu_f32(x)?, cpu_f32(dy)?, rows, cols, eps);
85            Ok((
86                f32_tensor(Storage::Cpu(CpuStorage::F32(dx.into())), x.shape().clone()),
87                f32_tensor(Storage::Cpu(CpuStorage::F32(dg.into())), pshape.clone()),
88                f32_tensor(Storage::Cpu(CpuStorage::F32(db.into())), pshape),
89            ))
90        }
91        Storage::Wgpu(_) => {
92            let (dx, dg, db) = gpu::ops::layernorm_bwd(
93                gpu_storage(x)?,
94                gpu_storage(gamma)?,
95                gpu_storage(dy)?,
96                rows,
97                cols,
98                eps,
99            );
100            Ok((
101                f32_tensor(Storage::Wgpu(dx), x.shape().clone()),
102                f32_tensor(Storage::Wgpu(dg), pshape.clone()),
103                f32_tensor(Storage::Wgpu(db), pshape),
104            ))
105        }
106    }
107}
108
109/// Column sums over all leading dims: `[.., cols]` -> `[cols]`. Bias gradients.
110pub fn sum_rows(x: &Tensor) -> Result<Tensor> {
111    let (rows, cols) = last_dim_rows(x)?;
112    let storage = match x.storage() {
113        Storage::Cpu(_) => Storage::Cpu(CpuStorage::F32(
114            cpu::sum_rows(cpu_f32(x)?, rows, cols).into(),
115        )),
116        Storage::Wgpu(_) => Storage::Wgpu(gpu::ops::sum_rows(gpu_storage(x)?, rows, cols)),
117    };
118    Ok(f32_tensor(storage, Shape::new(&[cols])))
119}
120
121/// `dst[ids[r]] += src[r]`, in place. Embedding-table gradient scatter.
122pub fn scatter_add_rows(dst: &mut Tensor, ids: &Tensor, src: &Tensor) -> Result<()> {
123    same_device(&[dst, ids, src])?;
124    let (vocab, c) = match dst.shape().dims() {
125        [v, c] => (*v, *c),
126        _ => return Err(ForgeError::Shape("scatter_add dst must be rank 2".into())),
127    };
128    let t = ids.shape().numel();
129    if src.shape().dims() != [t, c] {
130        return Err(ForgeError::Shape(format!(
131            "scatter_add src {} != [{t}, {c}]",
132            src.shape()
133        )));
134    }
135    match (&mut dst.storage, src.storage()) {
136        (Storage::Cpu(CpuStorage::F32(d)), Storage::Cpu(CpuStorage::F32(s))) => {
137            let ids = cpu_u32(ids)?;
138            if let Some(&bad) = ids.iter().find(|&&id| id as usize >= vocab) {
139                return Err(ForgeError::Shape(format!(
140                    "scatter_add id {bad} >= {vocab}"
141                )));
142            }
143            let d: &mut Vec<f32> = std::sync::Arc::make_mut(d);
144            cpu::scatter_add_rows(d, ids, s, c);
145            Ok(())
146        }
147        (Storage::Wgpu(d), Storage::Wgpu(s)) => {
148            gpu::ops::scatter_add_rows(d, gpu_storage(ids)?, s, t, c, vocab * c);
149            Ok(())
150        }
151        _ => Err(ForgeError::Device(
152            "scatter_add expects f32 tensors on one device".into(),
153        )),
154    }
155}
156
157/// Per-row NLL: `out[r] = -log(probs[r, ids[r]])`.
158pub fn gather_nll(probs: &Tensor, ids: &Tensor) -> Result<Tensor> {
159    same_device(&[probs, ids])?;
160    let (rows, cols) = last_dim_rows(probs)?;
161    if ids.shape().numel() != rows {
162        return Err(ForgeError::Shape("gather_nll ids length".into()));
163    }
164    let storage = match probs.storage() {
165        Storage::Cpu(_) => Storage::Cpu(CpuStorage::F32(
166            cpu::gather_nll(cpu_f32(probs)?, cpu_u32(ids)?, cols).into(),
167        )),
168        Storage::Wgpu(_) => Storage::Wgpu(gpu::ops::gather_nll(
169            gpu_storage(probs)?,
170            gpu_storage(ids)?,
171            rows,
172            cols,
173        )),
174    };
175    Ok(f32_tensor(storage, Shape::new(&[rows])))
176}
177
178/// Cross-entropy backward: dlogits = (probs - onehot(ids)) * scale.
179pub fn ce_bwd(probs: &Tensor, ids: &Tensor, scale: f32) -> Result<Tensor> {
180    same_device(&[probs, ids])?;
181    let (rows, cols) = last_dim_rows(probs)?;
182    if ids.shape().numel() != rows {
183        return Err(ForgeError::Shape("ce_bwd ids length".into()));
184    }
185    let storage = match probs.storage() {
186        Storage::Cpu(_) => Storage::Cpu(CpuStorage::F32(
187            cpu::ce_bwd(cpu_f32(probs)?, cpu_u32(ids)?, rows, cols, scale).into(),
188        )),
189        Storage::Wgpu(_) => Storage::Wgpu(gpu::ops::ce_bwd(
190            gpu_storage(probs)?,
191            gpu_storage(ids)?,
192            rows,
193            cols,
194            scale,
195        )),
196    };
197    Ok(f32_tensor(storage, probs.shape().clone()))
198}
199
200/// Inverted dropout with a deterministic counter RNG (identical masks on
201/// both backends). Apply the same (p, seed) to dy for the backward pass.
202pub fn dropout(x: &Tensor, p: f32, seed: u32) -> Result<Tensor> {
203    if !(0.0..1.0).contains(&p) {
204        return Err(ForgeError::Shape(format!("dropout p {p} outside [0, 1)")));
205    }
206    if p == 0.0 {
207        return Ok(x.clone());
208    }
209    // Compute the keep-scale once on the CPU and hand the same bits to both
210    // backends: GPU division isn't guaranteed correctly-rounded (unlike
211    // multiplication), so an independent 1.0/(1.0-p) on the GPU can be a
212    // few ULP off from the CPU's, breaking bit-for-bit parity.
213    let scale = 1.0f32 / (1.0 - p);
214    let n = x.shape().numel();
215    let storage = match x.storage() {
216        Storage::Cpu(_) => Storage::Cpu(CpuStorage::F32(
217            cpu::dropout(cpu_f32(x)?, p, scale, seed).into(),
218        )),
219        Storage::Wgpu(_) => Storage::Wgpu(gpu::ops::dropout(gpu_storage(x)?, n, p, scale, seed)),
220    };
221    Ok(f32_tensor(storage, x.shape().clone()))
222}
223
224/// split_heads backward for one of q/k/v (`which` in 0..3):
225/// d [h, t, hd] -> [t, 3c] with the other thirds zero.
226pub fn unsplit_head(d: &Tensor, which: usize) -> Result<Tensor> {
227    let (h, t, hd) = match d.shape().dims() {
228        [h, t, hd] => (*h, *t, *hd),
229        _ => return Err(ForgeError::Shape("unsplit_head needs [h, t, hd]".into())),
230    };
231    let c = h * hd;
232    let shape = Shape::new(&[t, 3 * c]);
233    let storage = match d.storage() {
234        Storage::Cpu(_) => Storage::Cpu(CpuStorage::F32(
235            cpu::unsplit_head(cpu_f32(d)?, t, c, h, which).into(),
236        )),
237        Storage::Wgpu(_) => Storage::Wgpu(gpu::ops::unsplit_head(gpu_storage(d)?, t, c, h, which)),
238    };
239    Ok(f32_tensor(storage, shape))
240}
241
242/// merge_heads backward: dy [t, c] -> [h, t, hd].
243pub fn unmerge_heads(dy: &Tensor, h: usize) -> Result<Tensor> {
244    let (t, c) = match dy.shape().dims() {
245        [t, c] => (*t, *c),
246        _ => return Err(ForgeError::Shape("unmerge_heads needs [t, c]".into())),
247    };
248    if c % h != 0 {
249        return Err(ForgeError::Shape(format!("unmerge_heads c={c} % h={h}")));
250    }
251    let shape = Shape::new(&[h, t, c / h]);
252    let storage = match dy.storage() {
253        Storage::Cpu(_) => Storage::Cpu(CpuStorage::F32(
254            cpu::unmerge_heads(cpu_f32(dy)?, t, c, h).into(),
255        )),
256        Storage::Wgpu(_) => Storage::Wgpu(gpu::ops::unmerge_heads(gpu_storage(dy)?, t, c, h)),
257    };
258    Ok(f32_tensor(storage, shape))
259}
260
261/// Sum of squares of all elements (for the global gradient norm).
262/// Training-path op with a sync readback — WebGPU tensors are native-only
263/// here (browser training is out of scope for 1.0).
264pub fn sumsq(x: &Tensor) -> Result<f32> {
265    match x.storage() {
266        Storage::Cpu(_) => Ok(cpu_f32(x)?.iter().map(|&v| v * v).sum()),
267        #[cfg(not(target_arch = "wasm32"))]
268        Storage::Wgpu(s) => {
269            let (partials, groups) = gpu::ops::sumsq_partials(s, x.shape().numel());
270            let bytes = partials.ctx.readback(&partials.buf, 0, groups * 4)?;
271            let vals: Vec<f32> = bytemuck::pod_collect_to_vec(&bytes);
272            Ok(vals.iter().sum())
273        }
274        #[cfg(target_arch = "wasm32")]
275        Storage::Wgpu(_) => Err(crate::error::ForgeError::Wgpu(
276            "sumsq readback unavailable on wasm32 (training is native-only)".into(),
277        )),
278    }
279}
280
281/// y = x * alpha.
282pub fn scale(x: &Tensor, alpha: f32) -> Result<Tensor> {
283    let n = x.shape().numel();
284    let storage = match x.storage() {
285        Storage::Cpu(_) => Storage::Cpu(CpuStorage::F32(
286            cpu_f32(x)?
287                .iter()
288                .map(|&v| v * alpha)
289                .collect::<Vec<_>>()
290                .into(),
291        )),
292        Storage::Wgpu(_) => Storage::Wgpu(gpu::ops::scale(gpu_storage(x)?, n, alpha)),
293    };
294    Ok(f32_tensor(storage, x.shape().clone()))
295}
296
297/// AdamW step with decoupled weight decay, updating param/m/v in place.
298/// `step` is 1-based.
299#[allow(clippy::too_many_arguments)]
300pub fn adamw(
301    param: &mut Tensor,
302    grad: &Tensor,
303    m: &mut Tensor,
304    v: &mut Tensor,
305    lr: f32,
306    beta1: f32,
307    beta2: f32,
308    eps: f32,
309    weight_decay: f32,
310    step: u32,
311) -> Result<()> {
312    if param.shape() != grad.shape() || param.shape() != m.shape() || param.shape() != v.shape() {
313        return Err(ForgeError::Shape("adamw shape mismatch".into()));
314    }
315    let n = param.shape().numel();
316    match (&mut param.storage, grad.storage()) {
317        (Storage::Cpu(CpuStorage::F32(p)), Storage::Cpu(CpuStorage::F32(g))) => {
318            let (Storage::Cpu(CpuStorage::F32(ms)), Storage::Cpu(CpuStorage::F32(vs))) =
319                (&mut m.storage, &mut v.storage)
320            else {
321                return Err(ForgeError::Device("adamw state device mismatch".into()));
322            };
323            let p: &mut Vec<f32> = std::sync::Arc::make_mut(p);
324            let ms: &mut Vec<f32> = std::sync::Arc::make_mut(ms);
325            let vs: &mut Vec<f32> = std::sync::Arc::make_mut(vs);
326            cpu::adamw(p, g, ms, vs, lr, beta1, beta2, eps, weight_decay, step);
327            Ok(())
328        }
329        (Storage::Wgpu(p), Storage::Wgpu(g)) => {
330            let (Storage::Wgpu(ms), Storage::Wgpu(vs)) = (&m.storage, &v.storage) else {
331                return Err(ForgeError::Device("adamw state device mismatch".into()));
332            };
333            gpu::ops::adamw(p, g, ms, vs, n, lr, beta1, beta2, eps, weight_decay, step);
334            Ok(())
335        }
336        _ => Err(ForgeError::Device(
337            "adamw expects f32 tensors on one device".into(),
338        )),
339    }
340}
forge/ops/train.rs

forge/ops/
train.rs
