MoE环游记:2、深入负载均衡

在上一篇文章中,我们介绍了MoE的一个几何诠释,旨在通过Dense模型的最佳逼近出发来推导和理解MoE。同时在文末我们也说了,给出MoE的计算公式仅仅是开始,训练一个实际有效的MoE模型还有很多细节补,比如本文要讨论的负载均衡(Load Balance)问题。 负载均衡,即"不患寡而患不均",说白了就是让每个Expert都在干活,并且都在干尽可能一样多的活,避免某些Expert浪费算力。负载均衡既是充分利用训练算力的需求,也是尽可能发挥MoE大参数量潜力的需求。 问题分析 我们知道,MoE的基本形式是 $$ \boldsymbol{y} = \sum_{i\in \mathop{\text{argtop}}_k \boldsymbol{\rho}} \rho_i \boldsymbol{e}_i $$ 对于传统MoE,$\boldsymbol{\rho}$是一个概率分布(Router),$\boldsymbol{e}_i=\boldsymbol{v}_i$,$\boldsymbol{v}_i$是一个小型FFN(Expert)的输出;而对于我们上一篇推导的几何MoE,$\boldsymbol{\rho}$没有归一化的要求,它预测的是Expert的模长,而$\boldsymbol{e}_i=\boldsymbol{v}_i/\Vert\boldsymbol{v}_i\Vert$预测的是Expert的方向。 不管哪种格式的MoE,实际表现都差不多,只是理解视角的不同。但要注意,虽然MoE的公式给人的感觉是"每遇到一个Token,就去找相应的Expert来计算",但实际训练时其实是反过来的:先给每个Expert分配好相应的算力,然后将Token分配(Route)到所属的Expert中并行计算,这也就为什么负责打分的$\boldsymbol{\rho}$被称为Router。 这样一来,如果Expert的分配不均衡,就可能出现如下局面:某些Expert(Dead Expert)几乎一直闲置,浪费算力;某些Expert要处理的Token太多,根本忙不过来,只能Token Drop(即放弃处理部分Token)。从理论上来说,出现Dead Expert意味着MoE没有达到预期的参数量,即花了大参数量的显存,结果只训出来小参数量的效果。 所以,不管是从训练还是性能角度看,我们都希望保证Expert的负载均衡。 辅助损失(Auxiliary Loss) 促进负载均衡的常规思路是添加与之相关的损失函数,我们通常称之为"Aux Loss(Auxiliary Loss)",目前主流用的Aux Loss最早可以追溯到2020年的《GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding》。 介绍Aux Loss之前,我们需要先引入一些新概念。首先,我们已经提到对于一般的MoE来说,$\boldsymbol{\rho}$未必是概率分布,我们将归一化的$\boldsymbol{\rho}$记为$\boldsymbol{p}=[p_1,p_2,\cdots,p_n]$,以及它Top-$k$版为$\boldsymbol{f}=[f_1,f_2,\cdots,f_n]$,其中 $$ p_i = \frac{\rho_i}{\sum_{i=1}^n \rho_i},\qquad f_i = \begin{cases}1/k, & i\in \mathop{\text{argtop}}_k \boldsymbol{\rho} \ 0, & i\not\in \mathop{\text{argtop}}_k \boldsymbol{\rho}\end{cases} $$ 接着我们定义$\boldsymbol{P}=\mathbb{E}[\boldsymbol{p}],\boldsymbol{F}=\mathbb{E}[\boldsymbol{f}]$,这里的$\mathbb{E}$是指对所有样本的所有Token做平均。不难看出,$\boldsymbol{F}$就是Expert当前的负载分布,而$\boldsymbol{P}$则相当于$\boldsymbol{F}$的一个光滑近似。 有了这些记号,我们就可以写出Aux Loss为: $$ \mathcal{L}_{\text{aux}} = \boldsymbol{F}\cdot \boldsymbol{P} = \sum_{i=1}^n F_i P_i \tag{1} $$ ...

August 10, 2025 · 5 min · 1055 words · Me

MoE环游记:1、从几何意义出发

