Star85
368
代码
代码
Issues55
Pull Requests147
项目讨论
Wiki
分析
Star85
368
已合并
moe unpermute optim and update moe readme #559
jializheng创建于 2024年7月18日
moe unpermute optim and update moe readme #559
已合并
jializheng创建于 2024年7月18日
从refs/pull/559/head合入到master
共 12 个文件变更+253-142
MREADME.md+33-36
@@ -84,42 +84,39 @@ MindSpeed 是针对华为昇腾设备的大模型加速库。
84 84 
85# 特性介绍85# 特性介绍
86 86 
87-| 特性 | 介绍 |87+| 特性 | 介绍 |
88-|--------------------------------| ----- |88+|------------------------------|-----------------------------------------------------------|
89-| Megatron 数据并行 | [link](docs/features/data-parallel.md) |89+| Megatron 数据并行 | [link](docs/features/data-parallel.md) |
90-| Megatron 张量并行 | [link](docs/features/tensor-parallel.md) |90+| Megatron 张量并行 | [link](docs/features/tensor-parallel.md) |
91-| Megatron 流水并行 | [link](docs/features/pipeline-parallel.md) |91+| Megatron 流水并行 | [link](docs/features/pipeline-parallel.md) |
92-| Megatron 虚拟流水并行 | [link](docs/features/virtual-pipeline-parallel.md) |92+| Megatron 虚拟流水并行 | [link](docs/features/virtual-pipeline-parallel.md) |
93-| Megatron 序列并行 | [link](docs/features/sequence-parallel.md) |93+| Megatron 序列并行 | [link](docs/features/sequence-parallel.md) |
94-| Megatron 重计算 | [link](docs/features/recomputation.md) |94+| Megatron 重计算 | [link](docs/features/recomputation.md) |
95-| Megatron 分布式优化器 | [link](docs/features/distributed-optimizer.md) |95+| Megatron 分布式优化器 | [link](docs/features/distributed-optimizer.md) |
96-| Megatron 异步DDP | [link](docs/features/async-ddp.md) |96+| Megatron 异步DDP | [link](docs/features/async-ddp.md) |
97-| Megatron 权重更新通信隐藏 | [link](docs/features/async-ddp-param-gather.md) |97+| Megatron 权重更新通信隐藏 | [link](docs/features/async-ddp-param-gather.md) |
98-| Ascend DeepSpeed MoE | [link](docs/features/moe.md) |98+| Megatron Mcore MoE | [link](docs/features/megatron_moe/megatron-moe.md) |
99-| Ascend Mask归一 | [link](docs/features/generate-mask.md) |99+| Ascend DeepSpeed MoE | [link](docs/features/deepspeed_moe/deepspeed-moe.md) |
100-| Ascend 内存碎片优化 | [link](docs/features/memory-fragmentation.md) |100+| Ascend Mask归一 | [link](docs/features/generate-mask.md) |
101-| Ascend 自适应选择重计算 | [link](docs/features/adaptive-recompute.md) |101+| Ascend 内存碎片优化 | [link](docs/features/memory-fragmentation.md) |
102-| Ascend 激活函数重计算 | [link](docs/features/activation-function-recompute.md) |102+| Ascend 自适应选择重计算 | [link](docs/features/adaptive-recompute.md) |
103-| Ascend 计算通信并行优化 | [link](docs/features/communication-over-computation.md) |103+| Ascend 激活函数重计算 | [link](docs/features/activation-function-recompute.md) |
104-| Ascend BF16 参数副本复用 | [link](docs/features/reuse-fp32-param.md) |104+| Ascend 计算通信并行优化 | [link](docs/features/communication-over-computation.md) |
105-| Ascend rms_norm 融合算子 | [link](docs/features/rms_norm.md) |105+| Ascend BF16 参数副本复用 | [link](docs/features/reuse-fp32-param.md) |
106-| Ascend swiglu 融合算子 | [link](docs/features/swiglu.md) |106+| Ascend rms_norm 融合算子 | [link](docs/features/rms_norm.md) |
107-| Ascend rotary_embedding 融合算子 | [link](docs/features/rotary-embedding.md) |107+| Ascend swiglu 融合算子 | [link](docs/features/swiglu.md) |
108-| Ascend MoE token重排性能优化 | [link](docs/features/token-rearrange.md) |108+| Ascend rotary_embedding 融合算子 | [link](docs/features/rotary-embedding.md) |
109-| Ascend MoE dropless性能优化 | [link](docs/features/efficient-moe-deepspeed-moe.md) |109+| Ascend flash attention 适配 | [link](docs/features/flash-attention.md) |
110-| Ascend MoE 负载感知内存均衡 | [link](docs/features/efficient-moe-megatron-moe.md) |110+| Ascend nano-pipe流水线并行 | [link](docs/features/nanopipe-pipeline-parallel.md) |
111-| Ascend MoE Grouped GEMM | [link](docs/features/efficient-moe-megatron-moe.md) |111+| Ascend MLP 通信隐藏 | [link](docs/features/pipeline-experts.md) |
112-| Ascend flash attention 适配 | [link](docs/features/flash-attention.md) |112+| Ascend 重计算流水线独立调度 | [link](docs/features/recompute_independent_pipelining.md) |
113-| Ascend nano-pipe流水线并行 | [link](docs/features/nanopipe-pipeline-parallel.md) |113+| Ulysses 长序列并行 | [link](docs/features/ulysses-context-parallel.md) |
114-| Ascend MLP 通信隐藏 | [link](docs/features/pipeline-experts.md) |114+| Ring Attention 长序列并行 | [link](docs/features/ring-attention-context-parallel.md) |
115-| Ascend 重计算流水线独立调度 | [link](docs/features/recompute_independent_pipelining.md) |115+| 【Prototype】混合长序列并行 | [link](docs/features/hybrid-context-parallel.md) |
116-| Ulysses 长序列并行 | [link](docs/features/ulysses-context-parallel.md) |116+| 【Prototype】Ascend MC2 | [link](docs/features/mc2.md) |
117-| Ring Attention 长序列并行 | [link](docs/features/ring-attention-context-parallel.md) |117+| 【Prototype】alibi | [link](docs/features/alibi.md) |
118-| 【Prototype】混合长序列并行 | [link](docs/features/hybrid-context-parallel.md) |118+| 【Prototype】PP自动并行 | [link](docs/features/automated-pipeline.md) |
119-| 【Prototype】Ascend MC2 | [link](docs/features/mc2.md) |119+| 【Prototype】其他昇腾亲和优化 | 暂无 |
120-| 【Prototype】alibi | [link](docs/features/alibi.md) |
121-| 【Prototype】PP自动并行 | [link](docs/features/automated-pipeline.md) |
122-| 【Prototype】其他昇腾亲和优化 | 暂无 |
123 120 
124# 自定义算子121# 自定义算子
125 122 
Rdocs/features/efficient-moe-deepspeed-moe.md→docs/features/deepspeed_moe/deepspeed-moe-efficient-moe.md+6-7
@@ -1,5 +1,4 @@
1-# Efficient-MOE-DeepSpeed-MOE1+# MoE token dropless性能优化
2-## 1. MOE token dropless性能优化
3 2 
4### 问题分析3### 问题分析
5 4 
@@ -13,7 +12,7 @@
13 12 
141. MoE模型训练过程中capacity具有一定的连续性，维护一个滑动窗口来保存近几次统计的capacity来预估下一个窗口的capacity。131. MoE模型训练过程中capacity具有一定的连续性，维护一个滑动窗口来保存近几次统计的capacity来预估下一个窗口的capacity。
152. 在Gate中不再每一个step都直接进行all reduce统计全局最大capacity，而是各个进程先判断当前的capacity能否满足不丢弃token，通过reduce统计全局的判断信息，若都能满足则无需进行all reduce通信，否则进行all reduce通信取得实际max capacity.142. 在Gate中不再每一个step都直接进行all reduce统计全局最大capacity，而是各个进程先判断当前的capacity能否满足不丢弃token，通过reduce统计全局的判断信息，若都能满足则无需进行all reduce通信，否则进行all reduce通信取得实际max capacity.
16-<p align="center"> <img src="../../sources/images/moe_dynamic_padding_a.png" height="300px" width="600px"></p>15+<p align="center"> <img src="../../../sources/images/moe_dynamic_padding_a.png" height="300px" width="600px"></p>
17 16 
18### 使用场景17### 使用场景
19 18 
@@ -36,16 +35,16 @@
36训练模型：Mixtral(4层)35训练模型：Mixtral(4层)
37 36 
38精度对比图如下：37精度对比图如下：
39-<p align="center"> <img src="../../sources/images/moe_dynamic_padding_b.png" height="300px" width="800px"></p>38+<p align="center"> <img src="../../../sources/images/moe_dynamic_padding_b.png" height="300px" width="800px"></p>
40 39 
41top2 多种并行方式 提速效果：40top2 多种并行方式 提速效果：
42-<p align="center"> <img src="../../sources/images/moe_dynamic_padding_c.png" height="400px" width="800px"></p>41+<p align="center"> <img src="../../../sources/images/moe_dynamic_padding_c.png" height="400px" width="800px"></p>
43 42 
44top1 多种并行方式 提速效果：43top1 多种并行方式 提速效果：
45-<p align="center"> <img src="../../sources/images/moe_dynamic_padding_d.png" height="400px" width="800px"></p>44+<p align="center"> <img src="../../../sources/images/moe_dynamic_padding_d.png" height="400px" width="800px"></p>
46 45 
47同时开启此优化减少显存占用3%：46同时开启此优化减少显存占用3%：
48-<p align="center"> <img src="../../sources/images/moe_dynamic_padding_e.png" height="500px" width="800px"></p>47+<p align="center"> <img src="../../../sources/images/moe_dynamic_padding_e.png" height="500px" width="800px"></p>
49 48 
50 49 
51 50 
Rdocs/features/token-rearrange.md→docs/features/deepspeed_moe/deepspeed-moe-token-rearrange.md+1-1
@@ -1,4 +1,4 @@
1-# token重排性能优化1+# Token 重排性能优化
2 2 
3## 问题分析3## 问题分析
4 4 
Rdocs/features/moe.md→docs/features/deepspeed_moe/deepspeed-moe.md+82-70
@@ -1,70 +1,82 @@
1-# Ascend DeepSpeed MoE 相关特性1+# Ascend DeepSpeed MoE 相关特性
2- 2+ 
3-## 整体方案：3+## 整体方案：
4- 4+ 
5-Mixture of Expert（MoE）是指混合专家模型功能。在大模型训练中使用该功能可以将常规的稠密大模型变成稀疏的MoE大模型，在计算量不显著增加的情况下大幅提升模型的参数量。5+Mixture of Expert（MoE）是指混合专家模型功能。在大模型训练中使用该功能可以将常规的稠密大模型变成稀疏的MoE大模型，在计算量不显著增加的情况下大幅提升模型的参数量。
6- 6+ 
7-通过使用专家并行（Expert Parallelism，EP），把专家分配到多个NPU上，减轻单个NPU的显存压力，也就是说专家并行（Expert Parallelism，EP），对全量专家进行分组。7+通过使用专家并行（Expert Parallelism，EP），把专家分配到多个NPU上，减轻单个NPU的显存压力，也就是说专家并行（Expert Parallelism，EP），对全量专家进行分组。
8-如图所示，一个包含6个专家的MoE模型在EP=2时的专家分布情况。可以把专家并行理解成模型并行的一种形态（模型被切分成多份），但是输入的数据又是不同的（DP），因此token在经过Router之后，可能会选中别的卡上的专家，此时就需要将这些token发送过去，即EP进程组内需要通过All2All通信交换token。值得注意的是，该MoE模型在token选择专家时，如果超过容量会drop掉token。8+如图所示，一个包含6个专家的MoE模型在EP=2时的专家分布情况。可以把专家并行理解成模型并行的一种形态（模型被切分成多份），但是输入的数据又是不同的（DP），因此token在经过Router之后，可能会选中别的卡上的专家，此时就需要将这些token发送过去，即EP进程组内需要通过All2All通信交换token。值得注意的是，该MoE模型在token选择专家时，如果超过容量会drop掉token。
9- 9+ 
10-![输入图片说明](moe.png)10+![输入图片说明](moe.png)
11- 11+ 
12-## 特性背景：12+## 特性背景：
13- 13+ 
14-支持moe模型及相关特性的兼容和适配，包含MoE基础模型、MoE适配序列并行（SP）、MoE适配长序列（CP）、MoE token重排性能优化。14+支持moe模型及相关特性的兼容和适配，包含MoE基础模型、MoE适配序列并行（SP）、MoE适配长序列（CP）、MoE token重排性能优化。
15- 15+ 
16-1.Mindspeed新增MoE混合专家模型（Mixtral 8*7B），支持使用MoE模型进行训练。16+1.Mindspeed新增MoE混合专家模型（Mixtral 8*7B），支持使用MoE模型进行训练。
17- 17+ 
18-2.MoE支持序列并行(sequence parallel)，支持MoE与序列并行同时开启，减少MoE模块计算，提升MoE模块的训练性能。18+2.MoE支持序列并行(sequence parallel)，支持MoE与序列并行同时开启，减少MoE模块计算，提升MoE模块的训练性能。
19- 19+ 
20-3.MoE适配长序列(context parallel)特性，支持MoE和CP特性同时开启。20+3.MoE适配长序列(context parallel)特性，支持MoE和CP特性同时开启。
21- 21+ 
22-4.MoE适配分布式优化器特性，支持MoE和分布式优化器同时开启，降低内存，减少OOM风险。22+4.MoE适配分布式优化器特性，支持MoE和分布式优化器同时开启，降低内存，减少OOM风险。
23- 23+ 
24-5.MoE token重排性能优化，减少token选择专家的gate计算量，提升训练性能。24+5.MoE token重排性能优化，减少token选择专家的gate计算量，提升训练性能。
25- 25+ 
26-## 使用场景26+## 使用场景
27- 27+ 
28-在需要处理大规模数据集和复杂任务的情况下，使用基于 moe 结构的大模型，以及其他SP、CP等特性。28+在需要处理大规模数据集和复杂任务的情况下，使用基于 moe 结构的大模型，以及其他SP、CP等特性。
29- 29+ 
30-### 使用建议：30+### 使用建议：
31- 31+ 
32-MoE+SP 建议默认开启。开启SP需要同时开启TP。32+MoE+SP 建议默认开启。开启SP需要同时开启TP。
33- 33+ 
34-MoE+cp 建议处理长序列时开启。34+MoE+cp 建议处理长序列时开启。
35- 35+ 
36-MoE+分布式优化器 建议默认开启。36+MoE+分布式优化器 建议默认开启。
37- 37+ 
38-Token重排优化 建议默认开启。38+Token重排优化 建议默认开启。
39- 39+ 
40- 40+ 
41-## 使用方法41+## 使用方法
42- 42+ 
43-MoE特性基础功能:43+MoE特性基础功能:
44-| 重要参数| 参数说明 |44+ 
45-| ---- | ---- |45+| 重要参数| 参数说明 |
46-|--moe-model-type deepspeed_moe |使用mixtral模型 |46+| ---- | ---- |
47-|--num_experts [int] |专家数 |47+|--moe-model-type deepspeed_moe |使用mixtral模型 |
48-|--expert_model_parallel_size [int] |专家并行 |48+|--num_experts [int] |专家数 |
49- 49+|--expert_model_parallel_size [int] |专家并行 |
50- 50+ 
51-MoE支持序列并行:51+ 
52-|重要参数| 参数说明 | 注意事项 |52+MoE支持序列并行:
53-| ---- | ---- | ---- |53+ 
54-|--sequence-parallel |开启SP | MoE场景开启SP，由于影响数据分发，节省内存的同时可能引起性能波动|54+|重要参数| 参数说明 | 注意事项 |
55- 55+| ---- | ---- | ---- |
56-MoE适配长序列:56+|--sequence-parallel |开启SP | MoE场景开启SP，由于影响数据分发，节省内存的同时可能引起性能波动|
57-| 重要参数| 参数说明 |57+ 
58-| ---- | ---- |58+MoE适配长序列:
59-|--context-parallel-algo megatron_cp_algo |配置CP算法|59+ 
60- 60+| 重要参数| 参数说明 |
61- 61+| ---- | ---- |
62-MoE适配分布式优化器:62+|--context-parallel-algo megatron_cp_algo |配置CP算法|
63-|重要参数| 参数说明 |63+ 
64-| ---- | ---- |64+ 
65-|--use-distributed-optimizer |开启分布式优化器特性|65+MoE适配分布式优化器:
66- 66+ 
67-MoE特性token重排优化:67+|重要参数| 参数说明 |
68-|重要参数| 参数说明 |68+| ---- | ---- |
69-| ---- | ---- |69+|--use-distributed-optimizer |开启分布式优化器特性|
70-|--enable-token-rearrange-opt |开启token重排|70+ 
71+MoE特性token重排优化:
72+ 
73+|重要参数| 参数说明 |
74+| ---- | ---- |
75+|--enable-token-rearrange-opt |开启token重排|
76+ 
77+## 特性优化
78+ 
79+| 特性 | 介绍 |
80+|------------------------------------|------------------------------------------------------|
81+| Ascend DeepSpeed MoE token 重排性能优化 | [link](deepspeed-moe-token-rearrange.md) |
82+| Ascend DeepSpeed MoE dropless 性能优化 | [link](deepspeed-moe-efficient-moe.md) |
Rdocs/features/moe.png→docs/features/deepspeed_moe/moe.png+0-0
文件重命名但无更改。
Adocs/features/megatron_moe/megatron-moe-adaptive-recompute-activation.md+20-0
@@ -0,0 +1,20 @@
1+# MoE 负载感知内存均衡算法
2+ 
3+## 问题分析
4+ 
5+MoE在预训练前期负载均衡 aux_loss 未起作用时，token 在专家层的分配不均会导致全局内存负载不均衡问题，并引入大量碎片内存，导致显存波动巨大，并呈现逐步递增的趋势，大集群训练下更容易出现OOM。
6+ 
7+## 优化方案
8+ 
9+根据模型设定参数（DP/TP/SeqLength/NumExpert等），设定token分布不均的判定阈值，当超过该阈值后执行重计算，避免产生内存不均衡的激活值。
10+ 
11+## 使用限制
12+ 
13+1. 使用时**建议**开启`--moe-router-load-balancing-type aux_loss`，这样会使得训练过程中token分布快速趋向于平衡。
14+2. 开启全局重计算后该功能无效。
15+ 
16+## 使用方法
17+ 
18+- 必选参数：脚本中加入`--moe-adaptive-recompute-activation`即可。
19+ 
20+- 可选参数：如果希望调节判定阈值，则修改`--moe-adaptive-recompute-activation-scale`即可，默认值为2.0，表示阈值为平均分配在每个专家上的序列的两倍长度。
Adocs/features/megatron_moe/megatron-moe-allgather-dispatcher.md+32-0
@@ -0,0 +1,32 @@
1+# Allgather Dispatcher 分支优化
2+ 
3+## 问题分析
4+### 1. gather & scatter 算子替换
5+ 
6+在 Megatron MoE 中的 Allgather 分支，存在使用 gather/scatter 操作。gather/scatter 功能为沿 dim 轴根据索引逐元素进行取值/赋值操作，此操作会有大量的随机地址，对性能造成巨大影响。
7+ 
8+在 Megatron MoE 中对 gather/scatter 的调用主要是以下调用方式，通过对 index 做 expand 操作对维度进行扩展，再通过扩展后 index 对 hidden_states 进行逐元素取值/赋值。
9+ 
10+```python
11+self.global_local_map = global_local_map.view(-1, 1).expand(-1, hidden_states.shape[-1])
12+local_hidden_states = torch.gather(global_hidden_states, 0, self.global_local_map)
13+```
14+ 
15+### 2. 异步通信
16+在 Allgather dispatcher 分支中，会 permutate 函数开头分别对 hidden_states、max_ind、max_prob 三个数据做 allgather 通信，此些操作为串行操作，但各计算任务之间并非串行依赖关系。
17+ 
18+ 
19+## 解决方案
20+### 1. gather & scatter 算子替换
21+由于 index 是通过 expand 进行扩展的，因此它的每一行中的内容都是一致，而我们没有必要使用 gather/scatter 进行逐元素的操作，可通过 index 算子以及 indexput 算子进行逐行操作，对 gather/scatter 进行等价替换。
22+ 
23+ 
24+### 2. 异步通信
25+通过对通信任务进行重新排序，并使用 async=True 参数进行异步下发，达到计算和通信并行的目的。
26+ 
27+## 使用场景
28+ 
29+在使用 mcore MoE 的场景下，开启了 `--moe-token-dispatcher-type allgather`。
30+ 
31+## 使用方法
32+开启参数 `--moe-permutation-async-comm`。
Adocs/features/megatron_moe/megatron-moe-alltoall-dispatcher.md+25-0
@@ -0,0 +1,25 @@
1+# Alltoall Dispatcher 分支优化
2+ 
3+## 问题分析
4+### 1. repeat_interleave 并行
5+在 Alltoall dispatcher 分支中，调用了 repeat_interleave 算子，此算子只使用了单个 block dim 在单个下发流上进行串行计算，且耗时较长，算子的输出也是在 alltoall、permute、alltoallv 之后才用到。
6+ 
7+### 2. 计算通信并行
8+在 alltoall 分支中的 permutation 函数最后会进行 allgather 操作对所有 tokens 被切分的 H 维进行补全，然后再对数据分块进行专家计算。此项操作为串行操作，但各专家间的 tokens 并没有存在依赖关系，可修改为并行操作。
9+ 
10+ 
11+## 解决方案
12+### 1. repeat_interleave 并行
13+通过新建一条下发流，将 repeat_interleave 算子调用分到新的流上，在 block dim 资源充足的情况下，可进行两个算子的并行计算，节省耗时。
14+ 
15+### 2. 计算通信并行
16+可按照每个专家需要的 tokens 进行切分，然后逐个对 tokens 进行 allgather 通信 + 专家计算，由于第一个专家计算只依赖第一个通信，专家之间无依赖关系，因此在做第一个专家计算的时候可同步进行第二专家的通信，达到计算和通信并行。
17+ 
18+## 使用场景
19+在使用 mcore MoE 的场景下，开启了 `--moe-token-dispatcher-type alltoall`。
20+ 
21+## 使用方法
22+开启参数 `--moe-permutation-async-comm`。
23+ 
24+## 场景限制
25+由于开启 `--moe-grouped-gemm` 后，专家计算被单一算子合并，因此计算通信并行优化会失效。
Rdocs/features/efficient-moe-megatron-moe.md→docs/features/megatron_moe/megatron-moe-gmm.md+10-27
@@ -1,35 +1,16 @@
1-# Efficient-MOE-Megatron-MOE1+# Megatron MoE Grouped GeMM
2-## 1. MoE 负载感知内存均衡算法
3 2 
4-### 问题分析3+## 问题分析
5 4 
6-MOE在预训练前期负载均衡aux_loss未起作用时，token在专家层的分配不均会导致全局内存负载不均衡问题，并引入大量碎片内存，导致显存波动巨大，并呈现逐步递增的趋势，大集群训练下更容易出现OOM。5+针对MoE单卡多专家计算，存在细碎的专家计算操作与通信，通过Grouped GeMM算子对多专家计算进行合并，提升MoE单卡多专家训练性能。
7 6 
8-### 优化方案7+## 解决方案
9 8 
10-根据模型设定参数（DP/TP/SeqLength/NumExpert等），设定token分布不均的判定阈值，当超过该阈值后执行重计算，避免产生内存不均衡的激活值。9+通过调用 gmm 融合算子，对多个专家计算进行融合，达到加速效果。
11 10 
12-### 使用限制11+## 使用方法
13 12 
14-1. 使用时**建议**开启`--moe-router-load-balancing-type aux_loss`，这样会使得训练过程中token分布快速趋向于平衡。13+设置`--moe-grouped-gemm`: 表示开启Grouped GeMM计算。
15-2. 开启全局重计算后该功能无效。
16- 
17-### 使用方法
18- 
19-- 必选参数：脚本中加入`--moe-adaptive-recompute-activation`即可。
20- 
21-- 可选参数：如果希望调节判定阈值，则修改`--moe-adaptive-recompute-activation-scale`即可，默认值为2.0，表示阈值为平均分配在每个专家上的序列的两倍长度。
22- 
23- 
24-## 2. MoE Group GEMM
25- 
26-### 问题分析
27- 
28-针对MoE单卡多专家计算，存在细碎的专家计算操作与通信，通过Group GEMM算子对多专家计算进行合并，提升MoE单卡多专家训练性能。
29- 
30-### 使用方法
31- 
32-设置`--moe-grouped-gemm`: 表示开启Group-GEMM计算。
33 14 
34## 效果说明15## 效果说明
35 16 
@@ -38,9 +19,10 @@ MOE在预训练前期负载均衡aux_loss未起作用时，token在专家层的
38- EP变小导致单卡专家数量增大 & DeepSeek MoE专家数量较多等场景。19- EP变小导致单卡专家数量增大 & DeepSeek MoE专家数量较多等场景。
39- DeepSeek MoE finegrained expert单个专家较小 & FFN规模不大 & TP变大导致单卡切分的计算变小。20- DeepSeek MoE finegrained expert单个专家较小 & FFN规模不大 & TP变大导致单卡切分的计算变小。
40 21 
41-1. 随着FFN规模提升，计算不再细碎，单专家计算效率提升，Group GEMM收益变小。22+1. 随着FFN规模提升，计算不再细碎，单专家计算效率提升，Grouped GeMM 收益变小。
42 23 
43表1：grok模型FFN大小和性能加速对比24表1：grok模型FFN大小和性能加速对比
25+ 
44|ffn_hidden_size| 32768 | 16384| 8192| 4096|26|ffn_hidden_size| 32768 | 16384| 8192| 4096|
45|--|--|--|--|--|27|--|--|--|--|--|
46|baseline|2280|1780|1537|1446|28|baseline|2280|1780|1537|1446|
@@ -51,6 +33,7 @@ MOE在预训练前期负载均衡aux_loss未起作用时，token在专家层的
512. TP越大，EP越小，收益更大。332. TP越大，EP越小，收益更大。
52 34
53表2：Mixtral8*7B模型配置不同性能收益35表2：Mixtral8*7B模型配置不同性能收益
36+ 
54|配置| tp4 ep2 16expert | tp4 ep2 8expert | tp2 ep4 16expert| tp2 ep4 8expert|37|配置| tp4 ep2 16expert | tp4 ep2 8expert | tp2 ep4 16expert| tp2 ep4 8expert|
55|--|--|--|--|--|38|--|--|--|--|--|
56|baseline|27969|20127|11976|13981|39|baseline|27969|20127|11976|13981|
Adocs/features/megatron_moe/megatron-moe.md+10-0
@@ -0,0 +1,10 @@
1+# Megatron mcore MoE 相关特性
2+ 
3+## 特性优化
4+ 
5+| 特性 | 介绍 |
6+|----------------------------------------------------------|-------------------------------------------------------|
7+| Ascend Megatron MoE 负载感知内存均衡算法 | [link](megatron-moe-adaptive-recompute-activation.md) |
8+| Megatron MoE GMM | [link](megatron-moe-gmm.md) |
9+| 【Prototype】Ascend Megatron MoE Allgather Dispatcher 性能优化 | [link](megatron-moe-allgather-dispatcher.md) |
10+| 【Prototype】Ascend Megatron MoE Alltoall Dispatcher 性能优化 | [link](megatron-moe-alltoall-dispatcher.md) |
Mmindspeed/core/transformer/moe/moe_utils.py+32-0
@@ -11,3 +11,35 @@ def permute(tokens, indices, topk: int = 1):
11 sorted_indices = torch.sort(flatten_indices.float(), stable=True)[1]11 sorted_indices = torch.sort(flatten_indices.float(), stable=True)[1]
12 permuted_tokens = tokens.index_select(0, sorted_indices // topk)12 permuted_tokens = tokens.index_select(0, sorted_indices // topk)
13 return permuted_tokens, sorted_indices13 return permuted_tokens, sorted_indices
14+ 
15+ 
16+def unpermute(
17+ permuted_tokens: torch.Tensor,
18+ sorted_indices: torch.Tensor,
19+ probs: torch.Tensor = None,
20+ padded_mode: bool = False,
21+ restore_shape: torch.Size = None,
22+):
23+ if padded_mode:
24+ return unpermute_with_padded_tokens(
A
Aaijgnem12024年8月5日

这里unpermute_with_padded_tokens没有定义

likedislike
25+ permuted_tokens, sorted_indices, probs, restore_shape=restore_shape
26+ )
27+ 
28+ assert sorted_indices.numel() == permuted_tokens.size(0)
29+ if probs is not None:
30+ # Unpermute and merge the tokens with their probabilities
31+ num_unpermuted_tokens = probs.numel()
32+ topk = probs.size(1)
33+ else:
34+ # Unpermute the tokens without merge
35+ num_unpermuted_tokens = permuted_tokens.size(0)
36+ topk = 1
37+ 
38+ sorted_indices = torch.argsort(sorted_indices.float()).int()
39+ unpermuted_tokens = permuted_tokens.index_select(0, sorted_indices)
40+ unpermuted_tokens = unpermuted_tokens.reshape(-1, topk, permuted_tokens.size(-1))
41+ if probs is not None:
42+ unpermuted_tokens = unpermuted_tokens * probs.unsqueeze(-1)
43+ unpermuted_tokens = unpermuted_tokens.sum(dim=1)
44+ 
45+ return unpermuted_tokens
Mmindspeed/megatron_adaptor.py+2-1
@@ -399,12 +399,13 @@ def mcore_moe_adaptation(pm, args):
399 if hasattr(args, 'moe_token_dispatcher_type') and args.moe_token_dispatcher_type == 'alltoall':399 if hasattr(args, 'moe_token_dispatcher_type') and args.moe_token_dispatcher_type == 'alltoall':
400 from .core.transformer.moe.token_dispatcher import preprocess, alltoall_token_permutation400 from .core.transformer.moe.token_dispatcher import preprocess, alltoall_token_permutation
401 from .core.transformer.moe.experts import sequential_mlp_forward401 from .core.transformer.moe.experts import sequential_mlp_forward
402- from .core.transformer.moe.moe_utils import permute402+ from .core.transformer.moe.moe_utils import permute, unpermute
403 pm.register_patch('megatron.core.transformer.moe.token_dispatcher.MoEAlltoAllTokenDispatcher.preprocess', preprocess)403 pm.register_patch('megatron.core.transformer.moe.token_dispatcher.MoEAlltoAllTokenDispatcher.preprocess', preprocess)
404 pm.register_patch('megatron.core.transformer.moe.token_dispatcher.MoEAlltoAllTokenDispatcher.token_permutation',404 pm.register_patch('megatron.core.transformer.moe.token_dispatcher.MoEAlltoAllTokenDispatcher.token_permutation',
405 alltoall_token_permutation)405 alltoall_token_permutation)
406 pm.register_patch('megatron.core.transformer.moe.experts.SequentialMLP.forward', sequential_mlp_forward)406 pm.register_patch('megatron.core.transformer.moe.experts.SequentialMLP.forward', sequential_mlp_forward)
407 pm.register_patch('megatron.core.transformer.moe.moe_utils.permute', permute)407 pm.register_patch('megatron.core.transformer.moe.moe_utils.permute', permute)
408+ pm.register_patch('megatron.core.transformer.moe.moe_utils.unpermute', unpermute)
408 409 
409 if args.use_ascend_mc2:410 if args.use_ascend_mc2:
410 # MoE MLP not use mc2 linear411 # MoE MLP not use mc2 linear