Skip to content

Commit b36f5e0

Browse files
authored
Merge pull request #98 from 840691168/Yang-doc
Yang doc
2 parents 5223fd1 + 1292bfe commit b36f5e0

File tree

6 files changed

+333
-0
lines changed

6 files changed

+333
-0
lines changed
121 KB
Loading
29.6 KB
Loading
71.7 KB
Loading
76.4 KB
Loading

app/docs/ai/MoE/MOE-intro.md

Lines changed: 140 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,140 @@
1+
---
2+
title: "MOE 浅谈"
3+
description: "混合专家(Mixture of Experts, MoE)架构的简要介绍"
4+
tags: ["MoE", "AI"]
5+
---
6+
7+
# 混合专家(Mixture of Experts, MoE)架构
8+
9+
为了在不显著增加计算需求的情况下有效扩展模型参数规模,MoE 架构已经成为一种可行的解决方案。MoE 利用一组专门化的子模型和一个门控机制,动态地选择合适的“专家网络”来处理给定输入。这使得模型能够根据需求分配计算资源,这一概念被称为 **条件计算(conditional computation)**
10+
11+
MoE 架构已经被广泛应用到大语言模型(LLMs)中,使这些模型能够在参数规模显著扩大的同时,获得相应的能力提升。
12+
例如,Mixtral AI 提出的 **Mixtral-8x7B** 管实际激活的参数量仅有 130 亿,**在多个基准测试上表现优于或相当于 Llama-2-70B 和 GPT-3.5**
13+
14+
---
15+
16+
## 传统 MoE 架构
17+
18+
自 MoE 最早被引入 Transformer 架构以来,MoE 主要作为 **前馈网络 (FFN)** 的替代模块使用。通常情况下,MoE 层中的每个专家都直接复制了其所替换的 FFN 结构。然后配备有一个 Router 来训练具体交给哪个专家处理。
19+
20+
![](./MOE-intro.assets/img-20250920112106486.png)
21+
22+
MoE 主要应用于 FFN 层,而不是自注意力层,原因在于:
23+
24+
- **注意力层**:稀疏性较低,更适用于全局交互。
25+
- **FFN 层**:稀疏性较高,更具有领域特性。
26+
其中 DS-MoE 使用 Wikitext 作为任务时,发现:FFN 层的专家仅有 **20%** 被激活
27+
而注意力层激活率高达 **80%**。这种高利用率表明注意力层的核心通讯机制不适用于特异化的专家。反之具有稀疏特性的 FFN 层,具有完整多专家特异化的潜力。
28+
29+
![](./MOE-intro.assets/img-20250920112106518.png)
30+
31+
---
32+
33+
## Routing 机制:Dense MoE 与 Sparse MoE
34+
35+
![](./MOE-intro.assets/img-20250920112106554.png)
36+
37+
- **Dense MoE**
38+
- gate 对于输入 token 使用 **softmax** 路由机制,传递给每个专家一定权重。
39+
- 优点:训练稳定。
40+
- 缺点:每次都要计算所有专家,计算成本高。
41+
42+
- **Sparse MoE**
43+
- 使用 **Top-K** 路由机制,仅激活前 K 个权重最大的专家。
44+
- 优点:极大减少计算量,这是目前主流模型(如 GShard、Switch Transformer、Mixtral、DeepSeek-MoE)的策略。
45+
- 缺点:Router 训练变复杂,容易出现「热门专家被频繁使用,冷门专家学不到东西」的问题 → **路由坍塌**
46+
- 解决办法:训练中需要引入额外的 **负载均衡损失**
47+
48+
---
49+
50+
## 专家个数选择
51+
52+
**GLaM (Google, 2021)** 在实验中探索了不同专家数量与 gating 策略的组合:
53+
发现 **64 个专家(per layer)+ Top-2 gating** 在性能和计算效率之间达到了最佳平衡。
54+
Top-2 gating 能显著提升效果,相比单一专家更稳定。并且 64 专家的配置在 **zero-shot、one-shot、few-shot** 场景下均表现优异。 所以后续的很多 MoE 工作(如 Mixtral, DBRX, DeepSeekMoE)也基本采用 ≤64 专家的规模,这个设计在实际应用中也具有参考价值。
55+
56+
---
57+
58+
## MoE 与 PEFT
59+
60+
近期仍有不少工作专注于 PEFT(参数高效微调)。
61+
论文 [_Pushing Mixture of Experts to the Limit: Extremely Parameter Efficient MoE for Instruction Tuning_](https://arxiv.org/abs/2309.05444) 首次提出将 **LoRA 类型的 PEFT 方法和 MoE 框架结合**
62+
其主要理念为,不直接在整个大模型上加 LoRA,而是专门在 MoE 的 expert 模块里应用 LoRA。因为 MoE 的每个专家就是 FFN(MLP),它们是知识写入的关键位置。这样每次只动一小部分 LoRA experts,并且大大增强了这种架构的易扩展性。
63+
64+
![](./MOE-intro.assets/img-20250920112106588.png)
65+
66+
该方法的核心思想是利用 **低秩近似更新** 来避免高级算力的微调。
67+
68+
1. **输入 (Input → Embedding)**
69+
- 输入 token(字或子词)先经过 Embedding。
70+
- 这部分和普通 Transformer 一样。
71+
72+
2. **Multi-Head Attention**
73+
- 输入 embedding 进入多头注意力模块。
74+
- 在这里,Q、K、V 一切正常,完全没被 LoRAMoE 改动。
75+
- 输出再走 **Add & Norm**,结果传给 FFN。
76+
77+
3. **FFN → MoE (Expert 路由)**
78+
- 普通 Transformer 的 FFN 被换成 **LoRA + MoE 专家网络**
79+
- Router 根据输入选择若干专家,每个专家是 **LoRA 化(低秩适配)的模块**,而不是全量可训练的 FFN。
80+
- 冻结的部分(❄️)是预训练的大模型主干。
81+
- 火花(🔥)表示 LoRA Adapter(可训练参数,低秩矩阵)。
82+
- Router 输出的加权组合:
83+
84+
$$
85+
y = \sum_i \alpha_i \cdot Expert_i(x)
86+
$$
87+
88+
其中 $\alpha_i$ 是 Router 根据输入算出来的权重。
89+
90+
4. **输出 (Add & Norm → Residual)**
91+
- Router 混合后的专家输出,和残差连接一起进入 Add & Norm,继续往后层传。
92+
93+
---
94+
95+
#### LoRA 拆解
96+
97+
LoRA(Low-Rank Adaptation)的核心思想:
98+
99+
对一个大的线性层权重 $W \in \mathbb{R}^{d_{out} \times d_{in}}$,不去训练整个矩阵,而是加上一个低秩近似更新:
100+
101+
$$
102+
W' = W + \Delta W, \quad \Delta W = BA
103+
$$
104+
105+
- $A \in \mathbb{R}^{r \times d_{in}}, B \in \mathbb{R}^{d_{out} \times r}$
106+
- 秩 $r \ll d_{in}, d_{out}$,通常只取个位数到几十
107+
- $W$:冻结(❄️,预训练参数)
108+
- $A, B$:可训练(🔥,参数量大幅减少)
109+
110+
这样,一个输入向量 $x$ 经过 LoRA 线性层时:
111+
112+
$$
113+
Wx + BAx
114+
$$
115+
116+
等于 **原始主干输出 + 一个小的低秩修正**
117+
118+
回到该图,我们可以发现每个专家 $Expert_i$ 不是全新的大 FFN,而是 **一个 FFN 的 LoRA adapter 组合**
119+
120+
$$
121+
Expert_i(x) = B_i A_i x
122+
$$
123+
124+
- Router 对输入 hidden state 计算一个分布 $\alpha$,然后加权组合:
125+
126+
$$
127+
y = \sum_i \alpha_i \cdot Expert_i(x)
128+
$$
129+
130+
- 最终结果再加上主干(冻结的 FFN 权重输出):
131+
132+
$$
133+
y_{final} = W_{FFN}x + \sum_i \alpha_i \cdot B_i A_i x
134+
$$
135+
136+
---
137+
138+
作者:**Yang Lewis**
139+
非商业转载请标明出处。
140+
商业转载请联系作者:**840691168ly@gmail.com**
Lines changed: 193 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,193 @@
1+
---
2+
title: "算力需求指南"
3+
description: "如何计算训练大模型所需显存大小"
4+
date: "2025-09-20"
5+
tags:
6+
- compute-platforms
7+
---
8+
9+
# 大模型多卡训练笔记
10+
11+
## 1. 单位说明
12+
13+
1 GB = 1024 MB = 1024×1024 KB = 1024×1024×1024 Byte(字节) = 1024×1024×1024×8 Bit
14+
15+
**参数类型与大小:**
16+
17+
| 参数类型 | 字节数 (Byte) |
18+
| ----------- | ------------- |
19+
| FP32 | 4 |
20+
| FP16 / BF16 | 2 |
21+
| INT8 | 1 |
22+
| INT4 | 0.5 |
23+
24+
---
25+
26+
## 2. 大模型训练显存计算
27+
28+
单卡哪怕有 80 GB 显存,也扛不住数十亿、上百亿参数的大模型全量训练。
29+
权重只是一部分,梯度、优化器状态、激活都需要显存。
30+
31+
假设模型有 **N 参数量**(例如 2B = 20 亿):
32+
33+
1. **权重 W**
34+
- 存储方式:BF16 (2 Byte)
35+
- 显存:$W = N \times 2$ Byte
36+
- 例:2B 参数 → ≈ 4 GB
37+
38+
2. **梯度 G**
39+
- 存储方式:BF16 (2 Byte)
40+
- 显存:$G = N \times 2$ Byte
41+
- 例:2B 参数 → ≈ 4 GB
42+
43+
3. **优化器状态(Adam)**
44+
- 包含动量项 V、平方梯度 S,各自 FP32 (4 Byte)
45+
- 每份 ≈ 8 GB,合计 ≈ 16 GB
46+
其中:优化器显存还需要的开销(权重梯度在训练中可能会复制很多份儿)
47+
1. **权重 W**
48+
- 模型本身的参数,BF16/FP16 存 2 Byte。
49+
50+
2. **梯度 G**
51+
- 反向传播临时存储,BF16/FP16 2 Byte。
52+
53+
3. **优化器状态**(不同优化器差别大):
54+
- **SGD**:通常只需要梯度本身,**0份复制**
55+
- **SGDM(带动量的 SGD)**:需要一个动量向量(FP32,4 Byte)。**一份复制**
56+
- **Adam/AdamW**
57+
- **一阶动量 (V)**:FP32(4 Byte)。
58+
- **二阶动量 (S)**:FP32(4 Byte)。
59+
- 所以是 **2 份状态**
60+
61+
4. **Master 权重 ($W^A$)**
62+
- 混合精度训练时常见:虽然前向/反向用 BF16,但优化器更新需要 FP32 精度 → 所以再存一份 FP32 权重。
63+
64+
4. **Activations**
65+
- 依赖 batch size、seq_len、实现细节
66+
- 粗略估计:≈ 0.7–1.0 × 权重大小
67+
68+
---
69+
70+
## 3. 公式总表
71+
72+
- **普通 Adam 模式:**
73+
74+
$W + G + W^A + V + S + 0.7W ≈ 24.8$ GB (以 2B 参数为例)
75+
76+
- **DeepSpeed ZeRO-3 模式:**
77+
78+
$W + G + W^A + G^A + V + S + 0.7W ≈ 32.8$ GB
79+
80+
说明:ZeRO-3 显存更省,但通信和 I/O 开销更大。
81+
82+
---
83+
84+
## 4. 实际案例:Mixtral-8×7B
85+
86+
### 设定与常数
87+
88+
- 架构:`d_model ≈ 4096``ffn_dim ≈ 14336`,每层 **8 个专家****32 层**,SwiGLU(gate/up/down 三个线性层)。
89+
- 单个专家参数量:
90+
$4096×14336×2 + 14336×4096 = 176,160,768$ ≈ **1.76×10^8**
91+
→ BF16 权重 ≈ **352 MB/专家**
92+
- 一层 8 专家 ≈ **2.82 GB/层**
93+
- 32 层合计 ≈ **90 GB(仅专家权重)**
94+
→ 全专家全参训练在 **44 GB 显存**上不可能。
95+
96+
---
97+
98+
### 案例 A:Router-only
99+
100+
- 路由器参数:`d_model × n_experts ≈ 4k × 8 = 32k`
101+
- 全 32 层 ≈ **百万级参数**
102+
- 开销极小(MB 级),显存主要花在 **激活**
103+
- 在 44 GB 上完全可行,但改进有限
104+
105+
---
106+
107+
### 案例 B:部分层 × 部分专家
108+
109+
例:只训底部 **6 层**,每层 **2 专家**,同时训路由。
110+
111+
- 可训练参数数:
112+
$6 × 2 × 176,160,768 = 2.114B$
113+
- 权重(BF16):≈ 4.23 GB
114+
- 梯度(BF16):≈ 4.23 GB
115+
- Adam 状态(V+S,FP32):≈ 16.9 GB
116+
- Master 权重(FP32):≈ 8.46 GB
117+
- **合计(持久内存 + 梯度)**:≈ 33.8 GB
118+
- 再加冻结权重占位、激活开销,44 GB 卡上需:
119+
- `batch=1–2`
120+
- `seq_len ≤ 1024`
121+
- `use_cache=False`
122+
- `gradient_checkpointing=True`
123+
- 可行,但需要严格控制。
124+
125+
---
126+
127+
### 案例 C:4-bit 全模型 + LoRA
128+
129+
在专家 / 路由上挂 LoRA(r=16)。
130+
131+
- 单专家 LoRA 参数:
132+
$r × (4096+14336 + 4096+14336 + 14336+4096) = r × 55296$
133+
→ r=16 → 0.885M/专家
134+
- 每层 8 专家:7.08M
135+
- 32 层:226.6M LoRA 参数
136+
- 显存开销:
137+
- 权重 ≈ 0.45 GB
138+
- 梯度 ≈ 0.45 GB
139+
- Adam + Master ≈ 2.72 GB
140+
- 合计 ≈ 3.6 GB
141+
- 远低于 44 GB 显存
142+
143+
---
144+
145+
## 5. 并行方式
146+
147+
### 数据并行 (DP)
148+
149+
- 各 GPU 拷贝全模型,喂不同 batch,梯度汇总
150+
- 优点:简单
151+
- 缺点:显存浪费大
152+
153+
### 分布式数据并行 (DDP)
154+
155+
- 一卡一进程,梯度分桶同步
156+
- 优点:主流、稳定
157+
- 缺点:依然每卡全模型
158+
159+
### ZeRO 优化 (DeepSpeed)
160+
161+
- ZeRO-1:拆优化器状态
162+
- ZeRO-2:再拆梯度
163+
- ZeRO-3:连参数也拆
164+
- 优点:显存省
165+
- 缺点:通信复杂
166+
167+
### 模型并行
168+
169+
- **张量并行 (TP)**:矩阵切块
170+
- **流水并行 (PP)**:层切片,像传送带
171+
- **MoE 并行**:专家分散在不同卡,token 激活部分专家
172+
173+
---
174+
175+
## 6. 踩坑经验
176+
177+
- 显存碎片化:
178+
`PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True`
179+
(需在 `import torch` 前设置)
180+
- 通信库:
181+
NCCL > Gloo > MPI(除非特殊环境)
182+
- DDP:必须同步 random seed
183+
- eval:
184+
- eval_mb_size 可大
185+
- 训练 batch 小 + 梯度累积
186+
- 关掉 `model.config.use_cache`
187+
- device = "auto": HF 会自动根据你机器的显存大小,把模型的不同部分切片分布,对于大模型,比如 7B,在单卡 44 GB 上:通常 attention + embedding + 部分 FFN 会放 GPU,冻结的模块、或者用不到的专家 (MoE inactive experts) 可以被放到 CPU。这个参数对推理来说自动分配还是很好的,但是训练的话**需要梯度的参数**必须常驻 GPU,否则每次 forward/backward 都要把参数搬上来,通信成本爆炸。所以,**device_map=auto 对训练不总是安全**,因为它可能把你要训练的层塞到 CPU 去,导致速度慢甚至无法训练。
188+
189+
---
190+
191+
作者:**Yang Lewis**
192+
非商业转载请标明出处。
193+
商业转载请联系作者:**840691168ly@gmail.com**

0 commit comments

Comments
 (0)