MoE(Mixture of Experts)架构的流行自不必多说,近来火出圈的DeepSeek-V3便是MoE架构,传言GPT-5也是MoE架构,国内最近出的一些模型(Qwen3系列相关)也有不少用上了MoE。然而,虽然MoE的研究由来已久,但其应用长时间内都不愠不火,大致上是从去年初的《Mixtral of Experts》开始,MoE才逐渐吸引大家的注意力,其显著优点是参数量大,但训练和推理成本都显著低。 但同时MoE也有一些难题,如训练不稳定、负载不均衡、效果不够好等,这也是它早年没有流行起来的主要原因。不过随着这两年关注度的提升,这些问题在很大程度上已经得到解决,我们在接下来的介绍中会逐一谈到这些内容。 问题定义 我们知道,Transformer模型由Attention层和MLP层组成,MoE替换的是模型中MLP层。MLP层又分FFN(FeedForward Network)和GLU(Gated Linear Unit)两种,主流的是GLU,但简单起见我们还是以FFN为例:$$y=f(xW^{(A)})W^{(B)}$$其中$x\in\mathbb{R}^d$ 是输入向量(行向量),$W^{(A)}\in\mathbb{R}^{d\times{D}}$, $W^{(B)}\in\mathbb{R}^{D\times{d}}$ 是两个参数矩阵,$f$是Element-wise的激活函数,设$n$是一个能整除$D$的整数,那么上面的FFN可以用分块矩阵等价: $$ \begin{equation}\boldsymbol{y} = f\big(\boldsymbol{x}\begin{bmatrix}\boldsymbol{W}^{(A)}_1 & \boldsymbol{W}^{(A)}_2 & \cdots & \boldsymbol{W}^{(A)}_n\end{bmatrix}\big)\begin{bmatrix}\boldsymbol{W}^{(B)}_1 \\ \boldsymbol{W}^{(B)}_2 \\ \vdots \\ \boldsymbol{W}^{(B)}_n\end{bmatrix} = \sum_{i=1}^n \underbrace{f(\boldsymbol{x}\boldsymbol{W}^{(A)}_i)\boldsymbol{W}^{(B)}_i}_{\boldsymbol{v}_i}\end{equation} $$ 其中 $W^{(A)}_i = W^{(A)}_{[:,(i-1)c:ic]}$, $W^{(B)}_i = W^{(B)}_{[(i-1)c:ic,:]}$, $c= D/n$,这里的切片按照Python规则来。由此可见,FFN可以等价表示成n个向量 $\boldsymbol{v}_1,\boldsymbol{v}_2,\cdots,\boldsymbol{v}_n$ 之和,每个向量代表了一个小模型$f(\boldsymbol{x}\boldsymbol{W}^{(A)}_i)\boldsymbol{W}^{(B)}_i$的输出,每个小模型计算量相同,这些小模型就是MoE中的“Expert”。 MoE提出的问题是: 能否只挑k个向量的和来逼近n个向量的和呢?这样就可以将计算量降低到k/n了。 模长排序 要解决上述的问题,实质上是要解决低秩近似的问题,数学公式就是: $$\begin{equation}\mathop{\text{argmin}}_{\lambda_1,\lambda_2,\cdots,\lambda_n\in\{0,1\}}\left\Vert\sum_{i=1}^n \lambda_i \boldsymbol{v}_i - \sum_{i=1}^n\boldsymbol{v}_i\right\Vert^2\quad\text{s.t.}\quad \sum_{i=1}^n \lambda_i = k\end{equation}$$ 记$\gamma_i = 1 - \lambda_i$,那么它又可以写成: $$\begin{equation}\mathop{\text{argmin}}_{\gamma_1,\gamma_2,\cdots,\gamma_n\in\{0,1\}}\left\Vert\sum_{i=1}^n \gamma_i \boldsymbol{v}_i\right\Vert^2\quad\text{s.t.}\quad \sum_{i=1}^n \gamma_i = n - k\end{equation}$$ 这个问题的精确求解是比较困难的(NP Hard),但有一个简单的近似解:当$v_i$两两正交时,我们有 $$\begin{equation}\left\Vert\sum_{i=1}^n \gamma_i \boldsymbol{v}_i\right\Vert^2 = \sum_{i=1}^n \gamma_i^2 \Vert\boldsymbol{v}_i\Vert^2 = \sum_{i=1}^n \gamma_i \Vert\boldsymbol{v}_i\Vert^2\end{equation}$$ 上式最优解显然就是让模长$\Vert\boldsymbol{v}_i\Vert$最小的$n-k$个$\gamma_i$等于1,这又等价于说挑出模长最大的$k$个向量来逼近$n$个向量之和。当$v_i$不满足两两正交的条件时,我们依然用它来作为一个近似解。它的几何意义也很直观,模长越大的向量,在求和过程中越不容易被抵消,从而作用越突出。 ...

