模型合并技术实践：融合多个模型的智慧

def simple_average(models):
    """简单权重平均"""
    avg_state = {}
    for key in models[0].state_dict():
        avg_state[key] = sum(m.state_dict()[key] for m in models) / len(models)
    return avg_state

def slerp(t, v0, v1):
    """球面线性插值"""
    v0_norm = v0 / v0.norm()
    v1_norm = v1 / v1.norm()
    omega = torch.acos(torch.clamp(v0_norm @ v1_norm, -1, 1))
    so = torch.sin(omega)
    if so < 1e-6:
        return (1-t)*v0 + t*v1
    return torch.sin((1-t)*omega)/so * v0 + torch.sin(t*omega)/so * v1

def merge_slerp(model_a, model_b, t=0.5):
    """SLERP合并"""
    merged = {}
    for key in model_a.state_dict():
        merged[key] = slerp(t, model_a.state_dict()[key], model_b.state_dict()[key])
    return merged

def ties_merge(models, base_model, density=0.5):
    """TIES合并: Trim, Elect Sign, Disjoint Merge"""
    # 1. 计算每个模型相对于base的delta
    deltas = [m.state_dict() - base_model.state_dict() for m in models]
    
    # 2. Trim: 只保留每个delta中top-k的参数
    for delta in deltas:
        for key in delta:
            threshold = torch.quantile(delta[key].abs(), 1 - density)
            delta[key] = torch.where(delta[key].abs() > threshold, delta[key], 0)
    
    # 3. Elect Sign: 投票决定每个参数的符号
    merged = {}
    for key in base_model.state_dict():
        signs = sum(torch.sign(d[key]) for d in deltas)
        elected_sign = torch.sign(signs)
        
        # 4. Disjoint Merge: 只保留与选举符号一致的delta，取平均
        consistent = []
        for delta in deltas:
            mask = torch.sign(delta[key]) == elected_sign
            consistent.append(torch.where(mask, delta[key], 0))
        
        merged[key] = base_model.state_dict()[key] + sum(consistent) / max(1, sum(elected_sign != 0))
    
    return merged

def dare_merge(model_a, model_b, base_model, drop_rate=0.9):
    """DARE: Drop And REscale"""
    delta_a = model_a.state_dict() - base_model.state_dict()
    delta_b = model_b.state_dict() - base_model.state_dict()
    
    merged = {}
    for key in base_model.state_dict():
        # 随机丢弃大部分delta
        mask_a = (torch.rand_like(delta_a[key]) > drop_rate).float()
        mask_b = (torch.rand_like(delta_b[key]) > drop_rate).float()
        
        # 重新缩放
        dropped_a = delta_a[key] * mask_a / (1 - drop_rate)
        dropped_b = delta_b[key] * mask_b / (1 - drop_rate)
        
        # 合并
        merged[key] = base_model.state_dict()[key] + dropped_a + dropped_b
    
    return merged

class LayerWiseMerger:
    def merge(self, models, base_model):
        merged = {}
        
        for key in base_model.state_dict():
            # embedding层：使用平均
            if "embed" in key:
                merged[key] = self.average([m.state_dict()[key] for m in models])
            
            # attention层：使用TIES
            elif "attention" in key:
                merged[key] = self.ties([m.state_dict()[key] for m in models], base_model.state_dict()[key])
            
            # MLP层：使用DARE
            elif "mlp" in key:
                merged[key] = self.dare([m.state_dict()[key] for m in models], base_model.state_dict()[key])
            
            # 输出层：使用SLERP
            else:
                merged[key] = self.slerp(models[0].state_dict()[key], models[1].state_dict()[key])
        
        return merged

async def evaluate_merged(merged_model, benchmarks):
    """评估合并后的模型"""
    results = {}
    
    for bench in benchmarks:
        results[bench] = await run_benchmark(merged_model, bench)
    
    # 检查是否保持了原始模型的能力
    for capability in ["coding", "math", "reasoning", "writing"]:
        score = results.get(f"{capability}_bench", 0)
        if score < original_scores[capability] * 0.9:
            logger.warning(f"Capability {capability} degraded: {score} vs {original_scores[capability]}")
    
    return results