?一、引言

本文主要聚焦于大語言模型預訓練相關階段的技術和行業最新進展，其中包括常見的分布式策略、模型結構、常見的優化手段等?？紤]到篇幅原因，暫不包含后訓練、多模態等領域。

二、模型結構

2.1 概述

當前 LLM 基本上都是 Decoder-Only 的 Transformer 模型，只不過都會進行一些修改。比如對 Attention 的修改衍生出來 Softmax Attention 系列和 Linear Attention 系列。而對 FFN 的修改衍生出了 Dense 模型和 MoE 模型。這個章節我們對這些模型結構的修改進行簡單的總結。

2.2 Attention

2.2.1 MHA、MQA、GQA、MLA

對于 Softmax Attention 系列，目前主要的是 MHA、MQA、GQA 和 MLA，主要區別如下圖 Figure 3 所示：

• MHA：Transformer 模型的標準實現，每個 Attention Head 都有獨立的 Q、K、V。

早期模型用的比較多，現在很少使用，至少是比較大的模型都沒有使用。

Inference 階段的問題比較多，最主要的問題是 KV Cache 占比是幾種方案里最大的，并且在 Continuous Batching 階段 Attention 核心計算無法只用 Tensor Core，存在一定的局限性。

• MQA：所有 Attention Head 共享相同的 K 和 V（[2003.04641] MQA: Answering the Question via Robotic Manipulation）。

計算量和 KV Cache 都是最小的，也是對 Inference 階段最友好的。

對模型效果影響較大，所以很少使用。

GQA：Attention Head 進行分組，一組 Attention Head 共享相同的 K 和 V（[2305.13245] GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints）。

是 MHA 和 GQA 的折衷方案，KV Cache 大小處于 MHA 和 GQA 之間，計算效率也處于兩者之間?？紤]到 Tensor Core 要求矩陣乘法的 Shape 是 8 的整數倍，因此每組有 8 個 Head 的整數倍是最優的。如果每組中不到 8 個 Head，計算時可以 Padding，也能用上 Tensor Core，只是存在部分冗余計算。

是當前模型中最常見的的 Attention 方式。

• MLA：DeepSeek 2024 年 5 月在 DeepSeek V2（DeepSeek-V2: A Strong, Economical, and Efficient Mixture-of-Experts Language Model）中提出，并在 DeepSeek V3 （[2412.19437] DeepSeek-V3 Technical Report）中沿用。核心思路是每個 Head 都有獨立的 K 和 V，但它們可以投影和反投影到同樣的、共享的 Latent KV。

KV Cache 和 MQA 相當，明顯少于 MHA 和 GQA，但效果還不錯，和 MHA 相當。但是也會額外的增加一些計算量。DeepSeek 也開源了部分針對 MLA Inference 的代碼實現（FlashMLA: Efficient MLA decoding kernels）。

當前主要是 DeepSeek V2、V3 以及最新開源的?Kimi K2 模型（Kimi-K2 - a moonshotai Collection）中使用，其他模型并沒有跟進。

如下圖 Table 1 可以看出，MLA 的 KV Cache 需求雖然依然大于 MQA，但明顯優于 MHA 和 GQA，同時效果更好：

如下圖 Table 8 和 Table 9 所示，DeepSeek 團隊在 DeepSeek V2 的技術報告中進行過一些消融實驗，MLA 的效果優于 GQA，甚至優于 MHA：

2.2.2 MHA -> GQA/MLA

模型在預訓練階段已經確定好是 MHA，還是 GQA、MLA?？紤]到 GQA 和 MLA 在 Inference 階段的優勢，也就衍生了一系列將 MHA 轉為 GQA、MLA 或者將 GQA 轉為 MLA 的方案（PS：這些方案和預訓練沒有太大關系，更多是對 Inference 的幫助，不過這里也簡單介紹）。

如下圖 Table 1 所示，在 [2412.20677] Align Attention Heads Before Merging Them: An Effective Way for Converting MHA to GQA 中作者提出了將 MHA 轉換為 GQA 的方案。其在 LLaMA2 7B 上做實驗（共 32 個 Attention Head ）：

• 在 GQA-16 上能取得與 MHA 相當甚至更好的效果，但是 KV Cache 壓縮比更高的 GQA-8（節約 75% KV Cache）和 GQA-4（節約 87.5% KV Cache）精度下降較大。
• LLaMA2 模型預訓練數據比較少，模型訓練不夠充分，有效性有待商榷。
• 當前的 LLM 基本至少會默認采用 GQA，因此相應場景也就更少。

如下圖 Table 1 所示，在 [2502.14837] Towards Economical Inference: Enabling DeepSeek's Multi-Head Latent Attention in Any Transformer-based LLMs 中，作者進一步探索了 MHA/GQA 到 MLA 的轉換（PS：其實 GQA 可以看做 MHA 的特例，因此如果支持 MHA 轉 MLA，就很容易支持 GQA 轉 MLA）。

• 在模型比較小時損失比較大，在 LLaMA2 7B 上損失還可以接受。
• 上述評測還存在兩方面問題：

LLaMA2 模型沒有充分訓練，Baseline 比較低。

LLaMA2 里面的評估任務中 OBQA 提升比較明顯拉升了整體得分，在更多任務上的有效性有待商榷。

如下圖 Table 1 所示，在 [2502.07864] TransMLA: Multi-Head Latent Attention Is All You Need 中，作者同樣提出了 MHA 到 MLA 的轉換方案 TransMLA：

• 相比上述的 MHA2MLA，在無需訓練的方式中提升較明顯。而在加入一定訓練數據后表現相當。
• 和 MHA2MLA 有同樣問題，評測任務中考 OBQA 拉高得分，另外模型是 LLaMA2MLA。
• 如果與直接將 GQA 轉為 MLA 的消融實驗會更有說服力。 

2.2.3 Linear Attention

由于 Softmax Attention 的二次方特性，當序列比較長時，Attention 的開銷急劇增加。針對這種場景，Linear Attention 和 Sparse Attention 是兩種常見的流派。在 Linear Attention 中比較常見的是 Mamba 和 RWKV 系列，不過在常見的開源模型中并非主流，業內質疑的聲音也比較多。

如下圖 Table 3 所示，在 RWKV-7[2503.14456] RWKV-7 "Goose" with Expressive Dynamic State Evolution 中，作者進行了相應對比實驗：

• 針對小規模模型，基于 Linear Attention 的 RWKV7 模型確實能獲得與基于 Softmax Attention 相當的性能。
• 主要實驗聚焦在 3B 及以下模型，缺少 7B 及更大規模模型的實驗。

2.2.4 Hybrid Linear Attention

相較于全部采用 Linear Attention 的方式，也有不少工作采用 Softmax Attention 和 Linear Attention 混合的方案。其相比全部 Linear Attention 更加保守，能夠同時保留 Softmax Attention 的高精度以及 Linear Attention 在長序列的高效率。

這類工作中規模最大，影響力最大的是 MiniMax 的開源大模型 MiniMax-01 系列模型（[2501.08313] MiniMax-01: Scaling Foundation Models with Lightning Attention）。如下圖 Figure 3 所示，其最大的 MiniMax-01 W456A46 模型包含 80 個 Transformer Block，其中每 7 個 Linear Attention 接一個 Softmax Attention（也就是下圖中的 M=7）。并且 Softmax Attention 采用了 GQA，而 Linear Attention 也采用了定制化的高性能實現 Lighting Attention。

如下圖 Figure 8 所示，作者也對比了不同機制的訓練吞吐?？梢钥闯?，隨著序列長度增加：

• Softmax Attention 的效率逐漸降低，并且序列越長，下降比例越快。不過對于預訓練常見的 8K 序列長度，基本還可以接受（PS：其實對于長序列也可以考慮采用 Sample Packing Mask 來優化，只不過是需要考慮負載均衡的問題）。
• Linear Attention 的 Lightning、HGRN2、Mamba2 幾乎都能維持性能不降。（PS：也可以看出，對于 LLM 預訓練常見的 4K、8K 訓練長度，Linear Attention 沒有特別明顯的優勢）
• Hybrid-Lightning 雖然也會出現性能下降的問題，但是明顯好于 Softmax Attention。

騰訊同樣也提出了混合結構的 Hunyuan-TurboS W560A56 模型（[2505.15431] Hunyuan-TurboS: Advancing Large Language Models through Mamba-Transformer Synergy and Adaptive Chain-of-Thought），其中 Softmax Attention 與 Mamba2 Layer 的比例大約為 1:8（7 vs 57）。

如下圖 Table 2 所示，其效果也還不錯，與 DeepSeek-V3、Qwen3-235B-A22B 相當：

除了 MiniMax 和 Hunyuan 外，NVIDIA 也發表過基于 Mamba 的混合模型 Mamba-2-Hybrid 8B（[2406.07887] An Empirical Study of Mamba-based Language Models），其混合結構如下圖 Table 6 所示，共 28 個 Attention Layer，其中：

• 24 個是 Mamba2 的 Linear Attention
• 4 個是 Softmax Attention，并且采用了 GQA。

如下圖 Table 7 所示，在一些常見的任務上確實可以獲得還不錯的效果：