August 8, 2025 · 1 min · 146 words · Me

[RL4LLM] 异步RL框架: Slime

https://github.com/THUDM/slime 一个异步实现但是非完全异步的RL框架 总体架构 从源码模块划分,有三大核心模块: training(Megatron):主训练流程,负责模型参数更新。 rollout(SGLang + router):负责采样、奖励/验证生成,产生训练数据。 data buffer:桥接训练与采样,管理数据流、缓存与生成方式。 分布式调度:关于资源分配、actor启动、任务调度都由于Ray管理,支持异步训练和采样 插件机制:支持自定义buffer、模型、模型格式转换(mbridge) flowchart LR subgraph Ray[Ray 分布式调度] A1[Actor Group<br>训练 Actor] A2[Rollout Group<br>采样/生成 Actor] A3[Placement Group<br>资源分配] end subgraph Training[Training <Megatron>] T1[模型训练] T2[权重同步] T3[评估/保存] end subgraph Rollout[Rollout <SGLang+Router>] R1[采样/生成] R2[奖励模型] R3[过滤器] end subgraph Buffer[Data Buffer] B1[数据缓存] B2[数据流转] B3[Offload/Onload] end subgraph Plugins[插件机制] P1[Buffer 插件] P2[Model 插件] P3[mbridge 格式转换] end A1-->|训练数据|B1 A2-->|生成数据|B1 B1-->|数据流|A1 B1-->|数据流|A2 A1-->|权重同步|A2 A1-->|评估/保存|T3 A2-->|采样/奖励/过滤|R1 R1-->|奖励|R2 R1-->|过滤|R3 B1-->|插件扩展|P1 A1-->|模型扩展|P2 A1-->|格式转换|P3 A3-->|资源分配|A1 A3-->|资源分配|A2 各模块视角的关系图 slime/rollout 组件图 rollout 负责采样、奖励、过滤,支持多种采样/奖励/过滤策略。 ...

August 7, 2025 · 15 min · 3119 words · Me

[RL4LLM] 异步RL框架: Areal

https://github.com/inclusionAI/AReaL 纯异步RL方案 异步PPO训练调用流程 graph TD A[用户执行: examples/run_async_ppo.sh] --> B[training/main_async_ppo.py] B --> C[AsyncPPOMATHConfig配置解析] C --> D[training/utils.py: run_experiment] D --> E[Ray初始化] E --> F[exp_cfg.initial_setup] F --> G[AsyncRLExperimentConfig.initial_setup] G --> H[创建ExperimentConfig] H --> I[启动Workers] I --> J[MasterWorker] I --> K[ModelWorker] I --> L[GenerationServer] I --> M[GserverManager] I --> N[RolloutWorker] %% MasterWorker训练流程 J --> J1[MasterWorker._poll_async] J1 --> J2[FunctionExecutor.execute_step] J2 --> J3[执行数据流图遍历] J3 --> J4[发送训练请求到ModelWorker] %% ModelWorker处理流程 K --> K1[ModelWorker._poll] K1 --> K2[接收MasterWorker请求] K2 --> K3[处理训练/推理请求] K3 --> K4[执行模型前向/反向传播] %% Rollout流程 N --> N1[RolloutWorker._poll_async] N1 --> N2[load_next_data] N2 --> N3[allocate_new_rollout] N3 --> N4[agent.collect_trajectory] N4 --> N5[env.step计算奖励] N5 --> N6[推送数据到训练端] %% 生成服务器流程 L --> L1[GenerationServer._poll] L1 --> L2[启动SGLang子进程] L2 --> L3[处理生成请求] %% 生成服务器管理器 M --> M1[GserverManager._poll] M1 --> M2[HTTP服务线程] M2 --> M3[请求调度和权重更新] %% 数据流 N6 --> O[stream_dataset.py] O --> J4 %% 异步通信 J4 -.->|异步请求| K2 N3 -.->|HTTP请求| M2 M2 -.->|调度请求| L3 %% 权重更新 K4 --> P[参数更新] P --> Q[权重同步] Q --> M3 M3 --> R[更新生成服务器权重] style A fill:#e1f5fe style J fill:#f3e5f5 style K fill:#e8f5e8 style L fill:#fff3e0 style M fill:#fce4ec style N fill:#f1f8e9 用户入口到配置解析 examples/run_async_ppo.sh → training/main_async_ppo.py ...

