mergekit项目地址
Mergekit提供模型合并方法可以概况为四大类：基本线性加权、基于球面插值、基于任务向量 以及一些专业化方法，今天我们来刷下基于任务向量的ties合并方法，熟悉原理和代码。

Ties背景

首先明确：Ties方法建立在“任务向量”的概念之上，任务向量表示参考的微调模型和基本模型之间的区别，这种方法在mergekit中非常多，Ties只是其中一种变体。
TIES-Merging ——可将多个同构不同参数的模型合并为单个多任务模型。主要解决模型合并中的两个主要挑战：

1.模型参数冗余：它识别并消除特定任务模型中的冗余参数。这是通过关注微调期间所做的更改、识别前 k% 最重要的更改并丢弃其余的来实现的。
2.参数符号之间的不一致：当不同模型对同一参数提出相反的调整时，就会出现冲突。TIES-Merging 通过创建一个统一的符号向量来解决这些冲突，该符号向量代表所有模型中最主要的变化方向。

Ties 核心思想

稀疏化密度：通过仅保留一小部分最重要的参数（密度参数）并将其余参数重置为零，减少特定于任务的模型中的冗余。
符号共识： 通过根据累积幅度的最主要方向（正或负）创建统一的符号向量，解决不同模型之间的符号冲突。
不相交合并：对与统一符号向量对齐的参数值进行平均，不包括零值。

具体表现为当多个模型进行合并时候，稀疏化密度只保留最重要的权重高化，符号共识**保留多个模型再更新方向上一致的参数。

具体流程

通过分析代码后，可以汇总为以下步骤：

1. 计算任务向量——每个模型和base模型的差异，比较简单

   def get_task_vectors(
    weight_info: WeightInfo,
    base_model: ModelReference,
    tensors: ImmutableMap[ModelReference, torch.Tensor],
    tensor_parameters: ImmutableMap[ModelReference, ImmutableMap[str, Any]],
) -> Tuple[List[Dict[str, Any]], torch.Tensor]:
    keys = list(tensors.keys())
    base = tensors[base_model]

    parameter_name = weight_info.name

    res = []
    for model in keys:
        if model == base_model:
            continue

        x = tensors[model].to(base.dtype)
        if x.shape != base.shape:
            if weight_info.is_embed:
                x = x[: base.shape[0], : base.shape[1]]
                logging.warning(f"Using submatrix of {model}:{parameter_name}")
            else:
                logging.warning(
                    f"skipping {model}:{parameter_name} due to size mismatch"
                )
                continue

        delta = x - base
        del x
        del tensors[model]

        d = {}
        d["model"] = model
        d["delta"] = delta
        for p in tensor_parameters[model]:
            d[p] = tensor_parameters[model][p]
        res.append(d)
    return res, base

2. 稀疏化权重

# sparsify
        if self.method.sparsification_method:
            for tv_info in tvs:
                kwargs = {}
                if "gamma" in tv_info:
                    kwargs["gamma"] = tv_info["gamma"]

                if "epsilon" in tv_info:
                    kwargs["epsilon"] = tv_info["epsilon"]

                tv_info["delta"] = sparsify(
                    tv_info["delta"],
                    density=tv_info["density"],
                    method=self.method.sparsification_method,
                    rescale_norm=self.rescale_norm,
                    **kwargs,
                )

然后，如何计算权重幅度等等，在源码中sparsify具备四种稀疏化方法：

• magnitude: 基于权重大小的剪枝
• random: 随机剪枝
• magnitude_outliers: 去除极大值和极小值的剪枝
• della_magprune: 基于概率的渐进式剪枝方法

这里我们主要说ties 基于权重幅度剪枝—— 保留绝对值最大的k个元素，k=density*总元素数，简单高效

def magnitude(
    tensor: torch.Tensor, density: float, rescale_norm: Optional[RescaleNorm] = None
) -> torch.Tensor:
    """Masks out the smallest values, retaining a proportion of `density`."""
    if density >= 1:
        return tensor

    k = int(density * tensor.numel())  #计算保留元素数量

    assert k > 0, "not gonna zero out the whole tensor buddy"
    mask = torch.zeros_like(tensor)
    w = tensor.abs().view(-1)
    if w.device.type == "cpu":
        w = w.float()
    topk = torch.argsort(w, descending=True)[:k]  #对绝对值进行降序排序，获取前k大值的索引
    mask.view(-1)[topk] = 1  #将掩码中对应top-k索引的位置设为1

    res = rescaled_masked_tensor(tensor, mask, rescale_norm)  #调用辅助函数应用掩码并根据需要重新缩放，保持特定的范数特性，保证输出和反向传播计算的稳定性
    return res

