Model Merging：如何把多个专家模型合并成一个？

import torch

def task_vector(base_model, finetuned_model, scale=1.0):
    """计算任务向量"""
    return (finetuned_model - base_model) * scale

def merge_task_vectors(base_model, task_vectors):
    """将多个任务向量合并到基础模型"""
    merged = {k: v.clone() for k, v in base_model.items()}
    for delta in task_vectors:
        for k in merged:
            merged[k] += delta[k]
    return merged

# 例子：合并数学和代码两个专家
math_delta = task_vector(base, math_model, scale=0.7)
code_delta = task_vector(base, code_model, scale=0.5)
merged_weights = merge_task_vectors(base.state_dict(), [math_delta, code_delta])

def trim(delta, keep_ratio=0.2):
    """只保留 top K% 的参数"""
    flat = delta.flatten()
    threshold = torch.kthvalue(
        flat, int(len(flat) * (1 - keep_ratio))
    ).values
    mask = delta.abs() >= threshold.abs()
    return delta * mask

def elect_sign(deltas):
    """符号选举，解决方向冲突"""
    stacked = torch.stack(deltas, dim=0)  # [num_tasks, ...]
    signs = torch.sign(torch.sum(stacked, dim=0))  # 多数符号
    result = []
    for delta in deltas:
        aligned = delta * (signs != 0)  # 保留与多数一致的
        result.append(aligned * (aligned * signs >= 0))  # mask 冲突的
    return result

def disjoint_merge(base, deltas, scales):
    """不交并合并"""
    merged = base.clone()
    for delta, scale in zip(deltas, scales):
        merged += delta * scale
    return merged

def dare_drop(delta, drop_prob=0.9):
    """DARE：丢弃 + 重缩放"""
    mask = (torch.rand_like(delta) > drop_prob).float()
    dropped = delta * mask
    rescaled = dropped / (1 - drop_prob)
    return rescaled

def dare_then_ties(base, finetuned_models, drop_prob=0.9, keep_ratio=0.2):
    deltas = [ft - base for ft in finetuned_models]

    # Step 1: DARE 激进裁剪
    deltas = [dare_drop(d, drop_prob) for d in deltas]

    # Step 2: TIES 符号对齐
    deltas = elect_sign(deltas)

    # Step 3: 不交并合并
    scales = [1.0 / (1 - drop_prob)] * len(deltas)
    return disjoint_merge(base, deltas, scales)

# Fisher Information 加权 Median
# 对每个参数，选择使 Σ Fᵢ|θ - Δᵢ| 最小的 θ
# 其中 Fᵢ 是该参数的 Fisher 信息

# 优化目标：最小化合并模型与各专家隐藏状态的差异
loss = α_l × L2(hidden_states_merge, hidden_states_expert_l)
       + KL(logits_merge, logits_expert)
# α_l 是每层的可学习系数向量

pip install mergekit

# merge.yaml
models:
  - model: meta-llama/Llama-3-8B-Instruct
    parameters:
      weight: 1.0
  - model: meta-llama/Llama-3-8B-Code
    parameters:
      weight: 0.7
  - model: meta-llama/Llama-3-8B-Math
    parameters:
      weight: 0.5

merge_method: ties
base_model: meta-llama/Llama-3-8B

ties:
  trim_weight: 0.2
  rescale_weights: true

mergekit merge merge.yaml merged-model/ --output trust_remote_code=True

# 正确的 LoRA 合并流程
from peft import PeftModel

# 先将 LoRA 权重合并回基础模型
model = PeftModel.from_pretrained(base, lora_path)
model = model.merge_and_unload()
# 然后再对合并后的模型做 Task Vector 合并

# 评估脚本框架
def evaluate_merge(original_models, merged_model, tasks):
    results = {}
    for task_name, eval_fn in tasks.items():
        original_scores = [evaluate(m, eval_fn) for m in original_models]
        merged_score = evaluate(merged_model, eval_fn)
        forgetting = merged_score - max(original_scores)
        results[task_name] = {
            "merged": merged_score,
            "forgetting": forgetting  # 负值说明退化了
        }
    return results