kylä lähtee!

network-poc/node/src/burn_smollm/modules.rs

@@ -0,0 +1,59 @@
+use burn::module::{Module, Param};
+use burn::tensor::{backend::Backend, Tensor};
+
+#[derive(Module, Debug)]
+pub struct RmsNorm<B: Backend> {
+    pub weight: Param<Tensor<B, 1>>,
+    epsilon: f64,
+}
+
+impl<B: Backend> RmsNorm<B> {
+    pub fn new(size: usize, epsilon: f64, device: &B::Device) -> Self {
+        let weight = Param::from_tensor(Tensor::ones([size], device));
+        Self { weight, epsilon }
+    }
+
+    pub fn forward(&self, x: Tensor<B, 3>) -> Tensor<B, 3> {
+        // x: [batch, seq_len, dim]
+        // RMSNorm: x * weight / sqrt(mean(x^2) + eps)
+        let x_sq = x.clone().powf_scalar(2.0);
+        // mean over last dim, keeping dims for broadcast
+        let [b, s, d] = x_sq.dims();
+        let variance = x_sq.sum_dim(2).div_scalar(d as f32);
+        let norm = x.div(variance.add_scalar(self.epsilon).sqrt());
+
+        let w = self.weight.val().unsqueeze::<2>().unsqueeze::<3>().reshape([1, 1, d]);
+        norm * w
+    }
+}
+
+#[derive(Module, Debug)]
+pub struct Mlp<B: Backend> {
+    pub gate_proj: Param<Tensor<B, 2>>, // [in, intermediate]
+    pub up_proj: Param<Tensor<B, 2>>,   // [in, intermediate]
+    pub down_proj: Param<Tensor<B, 2>>, // [intermediate, out]
+}
+
+impl<B: Backend> Mlp<B> {
+    pub fn new(hidden_size: usize, intermediate_size: usize, device: &B::Device) -> Self {
+        Self {
+            gate_proj: Param::from_tensor(Tensor::zeros([hidden_size, intermediate_size], device)),
+            up_proj: Param::from_tensor(Tensor::zeros([hidden_size, intermediate_size], device)),
+            down_proj: Param::from_tensor(Tensor::zeros([intermediate_size, hidden_size], device)),
+        }
+    }
+
+    pub fn forward(&self, x: Tensor<B, 3>) -> Tensor<B, 3> {
+        // x: [batch, seq, hidden]
+        // gate = x @ gate_proj  -> [batch, seq, intermediate]
+        let gate = x.clone().matmul(self.gate_proj.val().unsqueeze());
+        let up = x.matmul(self.up_proj.val().unsqueeze());
+
+        // SiLU(gate) * up
+        let silu = gate.clone() * burn::tensor::activation::sigmoid(gate);
+        let intermediate = silu * up;
+
+        // intermediate @ down_proj -> [batch, seq, hidden]
+        intermediate.matmul(self.down_proj.val().unsqueeze())
+    }
+}