combs-models 0.2.2

Combs Engine model architecture registry (Llama family)
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

//! RMSNorm (Llama-style, no learnable bias).

use burn::tensor::{Tensor, backend::Backend};

use crate::precision::{to_f32, to_float};

/// `y = x / rms(x) * w` where `rms` is taken over the last dimension and
/// `eps` is added inside the square root.
///
/// The reduction runs in f32 for f16 stability (no-op in f32 builds).
pub fn rms_norm<B: Backend, const D: usize>(
    x: Tensor<B, D>,
    weight: Tensor<B, 1>,
    eps: f64,
) -> Tensor<B, D> {
    let out_dtype = x.dtype();
    let dims = x.dims();
    let hidden = dims[D - 1];

    let xf = to_f32(x);
    // mean(x^2) over the last dim, keeping rank for broadcasting.
    let mean_sq = xf.clone().powf_scalar(2.0).mean_dim(D - 1);
    let inv_rms = mean_sq.add_scalar(eps).sqrt().recip();

    let mut shape = [1usize; D];
    shape[D - 1] = hidden;
    let y = xf * inv_rms * to_f32(weight).reshape(shape);
    to_float(y, out_dtype)
}

/// Gemma-style RMSNorm: `y = x / rms(x) * (1 + w)` — the learnable weight
/// is zero-centered (HF `Gemma3RMSNorm`).
pub fn gemma_rms_norm<B: Backend, const D: usize>(
    x: Tensor<B, D>,
    weight: Tensor<B, 1>,
    eps: f64,
) -> Tensor<B, D> {
    rms_norm(x, weight.add_scalar(1.0), eps)
}

/// LayerNorm (SigLIP-style, learnable weight + bias):
/// `y = (x - μ) / sqrt(σ² + eps) * w + b`, statistics over the last dim.
pub fn layer_norm<B: Backend, const D: usize>(
    x: Tensor<B, D>,
    weight: Tensor<B, 1>,
    bias: Tensor<B, 1>,
    eps: f64,
) -> Tensor<B, D> {
    let dims = x.dims();
    let hidden = dims[D - 1];

    let out_dtype = x.dtype();
    let xf = to_f32(x);
    let mean = xf.clone().mean_dim(D - 1);
    let centered = xf - mean;
    let var = centered.clone().powf_scalar(2.0).mean_dim(D - 1);
    let inv_std = var.add_scalar(eps).sqrt().recip();

    let mut shape = [1usize; D];
    shape[D - 1] = hidden;
    let y = centered * inv_std * to_f32(weight).reshape(shape.clone())
        + to_f32(bias).reshape(shape);
    to_float(y, out_dtype)
}

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

    type TestBackend = burn::backend::NdArray<f32>;

    #[test]
    fn normalizes_rows_to_unit_rms() {
        let device = burn::tensor::Device::<TestBackend>::default();
        let x: Tensor<TestBackend, 2> = Tensor::from_data(
            TensorData::new(vec![3.0f32, 4.0, 1.0, -2.0, 0.5, 2.5], [2, 3]),
            &device,
        );
        let w: Tensor<TestBackend, 1> = Tensor::ones([3], &device);
        let y = rms_norm(x, w, 1e-6);
        // With unit weights, each row of y must have RMS == 1 (up to eps).
        let rms = y
            .powf_scalar(2.0)
            .mean_dim(1)
            .sqrt()
            .into_data()
            .to_vec::<f32>()
            .unwrap();
        for (i, r) in rms.iter().enumerate() {
            assert!((r - 1.0).abs() < 1e-4, "row {i} rms = {r}");
        }
    }

    #[test]
    fn applies_weight() {
        let device = burn::tensor::Device::<TestBackend>::default();
        let x: Tensor<TestBackend, 2> =
            Tensor::from_data(TensorData::new(vec![1.0f32, 2.0], [1, 2]), &device);
        let w: Tensor<TestBackend, 1> =
            Tensor::from_data(TensorData::new(vec![2.0f32, 2.0], [2]), &device);
        let y = rms_norm(x.clone(), w, 1e-6);
        let z = rms_norm(x, Tensor::ones([2], &device), 1e-6);
        let yv: Vec<f32> = y.into_data().to_vec().unwrap();
        let zv: Vec<f32> = z.into_data().to_vec().unwrap();
        for i in 0..2 {
            assert!((yv[i] - 2.0 * zv[i]).abs() < 1e-4);
        }
    }
}