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 · 147 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

[VeRL] Multi-Turn RL训练源码走读(2)

在 Part 1 中,我们介绍了 verl 的初始化过程,我们进一步介绍 verl 的训练过程,包括rollout部分、make experience部分以及training部分。 在 GRPO 中,单个 step 包含四个阶段:load data -> rollout -> make experience -> update model。区别于前一节的详述,本节会使用伪代码结合源码的方式进行阐述。 flowchart LR subgraph W2["Initialize"] WP[Process Data] --> A direction TB D1[Data Prepare] --> A A[TaskRunner] --> B1[RayPPOTrainer] B1 --> Workers subgraph Workers["Workers"] direction TB WA[ActorRolloutWorker] --> WD[FSDP Engine] WB[CriticWorker] --> WD WC[RewardModelWorker] --> WD WD --> WE[SGLang Engine] end Workers --> C1[Hybrid Engine] end subgraph W3["Train Loop"] direction TB E[DataLoader] --> RolloutBox subgraph RolloutBox["Rollout"] F1[Prepare Data] --> F2[SGLang Async Rollout] F2 --> F3[Multi-turn Chat Process] end RolloutBox --> ExpBox subgraph ExpBox["Make Experience"] G1[Recompute Log Probs] --> G2[Compute Reward] G2 --> G3[Compute Advantage] end ExpBox --> UpdateBox subgraph UpdateBox["Train The Model"] H1[Load FSDP Model Weight] --> H2[Compute Gradient] H2 --> H3[Weights Update] H3 --> H4[Sync Weights] end UpdateBox --> E end W2 --> W3 数据加载与预处理 verl 通过 DataProto 和 RLHFDataset 来实现数据处理。具体来说,在 main_ppo.py 中,我们观察这个函数: ...

August 3, 2025 · 27 min · 5743 words · Me

AI Infra:颠覆性创新,还是经典工程范式的华丽转身?

近期看到一些关于传统基础设施(Traditional Infrastructure)与人工智能基础设施(AI Infrastructure,尤其大模型领域)差异的评论。其核心观点直指两者间的巨大鸿沟:许多精于网络、计算、存储等传统领域的工程师,在面对GPU集群、KV Cache管理、3D并行等全新概念时,常感过往经验难以直接套用,甚至产生踏入一个全然不同技术体系的“割裂感”。 这些看法颇具代表性,精准捕捉了工程师初探AI Infra时的普遍印象:陌生、高门槛、范式迥异。本文旨在分享我对此的一些初步思考:AI Infra究竟是颠覆传统的全新体系,抑或是既有Infra经验在智能时代的一次深度演化? (免责声明:本文纯属个人观点,旨在抛砖引玉,欢迎指正谬误!) 我的核心论点:AI Infra并非平地起高楼,它实质上是传统Infra工程智慧在新场景下的重构与系统性延展。 表象差异:新术语与新挑战带来的“视觉冲击” 乍看之下,AI Infra与传统Infra确实分野明显: 核心任务不同: 传统Infra聚焦于处理海量Web请求(毫秒级、无状态)、保障数据持久化存储、实现分布式服务协调。而AI Infra(尤以大模型为甚)则围绕GPU驱动的模型训练/推理、KV Cache的高效管理、百亿/千亿级参数的分布式执行框架展开。 请求形态迥异: Web请求追求瞬时响应(毫秒级)、天然无状态。大模型(LLM)推理则常承载持续的会话交互(秒级乃至更长,随上下文窗口扩展而递增),需动态维护细粒度的Token级状态(KV Cache)。 技术栈迭代: 熟悉的Kubernetes + Docker堆栈旁,涌现出GPU硬件抽象、vLLM、DeepSpeed、FlashAttention、Triton、NCCL等专为AI设计、名号“高深”的组件。 由此观之,认为传统经验难以直接迁移,确有其表象依据。但这仅仅是“水面之上的冰山”,远非其底层基石。 本质共性:工程核心挑战的永恒回归 拨开“AI专属”的面纱,工程实践的核心命题依然如故:系统设计与资源调度的精妙艺术。 我们面临的,仍是那些传统Infra领域中反复锤炼的同类问题,只是约束条件和优化目标发生了变化: 资源调度: 核心资源从CPU/内存/磁盘IO,转向了更稀缺、更昂贵的GPU显存与算力。 负载处理: 承载对象从HTTP资源请求,变为密集的Prompt请求与大规模训练任务。 核心目标: 高效、稳定、低成本地协调跨节点资源的核心诉求丝毫未变。 概念的映射:经典范式的AI实践 这种延续性,清晰地体现在关键概念的对应关系上: 传统 Infra 概念 AI Infra 对应实践 核心思想应用 数据分片 (Data Sharding) 数据并行 (Data Parallelism) 数据集拆分,多副本并行处理 负载均衡 (Load Balancer) MoE Router (Mixture of Experts) 动态分配请求(Token)至专家网络,避免热点 操作系统分页 (OS Paging) vLLM KV Cache Paging 虚拟化显存空间,高效管理请求状态 以vLLM为例: 其核心创新在于将操作系统经典的内存管理机制(分页、交换),创造性地应用于管理LLM推理中关键的KV Cache状态。它如同为LLM定制了一个“显存操作系统”,管理“进程”(推理请求)和“内存页”(KV Cache Blocks),极致优化昂贵显存的利用率。这绝非凭空创造,而是经典系统原理在特定约束下的卓越应用。 ...

August 1, 2025 · 1 min · 200 words · Me