2.2.5 Sparse Attention

與 Linear Attention 類似，Sparse Attention 也是為了降低長序列的計算開銷，只是方向不同。Sparse Attention 的主要動機是：Attention Score 是高度稀疏化的，只關注權重比較高的 Score 可以大幅降低計算復雜度并維持相當的精度。

Sparse Attention 在長序列 Inference 場景使用非常多，而預訓練場景序列長度比較小，比較少使用。這里面比較早的工作是 Mistral 7B（[2310.06825] Mistral 7B）中使用的 Sliding Window Attention，是 Sparse Attention 的一種特例。如下圖 Figure 1 所示，其相當于每個 Token 只關注附近的 Token。

在這個領域，今年比較火的有兩個工作（PS：都是針對長序列場景），一個是 Moonshot AI 的 [2502.13189] MoBA: Mixture of Block Attention for Long-Context LLMs，另外一個是 DeepSeek 的 Hardware-Aligned and Natively Trainable Sparse Attention。它們在某些方面都促進了一些共識：

• 長序列場景（Long Prefill 或 Long Decoding），Attention 是高度稀疏化的，也是高度動態化的。
• 固定 Pattern 的稀疏化方式往往很難保持精度，可學習 Sparse Pattern 是通用化且高精度的有效方案。
• Token 粒度的稀疏化很難充分發揮 GPU 算力，Block 粒度稀疏化是精度和性能（稀疏度、計算量）的良好平衡，基于此的高效 Block Sparse Attention 也成為標配。
• 當前常見的 LLM 通常會采用 GQA，也要充分結合 GQA 的特性來設計稀疏化方案，不然可能會影響整體的稀疏化程度。
• 在進行 Block 選擇時并不需要使用 Block 內所有的 KV Cache，選擇一個代表性的“聚類中心”即可，比如取 Avg 或者 Max。
• 不要隨意永久性丟棄 Token，由于 LLM 的自回歸特性，很難推測在后續的生成中是不是一定不需要某個 Token。這也就是為什么在 NSA 和 MOBA 中并不會節約 KV Cache 的存儲空間。

如下圖 Figure 1 所示為 Moonshot AI MoBA 的主要原理：

如下圖 Figure 2 所示是 DeepSeek NSA 的主要原理：

2.3 MoE

2.3.1 粗粒度 MoE

2023 年的大語言模型還以 Dense 模型為主，2024 年初 Mistral AI 發布 Mixtral 8x7B MoE 模型（[2401.04088] Mixtral of Experts），引發了 MoE LLM 的熱潮。不過早期的 MoE 還是比較粗粒度的專家，比如 Mixtral 8x7B 只有 8 個專家，后續的 Mixtral 8x22B 也是只有 8 個專家。

早期的 MoE 模型訓練都比較保守，往往采用先訓練 Dense 模型，然后通過 Upcycling 的方式擴展到 MoE 模型，比如上述的 Mixtral 8x7B 是由 Mistral 7B Upcycling 而來。在昆侖萬維的 Skywork-MoE（[2406.06563] Skywork-MoE: A Deep Dive into Training Techniques for Mixture-of-Experts Language Models）中也對相應方案有所探討，并將 Skywork 13B Dense 模型擴展為 Skywork 146B 的 MoE 模型（16 個專家）。

2.3.2 細粒度 + 共享 MoE

DeepSeek 團隊在 DeepSeek MoE 的技術報告（[2401.06066] DeepSeekMoE: Towards Ultimate Expert Specialization in Mixture-of-Experts Language Models）中提出了細粒度專家+共享專家方案，并且在后續模型中一直延續。如下圖 Figure 2 所示，DeepSeek MoE 模型主要有兩點改進：

• 細粒度專家（Routed Expert）：常見的 MoE 模型中通常是 8 或 16 個專家，而這里會將一個大專家切分為 M 個小專家。比如原來從 16 個專家中選擇 Top 2 大概有 120 種可能；而同樣計算量的 64 個專家（M=4）中選擇 8 個，對應了 4,426,165,368 種可能。
• 共享專家（Shared Expert）：額外增加了 1 個或多個共享專家，用于捕獲通用知識，每個 Token 都會經過這些共享專家。

3 個模型的具體配置如下所示，需要說明的是，3 個模型中都未使用 GQA，而是使用的 MHA（PS：這個應該是 23 年的工作，正好是 MHA 往 GQA 過渡的階段）：

在 DeepSeek MoE 階段（24 年 1 月），細粒度 MoE 還沒被廣泛接受（當模型規模不大時，細粒度 MoE 對訓練性能影響較大）。直到 DeepSeek V2（DeepSeek-V2: A Strong, Economical, and Efficient Mixture-of-Experts Language Model）甚至是后續的 DeepSeek V3（[2412.19437] DeepSeek-V3 Technical Report），這種細粒度專家和共享專家的方式的方式才被廣泛接受并得到大規模使用。

2.3.3 Dense + MoE

除了 Softmax Attention 與 Linear Attention 的混合架構外，也有一些 Dense 模型和 MoE 模型的混合架構。這種方式最早出現在 Google 經典的 Switch Transformer模型中（[2101.03961] Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity），其同樣是 MoE 相關工作中非常經典的 Paper。如下圖 Figure 2 所示，單看架構圖很容易誤解為是一個純粹的 MoE 模型（每一個 Transformer Layer 都包含 MoE Block），一些非官方的代碼實現中也是如此。然而，實際上該模型是一個混合架構模型。

如下圖 Table 9 所示，其中提到 Expert Freq 為 1/2，表明 MoE Transformer Layer 和 Dense Transformer Layer 各占 1/2：

Google 在此后的 GLaM 模型（[2112.06905] GLaM: Efficient Scaling of Language Models with Mixture-of-Experts）和 ST-MoE 模型（[2202.08906] ST-MoE: Designing Stable and Transferable Sparse Expert Models）中都沿襲了這種方式。不過這些都是在 ChatGPT 之前的工作。

在后續的 LLM 中比較少出現這種情況。最常見的是 DeepSeek 系列 MoE 模型以及沿襲 DeepSeek V3 的 Kimi K2 模型：

• DeepSeek MoE 包含 28 層，第 1 層是 Dense Layer，后續 27 層為 MoE Layer。
• DeepSeek V2 包含 60 層，第 1 層是 Dense Layer，后續 59 層為 MoE Layer。
• DeepSeek V3 包含 61 層，前 3 層是 Dense Layer，后續 58 層為 MoE Layer。
• DeepSeek K2 包含 61 層，第 1 層是 Dense Layer，后續 60 層為 MoE Layer。

DeepSeek 在 DeepSeek MoE 的 Paper 中提到，采用這種方式主要是第 1 層的負載均衡狀態收斂很慢，因此將第 1 層保留為 Dense 層。Moonshot 在實現 K2 模型時同樣發現第 1 層的 MoE 的 Router 很難做到負載均衡，但不同的是第 2 層并沒有發現此現象，因此相比 DeepSeek V3，只是第 1 層使用 Dense 層，第 2 和 第 3 層都使用 MoE 層。

其實后續的很多其他工作也采用了類似的方案：

• 小紅書 dots.llm1：62 層，第 1 層 Dense，后 61 層 MoE。
• 百度 ERNIE 4.5：54 層，前 3 層 Dense，后 51 層 MoE。

2.4 其他改進

2.4.1 MTP

DeepSeek V3 模型同樣采用了 DeepSeek V2 的 MLA 以及細粒度專家+共享專家的 MoE 結構。除了模型增大之外，模型層面最主要的變化是引入 MTP（Multi-Token Prediction），這個思路在 Inference 的投機采樣中經常使用，只不過這里也可以幫助提升模型訓練的效果。具體來說：

• 其中 Main Model 就是標準的 Next Token Prediction。
• MTP Module 1 用于預測下下一個 Token，MTP Module 2 用于預測下下下一個 Token（與 LLM 推理中常見的多頭投機采樣思路一致）。
• MTP Module 中的輸入都包含兩個部分，一個是上一個 Module 的 Output Head 的輸入，以及上一個輸入 Token，并且其中的 Embedding Layer 和 Output Head 都是共享自 Main Model，只有新增的 RMSNorm + Linear Projection 和一個 Transformer Block。由于這里有兩個輸入分別經過 RMSNorm 后 Concat 到一起，因此需要一個額外的 Linear Projection 進行降維，保持維度一致。

MTP 策略主要用于提升 Main Model 的性能，因此在推理階段，可以直接舍棄 MTP Module，Main Model 仍能獨立且正常運行。此外，還可將這些 MTP Module 用于投機解碼，以進一步降低生成延遲。

三、分布式并行策略

3.1 概述

LLM 模型規模很大，單 GPU 甚至單節點往往無法放下，即使可以放下也可能因為顯存空間有限而限制各種優化策略的實施。為了解決這些問題并獲得最大吞吐，通常會采用各種混合分布式并行策略來優化，常見的有 DP（Data Parallelism）、TP（Tensor Parallelism）、PP（Pipeline Parallelism）、EP（Expert Parallelism）和 SP（Sequence Parallelism）。除此之外，也有許多優化方案試圖改進上述并行策略，以便獲得更優吞吐。

3.2 DP

