fix(sync): resolve overlapping omni-model tied word embedding logic

Author: Aegis-AI

Libraries/MLXVLM/Models/Gemma4VL.swift

@@ -517,7 +517,11 @@ public class Gemma4VL: Module, VLMModel, KVCacheDimensionProvider, LayerPartitio
         // This MUST be done because we explicitly allocated a separate lm_head linear layer!
         if processed["lm_head.weight"] == nil || config.tieWordEmbeddings {
             // Check both prefixed and flat keys to be robust against different sanitization outputs
-            let embedKeys = ["model.embed_tokens.weight", "embed_tokens.weight", "model.embedTokens.weight", "embedTokens.weight"]
+            let prefix = "language_model."
+            let embedKeys = [
+                "\(prefix)model.embed_tokens.weight", "\(prefix)embed_tokens.weight", "\(prefix)model.embedTokens.weight", "\(prefix)embedTokens.weight",
+                "model.embed_tokens.weight", "embed_tokens.weight", "model.embedTokens.weight", "embedTokens.weight"
+            ]
 
             for key in embedKeys {
                 if let embedWeights = processed[key] {
@@ -528,15 +532,21 @@ public class Gemma4VL: Module, VLMModel, KVCacheDimensionProvider, LayerPartitio
             }
 
             // Repeat for scales/biases if present (quantized models)
-            let scaleKeys = ["model.embed_tokens.scales", "embed_tokens.scales", "model.embedTokens.scales", "embedTokens.scales"]
+            let scaleKeys = [
+                "\(prefix)model.embed_tokens.scales", "\(prefix)embed_tokens.scales", "\(prefix)model.embedTokens.scales", "\(prefix)embedTokens.scales",
+                "model.embed_tokens.scales", "embed_tokens.scales", "model.embedTokens.scales", "embedTokens.scales"
+            ]
             for key in scaleKeys {
                 if let embedScales = processed[key] {
                     processed["lm_head.scales"] = embedScales
                     break
                 }
             }
 
-            let biasKeys = ["model.embed_tokens.biases", "embed_tokens.biases", "model.embedTokens.biases", "embedTokens.biases"]
+            let biasKeys = [
+                "\(prefix)model.embed_tokens.biases", "\(prefix)embed_tokens.biases", "\(prefix)model.embedTokens.biases", "\(prefix)embedTokens.biases",
+                "model.embed_tokens.biases", "embed_tokens.biases", "model.embedTokens.biases", "embedTokens.biases"
+            ]
             for key in biasKeys {
                 if let embedBiases = processed[key] {
                     processed["lm_head.biases"] = embedBiases

Tests/MLXLMTests/Gemma4Tests.swift

@@ -41,7 +41,7 @@ struct Gemma4Tests {
     @Test("Gemma 4 Forward Pass - Determinism & Shape")
     func testGemma4ForwardPass() throws {
         let config = makeTinyConfig()
-        let model = Gemma4Model(config)
+        let model = Gemma4ModelInternal(config)
 
         let input = MLXArray(0..<8).reshaped(1, 8)
         let output = model(input)
@@ -56,7 +56,7 @@ struct Gemma4Tests {
     @Test("PLE Multimodal Signal Integrity")
     func testPLESignalIntegrity() throws {
         let config = makeTinyConfig()
-        let model = Gemma4Model(config)
+        let model = Gemma4ModelInternal(config)
 
         let input = MLXArray(Int32(0)..<Int32(5)).reshaped(1, 5)
 
@@ -70,7 +70,7 @@ struct Gemma4Tests {
     @Test("Weight Sanitization - PLE Mapping")
     func testGemma4Sanitization() throws {
         let config = makeTinyConfig()
-        let model = Gemma4Model(config)
+        let model = Gemma4ModelInternal(config)
 
         var weights = [String: MLXArray]()
         weights["model.layers.0.per_layer_conditioning.scale"] = MLXArray.ones([config.hiddenSize, config.hiddenSizePerLayerInput])
@@ -85,7 +85,7 @@ struct Gemma4Tests {
     @Test("Audio Configuration Dependency Safety")
     func testAudioConfigSafety() throws {
         let config = makeTinyConfig()
-        let model = Gemma4Model(config)
+        let model = Gemma4ModelInternal(config)
         #expect(model.model.layers.count == config.hiddenLayers)
     }
 }