使用NVIDIA NeMo AutoModel加速Transformer微調

NVIDIA NeMo AutoModel是一個開源庫，作為NVIDIA NeMo框架的一部分，用於大規模構建自定義生成式AI模型。它基於HuggingFace Transformers v5構建，增加了專家並行（Expert Parallelism）、DeepEP融合全對全排程和TransformerEngine核心，並利用v5的動態權重載入將這些最佳化應用於廣泛且不斷增長的模型家族。其成果是：在微調MoE模型時，相比原生Transformers v5，訓練吞吐量提升3.4-3.7倍，GPU記憶體減少29-32%，且使用相同的from_pretrained() API——只需更改一行匯入程式碼，無需其他修改。

MoE模型的興起帶來了新的訓練效率挑戰：跨數百個專家路由令牌、將專家矩陣乘法融合為單個核心、跨GPU分片權重以及重疊通訊與計算，這些都需要通用庫開箱即用之外的基礎設施。Transformers v5引入了對MoE的一流支援，如專家後端、動態權重載入和分散式執行的張量並行計劃。此外，v5透過將PyTorch的DeviceMesh直接整合到from_pretrained()中，使分散式訓練成為一等公民。

NeMo AutoModel透過繼承AutoModelForCausalLM，並新增專家並行、DeepEP融合全對全排程和TransformerEngine核心，在v5之上構建。DeepEP是v5尚未具備的部分：它將通訊與專家計算重疊。由於NeMo AutoModel利用v5的可逆權重轉換載入每個模型，它可以專注於這些可複用的核心操作，而非逐個模型的檢查點適配，同時save_pretrained()仍輸出標準HF檢查點，供vLLM和SGLang等工具載入。

效能測試展示了NeMo AutoModel在兩個場景下的優勢：跨16節點全微調千億級550B模型，以及單節點訓練兩個30B MoE模型。對於550B模型，Transformers v5因記憶體不足無法執行，而NeMo AutoModel透過專家並行（EP=64）將專家權重分片到各GPU，使全微調成為可能，達到每GPU每秒815個令牌，峰值記憶體58.2 GiB。在單節點8×H100 GPU上，Qwen3-30B-A3B模型從v5的3,075 TPS/GPU提升至11,340 TPS/GPU（3.69倍），峰值記憶體從68.2 GiB降至48.1 GiB（-29%）；Nemotron 3 Nano 30B A3B模型從4,583 TPS/GPU提升至15,421 TPS/GPU（3.36倍），峰值記憶體從62.1 GiB降至42.5 GiB（-32%）。

速度提升來自三個來源：專家並行降低記憶體壓力（EP=8將每GPU MoE記憶體佔用減少8倍）；DeepEP將通訊與計算融合，避免單獨的非阻塞通訊；TransformerEngine核心加速核心操作，如融合注意力、線性層和RMSNorm。NeMo AutoModel還提供了可配置的後端，透過BackendConfig選擇注意力、線性層、專家和排程器實現。總之，NeMo AutoModel為MoE模型微調提供了顯著的效能提升，同時保持與HuggingFace Transformers的API相容性。