rumus 0.3.1

A native-Rust deep learning framework with explicit memory safety and hardware acceleration
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129

// SPDX-License-Identifier: Apache-2.0 OR MIT
//! Stochastic Gradient Descent with optional momentum.
//!
//! When parameters are GPU-resident, dispatches a WGSL kernel —
//! zero D2H transfers.

use std::collections::HashMap;

use crate::autograd::{AutogradError, GradientStore};
use crate::nn::Parameter;
use crate::optim::Optimizer;
use crate::tensor::{ParamId, Tensor};

pub struct SGD {
    params: Vec<Parameter>,
    lr: f32,
    momentum: f32,
    velocity: HashMap<ParamId, Tensor>,
}

impl SGD {
    pub fn new(params: Vec<Parameter>, lr: f32, momentum: f32) -> Self {
        Self {
            params,
            lr,
            momentum,
            velocity: HashMap::new(),
        }
    }
}

impl Optimizer for SGD {
    fn step(&mut self, grads: &mut GradientStore) -> Result<(), AutogradError> {
        for param in &self.params {
            let raw_grad = grads
                .remove(param.grad_id())
                .ok_or(AutogradError::MissingGrad {
                    grad_id: param.grad_id(),
                })?;
            // Cast F16 gradients to F32 for master weight update.
            let grad = if raw_grad.dtype() != crate::tensor::DType::F32 {
                raw_grad.to_dtype(crate::tensor::DType::F32)
            } else {
                raw_grad
            };
            let numel = param.tensor.numel();

            // Ensure velocity buffer exists (needed for both GPU and CPU
            // when momentum > 0; also used as dummy for GPU path when
            // momentum == 0).
            let shape = param.tensor.shape().to_vec();
            let vel = self
                .velocity
                .entry(param.id)
                .or_insert_with(|| Tensor::new(vec![0.0; numel], shape));

            // ---- GPU path ----
            #[cfg(feature = "gpu")]
            if param.tensor.storage.is_gpu() {
                use crate::backend::gpu::{
                    compute as gpu_compute,
                    context::GpuContext,
                };

                let ctx = GpuContext::get().expect("GPU required for GPU param");

                grad.to_gpu();
                vel.to_gpu();

                let grad_buf = grad.storage.gpu_buffer();
                let vel_buf = vel.storage.gpu_buffer();
                let param_buf = param.tensor.storage.gpu_buffer();

                gpu_compute::sgd_step(
                    ctx,
                    &grad_buf,
                    &vel_buf,
                    &param_buf,
                    numel as u32,
                    self.lr,
                    self.momentum,
                );

                drop(grad_buf);
                drop(vel_buf);
                drop(param_buf);

                vel.storage.mark_gpu_dirty();
                param.tensor.storage.mark_gpu_dirty();
                param.tensor.storage.bump_version();
                continue;
            }

            // ---- CPU path ----
            let grad_guard = grad.storage.data();

            if self.momentum > 0.0 {
                {
                    let mut vel_data = vel.storage.data_write();
                    let mut param_data = param.tensor.storage.data_write();

                    for i in 0..numel {
                        vel_data[i] = self.momentum * vel_data[i] + grad_guard[i];
                        param_data[i] -= self.lr * vel_data[i];
                    }
                }
            } else {
                {
                    let mut param_data = param.tensor.storage.data_write();
                    for i in 0..numel {
                        param_data[i] -= self.lr * grad_guard[i];
                    }
                }
            }
            drop(grad_guard);
            param.tensor.storage.bump_version();
        }

        Ok(())
    }

    fn set_lr(&mut self, lr: f32) {
        self.lr = lr;
    }

    fn get_lr(&self) -> f32 {
        self.lr
    }
}