Pytorch分布式训练/多卡训练(六) —— Expert Parallelism (MoE的特殊策略)

hxxjxw2025-12-18 17:58

MoE 区别于传统的 Transformer 是其参数分布极为不均
复制代码
传统Transformer (175B GPT-3)：
├─ 96层，每层参数量相近
├─ Layer 1: 1.8B参数
├─ Layer 2: 1.8B参数
├─ ...
└─ Layer 96: 1.8B参数
特点：参数分布均匀 ✅

MoE模型 (Switch Transformer-1.6T)：
├─ Embedding: 1B参数
├─ Layer 1 (标准Attention): 1B参数
├─ Layer 2 (MoE层): 128B参数！⚠️
│   ├─ 64个Expert，每个2B参数
│   └─ 路由器：0.01B参数
├─ Layer 3 (标准Attention): 1B参数
├─ Layer 4 (MoE层): 128B参数！⚠️
├─ ...
└─ Layer 32 (MoE层): 128B参数
特点：参数分布极不均匀！❌
如果使用 ZeRO-3，当计算的时候all-gather 所有模型参数，如果 expert 数量很多，例如 64 个，那么总显存可能带 256G, 远超一个 GPU所能承载的

Expert Parallelism

核心思想

Expert并行 (EP)：

不同GPU负责不同的Expert

只加载被本GPU负责的Expert参数

Token根据路由结果发送到对应GPU

配置：EP=64（64个GPU，每个负责1个Expert）

参数分配：
├─ GPU 0: Expert 0 (2GB)
├─ GPU 1: Expert 1 (2GB)
├─ GPU 2: Expert 2 (2GB)
├─ ...
└─ GPU 63: Expert 63 (2GB)

每个GPU显存占用：

Expert参数：2GB（只有1个Expert）✅

共享层参数：1GB（Attention等）

激活值：根据收到的token数量