一、背景

DeepSeek V3/R1 模型的發(fā)布，以及 AI Infra 相關(guān)代碼庫的開源，對大模型從業(yè)者產(chǎn)生了不容忽視的影響。從短期來看，這些工作確實推動了業(yè)界對 AI Infra 建設(shè)的關(guān)注，至少促使人們開始重視算法與 Infra 的協(xié)同設(shè)計。這一變化也看似提升了 Infra 團隊在整個大模型生成鏈路的話語權(quán)，但也為相關(guān)從業(yè)者帶來了更大的學習壓力與追趕挑戰(zhàn)，甚至看到一些公司或團隊因而重新審視了原有的發(fā)展規(guī)劃。

近期，我依然保持閱讀了一些相關(guān)文章，但由于比較忙，一直沒來得及整理成文。在這些問題中，有一個我之前就有所關(guān)注，最近再次想起，然而看到相關(guān)討論很少，因此這里對其進行簡要梳理。這一問題聚焦于 DeepSeek 的模型設(shè)計選型，包括以下兩個方面：

• 問題一：為什么 DeepSeek V3 模型在預(yù)訓練中使用了 16 路 PP（Pipeline Parallelism），而其層數(shù)是 61，并不能被 16 整除，這一設(shè)計有何考量？
• 問題二：DeepSeek V3 模型的前 3 層為 Dense Layer，之后 58 層是 MoE Layer，這樣的混合結(jié)構(gòu)有何特殊的意義？

值得注意的是，DeepSeek 之前的模型也出現(xiàn)了類似的設(shè)計范式，比如：

• DeepSeek V2 包含 60 層，其中第 1 層是 Dense Layer，后續(xù) 59 層為 MoE Layer。
• DeepSeek MoE 包含 28 層，第 1 層是 Dense Layer，后續(xù) 27 層為 MoE Layer。

針對上述問題，筆者也有一些初步的個人猜測，比如：

• 關(guān)于層數(shù)與 PP 的關(guān)系：一種可能的解釋是，通過非均勻分配層數(shù)，可以優(yōu)化負載開銷，比如降低首尾 Stage 的內(nèi)存需求。
• 關(guān)于前幾層采用 Dense Layer：一個合理的假設(shè)是，前幾層更多關(guān)注低級且通用的特征，這些特征不太需要復(fù)雜的 MoE 機制，或者 MoE 中很難保持負載均衡。然而，這一推測尚缺乏實證支持，并且很多開源的 MoE 模型也并沒有采取類似策略。

在這個過程中，也發(fā)現(xiàn)了一些很有意思的現(xiàn)象：無論是 DeepSeek R1 還是 Grok-3 的 DeepSearch，在面對一些專業(yè)問題或者細節(jié)問題時，都很容易出現(xiàn)明顯的幻覺問題（hallucination）。仔細審查會發(fā)現(xiàn)其存在很多明顯的錯誤，這也凸顯了當前大模型在內(nèi)容生成方面可能存在的局限性.如下圖所示為 DeepSeek R1 的部分結(jié)果，紅框中的結(jié)果均存在不準確的地方：

如下圖所示為騰訊元寶中使用 DeepSeek-R1 的結(jié)果：

二、混合 MoE

2.1 Google Switch Transformer

Google Switch Transformer 論文（[2101.03961] Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity [1]）中對 MoE Layer 的分布介紹的非常隱晦，單看架構(gòu)圖很容易誤解為是一個純粹的 MoE 模型（每一個 Transformer Layer 都包含 MoE Block），一些非官方的代碼實現(xiàn)中也是如此。然而實際上并非如此：

如下圖 Table 1 的說明中提到，“with experts at every other feed-forward layer”，說明 MoE Transformer Layer 和 Dense Transformer Layer 是交替的：

如下圖 Table 9 所示，其中提到 Expert Freq 為 1/2，表明 MoE Transformer Layer 和 Dense Transformer Layer 各占 1/2，與上述的交替相對應(yīng)：

而在 Huggingface 的 Switch Transformer 實現(xiàn)中（Huggingface configuration_switch_transformers.py [2]），包含 num_sparse_decoder_layers 參數(shù)，表示每隔多少層有一個 MoE Layer，也可以間接支持上述結(jié)構(gòu)，如下所示：

2.2 Google GLaM

Google 在其 GLaM 模型（[2112.06905] GLaM: Efficient Scaling of Language Models with Mixture-of-Experts [3]）中依然延續(xù)了 Dense Layer 和 MoE Layer 交錯的方式。如下圖 Figure 2 所示，這次的架構(gòu)圖很清晰：