August 7, 2025 · 23 min · 4872 words · Me

昇腾超节点CloudMatrix384论文拆解

6.19发布的CloudMatrix384论文拆解,从宏观到基础概念 核心指标和计算方式 TPOT (Time Per Output Token) 公式: $$TPOT= \frac{Decode总耗时}{生成Token数量}$$ 测量方式: 从第一个输出Token开始计时,到生成结束(含MoE通信/KV读取) 为什么重要: 直接决定用户体验(如Chatbot响应速度),论文要求 <50ms(严格模式<15ms) 深层意义: 反映系统通信+计算综合能力,EP320下TPOT=42ms证明UB网络突破MoE通信墙 计算效率 (Tokens/s per TFLOPS) 公式: $$计算效率=\frac {吞吐量(tokens/s)} {NPU峰值算力(TFLOPS)}$$​ 论文数据: 阶段 值 对比基准 Prefill 4.45 超NVIDIA H100+SGLang(3.8) Decode 1.29 超NVIDIA H800+DeepSeek(0.9) 为什么重要: 揭示硬件利用率,1.0以上表明软硬件协同极致优化 深层意义: Decode阶段1.29 → 昇腾910的Cube引擎利用率达 86%(传统GPU仅60%) 缓存访问延迟 (KV Cache Access Latency) 公式: $$延迟=TMMU_{查询}+TUB_{传输}+TDRAM_{读取}​$$ 论文数据: 场景 延迟 对比传统 本地HBM命中 0.2μs - 远程DRAM访问(UB) 1.5μs >10μs (PCIe+IB) 为什么重要: 长上下文推理中70%时间花在KV缓存访问 深层意义: UB统一内存将远程访问性能提升至近本地水平,支撑百万Token上下文。 专家并行扩展性 (EP Degree) 定义:单个MoE层可分布的专家数量 论文突破:EP320(每个昇腾Die托管1个专家) 支撑公式: $$可扩展性=\frac {UB总带宽}{单个专家通信需求}$$ $$EPmax=\frac {384×392GB/s} {8B/token×10^6token/s}=320$$ 为什么重要: EP>100时传统网络崩溃,EP320证明UB突破通信可扩展性极限 INT8量化收益 公式:$$ 加速比=\frac {FP16吞吐}{INT8吞吐}×精度保持率$$ 论文数据: 吞吐提升:1.8倍 精度损失:<0.5%(16个基准测试) 为什么重要: Decode阶段内存带宽减少50%,解决NPU的“内存墙”问题 QA辅助理解 为什么用TPOT而非QPS? TPOT剥离Batch Size影响,纯粹衡量单次生成效率 更直观反映SLA(用户感知的延迟) 为什么强调计算效率而非绝对吞吐? 排除工艺优势(7nm vs 5nm),聚焦架构创新价值 1.29 tokens/s/TFLOPS → 证明UB+LEP设计优于NVLink+GPU 为什么测量远程DRAM访问延迟? 验证内存池化的实际效果,这是打破“内存墙”的核心 1.5μs延迟 → 实现“全集群如单机”的硬件基础 超节点架构 三级网络平面的物理隔离 硬件隔离原理 ...

August 7, 2025 · 6 min · 1211 words · Me