ferrotensor_backend/engine/
training.rs

1//! Losses, optimisers and gradient utilities.
2
3use super::{groups64, uniform, ComputeEngine};
4use crate::{UnifiedTensor, WgpuContext};
5
6/// Hyper-parameters for [`ComputeEngine::adamw_update`]; laid out to match
7/// the WGSL uniform.
8#[repr(C)]
9#[derive(Copy, Clone, bytemuck::Pod, bytemuck::Zeroable)]
10pub struct AdamParams {
11    /// Learning rate.
12    pub lr: f32,
13    /// First-moment decay.
14    pub beta1: f32,
15    /// Second-moment decay.
16    pub beta2: f32,
17    /// Denominator epsilon.
18    pub eps: f32,
19    /// Decoupled weight decay.
20    pub weight_decay: f32,
21    /// 1-based step number, used for bias correction.
22    pub t: f32,
23    /// Padding to a 16-byte multiple; set to zero.
24    pub _pad: [f32; 2],
25}
26
27impl ComputeEngine {
28    /// In-place SGD step: `weight -= lr * grad`.
29    pub fn sgd_update(
30        &self,
31        ctx: &WgpuContext,
32        weight: &UnifiedTensor,
33        grad: &UnifiedTensor,
34        lr: f32,
35    ) {
36        let lr_buf = uniform(ctx, bytemuck::cast_slice(&[lr]));
37        self.run(
38            ctx,
39            &self.sgd,
40            &[(0, weight.buffer()), (1, grad.buffer()), (2, &lr_buf)],
41            groups64(weight.shape.iter().product()),
42        );
43    }
44
45    /// In-place AdamW step, updating `weight` and the moment buffers `m`
46    /// and `v`.
47    pub fn adamw_update(
48        &self,
49        ctx: &WgpuContext,
50        weight: &UnifiedTensor,
51        grad: &UnifiedTensor,
52        m: &UnifiedTensor,
53        v: &UnifiedTensor,
54        params: &AdamParams,
55    ) {
56        let params_buf = uniform(ctx, bytemuck::bytes_of(params));
57        self.run(
58            ctx,
59            &self.adamw,
60            &[
61                (0, weight.buffer()),
62                (1, grad.buffer()),
63                (3, m.buffer()),
64                (4, v.buffer()),
65                (5, &params_buf),
66            ],
67            groups64(weight.shape.iter().product()),
68        );
69    }
70
71    /// Per-element squared error `out = 0.5 * (pred - target)^2` (not reduced).
72    pub fn mse_loss(
73        &self,
74        ctx: &WgpuContext,
75        pred: &UnifiedTensor,
76        target: &UnifiedTensor,
77        out: &UnifiedTensor,
78    ) {
79        self.binary(ctx, &self.mse_loss, pred, target, out);
80    }
81
82    /// Gradient of [`mse_loss`](Self::mse_loss): `grad_in = pred - target`.
83    pub fn mse_loss_backward(
84        &self,
85        ctx: &WgpuContext,
86        pred: &UnifiedTensor,
87        target: &UnifiedTensor,
88        grad_in: &UnifiedTensor,
89    ) {
90        self.run(
91            ctx,
92            &self.mse_loss_bw,
93            &[
94                (0, pred.buffer()),
95                (1, target.buffer()),
96                (3, grad_in.buffer()),
97            ],
98            groups64(pred.shape.iter().product()),
99        );
100    }
101
102    /// Cross-entropy loss of one logits vector: `loss_out[0] = -log(softmax(logits)[target])`.
103    pub fn cross_entropy_loss(
104        &self,
105        ctx: &WgpuContext,
106        logits: &UnifiedTensor,
107        target: u32,
108        loss_out: &UnifiedTensor,
109    ) {
110        let vocab_size = logits.shape.iter().product::<usize>() as u32;
111        let params = uniform(ctx, bytemuck::cast_slice(&[target, vocab_size, 0, 0]));
112        self.run(
113            ctx,
114            &self.ce_loss,
115            &[(0, logits.buffer()), (1, loss_out.buffer()), (2, &params)],
116            (1, 1, 1),
117        );
118    }
119
120    /// Cross-entropy loss backward: gradient = softmax(logits) - one_hot(target)
121    pub fn cross_entropy_backward(
122        &self,
123        ctx: &WgpuContext,
124        logits: &UnifiedTensor,
125        target: u32,
126        grad_out: &UnifiedTensor,
127    ) {
128        let vocab_size = logits.shape.iter().product::<usize>() as u32;
129        let params = uniform(ctx, bytemuck::cast_slice(&[target, vocab_size, 0, 0]));
130        self.run(
131            ctx,
132            &self.ce_loss_bw,
133            &[(0, logits.buffer()), (1, grad_out.buffer()), (2, &params)],
134            groups64(vocab_size as usize),
135        );
136    }
137
138    /// Compute squared elements (for gradient norm computation)
139    pub fn gradient_norm_squared(
140        &self,
141        ctx: &WgpuContext,
142        input: &UnifiedTensor,
143        out: &UnifiedTensor,
144    ) {
145        self.run(
146            ctx,
147            &self.grad_norm,
148            &[(0, input.buffer()), (1, out.buffer())],
149            groups64(input.shape.iter().product()),
150        );
151    }
152
153    /// Clip gradients by scaling: grad *= min(1.0, max_norm / actual_norm)
154    pub fn clip_grad_by_norm(&self, ctx: &WgpuContext, grad: &UnifiedTensor, scale: f32) {
155        let params = uniform(ctx, bytemuck::cast_slice(&[scale]));
156        self.run(
157            ctx,
158            &self.grad_clip,
159            &[(0, grad.buffer()), (1, &params)],
160            groups64(grad.shape.iter().product()),
161        );
162    }
163
164    /// Gradient accumulation: accum_grad += new_grad
165    /// Essential for micro-batching on memory-constrained devices
166    pub fn gradient_accumulate(
167        &self,
168        ctx: &WgpuContext,
169        accum_grad: &UnifiedTensor,
170        new_grad: &UnifiedTensor,
171    ) {
172        self.run(
173            ctx,
174            &self.grad_accum,
175            &[(0, accum_grad.buffer()), (1, new_grad.buffer())],
176            groups64(accum_grad.shape.iter().product()),
177        );
178    }
179
180    /// Gradient scaling: grad *= scale_factor
181    /// Used to average gradients after accumulation
182    pub fn gradient_scale(&self, ctx: &WgpuContext, grad: &UnifiedTensor, scale: f32) {
183        let params = uniform(ctx, bytemuck::cast_slice(&[scale]));
184        self.run(
185            ctx,
186            &self.grad_scale,
187            &[(0, grad.buffer()), (1, &params)],
188            groups64(grad.shape.iter().product()),
189        );
190    }
191}
ferrotensor_backend/engine/training.rs

ferrotensor_backend/engine/
training.rs
