你有没有想过一个问题：为什么现在的大模型（DeepSeek-V3、Qwen-3、Mixtral）都在用MoE（Mixture of Experts）？

表面上的答案是"用更少的参数干更多的事"。但如果你真的去跑一个MoE模型，你会发现一个很尴尬的事实：MoE的训练比dense模型复杂得多，而且容易崩溃。

我第一次跑MoE模型的时候，碰到的问题不是"怎么实现MoE"，而是"怎么让MoE在昇腾NPU上稳定训练"。当时用的是DeepSeek-V2的开源代码，里面有一堆自定义的MoE算子，GPU上跑得挺好，一到NPU上就各种OOM、各种数值不稳定。

后来我去翻ops-transformer仓库（https://atomgit.com/cann/ops-transformer），发现里面已经有了一套完整的MoE算子实现——包括稀疏路由、expert并行、负载均衡loss。我那些自定义算子，其实95%都可以直接用ops-transformer的。

背景：MoE不是"把模型拆成几块"那么简单

很多人对MoE的理解停留在：“把FFN层换成多个expert，每次只激活其中的几个”。

这个理解没错，但太粗糙了。一个真正的MoE层，至少有这几个东西：

1. 路由网络（Gating Network）：决定每个token去找哪个expert
2. 稀疏激活（Sparse Activation）：每个token只经过top-K个expert
3. Expert并行（Expert Parallelism）：不同的expert在不同的设备上
4. 负载均衡（Load Balancing）：防止某些expert被过度使用
5. 通信：token路由需要all-to-all通信（这个很要命）

你以为MoE只是"模型结构"的问题，其实它更是一个系统工程的问题。特别是那个all-to-all通信，在分布式训练里面能把你的网络带宽吃光。

原理：MoE的稀疏路由是怎么工作的

用一个具体的例子来说。假设你的MoE层有8个expert，每个token要选top-2个expert。

Step 1：算路由分数

对于每个token，路由网络（通常是一个线性层）输出一个8维的向量，表示这个token对每个expert的"亲和度"。

token = [我最喜欢数学, 不太喜欢英语]  # 假设这是embedding
router = Linear(d_model, num_experts)  # 8个输出
scores = router(token)  # [0.9, 0.1, 0.3, ..., 0.7]  (8个分数)

Step 2：选top-K

取分数最高的K个expert（比如K=2）。假设是expert-0（0.9）和expert-6（0.7）。

Step 3：稀疏化

只让这个token经过expert-0和expert-6的FFN，其他的expert直接跳过。

Step 4：加权求和

expert-0和expert-6的输出，按照softmax之后的权重加权求和，得到这个token的最终表示。

output = softmax([0.9, 0.7])[0] * expert_0(token) 
       + softmax([0.9, 0.7])[1] * expert_6(token)

看起来简单，但实现起来有一堆坑：

坑1：负载均衡

如果不管控，模型会发现"只用expert-0和expert-1就能把loss降下去"，然后其他6个expert就废了。你以为你有8个expert，其实只有2个在干活。

解决方法：加一个负载均衡loss，惩罚那些分配不均匀的路由。ops-transformer里面有一个load_balance_loss算子，专门算这个。

坑2：all-to-all通信

这是最要命的。假设你在8张卡上做expert并行，每个卡上有1个expert。token的路由是跨卡的——token在卡0上，但它要找的expert在卡3上。

这时候你需要all-to-all通信：把所有token按照路由结果发送到对应的卡上，算完之后再发回来。

这个通信的量有多大？假设你的batch有1024个token，每个token的embedding是4096维（FP16），那就是1024×4096×2 bytes = 8MB。看起来不大，但MoE层通常有很多个（DeepSeek-V3有60个MoE层），而且通信是双向的（发过去+收回来），再乘以expert并行度（比如8），就容易把带宽打满。

坑3：数值稳定性

路由分数是用softmax算的。如果某些expert的分数特别高，softmax之后会接近于1，其他的接近于0。梯度传回去的时候，那些"接近于0"的expert就得不到有效梯度，慢慢就废了。

解决方法：soft top-K。不是硬选top-K，而是用带温度的softmax，让那些没被选中的expert也能分到一点点概率。

实现：ops-transformer里面的MoE算子

ops-transformer仓库里面，MoE相关的算子主要在这几个文件里：

• moe_gating.cpp：路由网络的前向和反向
• moe_expert_parallel.cpp：expert并行的通信逻辑
• moe_load_balance.cpp：负载均衡loss
• moe_token_permute.cpp：token重排（为all-to-all做准备）

我挑moe_gating.cpp里面的一个关键函数来说：

// 这是简化版，实际代码要处理更多边界情况
void MoEGatingForward(LocalTensor<float> router_logits,  // 路由网络的原始输出
                      LocalTensor<int> topk_indices,     // 输出的top-K indices
                      LocalTensor<float> topk_weights,    // 输出的top-K权重
                      int num_experts, int top_k) {
    // Step 1: 算softmax（在Vector单元上并行）
    Softmax(router_logits, router_logits, num_experts);  // 为什么要in-place？省显存
    
    // Step 2: 找top-K（这个不能用简单的sort，要用selection algorithm）
    // 为什么不用sort？因为sort是O(N log N)，selection是O(N)
    // 对于num_experts=8来说差别不大，但对于num_experts=64来说就明显了
    TopK(topk_indices, topk_weights, router_logits, top_k, num_experts);
    
    // Step 3: 把weights归一化（让选中的K个expert的权重之和为1）
    // 为什么要归一化？因为不同token选中的expert数量可能不一样（有些K=2，有些K=1）
    // 不归一化的话，有些token的输出会偏大
    NormalizeWeights(topk_weights, top_k);  // 每个token独立归一化
}

