Mergekit——任务向量合并算法Ties解析

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.nameres = []for model in keys:if model == base_model:continuex = 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")continuedelta = x - basedel xdel tensors[model]d = {}d["model"] = modeld["delta"] = deltafor p in tensor_parameters[model]:d[p] = tensor_parameters[model][p]res.append(d)return res, base

2. 稀疏化权重

# sparsifyif 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 tensork = 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索引的位置设为1res = 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 mergedinto the final model.For the methodology described in the TIES paper use 'sum'. For asimpler naive count of signs, use 'count'."""if mask_dtype is None:mask_dtype = delta.dtypesign = 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    # 转换为±1del sign_weightelif method == "count":majority_sign = (sign.sum(dim=0) >= 0).to(mask_dtype) * 2 - 1else: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.dtypemask = 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] = 1else:mixed_delta = weighted_deltas.sum(dim=0) #只保留被掩码选中的deltas并求和divisor = weights.sum(dim=0)divisor[divisor.abs() < 1e-8] = 1if 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        #参考模型1parameters: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.33weight:- filter: mlp  #对于MLP是取0.5 其他层是0value: 0.5- value: 0
merge_method: ties
base_model: TheBloke/Llama-2-13B-fp16      #选择这个模型作为基础模型
parameters:normalize: trueint8_mask: true
dtype: float16

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