LLM剪枝2026：结构化vs非结构化

def magnitude_prune(weight, sparsity=0.5):
    """幅度剪枝：将最小幅度的权重置零"""
    # 计算阈值
    threshold = torch.quantile(weight.abs().flatten(), sparsity)
    
    # 创建掩码
    mask = (weight.abs() > threshold).float()
    
    # 应用掩码
    return weight * mask, mask

def apply_2_4_sparsity(weight):
    """应用2:4稀疏模式"""
    output = weight.clone()
    # 将权重重排为 [N, 4] 的块
    blocks = output.view(-1, 4)
    
    for i in range(blocks.shape[0]):
        block = blocks[i].abs()
        # 保留最大的2个，其余置零
        top2_indices = block.topk(2).indices
        mask = torch.zeros(4, device=weight.device)
        mask[top2_indices] = 1
        blocks[i] *= mask
    
    return output.view_as(weight)

def sparsegpt_prune(layer, calibration_data, sparsity=0.5):
    """SparseGPT单层剪枝"""
    W = layer.weight.data  # [out, in]
    H = compute_hessian(layer, calibration_data)  # [in, in]
    H += torch.eye(H.shape[0]) * 0.01  # 正则化
    
    # 逐列处理
    for i in range(W.shape[1]):
        # 计算每个权重的"重要性"分数
        importance = W[:, i].abs() ** 2 / H[i, i]
        
        # 选择保留的权重
        n_keep = int(W.shape[0] * (1 - sparsity))
        keep_indices = importance.topk(n_keep).indices
        
        # 对非保留权重进行补偿
        mask = torch.zeros(W.shape[0], device=W.device)
        mask[keep_indices] = 1
        
        # 重建误差补偿
        err = (W[:, i] * (1 - mask)) / H[i, i]
        W[:, i+1:] -= err.unsqueeze(1) @ H[i, i+1:].unsqueeze(0)
        
        # 应用掩码
        W[:, i] *= mask
    
    layer.weight.data = W

def structured_prune_heads(model, heads_to_prune):
    """剪枝注意力头"""
    for layer_idx, head_indices in heads_to_prune.items():
        layer = model.layers[layer_idx]
        
        # 移除对应的Q/K/V投影参数
        for head in sorted(head_indices, reverse=True):
            start = head * layer.head_dim
            end = start + layer.head_dim
            
            # 拼接所有非剪枝头的参数
            layer.W_q = nn.Parameter(
                torch.cat([layer.W_q[:start], layer.W_q[end:]], dim=0)
            )
            layer.W_k = nn.Parameter(
                torch.cat([layer.W_k[:start], layer.W_k[end:]], dim=0)
            )
            layer.W_v = nn.Parameter(
                torch.cat([layer.W_v[:start], layer.W_v[end:]], dim=0)
            )
        
        layer.n_heads -= len(head_indices)

def layer_pruning(model, layers_to_remove):
    """移除指定的Transformer层"""
    remaining_layers = []
    for i, layer in enumerate(model.layers):
        if i not in layers_to_remove:
            remaining_layers.append(layer)
    
    model.layers = nn.ModuleList(remaining_layers)
    model.config.n_layers = len(remaining_layers)

def short_layer_distillation(student, teacher, training_data):
    """短层蒸馏：层映射对齐"""
    # 建立Student层到Teacher层的映射
    # 例如，Student的第i层对应Teacher的第2i层
    layer_mapping = {i: 2 * i for i in range(len(student.layers))}
    
    for batch in training_data:
        # Teacher和Student前向传播
        with torch.no_grad():
            teacher_outputs = teacher(batch, output_hidden_states=True)
        
        student_outputs = student(batch, output_hidden_states=True)
        
        # 对齐隐藏状态
        hidden_loss = 0
        for s_idx, t_idx in layer_mapping.items():
            s_hidden = student_outputs.hidden_states[s_idx]
            t_hidden = teacher_outputs.hidden_states[t_idx].detach()
            
            # 可能需要投影对齐维度
            hidden_loss += F.mse_loss(s_hidden, t_hidden)
        
        # 输出对齐
        output_loss = F.kl_div(
            F.log_softmax(student_outputs.logits, dim=-1),
            F.softmax(teacher_outputs.logits, dim=-1),
            reduction='batchmean'
        )
        
        total_loss = hidden_loss + output_loss
        total_loss.backward()
        optimizer.step()

def wanda_prune(weight, activations, sparsity=0.5):
    """Wanda: 基于权重×激活的剪枝"""
    # 计算每个权重的重要性分数
    activation_norm = activations.norm(dim=0)  # [in_features]
    importance = weight.abs() * activation_norm.unsqueeze(0)  # [out, in]
    
    # 全局选择保留的权重
    n_keep = int(weight.numel() * (1 - sparsity))
    threshold = importance.flatten().topk(n_keep).values.min()
    
    mask = (importance >= threshold).float()
    return weight * mask