fix rmsnorm for gemma

2024-02-26 14:09:48 -08:00
parent a2ff291608
commit c7919cf7fe
1 changed files with 1 additions and 4 deletions
--- a/Libraries/LLM/Gemma.swift
+++ b/Libraries/LLM/Gemma.swift
@@ -19,10 +19,7 @@ private class RMSNorm: Module, UnaryLayer {
    }

    func norm(_ x: MLXArray) -> MLXArray {
-        let S = 1.0 / sqrt(Float(x.dim(-1)))
-
-        let n = (x * S).square().sum(axis: -1, keepDims: true)
-        return rsqrt(n + eps)
+        x * rsqrt(x.square().mean(axis: -1, keepDims: true) + eps)
    }

    public func callAsFunction(_ x: MLXArray) -> MLXArray {