3.2.1 概述

DP 是最常用的并行策略，因為它與其他并行策略正交，實現簡單并且通信量相對不是很大，很容易擴展訓練的規模。但是 DP 也存在一個比較明顯的問題：在每個 DP 組內都有完整的模型、優化器狀態和梯度副本，導致內存開銷比較大。

3.2.2 DeepSpeed ZeRO

為了解決內存開銷大的問題，微軟提出了 ZeRO 相關工作（[1910.02054] ZeRO: Memory Optimizations Toward Training Trillion Parameter Models）。如下圖 Figure 1 所示，根據不同的程度充分將優化器狀態（os）、梯度（g）和模型參數（p）切分到所有的設備中，也就是不同的 DP 組中會存儲不同的優化器狀態、梯度和參數切片。

• ZeRO-1（Pos）：將優化器狀態切分到所有設備，每個設備還有全量的模型參數和梯度。優先切分優化器狀態是因為其占用內存更多，并且與 Forward 和 Backward 的反向傳播無關，只影響模型權重參數更新階段。（PS：由于 ZeRO-1 中每個設備只需要對應部分的平均梯度，而不像 DP 那樣需要梯度的 AllGather，因此總的通信量不變。也就是說，在極大降低顯存開銷的情況下并不會增加通信量，這也是為什么常見的并行方案中基本都會默認采用 ZeRO-1 或者叫 ZeRO-DP）
• ZeRO-2（Pos+g）：在 ZeRO-1 的基礎上進一步切分梯度，切分梯度也不影響 Forward 過程。
• ZeRO-3（Pos+g+p）：在 ZeRO-2 的基礎上進一步切分模型參數，會影響 Forward 階段，需要 AllGather 所有參數才能計算，會引入更多通信。采用 ZeRO-3 幾乎可以將內存需求降低到 1/N，其中 N 表示設備數，這里是 64。

3.2.3 PyTorch FSDP

與 DeepSpeed 的 ZeRO 優化方案類似，Meta 也提供了 PyTorch 原生支持的 FSDP V1（[2304.11277] PyTorch FSDP: Experiences on Scaling Fully Sharded Data Parallel）和 FSDP V2（[2411.00284] SimpleFSDP: Simpler Fully Sharded Data Parallel with torch.compile） 方案。早期的 Megatron-LM 框架不支持 FSDP，限制了 FSDP 的發展，最近半年 NVIDIA 在 Megatron-Core 里實現了相應的能力并進行了一系列優化，也能獲得很不錯的吞吐。隨著后續集群 Scale-Up 域的擴展（比如 NVL72），FSDP 也許會有更大的空間。

3.3 TP

很多時候模型在單一 GPU 無法放下，此時通常會采用 TP 切分，對于 Transformer 模型而言，最常見的是 NVIDIA 在 [1909.08053] Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism 提出的切分方式。

如下圖 （a）所示，MLP 層的兩個 FC 采用先列切（A，Column Parallelism），然后行切（B，Row Parallelism）的方案，這樣兩個 FC 之間不用通信：

如下圖（b）所示，由于每個 Head 的 Attention，Softmax 都是獨立的，因此可以采用按照 Head 的方式切分（等價于 Column Parallelism），然后對之后的 FC 采用行切分（B，Row Parallelism），這樣 Self-Attention 中間也不用通信：

然而，由于每個 Transformer Layer 的 Forward 和 Backward 都需要 2 次 AllReduce 操作，并且通信量大，為了避免 TP 通信成為瓶頸，通常會將 TP 切分到一個節點內，因為節點內的 8 個 GPU 可以充分利用 NVLink + NVSwitch 的高帶寬（PS：這也是為什么 TP 通常不會大于 8）。

如下圖 Figure 8 所示為 Megatron-LM 中一個 DP + TP 的混合分布式并行方案，總共采用 64 臺 8 GPU 機器，共 512 GPU。

• 每臺機器的 8 個 GPU 組成一個 Model Parallelism Group（TP），共 64 個 TP Group；每個 TP Group 內的 GPU 包含不同的模型參數，并且使用相同的訓練數據。
• 所有設備的同號 GPU（比如 GPU 1，9，...，505）組成一個 Data Parallelism Group（DP），共 8 個 DP Group；每個 DP Group 內的 GPU 都有相同的模型參數，但是使用不同的訓練數據。

當然，TP 也存在明顯的問題：TP 會對 Tensor 進行切分，從而可能降低矩陣計算的算術強度。對于比較大的 Dense LLM，由于 TP 通常不會很大，基本還能接受；但是對于比較小的模型，或者細粒度的 MoE 模型，其矩陣乘法的 Shape 本身比較小，TP 切分后對算術強度的影響比較大，會導致吞吐的明顯下降，無法充分發揮 GPU 的性能，因此在細粒度 MoE 模型的專家部分比較少采用 TP 并行。

3.3 PP

3.3.1 概述

PP 是另一種常見的模型并行策略，其同樣是將模型切分成不同的部分，只不過與 TP 的切分方式不同，其主要是將模型的不同層切分到不同的設備上。

如下圖 Figure 3 所示為使用 4 個設備進行 PP 訓練的執行過程。其每一行代表一個設備，藍色表示 Forward，綠色表示 Backward，Forward 和 Backward 中的數字指的是 Mini Batch 的 ID。由于是按層切分，并且同一時間只有 1 個 Mini Batch，每個設備都必須等待之前的設備執行完才能執行對應的 Stage，導致存在大量的 Bubble。

3.3.2 1F1B

實際上當設備 1 執行完 Mini Batch 1 的 Forward 之后便可以開始 Mini Batch 2 的 Forward，以此類推。在調度的過程中，系統中的每個設備都會有兩個選擇：

• 對某個 Mini Batch 執行 Forward，進而可以將 Mini Batch 傳遞到下一個設備。
• 對另一個 Mini Batch 執行 Backward，進而確保學習向前推進。

如果始終優先考慮 Forward，則會導致阻塞 Backward，模型也就無法學習和更新，因為只有 Backward 后才能執行權重更新；同樣，如果只考慮 Backward 優先調度，則會導致計算資源閑置，無法充分發揮算力。

為了避免上述問題，1F1B （1次 Forward，1次 Backward，[1806.03377] PipeDream: Fast and Efficient Pipeline Parallel DNN Training）調度機制應運而生。如下圖 Figure 8 所示，4 個設備，分成 4 個 Stage。在起始階段允許執行多個 Mini Batch 的 Forward，穩定后就保持 Forward 和 Backward 的交替執行，這樣可以保證 GPU 在穩定狀態下沒有空閑，并且始終繼續學習。

上述的 1F1B 過程并不需要 Forward 和 Backward 一樣長，實際上，Backward 總是大于 Forward（大約 2 倍），此時 1F1B 依然是有效的調度機制。

PP 相比 TP 來說，只用在相鄰的 PP Stage 間進行 P2P 通信即可，其通信次數、通信量相對較少，因此比較適合跨節點間通信；除此之外，從上圖也可以推測出，增加梯度累加的次數，也就是 Mini Batch 的數量，可以較大程度降低 Bubble 率，梯度累加次數越大，Bubble 率越小。

3.3.3 Interleaved 1F1B

采用 PP 最需要關注的問題就是 Bubble 率，需要盡可能的利用流水線機制讓所有設備處于工作狀態，提升整體吞吐。

NVIDIA 在 [2104.04473] Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM 中提出的 Interleaved-1F1B 方案正是為了降低 Bubble 率。

如下圖 Figure 4 所示為基于 1F1B 提出的 Interleaved 1F1B 調度方案。

• 上圖為 1F1B：假設模型有 16 層，每個 Device 有 4 層。

• 對應 K=4 個 Device，Micro Batch 個數為 M=8，也就是每 M=8 個 Micro Batch 進行一次同步梯度更新。
• 模型被分為 4 個 Stage，Device 1 包含 Layer (0,1,2,3)，Device 2 包含 Layer (4,5,6,7)，Device 3 包含 Layer(8,9,10,11)，Device 4 包含 Layer(12,13,14,15)。

• 下圖為?Interleaved 1F1B：
• 與標準 1F1B 的主要不同是層的切分方式。
• 模型被分為 8 個 Stage，Device 1 包含 Layer (0,1,8,9)，Device 2 包含 Layer (2,3,10,11)，Device 3 包含 Layer(4,5,12,13)，Device 4 包含 Layer(6,7,14,15)?？梢钥闯?，相當于將模型切分為 8 個 Stage，但是交替放在 4 個 Device 上，下圖中深色代表前 4 個 Stage（Layer 0-7），淺色代表后 4 個 Stage（Layer 8-15）。以此就可以實現更細力度的調度，減少 Bubble。

3.3.4 ZeroBubble

Sea AI-Lab 團隊在 [2401.10241] Zero Bubble Pipeline Parallelism 中進一步提出了 Zero Bubble 方案，可以進一步降低 Bubble 率。

如下圖 Figure 1 所示，ZeroBubble 中將 Backward 分成兩個部分，一部分計算輸入的梯度，一部分計算權重的梯度。這里計算輸入的梯度有明確的依賴關系，也是鏈式法則不斷傳遞的基礎；而計算權重的梯度卻沒有明確的依賴，甚至可以滯后很多。此外，三個紅色部分計算量相當，這也就是為什么之前 1F1B 或者 Interleaved-1F1B 中 Backward 的長度為 Forward 的 2 倍。

