Transformers 中的混合專家模型 (MoEs) 技術解析

Huggingface

3 天前

AI 生成摘要

本文探討了從稠密模型轉向使用混合專家模型（MoEs）的過程，並詳細介紹了 Hugging Face Transformers 函式庫如何透過重構權重加載與專家後端系統來優化效能。

Transformers 中的混合專家模型 (MoEs)

前言

在過去幾年中，擴展稠密語言模型（dense language models）推動了 LLM 的大部分進展。從早期的模型如原始的 ULMFiT（~30M 參數）或 GPT-2（1.5B 參數，當時被認為「太危險而無法發佈」🧌），到如今的千億級參數系統，其配方非常簡單：

更多的數據 + 更多的參數 = 更好的性能。

縮放法則（Scaling laws）強化了這一趨勢，但稠密縮放存在實際限制：

這正是混合專家模型 (MoEs) 登場的地方。

如果您已經熟悉 MoE 並想直接了解 Transformers 中的工程實現，可以直接跳轉到 Transformers 與 MoEs 章節。

從稠密到稀疏：什麼是 MoE？

混合專家模型保留了 Transformer 的骨幹，但將某些稠密的前饋層（feed-forward layers）替換為一組專家。一個「專家」並非針對特定主題的模組（例如「數學專家」、「代碼專家」），它僅僅是一個可學習的子網絡。對於每個 token，路由器（router）會選擇一小部分專家來處理它。

不同的 token 根據其隱藏表示（hidden representations）激活不同的專家。

模型容量取決於總參數，但推理速度取決於激活參數。

這是核心思想。

例如，以 gpt-oss-20b 為例。它總共有 21B 參數，但每個 token 僅從 32 個專家中選用 4 個激活專家。考慮到共享組件加上激活專家，該模型每個 token 使用約 3.6B 個激活參數。在內存頻寬約為 800 GB 的 M3 Ultra Mac 上運行，我們可以使用 bfloat16（每個參數佔 2 字節）估算生成速度為 ~ 800 / (3.6 * 2)。這得出大約每秒 111 個 token。我們實際獲得的性能數據約為 115 tok/s，這與粗略計算非常接近。

這種極快的速度證實了該模型運行起來大約像一個 3.6B 參數的模型，但它具有與 21B 參數模型相同的容量（或品質）。

（註：如果我們使用該模型原生 mxfp4 量化的算子，速度會更快）。

MoE 因以下原因而具有吸引力：

更好的計算效率

在固定的訓練 FLOP 預算下，MoE 通常優於稠密對手。

這意味著更快的迭代和更好的縮放效率。

天然的並行軸

專家在計算圖中提供了結構邊界。由於不同的 token 涉及不同的專家，我們可以跨專家進行並行化（我們稍後在專家並行中討論）。

行業採用

過去幾週發佈的近期主要開源 MoE 模型包括 Qwen 3.5、MiniMax M2、GLM-5 或 Kimi K2.5。

這一趨勢在 2025 年 1 月 DeepSeek R1 取得成功後加速發展，其建立在 DeepSeek V2 等早期系統之上。另一個早期的 MoE 是 2023 年 12 月發佈的 Mixtral-8x7B。

2-year timeline of MoE model addition in the transformers package

封閉實驗室也使用 MoE。長期以來一直傳聞 ChatGPT 使用稀疏架構，而開源的 gpt-oss 模型確實如此。

如果您想了解更多關於 MoE 的一般知識，我們強烈建議閱讀這篇部落格並觀看我們最近關於路由的 YouTube 影片。

Transformers 與 MoEs

生態系統中的大多數工具，包括模型加載、設備分配、量化和後端執行，最初都是為稠密模型設計的。MoE 挑戰了這些假設。

讓 MoE 成為 transformers 中的一等公民意味著重新設計部分加載流水線、執行模型和分佈式抽象，而不僅僅是添加新的模型類別。我們將重點介紹 transformers 庫如何演進以支持稀疏架構，涵蓋以下方面：

權重加載重構

AutoModelForCausalLM.from_pretrained("model_id") 下載並將模型權重加載到 PyTorch 模型中。對於稠密模型，加載相對簡單，檢查點（checkpoint）中的每個張量與運行時模組中的參數一一對應。

對於 MoE，情況更為複雜。在大多數 MoE 檢查點中，每個專家都是獨立序列化的。如果您查看 DeepSeek-V3 檢查點索引，您會看到如下鍵值：

每個專家都有自己的一組權重矩陣，以 DeepSeek-V3 為例，基本上是 256 個（總共 0 到 255）並排保存的小型前饋網絡。然而，在運行時，GPU 執行優化過的算子。現代 MoE 算子（如 grouped GEMM 和融合 MoE 實現）旨在通過單次操作處理所有專家，而不是逐個循環處理。

為了高效實現這一點，它們要求專家權重被打包成單個連續的張量。

因此我們面臨不匹配：

系統性地彌合這一差距正是權重加載重構所實現的。

隨著通用 WeightConverter 的引入，思維模型從：

