Merge pull request #21 from SharpAI/fix/gemma4-moe-unhandled-keys

solderzzc · web-flow · commit 7fb05f7dbe5b · 2026-04-16T10:06:46.000-07:00
Fix: Gemma 4 MoE Loading Failure (gemma4_text)
diff --git a/Libraries/MLXLLM/Models/Gemma4Text.swift b/Libraries/MLXLLM/Models/Gemma4Text.swift
@@ -33,6 +33,10 @@ public struct Gemma4TextConfiguration: Codable, Sendable {
     var attentionKeqV: Bool = false
     var finalLogitSoftcapping: Float = 30.0
     var useDoubleWideMlp: Bool = true
+    var enableMoEBlock: Bool = false
+    var numExperts: Int?
+    var topKExperts: Int?
+    var moeIntermediateSize: Int?
     var layerTypes: [String] = []
     var tieWordEmbeddings: Bool = true
 
@@ -66,6 +70,10 @@ public struct Gemma4TextConfiguration: Codable, Sendable {
         case attentionKeqV = "attention_k_eq_v"
         case finalLogitSoftcapping = "final_logit_softcapping"
         case useDoubleWideMlp = "use_double_wide_mlp"
+        case enableMoEBlock = "enable_moe_block"
+        case numExperts = "num_experts"
+        case topKExperts = "top_k_experts"
+        case moeIntermediateSize = "moe_intermediate_size"
         case layerTypes = "layer_types"
         case tieWordEmbeddings = "tie_word_embeddings"
         case ropeParameters = "rope_parameters"
@@ -110,6 +118,14 @@ public struct Gemma4TextConfiguration: Codable, Sendable {
             try container.decodeIfPresent(Float.self, forKey: .finalLogitSoftcapping) ?? 30.0
         self.useDoubleWideMlp =
             try container.decodeIfPresent(Bool.self, forKey: .useDoubleWideMlp) ?? true
+        self.enableMoEBlock =
+            try container.decodeIfPresent(Bool.self, forKey: .enableMoEBlock) ?? false
+        self.numExperts =
+            try container.decodeIfPresent(Int.self, forKey: .numExperts)
+        self.topKExperts =
+            try container.decodeIfPresent(Int.self, forKey: .topKExperts)
+        self.moeIntermediateSize =
+            try container.decodeIfPresent(Int.self, forKey: .moeIntermediateSize)
         if let decoded = try container.decodeIfPresent([String].self, forKey: .layerTypes) {
             self.layerTypes = decoded
         } else {
@@ -374,6 +390,89 @@ private class Gemma4MLP: Module {
     }
 }
 
+// MARK: - MoE Router
+
+private class Gemma4TextRouter: Module {
+    let topKExperts: Int
+    let rootSize: Float
+
+    @ModuleInfo(key: "norm") var norm: RMSNormNoScale
+    @ModuleInfo(key: "proj") var proj: Linear
+    @ModuleInfo(key: "scale") var scale: MLXArray
+    @ModuleInfo(key: "per_expert_scale") var perExpertScale: MLXArray
+
+    init(_ config: Gemma4TextConfiguration) {
+        guard let numExperts = config.numExperts, let topKExperts = config.topKExperts else {
+            fatalError("Gemma4 MoE router requires numExperts and topKExperts")
+        }
+
+        self.topKExperts = topKExperts
+        self.rootSize = pow(Float(config.hiddenSize), -0.5)
+
+        self._norm.wrappedValue = RMSNormNoScale(eps: config.rmsNormEps)
+        self._proj.wrappedValue = Linear(config.hiddenSize, numExperts, bias: false)
+        self._scale.wrappedValue = MLXArray.ones([config.hiddenSize])
+        self._perExpertScale.wrappedValue = MLXArray.ones([numExperts])
+        super.init()
+    }
+
+    func callAsFunction(_ x: MLXArray) -> (MLXArray, MLXArray) {
+        var x = norm(x)
+        x = x * MLXArray(rootSize, dtype: x.dtype)
+        x = x * scale.asType(x.dtype)
+
+        let expertScores = proj(x)
+        let routerProbabilities = MLX.softmax(expertScores, axis: -1, precise: true)
+
+        let topKIndices = MLX.argPartition(-expertScores, kth: topKExperts - 1, axis: -1)[
+            .ellipsis, ..<topKExperts,
+        ]
+        var topKWeights = MLX.takeAlong(routerProbabilities, topKIndices, axis: -1)
+        topKWeights = topKWeights / MLX.sum(topKWeights, axis: -1, keepDims: true)
+        topKWeights = topKWeights * perExpertScale[topKIndices].asType(topKWeights.dtype)
+        return (topKIndices, topKWeights)
+    }
+}
+
+// MARK: - MoE Experts
+
+private class Gemma4TextExperts: Module {
+    @ModuleInfo(key: "switch_glu") var switchGLU: SwitchGLU
+
+    init(_ config: Gemma4TextConfiguration) {
+        guard let numExperts = config.numExperts,
+            let moeIntermediateSize = config.moeIntermediateSize
+        else {
+            fatalError("Gemma4 MoE experts require numExperts and moeIntermediateSize")
+        }
+
+        self._switchGLU.wrappedValue = SwitchGLU(
+            inputDims: config.hiddenSize,
+            hiddenDims: moeIntermediateSize,
+            numExperts: numExperts,
+            activation: geluApproximate,
+            bias: false
+        )
+        super.init()
+    }
+
+    func callAsFunction(
+        _ x: MLXArray, topKIndices: MLXArray, topKWeights: MLXArray
+    ) -> MLXArray {
+        let batch = x.dim(0)
+        let length = x.dim(1)
+        let hidden = x.dim(2)
+        let topK = topKIndices.dim(-1)
+
+        let expertOutput = switchGLU(
+            x.reshaped(batch * length, hidden),
+            topKIndices.reshaped(batch * length, topK)
+        )
+        let weights = topKWeights.reshaped(batch * length, topK, 1).asType(expertOutput.dtype)
+        return (expertOutput * weights).sum(axis: -2).reshaped(batch, length, hidden)
+    }
+}
+
 // MARK: - Decoder Layer
 
 private class Gemma4DecoderLayer: Module {
@@ -388,6 +487,11 @@ private class Gemma4DecoderLayer: Module {
     @ModuleInfo(key: "post_attention_layernorm") var postAttentionLayernorm: RMSNorm
     @ModuleInfo(key: "pre_feedforward_layernorm") var preFeedforwardLayernorm: RMSNorm
     @ModuleInfo(key: "post_feedforward_layernorm") var postFeedforwardLayernorm: RMSNorm
+    @ModuleInfo(key: "router") var router: Gemma4TextRouter?
+    @ModuleInfo(key: "experts") var experts: Gemma4TextExperts?
+    @ModuleInfo(key: "post_feedforward_layernorm_1") var postFeedforwardLayernorm1: RMSNorm?
+    @ModuleInfo(key: "post_feedforward_layernorm_2") var postFeedforwardLayernorm2: RMSNorm?
+    @ModuleInfo(key: "pre_feedforward_layernorm_2") var preFeedforwardLayernorm2: RMSNorm?
 
     // Per-layer input (PLE) gating
     @ModuleInfo(key: "per_layer_input_gate") var perLayerInputGate: Linear?
@@ -415,6 +519,17 @@ private class Gemma4DecoderLayer: Module {
         self._postFeedforwardLayernorm.wrappedValue = RMSNorm(
             dimensions: config.hiddenSize, eps: config.rmsNormEps)
 
+        if config.enableMoEBlock {
+            self._router.wrappedValue = Gemma4TextRouter(config)
+            self._experts.wrappedValue = Gemma4TextExperts(config)
+            self._postFeedforwardLayernorm1.wrappedValue = RMSNorm(
+                dimensions: config.hiddenSize, eps: config.rmsNormEps)
+            self._postFeedforwardLayernorm2.wrappedValue = RMSNorm(
+                dimensions: config.hiddenSize, eps: config.rmsNormEps)
+            self._preFeedforwardLayernorm2.wrappedValue = RMSNorm(
+                dimensions: config.hiddenSize, eps: config.rmsNormEps)
+        }
+
         if hiddenSizePerLayerInput > 0 {
             self._perLayerInputGate.wrappedValue = Linear(
                 config.hiddenSize, hiddenSizePerLayerInput, bias: false)
@@ -446,8 +561,26 @@ private class Gemma4DecoderLayer: Module {
         var out = residual + postAttn
 
         let residual2 = out
-        out = preFeedforwardLayernorm(out)
-        out = mlp(out)
+        if let router, let experts,
+            let postFeedforwardLayernorm1,
+            let postFeedforwardLayernorm2,
+            let preFeedforwardLayernorm2
+        {
+            // MoE: dual dense + sparse feedforward
+            var dense = preFeedforwardLayernorm(out)
+            dense = mlp(dense)
+            dense = postFeedforwardLayernorm1(dense)
+
+            let (topKIndices, topKWeights) = router(out)
+            var sparse = preFeedforwardLayernorm2(out)
+            sparse = experts(sparse, topKIndices: topKIndices, topKWeights: topKWeights)
+            sparse = postFeedforwardLayernorm2(sparse)
+
+            out = dense + sparse
+        } else {
+            out = preFeedforwardLayernorm(out)
+            out = mlp(out)
+        }
         out = postFeedforwardLayernorm(out)
         out = residual2 + out
 
@@ -675,6 +808,34 @@ public class Gemma4TextModel: Module, LLMModel, KVCacheDimensionProvider {
             {
                 continue
             }
+
+            // MoE expert weight remapping: fused HF tensors → SwitchGLU layout
+            if k.hasSuffix(".experts.down_proj") {
+                sanitized[
+                    k.replacingOccurrences(
+                        of: ".experts.down_proj",
+                        with: ".experts.switch_glu.down_proj.weight"
+                    )
+                ] = v
+                continue
+            }
+            if k.hasSuffix(".experts.gate_up_proj") {
+                let mid = v.dim(-2) / 2
+                sanitized[
+                    k.replacingOccurrences(
+                        of: ".experts.gate_up_proj",
+                        with: ".experts.switch_glu.gate_proj.weight"
+                    )
+                ] = v[.ellipsis, ..<mid, 0...]
+                sanitized[
+                    k.replacingOccurrences(
+                        of: ".experts.gate_up_proj",
+                        with: ".experts.switch_glu.up_proj.weight"
+                    )
+                ] = v[.ellipsis, mid..., 0...]
+                continue
+            }
+
             sanitized[k] = v
         }
         return sanitized
diff --git a/Tests/MLXLMTests/Gemma4Tests.swift b/Tests/MLXLMTests/Gemma4Tests.swift
@@ -97,5 +97,59 @@ extension MLXTestingSuite {
         #expect(!sum.isNaN)
         #expect(!sum.isInfinite)
     }
+
+    /// Create a minimal test configuration for Gemma 4 Text MoE
+    private func makeTinyTextMoEConfigData() -> Data {
+        let json = """
+        {
+            "model_type": "gemma4_text",
+            "hidden_size": 64,
+            "num_hidden_layers": 2,
+            "intermediate_size": 128,
+            "num_attention_heads": 4,
+            "head_dim": 16,
+            "global_head_dim": 64,
+            "rms_norm_eps": 1e-6,
+            "vocab_size": 100,
+            "num_key_value_heads": 2,
+            "rope_traditional": false,
+            "sliding_window": 128,
+            "sliding_window_pattern": 1,
+            "max_position_embeddings": 512,
+            "num_kv_shared_layers": 0,
+            "use_double_wide_mlp": false,
+            "tie_word_embeddings": true,
+            "hidden_size_per_layer_input": 32,
+            "vocab_size_per_layer_input": 10,
+            "final_logit_softcapping": 30.0,
+            "enable_moe_block": true,
+            "num_experts": 4,
+            "top_k_experts": 2,
+            "moe_intermediate_size": 128,
+            "attention_k_eq_v": false
+        }
+        """
+        return json.data(using: .utf8)!
+    }
+
+    @Test("Gemma 4 Text MoE Instantiation & Forward Pass")
+    func testGemma4TextMoEInstantiationAndForward() throws {
+        let data = makeTinyTextMoEConfigData()
+        let config = try JSONDecoder().decode(Gemma4TextConfiguration.self, from: data)
+        let model = Gemma4TextModel(config)
+        #expect(model.vocabularySize == 100)
+        
+        // This validates that the conditional MoE logic and SwitchGLU layer initialize properly
+        // without crashing, proving we correctly load gemma4_text active MoEs.
+        let input = MLXArray(0..<8).reshaped(1, 8)
+        let output = model(input, cache: nil)
+
+        // Ensure dimensionality is correct
+        #expect(output.shape == [1, 8, model.vocabularySize])
+        
+        let sum = output.sum().item(Float.self)
+        #expect(!sum.isNaN)
+        #expect(!sum.isInfinite)
+    }
 }
 }