可以看出，ZeroBubble 的方式為降低 Bubble 率提供了更多的可能，然而也有一定局限性。首先，當梯度累積次數比較多時，Bubble 率本身不大，提升的空間也就比較有限；此外，通常優化方案中會將上述兩個梯度的計算放在一個 Kernel 里，ZeroBubble 會將其變成兩個 Kernel，有可能導致效率的降低。

3.3.5 非均勻 PP

由于 LLM 模型中除了 Transformer Block 外還有一些其他組件，比如還有開始的 Word Embedding，結尾的 LM Head 以及 Loss 計算，對于 DeepSeek 模型還有 MTP Layer，如果按照 Transformer Block 平均切分會存在計算的負載不均衡；此外，由于 1F1B 的機制問題，PP Stage 越靠前的部分就會占用越多的 GPU 顯存，導致顯存的不均衡。

對于計算負載不均的問題，思路也比較簡單，既然首尾 Stage 的計算負載更大，那么讓首尾的層數少一些即可。比如：

• 智譜在 GLM-130B 模型（[2210.02414] GLM-130B: An Open Bilingual Pre-trained Model）時就提出將 72 層 Transformer Block 變成 70 層，PP 為 8，中間各 9 層，起始各 8 層，這樣可以降低首尾 Stage 的壓力。
• Meta 在 LLaMA3 405B 模型（[2407.21783] The Llama 3 Herd of Models）中同樣采用了類似的方案，共包含 126 層，同樣是首尾 Stage 少 1 層。
• 昆侖萬維在 Skywork-MoE 也類似，24 層，由 [6, 6, 6, 6] 切分變為 [5, 5, 5, 5, 4]。

對于顯存開銷不均的問題，通常會使用重計算和 Offloading 機制來緩解，不過也需要注意重計算或者 Offloading 的粒度，通常不會對整個 Transformer Layer 進行，而是針對個別計算，以便盡可能降低重計算等帶來的額外開銷。

3.3.6 DeepSeek DualPipe

對于 DeepSeek V3 而言，跨節點 EP 引入的通信開銷導致計算與通信比約為 1:1，效率很低。為了應對這一挑戰，作者設計了一種創新的流水線并行算法 DualPipe。

DualPipe 的核心思想是：將一對獨立的 Forward 與 Backward Chunk 內的計算與通信進行 Overlap。特別地，對于 Backward Chunk，借鑒 ZeroBubble，將 Attention 與 MLP 的 Backward 分為兩部分：Backward for Input 及 Backward for Weight。

如下圖 Figure 4 所示，針對一對 Forward 與 Backward Chunk，重新排列這些組件，并手動調整 GPU SM 在通信與計算間的分配比例。在此 Overlap 策略下，能夠確保 All2All 和 PP 通信在執行過程中完全隱藏，其中：

• 橙色表示 Forward
• 綠色表示 Backward for Input
• 藍色表示 Backward for Weight
• 紫色表示 PP 通信
• 紅色表示 Barrier 同步

完整的 DualPipe 調度如下圖 Figure 5 所示，其采用雙向 PP 調度，同時從流水線兩端輸入 Micro Batch，使得大部分通信得以完全 Overlap（PS：8PP，雙向 20 Micro Batch，反方向 10-19 的 10 個 Micro Batch 并沒有列出來，因此我們用紅色 10-19 補充了部分 Micro Batch）。這種 Overlap 還確保了隨著模型進一步擴展，只要保持恒定的計算與通信比，仍可在跨節點部署細粒度專家的同時，實現近乎零的 All2All 通信開銷。

PS：正常來說是無法實現雙向 PP 調度的，主要是因為 Forward 執行順序是從前往后，比如從 Layer 0,1,2,...,14,15，而 Backward 執行順序是從后往前，比如 Layer 15,14,...,2,1,0。而常見 PP 中的 Layer 只會在某一個 PP Stage，比如 8 PP，那么：

• Stage 0 上有 Layer 0 和 1 的權重
• Stage 1 上有 Layer 2 和 3 權重
• Stage 7 上有 Layer 14 和 15 的權重
• Forward 的順序也只能從 Stage 0 到 Stage 7，不能從 Stage 7 到 Stage 0。

而 DeepSeek V3 的雙向 PP 調度中，還是 8 PP 為例：

• Stage 0 上有 Layer 0, 1 以及 Layer 14, 15 的權重
• Stage 1 上有 Layer 2, 3 以及 Layer 12, 13 的權重
• Stage 7 上有 Layer 14, 15 以及 Layer 0, 1 的權重
• 相當于有 2 份相同的模型副本，Forward 的順序可以從 Stage 0 到 7，也可以從 Stage 7 到 0。

3.3.7 NVIDIA Merged FWD-BWD

DeepSeek 的 DualPipe 會導致靜態顯存翻倍，此外仍然存在較高的 Bubble 率。針對上述問題，NVIDIA 也提出利用奇&偶 Micro-Batch 實現 Overlap 的方案，已在 Megatron-LM 中集成，如下圖所示：

此外，華為在 [2505.04519] Pangu Ultra MoE: How to Train Your Big MoE on Ascend NPUs 中提出了和 NVIDIA 類似的方案，核心思路也是利用 Micro Batch 間的獨立性，以 Forward 計算掩蓋 Backward 通信（反之亦然）。

3.4 EP

3.4.1 概述

對于 Dense LLM 而言，采用 DP（Zero-1） + TP + PP 基本都能獲得不錯的吞吐。而對于 MoE 模型，尤其是細粒度 MoE，TP 已經不再合適，一般都會引入 EP 策略。這里通常會涉及幾個方面的問題：

• 在 MoE 之前引入 All2All Dispatch，在 MoE 之后引入 All2All Combine 操作，通信開銷比較大。
• Dispatch 和 Combine 處存在比較多小的 Kernel，效率比較低。
• EP 中不同設備可能存在負載不均的問題。

3.4.2 DeepEP

MoE 模型 EP 的主要挑戰就是高性能的 All2All 實現。DeepSeek V3 中，為了確保 DualPipe 具有足夠的計算性能，定制了高效的跨節點 All2All 通信 Kernel（包括 Dispatching 和 Combining），以節省專用于通信的 SM 數量。Kernel 的實現也與 MoE Gating 算法及集群的網絡拓撲共同設計。

• 為了有效利用 IB 和 NVLink 的不同帶寬，將每個 Token 限制為最多被發送到 4 個節點，從而減少 IB 流量。
• 對于每個 Token，在做出路由決策時，首先通過 IB 傳輸到其目標節點上具有相同節點內索引的 GPU。一旦到達目標節點，將努力確保它通過 NVLink 立即轉發到承載目標專家的特定 GPU，而不會被隨后到達的 Token 阻塞。（PS：比如說，節點 A 上 GPU 0 的 Token 要發送到節點 B 上的 GPU 3，則對應的路徑為：節點 A GPU 0 -> 節點 B GPU 0 -> 節點 B GPU 3。這樣做是因為高性能 GPU 訓練集群往往會采用軌道優化，同號 GPU 在一個 Leaf Switch 下，因此可以利用高速的 NVLink 來代替從 Leaf Switch 到 Spine Switch 的流量，從而降低 IB 通信時延，并且減少 Leaf Switch 和 Spine Switch 之間的流量，這也是 NVIDIA PXN 的核心思路）

DeepSeek 隨后也開源了相應實現 DeepEP（GitHub - deepseek-ai/DeepEP: DeepEP: an efficient expert-parallel communication library），是專為 MoE 和 EP（Expert Parallelism, EP）設計的通信庫。提供了一系列優化的通信 Kernel，實現了以下能力：

• 高度優化的 All2All 通信，適合 MoE 模型 2 個主要過程：

Dispatch：將 Token 發送給專家。

Combine：從專家接收處理過的 Token 過程。

• 同時支持不同的通信類型：

節點內（intra-node）：可以使用 NVLink + NVSwitch 通信。

節點間（inter-node）：可以使用 RDMA 通信。

• 針對不同場景的 Kernel：

常規（高吞吐） Kernel（Normal Kernel）：針對 Training 和 Inference Prefill。節點內 NVLink + 節點間 RDMA 通信。

低時延 Kernel（Low-Latency Kernel）：針對 Inference Decoding。使用純 RDMA 通信來最小化時延。

• 原生支持 FP8，減少數據傳輸需求，相比 FP16 通信量減半。
• 靈活的 GPU 資源（SM）控制，支持計算和通信的 Overlap。

除此之外，如果模型不是非常大，將 EP 的 All2All 放在一個節點內，使用 NVLink + NVSwitch 的高帶寬通信也是常見的優化方案。

3.4.3 專家負載均衡

專家負載均衡是 MoE 模型另外一個非常常見的問題。其不僅影響預訓練階段，還會影響 Inference 階段的性能。針對這個問題，通常會采用負載均衡損失來解決，由于 DeepSeek 系列模型的負載均衡比較典型，這里以 DeepSeek 系列為主。

