teeny_kernels/nn/optim/
sgd.rs

1/*
2 * Copyright (c) 2026 Teenygrad.
3 *
4 * Licensed under the Apache License, Version 2.0 (the "License");
5 * you may not use this file except in compliance with the License.
6 * You may obtain a copy of the License at
7 *
8 *   http://www.apache.org/licenses/LICENSE-2.0
9 *
10 * Unless required by applicable law or agreed to in writing, software
11 * distributed under the License is distributed on an "AS IS" BASIS,
12 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13 * See the License for the specific language governing permissions and
14 * limitations under the License.
15 */
16
17#![allow(non_snake_case)]
18
19use teeny_macros::kernel;
20use teeny_triton::triton::{
21    types::{AddOffsets, Comparison},
22    *,
23};
24
25/// SGD step (no momentum).
26///
27/// ```text
28/// p = p - lr * (g + weight_decay * p)
29/// ```
30///
31/// Grid: `[ceil(n_elements / BLOCK_SIZE), 1, 1]`.
32#[kernel]
33pub fn sgd_step<T: Triton, const BLOCK_SIZE: i32>(
34    params_ptr: T::Pointer<f32>,
35    grad_ptr: T::Pointer<f32>,
36    n_elements: i32,
37    lr: f32,
38    weight_decay: f32,
39) where
40    T::I32Tensor: types::Tensor<i32, 1>,
41    T::I32Tensor: Comparison<i32, BoolTensor = T::BoolTensor>,
42    T::Pointer<f32>: AddOffsets<i32, 1, T::I32Tensor, Output = T::Tensor<T::Pointer<f32>>>,
43{
44    let pid = T::program_id(Axis::X);
45    let block_start = pid * BLOCK_SIZE;
46    let offsets = T::arange(0, BLOCK_SIZE) + block_start;
47    let mask = offsets.lt(n_elements);
48
49    let p = T::load(
50        params_ptr.add_offsets(offsets),
51        Some(mask),
52        None,
53        &[],
54        None,
55        None,
56        None,
57        false,
58    );
59    let g = T::load(
60        grad_ptr.add_offsets(offsets),
61        Some(mask),
62        None,
63        &[],
64        None,
65        None,
66        None,
67        false,
68    );
69
70    let lr_t = T::full(&[BLOCK_SIZE], lr);
71    let wd_t = T::full(&[BLOCK_SIZE], weight_decay);
72
73    let p_new = p - lr_t * (g + wd_t * p);
74    T::store(
75        params_ptr.add_offsets(offsets),
76        p_new,
77        Some(mask),
78        &[],
79        None,
80        None,
81    );
82}
83
84/// SGD step with momentum (non-Nesterov).
85///
86/// ```text
87/// buf = momentum * buf + (1 - dampening) * (g + weight_decay * p)
88/// p   = p - lr * buf
89/// ```
90///
91/// Grid: `[ceil(n_elements / BLOCK_SIZE), 1, 1]`.
92#[kernel]
93pub fn sgd_momentum_step<T: Triton, const BLOCK_SIZE: i32>(
94    params_ptr: T::Pointer<f32>,
95    grad_ptr: T::Pointer<f32>,
96    buf_ptr: T::Pointer<f32>,
97    n_elements: i32,
98    lr: f32,
99    momentum: f32,
100    dampening: f32,
101    weight_decay: f32,
102) where
103    T::I32Tensor: types::Tensor<i32, 1>,
104    T::I32Tensor: Comparison<i32, BoolTensor = T::BoolTensor>,
105    T::Pointer<f32>: AddOffsets<i32, 1, T::I32Tensor, Output = T::Tensor<T::Pointer<f32>>>,
106{
107    let pid = T::program_id(Axis::X);
108    let block_start = pid * BLOCK_SIZE;
109    let offsets = T::arange(0, BLOCK_SIZE) + block_start;
110    let mask = offsets.lt(n_elements);
111
112    let p = T::load(
113        params_ptr.add_offsets(offsets),
114        Some(mask),
115        None,
116        &[],
117        None,
118        None,
119        None,
120        false,
121    );
122    let g = T::load(
123        grad_ptr.add_offsets(offsets),
124        Some(mask),
125        None,
126        &[],
127        None,
128        None,
129        None,
130        false,
131    );
132    let buf = T::load(
133        buf_ptr.add_offsets(offsets),
134        Some(mask),
135        None,
136        &[],
137        None,
138        None,
139        None,
140        false,
141    );
142
143    let lr_t = T::full(&[BLOCK_SIZE], lr);
144    let mu_t = T::full(&[BLOCK_SIZE], momentum);
145    let damp_t = T::full(&[BLOCK_SIZE], 1.0_f32 - dampening);
146    let wd_t = T::full(&[BLOCK_SIZE], weight_decay);
147
148    let g_eff = g + wd_t * p;
149    let buf_new = mu_t * buf + damp_t * g_eff;
150    let p_new = p - lr_t * buf_new;
151
152    T::store(
153        params_ptr.add_offsets(offsets),
154        p_new,
155        Some(mask),
156        &[],
157        None,
158        None,
159    );
160    T::store(
161        buf_ptr.add_offsets(offsets),
162        buf_new,
163        Some(mask),
164        &[],
165        None,
166        None,
167    );
168}
169
170/// SGD step with Nesterov momentum.
171///
172/// ```text
173/// buf   = momentum * buf + (1 - dampening) * (g + weight_decay * p)
174/// g_nes = (g + weight_decay * p) + momentum * buf
175/// p     = p - lr * g_nes
176/// ```
177///
178/// Grid: `[ceil(n_elements / BLOCK_SIZE), 1, 1]`.
179#[kernel]
180pub fn sgd_nesterov_step<T: Triton, const BLOCK_SIZE: i32>(
181    params_ptr: T::Pointer<f32>,
182    grad_ptr: T::Pointer<f32>,
183    buf_ptr: T::Pointer<f32>,
184    n_elements: i32,
185    lr: f32,
186    momentum: f32,
187    dampening: f32,
188    weight_decay: f32,
189) where
190    T::I32Tensor: types::Tensor<i32, 1>,
191    T::I32Tensor: Comparison<i32, BoolTensor = T::BoolTensor>,
192    T::Pointer<f32>: AddOffsets<i32, 1, T::I32Tensor, Output = T::Tensor<T::Pointer<f32>>>,
193{
194    let pid = T::program_id(Axis::X);
195    let block_start = pid * BLOCK_SIZE;
196    let offsets = T::arange(0, BLOCK_SIZE) + block_start;
197    let mask = offsets.lt(n_elements);
198
199    let p = T::load(
200        params_ptr.add_offsets(offsets),
201        Some(mask),
202        None,
203        &[],
204        None,
205        None,
206        None,
207        false,
208    );
209    let g = T::load(
210        grad_ptr.add_offsets(offsets),
211        Some(mask),
212        None,
213        &[],
214        None,
215        None,
216        None,
217        false,
218    );
219    let buf = T::load(
220        buf_ptr.add_offsets(offsets),
221        Some(mask),
222        None,
223        &[],
224        None,
225        None,
226        None,
227        false,
228    );
229
230    let lr_t = T::full(&[BLOCK_SIZE], lr);
231    let mu_t = T::full(&[BLOCK_SIZE], momentum);
232    let damp_t = T::full(&[BLOCK_SIZE], 1.0_f32 - dampening);
233    let wd_t = T::full(&[BLOCK_SIZE], weight_decay);
234
235    let g_eff = g + wd_t * p;
236    let buf_new = mu_t * buf + damp_t * g_eff;
237    let g_nes = g_eff + mu_t * buf_new;
238    let p_new = p - lr_t * g_nes;
239
240    T::store(
241        params_ptr.add_offsets(offsets),
242        p_new,
243        Some(mask),
244        &[],
245        None,
246        None,
247    );
248    T::store(
249        buf_ptr.add_offsets(offsets),
250        buf_new,
251        Some(mask),
252        &[],
253        None,
254        None,
255    );
256}
teeny_kernels/nn/optim/sgd.rs

teeny_kernels/nn/optim/
sgd.rs
