DeepSeek-v3在训练和推理方面的优化

1. 基础架构：MLA，大幅减少了KV cache大小。（计算量能不能减少？）

2. 基础架构：MoE，同等参数量（模型的”能力“）下，训练、推理的计算量大幅减少。

3. MoE的load-balance

训练中，边训练，边调整每个MoE的路由权重。负载高的减少权重，负载低的增加权重。（传统方法是将负载均衡情况作为附加旁路loss。缺点是影响模型训练的主目标）。

训练中，用旁路loss来鼓励句子中命中expert更均匀。

限制每个token最多和M个GPU上的experts进行通信。

4. Multi-Token Prediction (MTP)

推理的时候就是Speculative Decoding。可以一次推理多个tokens（第1次全量推理，后面K次用很小的模型链式推理）。

训练的时候，一次可训练整个一个句子。可视为一张网络，用所有推理步的loss加和，作为总loss，更新一次模型。

把“未来”考虑到训练里，可以让模型不再“短视”。

推理时，既可以扔掉MTP Module使用token-by-token老式推理，又可以利用MTP Module使用speculative decoding推理。

5. DualPipe

总体架构：16路Pipeline并行，64路EP并行，Zero-1 DP并行。（如何布局在2048张卡上的？）

如果不做任何优化，MoE会导致训练集群的计算:通信=1:1.

forward阶段的batch，和backward阶段的batch，可同时调度到同一张卡上，从而overlap通信和计算：

想象有一排工人（GPU），面前有2条传送带(DualPipe)，各自往相反的方向传送，每个工人可对1~2条传送带上的东西（任务）进行加工处理，如果，2个东西一个是计算任务一个是通信任务，则可并行处理之。

优点：1. overlap通信和计算。2.减少了bubble。

缺点：每张卡要保存2段模型（被EP路数分片多所化解）。新增少量要缓存的activation。

我的疑问：更新模型的同时还在做前向操作，如何确保同一个样本使用同一版本的模型参数呢？

6. All-to-All通信优化

NVLink带宽：160GB/s;  IB带宽：50GB/s

限制每个token只能被dispatch到最多4台机器(node)，减少IB通信量。

先去重，后通过IB网传给各个node，再在各个node内部通过NVLink传给需要的GPU。好处1：去重，避免同一个token的激活在IB上传给同一个node多次，减少了跨节点通信量。好处2：形成流水线，目的地node的NVLink在多播第N个token时，IB可以同时在往目的地node发送第N+1个token。

用于通信的SM，20个就够了（一张GPU卡上一般有100多个SM)。

用于通信的warp数目，动态可调的。用PTX指令精细制作，减少L2-cache使用和对其他SM的干扰。

7. 显存优化

激活缓存：对计算量小的操作（例如RMSNorm)，只保存input，output在backward时重新计算。

EMA参数放至CPU memory，异步放。

Token embedding层和output head层，放到了同一张卡上，可以share。share参数，share梯度存储。

8. FP8

开源大模型里首次使用FP8混合精度训练的。

FP8和INT8一样快；FP8比FP16(或BF16)要快1倍。

FP8原理类似INT8，先要统计矩阵数值，拿到scaling factor，矩阵绝对值的最大值对准FP8最大值；FP8乘法，累加到FP16或FP32的结果里。最后再用scaling factor恢复到FP16或FP32。

FP8混合精度训练架构：

大部分计算密集任务，用FP8来计算。小部分需要精度高才行的任务，保持BF16或FP32计算。