DeepSeek MoE 中采用了 2 種負載均衡損失：

• 專家級負載損失（Expert-Level Balance Loss）：讓各個專家的負載平均，盡量避免都路由到少數專家。
• 設備級負載損失（Device-Level Balance Loss）：強制讓各個專家負載平均可能影響效果。由于每個設備（GPU）包含多個專家，因此讓不同設備上的計算量盡量均衡同樣可以一定程度避免負載不均的問題。

DeepSeek V2 中采用了 3 種輔助負載均衡損失：

• 專家級負載損失（Expert-Level Balance Loss）：同 DeepSeek MoE。
• 設備級負載損失（Device-Level Balance Loss）：同 DeepSeek MoE。
• 通信負載損失（Communication Balance Loss）：設備限制路由可以保證每個設備發送的通信量是有界的，但如果某個設備接收的 Token 比其他設備多，實際的通信效率也會受到影響。為此，引入通信負載損失，正是為了緩解不同設備接收 Token 不均衡的問題。保證設備之間均衡交換信息，促進高效通信。

DeepSeek V2Token Dropping 策略：負載均衡損失可以促進均衡，但是無法嚴格保證。為了進一步減少負載不均導致的計算資源浪費，DeepSeek V2 中額外引入了設備級 Token 丟棄策略（Device Level Token Dropping Strategy）。首先計算每個設備的平均計算預算，也就意味著每個設備的容量因子等同于 1.0，然后在每個設備上丟棄親和力得分最低的 Token，直到達到計算預算需求；除此之外，確保約 10% 的訓練序列中的 Token 永遠不會被丟棄。根據效率需求，可以在 Inference 過程中靈活配置是否丟棄 Token，并始終保證 Training 和 Inference 的一致性。

DeepSeek V3 中進一步對專家負載均衡策略進行了調整：

• 無需輔助損失的負載均衡策略（Auxiliary-Loss-Free Load Balancing Strategy）：來自 DeepSeek 2024 年的論文（[2408.15664] Auxiliary-Loss-Free Load Balancing Strategy for Mixture-of-Experts），具體來說，其通過動態更新每個專家的偏置（b）來維持專家的負載均衡，而不會引入額外的干擾梯度。
• 補充的序列級輔助損失（Complementary Sequence-Wise Auxiliary Loss）：盡管 DeepSeek-V3 主要依賴于無輔助損失的策略來實現負載平衡，但為了防止在任何單一序列中出現極端的不平衡，作者采用一種補充的序列級均衡損失。這種序列級均衡損失的目的是鼓勵每個序列中的專家負載更加均衡，避免負載集中在少數專家上，從而提高模型的效率和公平性。
• 節點限制路由（Node-Limited Routing）：與 DeepSeek-V2 所采用的設備限制路由類似，DeepSeek-V3 同樣使用了一種約束路由機制以控制訓練過程中的通信開銷。簡而言之，確保每個 Token 最多被發送至 M 個節點，這些節點的選擇依據是分布在各節點上的專家中，其親和度得分最高的 Kr/M 項之和。在此約束下，MoE 訓練框架幾乎能夠實現計算與通信的完全重疊。

DeepSeek V3No Token-Dropping：得益于高效的負載均衡策略，DeepSeek-V3 在整個訓練過程中保持了良好的負載平衡。因此，DeepSeek-V3 在訓練期間未丟棄任何 Token。此外，作者還實施了特定的部署策略以確保推理過程中的負載均衡，故 DeepSeek-V3 在推理階段同樣不會丟棄 Token。

3.4.4 MoE 計算優化 —— 設備內

當前常見粗粒度 MoE 模型的專家數通常是 8-32，而細粒度 MoE 的專家數通常可以達到 64-256，在 Kimi K2 中更是高達 384。在 Inference 階段經常會采用超大 EP 的方式以盡可能提升計算強度；而在 Training 階段，不像 Inference 的 Decoding 有那么明顯的 Memory Bound 問題，因此通常 EP 不會特別大。此時，一個設備上會存在多個專家，相應的計算問題就是經典的 Grouped GEMM 計算問題，由于每個專家的 Token 數可能不同，相應矩陣計算的 Shape 也可能不等。

Grouped GEMM 有幾種常見的解法，在不同 Shape 下它們的性能也會差距很大（如下圖 Figure 1 是非常早的一個圖，一定程度上可以說明這個問題）：

• For 循環串行調度：效率最低，但是當 Shape 非常大時可能也不會特別差。
• Multi-Stream 調度：多個 Stream 同時執行，在 Shape 不是特別大時可以更充分的利用 GPU 資源，但是 Kernel Launch 開銷無法避免。
• Batched GEMM：對于同 Shape 的多個 GEMM 計算，cuBLAS 提供了 BatchedGEMM 的 API（cublasgemmgroupedbatchedex），可以直接使用，通常能獲得很不錯的性能。
• Grouped GEMM：對于不同 Shape 的多個 GEMM 計算，CUTLASS 和 DeepSeek 開源的 DeepEP（DeepGEMM: clean and efficient FP8 GEMM kernels with fine-grained scaling）也提供了相應方案。這種方式通常能獲得很不錯的性能。

如下圖 Table 1 所示，小紅書團隊在 [2506.05767] dots.llm1 Technical Report 中也提到了相關優化方案，通過優化 GroupedGEMM 獲得相應算子 Forward 14%，Backward 8% 的提升：

3.4.5 MoE 負載優化 —— 設備間

MoE 負載均衡損失可以盡可能的降低負載不均的問題，但是依然無法嚴格保證。此外，Grouped GEMM 可以優化設別內的計算效率，但是無法緩解節點間負載不均的問題。

針對上述問題，華為在 Pangu Ultra MoE（[2505.04519] Pangu Ultra MoE: How to Train Your Big MoE on Ascend NPUs）中提出了動態重排的方案。如下圖 Figure 11 所示，如果某個 Device 上被路由到的 Token 數都比較少（Device 0），則相較于路由到比較多 Token 的 Device 會出現計算的 Bubble（Device 1）。通過 Planner 和 Executor 的協同合作，動態調整 Device 上 Expert 的排布，可以讓不同 Device 上的負載盡可能均衡，從而提升整體的利用率。

華為在 Pangu Pro MoE（[2505.21411] Pangu Pro MoE: Mixture of Grouped Experts for Efficient Sparsity）中還采用了專家分組的方式，在專家選擇時，讓每個組都選擇固定數量的專家，這樣可以保證每組的負載是均衡的，但會降低專家的可組合數，也就是降低專家組合的空間。（DeepSeek 中也有類似的方案）

其實 DeepSeek V3 在 Inference 階段也會采用類似但稍有不同的負載均衡方案，比如針對共享專家、高負載專家采用冗余部署，并動態重排的方案來盡可能的實現負載的均衡。

3.4.5 Permute & UnPermute 算子優化

在 MoE 模型中另外一個容易影響吞吐的地方就是 MoE 相關的 Permute 和 Unpermute 操作，此處會存在很多較小的 Kernel，如果不進行優化也會一定程度上影響性能。NVIDIA 在 GitHub - NVIDIA/TransformerEngine 中也提供了相應 Kernel 融合的優化，具體可以搜索 moe_permute 和 moe_unpermute 相關實現。

3.5 CP & SP

在 Transformer 模型中，自注意力機制的內存需求、計算量與序列長度成二次方關系，導致序列比較長時可能存在明顯瓶頸。為此，也有一系列工作嘗試解決這個問題，比如：

• [2105.13120] Sequence Parallelism: Long Sequence Training from System Perspective：核心是將輸入序列分為不同的 Chunk，每個 Chunk 都使用獨立的設備處理，而 Attention 部分采用 Ring Attention 的實現。
• [2205.05198] Reducing Activation Recomputation in Large Transformer Models：NVIDIA 在 Megatron 中也提供了相應實現，主要思路是在 TP 的基礎上，將原來未切分的部分激活進一步切分，從而降低內存開銷。
• [2309.14509] DeepSpeed Ulysses: System Optimizations for Enabling Training of Extreme Long Sequence Transformer Models：微軟進一步解決了序列并行中通信量比較大的問題，同樣是按序列切分，不過在 Attention 計算時通過 All2All 將同一個 Head 的序列匯聚在一起。這里也有個約束條件，Head 數需要是序列并行度的整數倍。
• [2310.03294] DISTFLASHATTN: Distributed Memory-efficient Attention for Long-context LLMs Training：進一步解決了長序列時 Ring-Attention 切分方式導致的負載不均衡問題。整體思路還是對調度進行重排，減少 Bubble。
• [2405.07719] USP: A Unified Sequence Parallelism Approach for Long Context Generative AI：結合了上述 Ulysses 和 Ring-Attention 的優勢。

在 LLM 的預訓練場景中通常序列長度不會超過 8K，因此沒有太大必要使用序列并行，在序列長度擴展或視頻場景中比較需要，這里不再展開。

四、常見優化方案

4.1 FP8 訓練

4.1.1 概述

