flodl 0.5.2

floDl — a flow-graph deep learning framework built on libtorch
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

use crate::autograd::Variable;
use crate::tensor::{Device, DType, Result, Tensor, TensorOptions};

use super::parameter::Parameter;
use super::Module;

/// Root Mean Square Layer Normalization.
///
/// Normalizes the last dimension without centering (no bias subtraction),
/// making it cheaper than LayerNorm. Used in LLaMA, Gemma, and other
/// modern architectures.
///
/// Formula: `y = x / sqrt(mean(x^2) + eps) * weight`
pub struct RMSNorm {
    pub weight: Parameter,
    eps: f64,
}

impl RMSNorm {
    /// Create an RMSNorm normalizing over the last `size` elements on CPU.
    pub fn new(size: i64) -> Result<Self> {
        Self::on_device(size, Device::CPU)
    }

    /// Create an RMSNorm on a specific device.
    pub fn on_device(size: i64, device: Device) -> Result<Self> {
        let opts = TensorOptions { dtype: DType::Float32, device };
        let weight = Variable::new(Tensor::ones(&[size], opts)?, true);

        Ok(RMSNorm {
            weight: Parameter {
                variable: weight,
                name: "weight".into(),
            },
            eps: 1e-5,
        })
    }

    /// Set epsilon for numerical stability (default: 1e-5).
    pub fn eps(mut self, eps: f64) -> Self {
        self.eps = eps;
        self
    }
}

impl Module for RMSNorm {
    fn name(&self) -> &str { "rmsnorm" }

    fn forward(&self, input: &Variable) -> Result<Variable> {
        // RMS = sqrt(mean(x^2) + eps)
        let x2 = input.pow_scalar(2.0)?;
        let mean_x2 = x2.mean_dim(-1, true)?;
        let rms = mean_x2.add_scalar(self.eps)?.sqrt()?;
        let normalized = input.div(&rms)?;
        normalized.mul(&self.weight.variable)
    }

    fn parameters(&self) -> Vec<Parameter> {
        vec![self.weight.clone()]
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::tensor::test_device;

    #[test]
    fn test_rmsnorm_shape() {
        let device = test_device();
        let norm = RMSNorm::on_device(4, device).unwrap();
        let x = Variable::new(
            Tensor::randn(&[2, 3, 4], crate::tensor::test_opts()).unwrap(),
            false,
        );
        let y = norm.forward(&x).unwrap();
        assert_eq!(y.shape(), vec![2, 3, 4]);
    }

    #[test]
    fn test_rmsnorm_unit_rms() {
        let device = test_device();
        let norm = RMSNorm::on_device(4, device).unwrap();
        // Input with known RMS: [1, 1, 1, 1] → RMS=1 → output ≈ [1, 1, 1, 1]
        let x = Variable::new(
            Tensor::ones(&[1, 4], crate::tensor::test_opts()).unwrap(),
            false,
        );
        let y = norm.forward(&x).unwrap().data().to_f32_vec().unwrap();
        for v in &y {
            assert!((v - 1.0).abs() < 1e-3);
        }
    }

    #[test]
    fn test_rmsnorm_gradient() {
        let device = test_device();
        let norm = RMSNorm::on_device(4, device).unwrap();
        let x = Variable::new(
            Tensor::randn(&[2, 4], crate::tensor::test_opts()).unwrap(),
            true,
        );
        let y = norm.forward(&x).unwrap();
        let loss = y.sum().unwrap();
        loss.backward().unwrap();

        let grad = x.grad().unwrap();
        assert_eq!(grad.shape(), vec![2, 4]);
        // Weight should also have gradients
        let w_grad = norm.weight.variable.grad().unwrap();
        assert_eq!(w_grad.shape(), vec![4]);
    }

    #[test]
    fn test_rmsnorm_parameters() {
        let norm = RMSNorm::new(8).unwrap();
        let params = norm.parameters();
        assert_eq!(params.len(), 1);
        assert_eq!(params[0].name, "weight");
    }
}