Skip to content

Commit 3557601

Browse files
author
Jiwoong
authored
fix: handle quantized tie-embedding for Gemma3Text (#372)
Gemma3TextModel.santize method can now handle quantized tie-embedding.
1 parent 55daf5b commit 3557601

File tree

1 file changed

+4
-2
lines changed

1 file changed

+4
-2
lines changed

Libraries/MLXLLM/Models/Gemma3Text.swift

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -362,8 +362,10 @@ public class Gemma3TextModel: Module, LLMModel {
362362
}
363363

364364
if processedWeights["lm_head.weight"] == nil {
365-
if let embedWeight = processedWeights["model.embed_tokens.weight"] {
366-
processedWeights["lm_head.weight"] = embedWeight
365+
["weight", "scales", "biases"].forEach { key in
366+
if let embedWeight = processedWeights["model.embed_tokens.\(key)"] {
367+
processedWeights["lm_head.\(key)"] = embedWeight
368+
}
367369
}
368370
}
369371
return processedWeights

0 commit comments

Comments
 (0)