NVIDIA GPU 從 Hopper/Ada Lovelace 開始支持 FP8 計算，其 FP8 算力通常是 FP16 算力的兩倍，而常見的訓練還是以 BF16 為主，利用 FP8 實現訓練加速也是一個值得探索的方向。此外，NVIDIA 從 Blackwell GPU 開始也進一步支持了 MX Format，可以原生支持細粒度的 FP8 甚至 FP4 量化，為 FP8 訓練，FP4 推理提供了更多可能。

微軟在 23 年就提出了使用 FP8 加速訓練的工作（[2310.18313] FP8-LM: Training FP8 Large Language Models）。NVIDIA 也在 Transformer Engine（Using FP8 with Transformer Engine）和 Megatron-LM 中提供了相應支持。

然而，LLM 預訓練的代價很高，對 FP8 訓練是否能真的獲得和 BF16 相當的性能也不得而知，因此，早期業內一直沒有真正的使用 FP8 做預訓練。

為了解決上述問題，零一萬物在 零一萬物面向萬卡集群的 AI Infra 建設 中提到了一個 Trick 的方法。如下圖所示，每隔一段時間就會 Load FP8 的 Checkpoint 并使用 BF16 進行訓練，驗證 Loss 是否和 FP8 訓練的 Loss 一致（PS：Loss 對齊就真的表示下游任務也能對齊嗎？）。如果出現不一致的情況，就會使用 BF16 的訓練代替 FP8，并在一段時間后繼續使用 FP8 訓練。最終獲得了 1.3x 的吞吐提升，不過并沒有說明這個提升是純粹的 FP8 相比 BF16 還是也包含了 BF16 的校驗預算。

4.1.2 DeepSeek FP8

DeepSeek V3 中另一個比較大的優化是 FP8 混合精度訓練，也是業內首個宣稱使用 FP8 進行端到端預訓練并進行了大量優化的工作。在 DeepSeek V3 的訓練中，大多數計算密集型操作以 FP8 執行，而少數關鍵操作則保留原始數據格式，以平衡訓練效率與數值穩定性。整體框架如下圖 Figure 6 所示，與線性算子相關的三個 GEMM 操作，包括 Forward（Fprop）、激活 Backward（Dgrad）和權重 Backward（Wgrad），接受 FP8 Tensor 作為輸入，并輸出 BF16 或 FP32 格式的結果，理論上使計算速度較原 BF16 方法提升一倍。此外，FP8 Wgrad GEMM 允許激活值以 FP8 存儲，供 Backward 使用，從而顯著降低內存消耗。

盡管 FP8 格式具有效率優勢，但某些算子對低精度計算比較敏感，仍需更高精度。同時，一些低成本算子也可采用更高精度，對整體訓練性能的影響較小。因此，對以下組件保持原始精度（如 BF16 或 FP32）：Embedding Module、輸出 Head、MoE 門控模塊、歸一化算子及 Attention 算子。為進一步保證數值穩定性，將主權重、權重梯度和優化器狀態以更高精度存儲。

DeepSeek V3 中還引入了一系列策略來提升 FP8 訓練的準確率，包括：

• 細粒度量化。
• 提升累加精度。
• 全部使用 E4M3 以獲得更高精度。
• 在線量化。

除此之外，DeepSeek V3 還通過將緩存的激活值和優化器狀態壓縮為低精度格式，進一步減少內存消耗和通信開銷。

4.2 CUDA Graph

CUDA Graph 首次出現在 CUDA 10 中，是 NVIDIA 在 CUDA 編程模型中引入的一種工作提交模型。允許將一系列 GPU 操作（如 Kernel、內存拷貝、Event 記錄等）按依賴關系組織成一個 Graph，并與其執行分離開來。換言之，Graph 描述了整個任務流程的靜態依賴關系，一旦定義完成，就可以多次重復執行，而無需每次都重新 Launch Kernel 和設置依賴。

在傳統的執行模型中，每次 Kernel Launch 都需要 CPU 配合執行一系列準備和調度操作，這對每個 Kernel 都是額外開銷。當 Kernel 執行時間很短時，這些 Launch 開銷可能成為整個任務的主要瓶頸（PS：這也是 Kernel Fusion 的一個好處）。CUDA Graph 通過提前定義工作流、預先實例化 Graph，將這些開銷挪至 Graph 的準備階段，大幅降低了每次執行時的 CPU 負擔。一旦定義完成，就可以多次重復執行，而無需每次都重新 Launch Kernel 和設置依賴。

如下圖所示，Launch Graph 的時間遠小于 A、B、C、D、E 這 5 個 Kernel 總的 Launch 時間。也就是在多個小 Kernel 按順序執行的場景中，用單次 Graph Launch 替代多次小 Kernel Launch，可以顯著減少 GPU 閑置等待時間，提高整體吞吐率。

實踐中，當 Kernel 執行時間較短（微秒級）時，Graph 可顯著減少調度開銷并提升性能。此外，CUDA Graph 將完整的計算流程呈現給驅動程序，使得驅動能夠針對整個流程進行優化（如更高效的線程塊調度），這是逐次提交無法輕易做到的。

PyTorch 和 Megatron-LM 都提供了對 CUDA Graph 的支持，在一些小規模細粒度的 MoE 模型訓練中會有比較明顯的幫助。

4.3 細粒度 Overlap

DeepSeek V3 中的通信、計算細粒度 Overlap 也是其軟硬協同設計的一個重要組成部分。然而 DeepSeek 方案是針對特定硬件環境、特定模型的高度定制化，在其他場景并不一定最優；此外，隨著模型、硬件的變化，分別定制化的成本比較高，因此也就衍生出一系列更通用的方案。

其實在 DeepSeek V3 之前已經有一些相關工作，但是沒有被廣泛使用，在 DeepSeek V3 之后，字節跳動也出現了一些相關方案，值得關注。

24 年北大提出過 Centauri 框架（[ASPLOS 24.04] Centauri: Enabling Efficient Scheduling for Communication-Computation Overlap in Large Model Training via Communication Partitioning），其構建了一個由三個固有抽象維度組成的切分空間：原語替換、拓撲感知組切分及工作負載切分。這些維度共同構成了一個全面的優化空間，用于高效 Overlap。為確定通信與計算的高效 Overlap，將混合并行訓練中的調度任務分解為 OP、Layer 和模型三個層次。如下圖 Figure 3 所示，通過通信切分和層次調度實現通信和計算的細粒度 Overlap。

字節也相應提出了 Flux（[2406.06858] FLUX: Fast Software-based Communication Overlap On GPUs Through Kernel Fusion），旨在通過依賴計算隱藏 GPU 間的通信時延。Flux 將通信和計算操作分解為更細粒度的操作，并進一步融合成更大的 Kernel，從而在不損害 Kernel 效率的前提下有效隱藏通信。

25 年，字節又在 Flux 的基礎上提出 TileLink（[2503.20313] TileLink: Generating Efficient Compute-Communication Overlapping Kernels using Tile-Centric Primitives），旨在高效編譯并生成計算-通信 Overlap 執行的 Kernel，并且聚焦于層內并行。TileLink 由前端（Frontend）和后端（Backend）構成：

• 在前端，系統通過以 Tile 為中心的原語將通信與計算的設計空間解耦并建立關聯。
• 在后端，將這些原語轉換為底層指令，整合通信與計算組件以實現 Overlap 執行。

在 TileLink 之后，字節很快又發布了 Triton-distributed（Programming Overlapping Kernels on Distributed AI Systems with the Triton Compiler），其作為對 Triton 編譯器的擴展方案，以更好的支持分布式 AI 工作負載的 Overlap 優化。其首先將符合 OpenSHMEM（NVSHMEM）標準的通信原語集成到編譯器中，使得能夠通過更高層次的 Python 編程模型使用這些原語。其次，還闡述了如何借助編譯器實現計算、內存訪問與通信的復雜聯合優化，并重點介紹了如何利用 Overlap 技術隱藏時延。

4.4 Attention Mask

在 Attention 模塊中的 Attention Mask 是實現序列中 Token 之間信息融合的關鍵模塊，也是實現各種業務場景的關鍵所在。

• 在 LLM 預訓練中，由于 Attention 計算占比不高，并且基本都是 Causal Mask，因此這塊相應的關注比較少，基本上用 FlashAttention 就能獲得非常高的性能。
• 在 LLM 后訓練或者多模態等復雜場景，當序列比較長時，Attention 的占比變大，Attention 相應的優化問題也需要關注。

如下圖所示就是一系列常見的 Attention Mask，比如：

• （1）Causal：Decoder 的標準 Mask。
• （2）Sliding Window：滑動窗口 Attention 常用的 Mask。
• （3）Causal Document Mask：SFT 等后訓練常用的 Sample Packing Mask。
• （9）Prefix LM Cache Mask：一些多模態場景會用的 Mask。
• （12）Random Eviction Mask：有點類似 Tree Attention Mask，在投機采樣，推薦等場景比較常見。

Attention 部分最經典的優化就是 FlashAttention 系列工作，包括 V1、V2、V3，隨后也有其他相關優化，比如：

• PyTorch 官方的 FlexAttention 優化：FlexAttention: The Flexibility of PyTorch with the Performance of FlashAttention。
• 百度的 FlashMask 優化：[2410.01359] FlashMask: Efficient and Rich Mask Extension of FlashAttention
• FlashInfer 的 Sparse Attention：flashinfer.sparse