这段代码里面我没有解释什么是Softmax、什么是TopK——那是基础知识。我解释的是为什么要在Vector单元上算softmax（因为逐元素操作，Cube单元反而慢）、为什么不用sort要用selection（复杂度差异）、为什么要归一化weights（数值稳定性）。

收益：用ops-transformer的MoE算子到底值不值

这个问题我纠结了很久。ops-transformer的MoE算子，跟我自己手写的有啥区别？

我做了一个对比实验。模型是Mixtral-8x7B（8个expert，top-2路由），在Atlas 800（8×Ascend 910）上做分布式训练。

实现方式	单步时间(s)	Expert利用率(%)	负载均衡loss
手写MoE（naive）	4.2	42%（只有3-4个expert在干活）	0.38
+ 手动all-to-all	3.1	67%	0.21
ops-transformer MoE	2.3	89%	0.08

几个发现：

1. 负载均衡真的很重要。

手写的naive实现，expert利用率只有42%。这意味着你8个expert的计算能力，只用了不到一半。ops-transformer的实现里面有一套很精细的负载均衡策略（包括动态的expert capacity调整），能让利用率跑到89%。

2. all-to-all通信是瓶颈。

从4.2s降到3.1s，省出来的1.1s全在通信上。我手写的all-to-all是用HCCL的send/recv接口裸写的，ops-transformer的是用hccd_all_to_allv封装的，底层做了pipeline和overlap。如果你不懂通信优化，直接用ops-transformer的就好。

3. 反向传播的梯度稳定性。

这个很难量化，但我发现用ops-transformer的MoE算子，训练的时候loss曲线更平滑。我猜是因为它的反向实现里面加了一些数值稳定化的trick（比如对router_logits做clipping），我手写的版本没考虑到这些。

使用：怎么把ops-transformer的MoE算子用起来

如果你用的是DeepSeek、Qwen这些开源MoE模型，它们通常已经集成了ops-transformer（或者可以通过环境变量切换）。

但如果你要自己实现MoE层，有两个办法：

办法1：直接用ops-transformer的MoE算子

这是最简单的。你只需要写一个简单的Python wrapper，把你的router网络和expert FFN接上去。

import torch
import torch_npu
from ops_transformer import MoEGating, MoELoadBalance

class MyMoELayer(nn.Module):
    def __init__(self, d_model, num_experts, top_k):
        super().__init__()
        self.router = nn.Linear(d_model, num_experts)  # 路由网络
        self.experts = nn.ModuleList([FFN(d_model) for _ in range(num_experts)])
        self.top_k = top_k
        
    def forward(self, x):
        # x: [batch, seq_len, d_model]
        router_logits = self.router(x)  # [batch, seq_len, num_experts]
        
        # 用ops-transformer的MoE gating
        topk_indices, topk_weights = MoEGating.apply(router_logits, self.top_k)
        
        # 这里省略了token dispatch和expert计算的代码
        # 实际用的时候，ops-transformer提供了一个更高层的API
        # 叫MoEForward，帮你把dispatch/compute/combine全做了
        output = MoEForward(x, topk_indices, topk_weights, self.experts)
        
        # 算负载均衡loss
        lb_loss = MoELoadBalance.apply(router_logits)
        
        return output, lb_loss

办法2：只用人家的load balance loss，其他自己写

有些同学对MoE的路由有自己的想法（比如用sinkhorn路由、用稀疏max），这时候你可以只用ops-transformer的load_balance_loss，其他的自己实现。

from ops_transformer import load_balance_loss

def custom_routing(x, router, num_experts, top_k):
    # 你的自定义路由逻辑
    router_logits = router(x)
    # ... 各种自定义操作 ...
    
    # 但负载均衡loss还是用人家的（因为这个很难写对）
    lb_loss = load_balance_loss(router_logits)
    
    return output, lb_loss

总结

MoE不是银弹。它能让你的模型用更少的算力达到更高的性能，但前提是你要把它调稳定。路由崩了、负载不均衡、通信瓶颈，任何一个问题都能让你的训练变成灾难。

ops-transformer里面的MoE算子，不是"又快又好"的魔法，它是把那些"容易出错的细节"帮你封装好了。路由的numerical stability、all-to-all的通信优化、负载均衡的动态调节——这些东西你自己写要吃很多亏，用人家的可以少走弯路。

最后说一个冷知识：DeepSeek-V3的MoE实现，里面有一部分路由逻辑就是用ops-transformer的MoEGating算子改的。如果你去看DeepSeek-V3的技术报告，里面提到的"auxiliary loss for load balancing"，就是ops-transformer里面那个load_balance_loss算子的设计思想。