MergeKit项目中Gemma2模型合并时的参数异常问题分析

在开源项目MergeKit的使用过程中，用户报告了一个有趣的技术现象：当合并两个Gemma2 9B模型时，结果模型意外获得了额外的10亿参数。经过技术分析，我们发现这是由于模型权重文件中存在重复的lm_head.weight张量导致的。

问题现象

用户在使用MergeKit工具合并两个Gemma2 9B模型时，发现合并后的模型参数数量从预期的90亿增加到了100亿。类似现象也在Gemma2 2B模型的合并过程中被观察到。经过检查，发现合并后的模型权重文件中存在重复的lm_head.weight张量。

技术原理

这种现象的根本原因在于Gemma2模型架构的特殊设计。与许多其他LLM模型类似，Gemma2采用了共享权重机制——模型的输出层(lm_head)权重实际上是输入嵌入层权重的转置。这种设计既能减少模型参数数量，又能保持嵌入空间的一致性。

在模型合并过程中，MergeKit可能错误地保留了原始模型中的lm_head.weight张量，而实际上这个张量应该被移除，因为它已经作为嵌入权重的一部分存在于模型中。

解决方案

要解决这个问题，需要手动从合并后的模型文件中移除重复的lm_head.weight张量。具体步骤如下：

1. 编辑model.safetensors.index.json文件，移除其中对lm_head.weight的引用
2. 使用Python脚本处理safetensors文件，移除实际的张量数据

from safetensors import safe_open
from safetensors.torch import save_file

input_file = "model-00001-of-00002.safetensors"
output_file = "fixed-model-00001-of-00002.safetensors"

tensors = {}
with safe_open(input_file, framework="pt", device="cpu") as f:
    for key in f.keys():
        if key != "lm_head.weight":
            tensors[key] = f.get_tensor(key)

save_file(tensors, output_file)

模型架构设计的启示

这个问题揭示了LLM模型设计中权重共享机制的重要性。Gemma2模型采用的这种输出层与嵌入层权重共享的设计，与Command-R等模型类似，是一种常见的参数优化技术。开发者在进行模型合并时，需要特别注意这种特殊的架构设计，避免参数重复计算。

结论

MergeKit作为模型合并工具，在处理具有特殊权重共享机制的模型时需要特别注意。开发者在使用时应当：

1. 了解原始模型的架构特点
2. 检查合并后的参数数量是否符合预期
3. 必要时手动处理权重文件中的重复张量

这个问题不仅是一个工具使用问题，更反映了现代LLM模型架构设计的复杂性，提醒我们在模型操作时需要更深入地理解其内部机制。