五、業內常見模型

5.1 Meta LLaMA

Meta LLaMA 3.1 405B（[2407.21783] The Llama 3 Herd of Models）是開源的最大的 Dense 模型：

• 總參數量 405B。
• 采用 GQA，比例為 128/8。
• 共 126 層。
• 詞表 128,000。
• 15T Token。

如下圖 Figure 5 所示為 405B 模型對應的分布式排布，其 MFU 在 38%-43% 之間（未使用 FP8）：

• TP 始終是 8，保持在一個節點內，充分利用 NVLink + NVSwitch。
• 在 8K 序列長度時不用 CP，在 128K 時使用 16 CP。
• 始終使用 16 PP。
• 根據預算 GPU 數調整相應的 DP 數和 Batch Size。

主要提到的優化就是對 VPP 的一些改進，以便解決內存和計算的不均衡，比如首尾 Stage 都少一層，主動釋放一些不需要的 Tensor 等：

5.2 DeepSeek

如下圖所示，今年上半年我們曾總結過 DeepSeek 相關工作中的關鍵技術點，這里不再贅述，詳細內容可以參考：???綜述：DeepSeek Infra/V1/MoE/V2/V3/R1 & 開源關鍵技術??。其中 DeepSeek V3 預訓練 MFU 預估在 38.2% 左右。

5.3 阿里 Qwen3 MoE

阿里之前也開源了其最新的 Qwen3 系列模型（[2505.09388] Qwen3 Technical Report），其中最大的模型為 Qwen3 W235A22 模型，具體配置如下：

• 總參數量 235B，激活參數 22B。
• 采用 GQA，比例為 64/4。
• 128 個路由專家，每個 Token 激活 8 個專家。
• 共 94 層。
• 詞表 151,669。
• 預訓練序列長度 4K。
• 預訓練 36T Token，119 種語言和方言。

Qwen3 雖然提供了技術報告，但并沒有提供使用的 GPU 資源量、MFU 以及相關的訓練優化手段。

5.4 騰訊 Hunyuan-TurboS

我們前面提到過，騰訊的 Hunyuan-TurboS W560A56 模型（[2505.15431] Hunyuan-TurboS: Advancing Large Language Models through Mamba-Transformer Synergy and Adaptive Chain-of-Thought）采用了 Softmax Attention 與 Mamba2 混合的架構：

• 總參數量 560B，激活參數 56B。
• 采用 GQA，比例為 64/8。
• 預訓練 16T Token。
• 1 個共享專家，32 個路由專家，每個 Token 激活 2 個路由專家。

不過并沒有提供預訓練 Infra 優化相關部分，沒有介紹使用的資源量、MFU 及相應的優化方案。

5.5 小紅書 dots.llm1

5.5.1 模型配置

小紅書于 2025.06.06 發布并開源自研大模型 dots.llm1 W142A14，并提供相應技術報告（[2506.05767] dots.llm1 Technical Report）。對應的模型配置為：

• 總參數量 142B，激活參數 14B。
• 采用 MHA，32 Head。（PS：在大模型里依然采用 MHA 是比較少見的）
• 第 1 層是 Dense 層。
• 2 個共享專家，128 個路由專家，每個 Token 激活 6 個共享專家。
• 共 62 層。
• 預訓練 11.2T Token。
• 預訓練序列長度 8K。

5.5.2 預訓練優化

論文中針對預訓練主要提了兩個優化手段，在前面已經介紹過：

• 和 NVIDIA 合作的類似 DualPipe 的 Interleaved-1F1B 優化。
• Grouped GEMM 優化。

根據公式可以大概推出 DeepSeek V2 預訓練的 MFU：

MFU = (Token 數 * Ctoken) / (訓練 GPU 小時數 * GPU FLOPs * 3600)

如下圖 Table 4 所示 Qwen2.5 72B MFU 大約為 35.69%，dots.llm1 W142A14 MFU 大約為 18.15%：

Qwen2.5 72B：(1T*72B*6) / (340K*989T*3600) = 35.69%

dots.llm1 W142A14：(1T*14B*6) / (130K*989T*3600) = 18.15%

5.6 百度 ERNIE 4.5

5.6.1 模型配置

百度于 2025.06.28 也開源了自研的 ERNIE 4.5 系列模型，并提供了詳細的技術報告（ERNIE 4.5 Technical Report）。其中最主要的 LLM 為 ERNIE-4.5-300B-A47B：

• 總參數量 300B，激活參數 47B。
• 采用 GQA，比例為 64/8。
• 前 3 層是 Dense 層。
• 64 個路由專家，每個 Token 激活 8 個專家。
• 共 54 層。
• 預訓練序列長度 4K。

5.6.2 預訓練優化

W300A47 LLM 預訓練的分布式策略為 21DP（ZeRO-1）、8EP（Attention 8TP）、12PP，共 2016 H800 GPU，Global Batch Size 15120，GPU 之間使用 RoCE NIC，MFU 達到 47%。此外，ERNIE 也采用了一系列的優化措施。

機內 EP：從上可以看出，其訓練使用 8EP，可以將 EP 的 All2All 全部放在機內，進而降低跨機 All2All 的壓力。

All2All 內存優化：如下圖 Figure 9a，傳統 MoE 實現在第二次 All2All 后應用 Gating 概率乘法算子。該方法需保留第二次 All2All 的輸出 Tensor 以供 Backward，造成顯著內存壓力。如下圖 Figure 9b，作者提出將 Gating 概率乘法算子重新放于專家計算模塊內部。這一架構改進使得第二次 All2All 輸出 Tensor 在使用后可立即釋放。盡管通過概率置換和額外輕量級 All2All 操作引入了微小開銷，但該優化顯著降低了峰值內存使用量，并消除了 Backward 過程中的大量重計算。

VPP 優化：至于 PP 方案，當梯度累積比較少時，比如小于 PP-Degree 時，采用 1F1B；當梯度累積比較大時，采用 VPP（Interleaved 1F1B）?？紤]到最后一個 PP Stage 還包含損失函數相關計算，也會占用較高內存，因此對其進行優化，一旦最后一個 PP Stage 的 Forward 完成，立即啟動 Backward 計算并釋放損失函數的激活內存，這樣最后一個 PP Stage 最多只保留單個 VPP 階段的激活內存。如下圖 Figure 11 中的紅框所示。

FP8 混合精度訓練：采用和 DeepSeek V3 類似的量化策略，使用 E4M3 FP8 數據格式及在線量化方法，對權重實施 Block-wise 量化，對激活實施 Tile-wise 量化。如下圖 Figure 12 所示為 ERNIE 4.5 的 FP8 混合精度訓練策略。

• FP8 訓練中的細粒度內存優化：FP8 訓練可以降低內存開銷，因此可以降低重計算以提升吞吐。MoE 中主要激活內存消耗來自 Up-gate Linear、Down-gate Linear、Down Linear、SwiGLU 以及 Gate 概率乘法的輸入激活。

• 對于 Up-gate Linear，保留其 FP8 輸入激活值 XFP8 而非 BF16 張量 XBF16 用于 Backward。在 Backward 中，需要利用轉置后 XBF16 的 FP8 量化版本來計算權重梯度。因此，在權重梯度計算階段，需要對 XFP8 執行反量化-轉置-再量化操作。該策略可降低 Up-gate Linear 的內存占用，并使得第一個 All2All 操作使用 FP8 精度執行以節約通信開銷。這是內存占用、通信成本與計算精度之間的權衡方案，實驗表明該方法能保持與基線實現相同的收斂速率。
• 對于 Down Linear，有兩種優化方案：（1）保留 Up-gate Linear 的 BF16 輸出張量；（2）利用上述 XFP8 張量重新計算 Up-gate Linear 以生成其 BF16 輸出張量。這兩種方法均需要對 SwiGLU 激活函數和 Gate 概率乘法算子執行輕量級重計算，從而節約這兩個算子輸入張量的內存占用。

• FP8 算子融合優化：通過算子融合降低數據移動開銷并提升計算強度，具體包括：(1) Forward 中 Permutation 操作與 FP8 量化的融合；(2) Forward 與 Backward 中 SwiGLU、Gate 概率乘法及 FP8 量化的三重融合。
• FP8 通信優化與 Overlap：Forward 階段采用 FP8 精度執行第一個 All2All 以降低 BF16 通信成本；Backward 階段將第二次 All2All 與 Up-gate Linear 權重梯度計算進行 Overlap。

重計算優化：為了制定最優的重計算策略，對模型中的每個算子進行精細化分析，系統評估了內存占用與計算時間的權衡。選擇性的對性價比最高的算子（能以最小運行時代價換取顯著內存節省的算子）實施算子級重計算，最終設計出能最大化訓練效率的最優重計算方案。

Paddle 框架原生容錯系統：為了實現快速識別硬件故障等異常，并做到快速恢復，設計了框架原生的容錯系統。如下圖 Figure 13 所示，包括以下幾個關鍵組件：