此外，對應(yīng)的層數(shù)也是標準的 4 或 8 的整數(shù)倍：

2.3 Google ST-MoE

Switch Transformer 的作者在 [2202.08906] ST-MoE: Designing Stable and Transferable Sparse Expert Models [4] 中進一步提出了 ST-MoE 模型，其相比 Switch Transformer，其中 MoE Layer 的占比更低，只有 1/4，并且在這篇論文中也給出了 Expert Freq 的說明：

為了支持 Switch Transformer 和 ST-MoE 模型，作者開源了對應(yīng)的 GitHub - google/flaxformer [5]。如下圖所示，其通過 sparse_layout 和 num_sparse_layers 可以提供各種靈活配置（也可以支持 Switch Transformer），比如：

• 只有前幾層是 MoE。
• 只有中間幾層是 MoE。
• Dense Transformer Layer 中交替插入 MoE。
• 只有最后幾層是 MoE。

現(xiàn)在想想，不知道為何 Google 一直在堅持這種交錯的范式，也許是早期的 MoE 模型還不成熟，專家負載均衡的問題還沒有被很好的解決，或者是對訓練穩(wěn)定性與計算效率的一種妥協(xié)。

2.4 Snowflake Arctic

在使用 Grok DeepSearch 查找相關(guān)工作時也有意外收獲，比如其提到了 Snowflake AI Research 發(fā)布的 Snowflake Arctic 模型（GitHub - Snowflake-Labs/snowflake-arctic [6]）。

如下圖所示，其所謂的混合是在一個 Transformer Layer 中既有常規(guī)的 FFN，還有一個額外的 MoE 分支，只不過 MoE 分支和 Attention+FFN 是并聯(lián)的方式。實際上，如果 MoE 的輸入和 FFN 的輸入相同（如紅線所示），對應(yīng)的就是 DeepSeek 的 Shared Expert + Sparse Expert 的方式。該模型的 Dense 分支總共包含 10B 參數(shù)量，殘差連接的 MoE 分支包含 128x3.66B 參數(shù)量，總計參數(shù)量達到 480B，其中通過 top-2 Gating 機制選定的活躍參數(shù)為 17B。

2.5 MixCon

還看到其他的隔層 MoE 的架構(gòu)，比如 MixCon: A Hybrid Architecture for Efficient and Adaptive Sequence Modeling [7]。MixCon 是一種新的混合序列建模的架構(gòu)，主要是結(jié)合了基于 Attention 機制的標準 Transformer Layer、Conba Layer 以及 MoE Block。其在長序列建模中可以獲得很高的計算效率。

2.6 MiniMax-01

與 MixCon 類似，其實 MiniMax 發(fā)布的開源大模型 [2501.08313] MiniMax-01: Scaling Foundation Models with Lightning Attention [8] 中也采用了混合架構(gòu)，只不過是 Attention 的混合。為了彌補 Linear Attention 長距離序列建模能力不足的問題，作者在每隔 M 層會采用一個標準的 Softmax Attention Block。

2.7 DeepSeekMoE

對于之前的疑問 “DeepSeek 中為什么在淺層使用 Dense Layer 以及使用多少個 Dense Layer 更好？”，在 DeepSeek V2、V3 以及其中使用的 MoE 負載均衡損失（[2408.15664] Auxiliary-Loss-Free Load Balancing Strategy for Mixture-of-Experts [9]）等 Paper 中并沒有找到具體答案。最終在 [2401.06066] DeepSeekMoE: Towards Ultimate Expert Specialization in Mixture-of-Experts Language Models [10] 看到，作者提到第一層的負載均衡狀態(tài)收斂很慢，因此保留第一層為 Dense Layer：

不過論文中針對這個問題并沒有太多的消融實驗，在 DeepSeek V3 中也是直接將 Dense Layer 從 1 層擴展到 3 層。

2.8 Router 負載實驗

在最近的一篇論文 [2503.05029] Continual Pre-training of MoEs: How robust is your router? [14] 中，作者發(fā)現(xiàn)，在多個 MoE 模型中都一致地出現(xiàn)了類似上述 DeepSeekMoE 的結(jié)論。也就是：在多個 MoE 模型的早期層都一致的出現(xiàn)了 MRI（Maximum Routing Imbalance，最大路由不平衡）高的情況。MRI 定義為在某個訓練迭代和 MoE 層中，單個專家處理的 Token 比例的最大值。MRI 的值越高，表示路由越不平衡，即某個專家被分配了更多的 Token，可能導致負載不均衡。如下圖所示：