3个GEMM(正向传播，反向计算Weight梯度，反向计算Activation梯度），都用FP8计算。正向传播是的激活值，以FP8进行缓存，供反向传播计算Weight梯度时使用，减少了显存占用量。

非计算密集任务，或精度敏感的操作，仍保持高精度：embedding, output head，MoE gating，normalization操作，attention操作。

为维持数值稳定性，以下仍保持高精度：主Wegiht，梯度，Optmizer状态(Adam的2个参数）。可用Zero来减少显存占用。

细粒度量化：激活的量化单元设为1*128，权重的量化单元设为128*128。和整个矩阵做量化单元相比，好处：减少量化误差。

如上图，分片矩阵的FP8乘法之后，结果再乘以各自的scale（2个），加和到结果矩阵里。

这个叫做microscaling，NVIDIA下一代GPU架构官宣会支持。

提升累加的精度：发现H800的FP8乘加，会累加到14位数上面，精度不够易造成误差，特别是累加的数多的时候。

解决：在Tensor Core上，每累加一定次数后，就将中间累加结果，传输至CUDA core上和FP32进行累加。

指数和尾数：

老方法一般是用E4M3做前向传播，用E5M2做反向传播。这里得益于细粒度的分片量化，减少了分片内的数值范围，因此全部使用E4M3。

在线量化：

老方法一般是calibration时统计激活值范围，确定好量化scale，实际推理时复用该scale（类似“静态量化”的概念）。这里为了保持精度，不预先统计，随用随统计（类似“动态量化”的概念）。（权重因为会边训练边更新，所以训练时Weight也是随用随统计scale）。

9. 低精度存储和通信

Adam Optimizer的和，用的BF16。累加用的Gradient，全量Weight，用的FP32。

大部分激活值，缓存FP8值。特例：1. attention之后Linear变换之前的激活，用E5M6共12位来缓存。2. 只缓存SwiGLU的输入激活值，其输出激活值在backward时临时计算，可进一步减少缓存存储量。

MoE通信：1. forward时，激活值的dispatch通信（第1个矩阵乘法之前的），先量化为FP8，再all-to-all通信。2. backward时，梯度值的dispatch通信（第2个矩阵乘法之后的），先量化为FP8，再all-to-all通信。3. forward时，激活值的combine通信（第2个矩阵乘法之后的），保持BF16，进行all-to-all通信。4. backward时，梯度值的combine通信（第1个矩阵乘法之前的），保持BF16，进行all-to-all通信。  3和4涉及到加和操作，为了保持精度，所以采用BF16。

10. 推理：

采用PD分离。Decoding阶段用了10倍于Prefilling阶段的卡。暗合了我的测试结论，相同的token长度，Decoding阶段的耗时大约是Prefilling阶段的10倍。

Prefilling:

4台*8卡=32卡。

attention层：4路TP，外面套8路DP。4路TP比8路TP的好处是通信开销减小。

MoE层：32路Expert并行。好处：每个Expert可以分到更多的tokens（如果每个DP是一份完整的MoE，则总的Expert数目变成DP倍，每个Expert分到的token就少了）。我猜测的另一个好处：节约显存。

MoE的all-to-all：和训练阶段类似，也是先去重后通过IB网传给各个node，再在各个node内部通过NVLink传给需要的GPU。

MoE的负载均衡：在线推理阶段，监测Expert负载情况，每10分钟调整一次。让Heavy的GPU和清闲的GPU，互换一些Expert。对Heavy的Expert，启动一个冗余副本Expert，分散一部分token流量。比例：256个expert里选32个流量最大的，启动32个副本。

计算通信overlap: 把1个大batch拆分成2个小batch，错开发射时间，力争将<MoE前后的2次All-to-All，attention&MoE计算>，重叠起来。

显存换通信：每个GPU上，多放一倍的Expert。token路由时，可以找更好的发送方案。

Decoding:

40台*8卡=320卡。

attention层：TP4+SP，外层套80路DP。

MoE层：EP320，每张卡上就放1个Expert。256个常规Expert&64个shared/冗余Expert。

MoE的all-to-all：使用IB的point-to-point通信，降低了延迟。IBGDA技术。

MoE的副本冗余Expert: 也会在线上根据监控负载，来动态调整。把流量少的下掉，上线流量多的。

计算通信overlap: attention的耗时更大。因此，仍然采用2个小batch并行，只是改为<attention计算，all-to-all&MoE计算>，重叠起来。SM分配：后者计算量更小，因此，后者的SM少分些，前者的SM多分些。