少算 = 算得对

§6.1SparseGPT / Wanda · 一发剪枝

# Wanda: 一行剪枝公式
def wanda_prune(W, X, sparsity=0.5):
    # W: [d_out, d_in],  X: [n_calib, d_in]
    act_norm = X.norm(p=2, dim=0)            # [d_in]
    importance = W.abs() * act_norm[None, :] # [d_out, d_in]
    # 按行选 top-(1-sparsity) 比例保留
    n_keep = int(W.shape[1] * (1 - sparsity))
    thr = importance.topk(n_keep, dim=1).values[:, -1:]
    mask = importance >= thr
    return W * mask

§6.22:4 / N:M 结构化稀疏 · 实际能加速

上面那种非结构化稀疏，理论 50% 但 GPU 上很难真加速—— 随机零分布在 SM 调度时跟 dense 没区别。 2:4 稀疏是 Ampere/Hopper 硬件支持： 每 4 个连续权重必须有 2 个零。 tensor core 直接 2× 算力（一个 cycle 处理 2 个非零）。

# 2:4 mask: 每 4 连续位置保留 magnitude 最大的 2 个
def make_24_mask(W):
    # W: [d_out, d_in], d_in % 4 == 0
    d_out, d_in = W.shape
    g = W.abs().reshape(d_out, d_in // 4, 4)
    top2 = g.topk(2, dim=-1).indices       # [d_out, d_in//4, 2]
    mask = torch.zeros_like(g, dtype=bool)
    mask.scatter_(-1, top2, True)
    return mask.reshape(d_out, d_in)

§6.3Mixture of Experts · 路由是两难

Router 的形状

# 教学版 top-K router: 一个小线性层 + softmax + top-K
class MoELayer(nn.Module):
    def __init__(self, D, N_experts, K=2):
        self.gate    = nn.Linear(D, N_experts, bias=False)
        self.experts = nn.ModuleList([FFN(D) for _ in range(N_experts)])
        self.K = K
    def forward(self, x):                  # x: [B, S, D]
        scores = self.gate(x)              # [B, S, N]
        probs = scores.softmax(dim=-1)
        # top-K 选择
        top_p, top_i = probs.topk(self.K, dim=-1)
        top_p = top_p / top_p.sum(dim=-1, keepdim=True)
        # 分发: 每个 expert 接收的 token 子集
        out = torch.zeros_like(x)
        for e_idx, expert in enumerate(self.experts):
            mask = (top_i == e_idx).any(dim=-1)
            if not mask.any(): continue
            weight = (top_p * (top_i == e_idx)).sum(dim=-1, keepdim=True)
            out += mask[..., None] * weight * expert(x)
        return out

三大工程难点

• 负载均衡——router 倾向于把所有 token 倒进同一个 expert， 需要 aux loss / capacity factor 控制；DeepSeek-V3 用 "auxiliary-loss-free" per-expert bias，避免 aux loss 拖累质量。
• 通信——all-to-all 把 token 发去对应 expert， 通信量随 MoE size 暴增；DeepEP 用 FP8 send + double buffer overlap。
• 稀疏激活的不规则计算—— 每个 expert 拿到的 token 数不同，需要"grouped GEMM"或者 padding。

§6.4Expert Parallelism · DeepSeek-V3 的工程范式

# auxiliary-loss-free balance: 每个 expert 维护一个 bias
class AuxFreeRouter(nn.Module):
    def __init__(self, D, N):
        self.gate = nn.Linear(D, N, bias=False)
        # bias 不是参数, 用 running avg 调整
        self.register_buffer('bias', torch.zeros(N))
    def forward(self, x):
        logits = self.gate(x) + self.bias        # 加 bias (不参与梯度)
        probs = logits.softmax(dim=-1)
        top_p, top_i = probs.topk(self.K, dim=-1)
        # update bias based on observed load
        with torch.no_grad():
            load = torch.bincount(top_i.flatten(), minlength=self.N).float()
            avg = load.mean()
            self.bias -= 0.01 * (load - avg).sign()  # over-load → bias 减小
        return top_p, top_i