三、Pipeline Parallelism & Layer

3.1 Meta LLaMA 3

Meta 的 LLaMA 3 405B（[2407.21783] The Llama 3 Herd of Models [11]） 包含 126 個 Transformer Layer。如果是 128 個 Layer，很容易將它們均分到不同的 PP Stage。比如，16 個 PP Stage，則每個 Stage 包含 8 個 Layer。

然而實際只有 126 個 Layer，這是因為默認的 Pipeline Parallelism 實現(xiàn)會導致資源消耗不均衡。由于 Word Embedding 和 Warmup Micro Batch 的存在，第一個 Stage 會消耗更多內(nèi)存；由于 LM Head 的存在，最后一個 Stage 也會多消耗一些內(nèi)存。

如下圖所示為標準的 1F1B PP：

• Device 1 上執(zhí)行 Micro Batch 1 的 Backward（橙色）之前，Micro Batch 1/2/3/4 Forward（藍色）產(chǎn)生的激活會一直占用內(nèi)存。占用空間最大。
• Device 3 上執(zhí)行 Micro Batch 1 的 Backward（橙色）之前，Micro Batch 1/2 Forward（藍色）產(chǎn)生的激活會一直占用內(nèi)存。
• Device 4 上執(zhí)行 Micro Batch 1 的 Backward（橙色）之前只有 Micro Batch 1 Forward（藍色）產(chǎn)生的激活。占用空間最小。

為了使 PP 的資源消耗更加均衡，作者從第一個 PP Stage 和最后一個 PP Stage 都減少了一個 Transformer Layer。

3.2 GLM-130B

如下圖所示，其實智譜 2022 年在訓練 GLM-130B 模型（[2210.02414] GLM-130B: An Open Bilingual Pre-trained Model [12]）時已經(jīng)采用過類似上述 Meta 的方案。具體來說，其會在第一個 PP 和最后一個 PP 時讓 Transformer Layer 少一點，以便更加的均衡（PP 為 8，共 70 層，中間各 9 層，起始各 8 層）。

3.3 PipeOffload

最近看到 Sea AI-Lab 發(fā)表了一篇 PipeOffload: Improving Scalability of Pipeline Parallelism with Memory Optimization [13]，其在 Pipeline Parallelism 中通過調(diào)度和 Offload 等機制來降低內(nèi)存開銷。實驗證明，每設(shè)備的激活內(nèi)存隨 Stage 總數(shù)有效減少，使 PP 成為比 TP（Tensor Parallelism）更具優(yōu)勢的選擇，在內(nèi)存消耗更低的情況下，最高可帶來 19% 的加速。

如下圖 Figure 5 所示，通過最簡單的調(diào)整 Interleaved 1F1B 中 Warmup 的調(diào)度（GIS），從而實現(xiàn)內(nèi)存的節(jié)約。其中紅框部分的幾個 Forward Stage 被延后處理，進而可以延緩對應(yīng)激活的生命周期：

如下圖 Figure 6 所示，進一步對 GIS 進行改進（GIS-H），可以進一步降低激活內(nèi)存的需求，不過相應(yīng)的 Bubble 率也會提升：

既然一些激活要很長一段時間之后才會被用到，那么將其 Offload 到 Host 內(nèi)存中就是一種很自然的想法。如下圖 Figure 12 所示為一種比較極端的 Case。由于最后一個 PP Stage 在計算完 Forward 后可以馬上計算 Backward，因此不需要 Offload，其他的 Stage 均可使用。不過需要在 Offload 開銷與內(nèi)存節(jié)約之間取得平衡。其中的箭頭標注了 Offload 對前序 Forward 計算的依賴，以及 Backward 計算對 Reload 的依賴。

四、參考鏈接

• ??https://arxiv.org/abs/2101.03961??
• ??https://github.com/huggingface/transformers/blob/v4.49.0/src/transformers/models/switch_transformers/configuration_switch_transformers.py??
• ??https://arxiv.org/abs/2112.06905https://arxiv.org/abs/2202.08906??
• ??https://github.com/google/flaxformer??
• ??https://github.com/Snowflake-Labs/snowflake-arctic??
• ??https://zhouchenlin.github.io/Publications/2024-ECAI-MixCon.pdf??
• ??https://arxiv.org/abs/2501.08313https://arxiv.org/abs/2408.15664??
• ??https://arxiv.org/abs/2401.06066??
• ??https://arxiv.org/abs/2407.21783??
• ??https://arxiv.org/abs/2210.02414??
• ??https://arxiv.org/pdf/2503.01328??