通过上述计算，可以保留所谓权重幅度最大的参数，完成稀疏化。
3. 权重应用

		deltas = torch.stack([tv["delta"] for tv in tvs], dim=0)

        weights = torch.tensor(
            [tv["weight"] for tv in tvs], dtype=deltas.dtype, device=deltas.device
        )
        while len(deltas.shape) > len(weights.shape):
            weights.unsqueeze_(-1)

        weighted_deltas = deltas * weights
       

4.符号共识
Ties提供两种方式计算符号
”sum“: 加权，考虑参数幅度
”count“，基于符号数量统计

def get_mask(
    delta: torch.Tensor,
    method: Literal["sum", "count"] = "sum",
    mask_dtype: Optional[torch.dtype] = None,
):
    """Returns a mask determining which delta vectors should be merged
    into the final model.

    For the methodology described in the TIES paper use 'sum'. For a
    simpler naive count of signs, use 'count'."""
    if mask_dtype is None:
        mask_dtype = delta.dtype

    sign = delta.sign().to(mask_dtype)  # 获取每个元素的符号（-1, 0, +1）

    if method == "sum":
        sign_weight = delta.sum(dim=0) # 沿模型维度求和
        majority_sign = (sign_weight >= 0).to(mask_dtype) * 2 - 1    # 转换为±1
        del sign_weight
    elif method == "count":
        majority_sign = (sign.sum(dim=0) >= 0).to(mask_dtype) * 2 - 1
    else:
        raise RuntimeError(f'Unimplemented mask method "{method}"')

    return sign == majority_sign  #生成bool mask

在Sum方法中也是ties的论文方法，考虑差异的幅度和方向：

对每个参数位置，计算所有模型差异的总和；如果总和≥0，多数符号为+1，否则为-1，这样幅度的差异大对结果影响较大；

在count方法中，对每个参数位置，统计正负号数量，正号多则多数符号为+1，否则为-1，忽略差异幅度，只考虑方向

计算更简单但可能不够精确。

5. 合并

回到主线代码，这里我们已经拿到了掩码，可以确定哪些参数变化(deltas)应该被合并

if self.method.consensus_method:
            mask_dtype = torch.int8 if self.int8_mask else base.dtype
            mask = get_mask(
                weighted_deltas,
                method=self.method.consensus_method,
                mask_dtype=mask_dtype,
            ) #拿到mask后
            mixed_delta = (weighted_deltas * mask).sum(dim=0) #直接对所有加权deltas求和
            divisor = (weights * mask).sum(dim=0) #计算有效权重的和（用于归一化）
            divisor[divisor == 0] = 1
        else:
            mixed_delta = weighted_deltas.sum(dim=0) #只保留被掩码选中的deltas并求和
            divisor = weights.sum(dim=0)
            divisor[divisor.abs() < 1e-8] = 1

        if self.normalize:  # 归一化
            mixed_delta /= divisor 

        if self.lambda_ != 1: #系数缩放
            mixed_delta *= self.lambda_

        return (base + mixed_delta).to(base.dtype) 合并

总结

Mergekit ties的配置参考

models:
  - model: psmathur/orca_mini_v3_13b        #参考模型1
    parameters:
      density: [1, 0.7, 0.1] # density gradient   这是稀疏化的密度列表 对应不同层的稀疏化成都
      weight: 1.0                                #权重值
  - model: garage-bAInd/Platypus2-13B         #参考模型2 其余同上
    parameters:
      density: 0.5                           
      weight: [0, 0.3, 0.7, 1] # weight gradient
  - model: WizardLM/WizardMath-13B-V1.0       #参考模型3 其余同上
    parameters:
      density: 0.33
      weight:
        - filter: mlp  #对于MLP是取0.5 其他层是0
          value: 0.5
        - value: 0
merge_method: ties
base_model: TheBloke/Llama-2-13B-fp16      #选择这个模型作为基础模型
parameters:
  normalize: true
  int8_mask: true
dtype: float16

这里当参数为List时候，Mergekit会进行映射到不同层，在处理每个参数时候根据位置选择List中的密度或者权重值，保留嵌入层 、中间层、末尾层对应不同密度稀疏，考虑使用线性插值来计算中间层的密度。