檢查點已經匹配我的運行時佈局；加載主要是逐鍵複製。

轉變為：

檢查點只是張量的序列化來源。加載是一個轉換流水線，將它們轉換為我們想要的運行時佈局。

使用 WeightConverter 進行動態權重加載

這次重構引入的核心抽象是通過 WeightConverter 進行動態權重加載。

WeightConverter 讓我們定義：

原始操作（切分 chunk、拼接 concatenate 等）是可以組合的。其中兩個對 MoE 特別有用：

MergeModulelist 將張量列表合併為單個張量。例如，您可以將 MergeModulelist 與 Concatenate 組合，以堆疊 MoE 中的專家並將其打包成一個張量。

SplitModulelist 將張量拆回張量列表。例如，您可以將堆疊的專家拆回單個專家。

張量的延遲實例化

重構不僅改進了存在的轉換類型，還改進了它們的調度方式。

加載器掃描一次檢查點鍵，將它們與轉換器模式匹配，並按轉換器對張量進行分組。一旦某個鍵被確定為所需，它就會被註冊為 future，並通過線程池實例化。轉換操作僅在其依賴項就緒時運行。例如，MergeModulelist 會等待某一層的所有專家都加載完畢。

這避免了重複掃描並降低了內存峰值。

基準測試：權重加載流水線改進

為了評估新權重加載流水線帶來的改進，我們對 transformers 的 v4 與 v5 版本進行了基準測試。重點在於大型 MoE 模型的加載速度，這通常是訓練和推理中的瓶頸。

我們使用以下配置測試了 v4 與 v5：

範例：

兩個相關的環境變量：

HF_ENABLE_PARALLEL_LOADING: 通過線程啟用並行分片加載。

HF_DEACTIVATE_ASYNC_LOAD: 禁用新的異步流水線（v5 的逃生艙口）。

結果

模型：Qwen/Qwen1.5-110B-Chat
GPU：1× A100 (80GB)

提速不僅僅是因為「更多線程」。

它是單次路由（Single-pass routing）、異步實例化（Async materialization）和轉換感知調度（Conversion-aware scheduling）的結合，共同避免了不必要的實例化和內存峰值，同時在加載時實現了專家打包和投影融合。

量化的位置

通過這次重構，我們現在可以先創建運行時模組結構，然後將權重轉換為該結構。我們現在可以選擇在轉換流水線中附加量化，使量化成為權重加載流水線本身的一部分。這至關重要，因為「按專家」量化只有在專家以可預測的打包佈局存在時才有意義。

這種端到端的流水線以前是不可能的，現在作為公開 API 提供給用戶。

專家後端 (Expert Backend)

一旦專家被打包成單個運行時張量，另一個問題出現了：

你如何高效地通過它們進行路由？

在混合專家模型中，每個 token 被路由到不同的專家。這意味著運行時必須將 token 分發到其選定的專家權重，高效執行投影，應用路由權重，然後收集並重新排序結果。

這就是專家後端系統（在 PR #42697 中引入）所解決的問題。專家後端引入了一種可插拔的執行架構，將專家計算與模型實現解耦。系統允許專家層在運行時動態選擇後端，而不是在每個 MoE 模型內部硬編碼一種分發策略。

這是通過裝飾器模式實現的：

裝飾器封裝了專家類，並自動將計算分發到選定的後端。

目前提供三種後端：

eager: 循環遍歷選定的專家並對每個專家應用投影。這用於正確性參考和調試。

batched_mm: 使用 torch.bmm API。這會為每個 token 複製選定的專家權重，並執行單次批次 GEMM。此後端非常適合內存充足的小批次、重 GPU 負載。

grouped_mm: 使用 torch._grouped_mm API。在這裡，我們按專家 ID 對 token 進行排序、分組，然後執行單次分組 GEMM。此後端在大批次或內存受限的設置中表現出色。

專家並行 (Expert Parallelism)

混合專家 (MoE) 模型可以擁有數千億個參數（遠超單個 GPU 的承載能力）。專家並行 (EP) 通過將專家分佈在多個設備上來解決這個問題。每個設備僅加載其分配的專家子集，為這些專家進行計算，然後參與結果聚合。這種方法在不增加計算成本的情況下將模型擴展到更大的參數數量，因為每個 token 僅激活少數專家。

專家並行通過 enable_expert_parallel 啟用：

啟動方式：

其中 N 整除專家總數，並可能與您節點中的 GPU 數量匹配。

當 enable_expert_parallel=True 時，模型從標準的張量並行 (TP) 方案切換到具有專門分片策略的專家並行 (EP) 方案。

EP 的核心組件在於：

GroupedGemmParallel: 這沿著專家維度 (dim=0) 拆分專家權重。在這裡，每個設備僅加載 num_experts / num_devices。

RouterParallel: 這將全局專家索引映射到本地索引，掩蓋未分配給當前 rank 的專家，確保每個設備僅使用其本地專家進行計算，並使用 all-reduce 合併跨設備的部分輸出。