DeepSeek-v3在训练和推理方面的优化
DeepSeek-v3在训练和推理方面的优化
DeepSeek-v3在训练和推理方面的优化主要集中在基础架构、MoE（Mixture of Experts）的负载均衡、Multi-Token Prediction、DualPipe并行、通信优化、显存优化、FP8混合精度训练、低精度存储和通信，以及推理阶段的优化等。这些内容对于了解深度学习模型的优化技术具有较高的参考价值。
1. 基础架构：MLA
大幅减少了KV cache大小。以1个token的1个layer为例：
- Grouped Query Attention是存group个K和group个V，LLama2-70B是group=8，维度=128；总共是28128；
- MLA是存1个KV中间量，维度=4128；再存1个K旋转中间量，维度=0.5128；总共是4*128；
2. 基础架构：MoE
同等参数量（模型的”能力“）下，训练、推理的计算量大幅减少。采用256个路由专家+64个共享专家的架构。
3. MoE的load-balance
训练中，边训练，边调整每个MoE的路由权重。负载高的减少权重，负载低的增加权重。同时，用旁路loss来鼓励句子中命中expert更均匀。限制每个token最多和M个GPU上的experts进行通信。
4. Multi-Token Prediction (MTP)
推理的时候就是Speculative Decoding。可以一次推理多个tokens（第1次全量推理，后面K次用很小的模型链式推理）。训练的时候，一次可训练整个一个句子。可视为一张网络，用所有推理步的loss加和，作为总loss，更新一次模型。
5. DualPipe
总体架构：16路Pipeline并行，64路EP并行，Zero-1 DP并行。如果不做任何优化，MoE会导致训练集群的计算:通信=1:1。
forward阶段的batch，和backward阶段的batch，可同时调度到同一张卡上，从而overlap通信和计算：
6. All-to-All通信优化
NVLink带宽：160GB/s; IB带宽：50GB/s。限制每个token只能被dispatch到最多4台机器(node)，减少IB通信量。先去重，后通过IB网传给各个node，再在各个node内部通过NVLink传给需要的GPU。
7. 显存优化
- 激活缓存：对计算量小的操作（例如RMSNorm)，只保存input，output在backward时重新计算。
- EMA参数放至CPU memory，异步放。
- Token embedding层和output head层，放到了同一张卡上，可以share。
8. FP8
开源大模型里首次使用FP8混合精度训练的。FP8和INT8一样快；FP8比FP16(或BF16)要快1倍。
9. 低精度存储和通信
Adam Optimizer的和，用的BF16。累加用的Gradient，全量Weight，用的FP32。大部分激活值，缓存FP8值。
10. 推理优化
采用PD分离。Decoding阶段用了10倍于Prefilling阶段的卡。Prefilling阶段使用32卡，Decoding阶段使用320卡。通过TP、EP、DP并行优化推理性能。
11. 预训练集群并行架构
256个路由专家，EP是64路，则每张卡放4个路由专家。DP是128，也就是每张卡放自己这个PP-stage的attention层的所有参数。