• TraceHang：通過并行度信息和通信記錄的協同分析，實現對 Hang 源頭的自動化診斷。TraceHang 可以精準定位 Hang 的根本原因，從而加速問題解決并最大限度減少停機時間。
• Online SDC Scanner：靜默數據損壞（Silent Data Corruption，SDC，這種情況 ECC、CRC 可能捕獲不到）因其隱蔽性特征對模型收斂構成重大威脅。利用 PP 中的 Bubble，采用固定輸入參數執行計算與通信操作，并將結果與基準真值進行實時比對驗證，識別出多個存在 SDC 的節點單元。
• Parallelized Warmup：PP Warmup 階段存在數據依賴性，因而初始化導致的性能退化被放大 P 倍（PP Stage 數量）。因此，采用了跨 PP Stage 的并行同步 Warmup 方案，將首個訓練 Step 的延遲降低到 1/P。
• Zero Cost Checkpoint (ZCC)：支持在每個訓練步驟保存檢查點，且不會對訓練吞吐量產生任何開銷，從而確保訓練中斷時不會丟失任何進度。首先是訓練中異步保存，其次是將異常節點 Checkpoint 通過 RDMA P2P 操作快速傳輸到健康節點，并與初始化 Overlap；如果異常節點內存不可訪問，則從持久化 Checkpoint 恢復。

5.7 Kimi K2

5.7.1 模型配置

Kimi K2 是 Moonshot AI 在 2025.07.11 新發布的模型（Kimi K2: Open Agentic Intelligence），其采用了和 DeepSeek V3 類似的模型，也獲得了很不錯的效果。

如下圖（??https://x.com/rasbt/status/1944056316424577525/photo/1??）所示為 Kimi K2 與 DeepSeek V3 模型的主要區別：

除此之外還有些需要關注的地方：

• Kimi K2 采用了自研的 MuonClip Optimizer。
• DeepSeek V3 增加了 Expert 分組以便降低通信開銷，同時強制負載均衡；Kimi K2 中去除了這一個限制，以便增加更多的專家組合空間。
• Kimi K2 參數量更大，但是激活更少，主要是因為 MoE 部分共享專家增多，但是激活專家并沒有增加；同時 Attention Head 減半，因此總的激活參數反而降低。

2025.07.22 Moonshot 發布了 Kimi K2 的技術報告（??https://github.com/MoonshotAI/Kimi-K2/blob/main/tech_report.pdf??），模型配置與上述一致：

5.7.2 預訓練優化

Kimi K2 在 H800 GPU 集群訓練，每個節點包含 8 個 H800 GPU，2T 內存，采用 NVLink + NVSwitch 高速互聯，并且包含 8 個 400 Gbps 的 RoCE 網卡。不過沒有提供 MFU 相關信息。

采用動態資源訓練，并保證資源數量是 32 個節點（256 H800）的整數倍，其 ZeRO-1 DP + 16PP + 16EP 的策略，訓練時直接擴展 DP 組即可。訓練時以 BF16 格式存儲模型參數，并采用 FP32 梯度累加，約需要 6TB 顯存，分散在 256 個 GPU 的模型并行組中。而對于優化器狀態，節點較多時分散在所有節點，節點較少時，則 Offload 到 CPU 內存。

通信和計算 Overlap：通過增加 Warmup Micro Batch 數量，實現 EP 中 All2All 與 Interleaved 1F1B 的計算 Overlap。不過 Interleaved 1F1B 把 PP 切的更碎，會引入更多的通信開銷，為了降低這一成本，同樣解耦了權重梯度重計算，使其能與 PP 并行通信 Overlap。

最小化 EP 并行：由于 Attention Head 減半，Attention 計算時間變少，為了更好的實現計算和通信的 Overlap，需要最小化 EP 耗時，因此采用了 16EP 的最小化 EP 并行策略，這也放寬了負載均衡約束。

細粒度重計算：對 LayerNorm、SwiGLU 和 MLA Up-proj、MoE Down-proj，已最小化計算開銷并最大化顯存節約。

FP8 存儲激活：對于 MoE Up-proj 和 SwiGLU 等不敏感激活采用 FP8-E4M3 存儲（ 1 x 128 Tile 的 FP32 Scale）。潛在性能下降風險，未采用 FP8 計算。

激活 Offload：如上圖 Figure 7 所示，采用了激活 Offload 的機制。

六、參考鏈接

1. MQA：[2003.04641] MQA: Answering the Question via Robotic Manipulation
2. GQA：[2305.13245] GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints
3. FlashMLA：FlashMLA: Efficient MLA decoding kernels
4. MHA to GQA：[2412.20677] Align Attention Heads Before Merging Them: An Effective Way for Converting MHA to GQA
5. GQA to MLA：[2502.14837] Towards Economical Inference: Enabling DeepSeek's Multi-Head Latent Attention in Any Transformer-based LLMs
6. TransMLA：[2502.07864] TransMLA: Multi-Head Latent Attention Is All You Need
7. RWKV-7：[2503.14456] RWKV-7 "Goose" with Expressive Dynamic State Evolution
8. MiniMax-01：[2501.08313] MiniMax-01: Scaling Foundation Models with Lightning Attention
9. Hunyuan-TurboS：[2505.15431] Hunyuan-TurboS: Advancing Large Language Models through Mamba-Transformer Synergy and Adaptive Chain-of-Thought
10. Mamba-2-Hybrid 8B：[2406.07887] An Empirical Study of Mamba-based Language Models
11. Mistral 7B：[2310.06825] Mistral 7B
12. MOBA：[2502.13189] MoBA: Mixture of Block Attention for Long-Context LLMs
13. NSA：Hardware-Aligned and Natively Trainable Sparse Attention
14. Mixtral 8x7B：[2401.04088] Mixtral of Experts
15. Skywork-MoE：[2406.06563] Skywork-MoE: A Deep Dive into Training Techniques for Mixture-of-Experts Language Models
16. Switch Transformer：[2101.03961] Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity
17. GlaM：[2112.06905] GLaM: Efficient Scaling of Language Models with Mixture-of-Experts
18. ST-MoE：[2202.08906] ST-MoE: Designing Stable and Transferable Sparse Expert Models
19. DeepSeek MoE：[2401.06066] DeepSeekMoE: Towards Ultimate Expert Specialization in Mixture-of-Experts Language Models
20. DeepSeek V2：[2405.04434] DeepSeek-V2: A Strong, Economical, and Efficient Mixture-of-Experts Language Model
21. DeepSeek V3：[2412.19437] DeepSeek-V3 Technical Report
22. ZeRO：[1910.02054] ZeRO: Memory Optimizations Toward Training Trillion Parameter Models
23. FSDP1：[2304.11277] PyTorch FSDP: Experiences on Scaling Fully Sharded Data Parallel
24. FSDP2：[2411.00284] SimpleFSDP: Simpler Fully Sharded Data Parallel with torch.compile
25. Megatron-LM：[1909.08053] Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism
26. PipeDream：[1806.03377] PipeDream: Fast and Efficient Pipeline Parallel DNN Training
27. Interleaved 1F1B：[2104.04473] Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM
28. ZeRO Bubble：[2401.10241] Zero Bubble Pipeline Parallelism
29. GLM 130B：[2210.02414] GLM-130B: An Open Bilingual Pre-trained Model
30. LLaMA 3：[2407.21783] The Llama 3 Herd of Models
31. Pangu Ultra MoE：[2505.04519] Pangu Ultra MoE: How to Train Your Big MoE on Ascend NPUs
32. DeepEP：GitHub - deepseek-ai/DeepEP: DeepEP: an efficient expert-parallel communication library
33. DeepGEMM：DeepGEMM: clean and efficient FP8 GEMM kernels with fine-grained scaling
34. Pangu Pro MoE：[2505.21411] Pangu Pro MoE: Mixture of Grouped Experts for Efficient Sparsity
35. 小紅書 dots.llm1：[2506.05767] dots.llm1 Technical Report
36. DeepSpeed Ulysses：[2309.14509] DeepSpeed Ulysses: System Optimizations for Enabling Training of Extreme Long Sequence Transformer Models
37. DistAttention：[2310.03294] DISTFLASHATTN: Distributed Memory-efficient Attention for Long-context LLMs Training
38. USP：[2405.07719] USP: A Unified Sequence Parallelism Approach for Long Context Generative AI
39. FP8-LM：[2310.18313] FP8-LM: Training FP8 Large Language Models
40. Centauri：Centauri: Enabling Efficient Scheduling for Communication-Computation Overlap in Large Model Training via Communication Partitioning
41. FLUX：[2406.06858] FLUX: Fast Software-based Communication Overlap On GPUs Through Kernel Fusion
42. TileLink：[2503.20313] TileLink: Generating Efficient Compute-Communication Overlapping Kernels using Tile-Centric Primitives
43. Triton-Distributed：Programming Overlapping Kernels on Distributed AI Systems with the Triton Compiler
44. FlexAttention：??https://pytorch.org/blog/flexattention/??
45. FlashMask：[2410.01359] FlashMask: Efficient and Rich Mask Extension of FlashAttention
46. Qwen3：[2505.09388] Qwen3 Technical Report
47. ERNIE 4.5：ERNIE 4.5 Technical Report
48. Kimi K2：Kimi K2: Open Agentic Intelligence?

本文轉載自?????AI閑談???????，作者：AI閑談