消除隱藏瓶頸：Unsloth 與 NVIDIA 如何加速 LLM 訓練

微調大型語言模型（LLM）仍然是現代 AI 中計算最密集的任務之一。雖然 NVIDIA GPU 為大規模平行運算而設計，但實際的訓練速度往往不是受限於核心運算的原始算術，而是受限於「黏合程式碼」——在繁重計算之間負責元資料管理與資料移動的部分。

在最近的合作中，Unsloth 與 NVIDIA 針對這些隱藏瓶頸進行優化，實現了約 25% 的整體訓練速度提升。此工作聚焦於三個主要領域：快取打包序列的元資料、實作雙緩衝檢查點重新載入，以及優化混合專家（MoE）路由。

1. 快取打包序列元資料

為了最大化 GPU 的使用率，開發者常使用「打包序列」，即將多個短樣本串接成單一長序列，以避免在填充標記上浪費計算資源。然而，這需要模型追蹤元資料——例如序列長度、累積偏移 (cu_seqlens) 與注意力遮罩——以確定每個原始序列的起始與結束位置。

傳統上，這些元資料會在 transformer 的每一層都重新建構。如果模型有 $L$ 層，系統就會執行 $L$ 次相同的帳務處理。這種重複的重建常常迫使裝置與主機同步，產生 GPU 與 CPU 的同步點，導致管線停滯。

Unsloth 為此可重用的元資料實作了快取。系統不再於每層重新建構打包序列資訊與 SDPA（Scaled Dot Product Attention）遮罩，而是為當前批次在每個裝置上快取這些結構。

在一次 Qwen3-14B QLoRA SFT 執行中，結果相當顯著：

前向傳播受益最大，因為它是最頻繁消耗重複元資料的階段。NVIDIA Blackwell GPU 的微基準測試顯示，單次元資料呼叫雖然很小（約 0.2 ms），但遮罩建構路徑每層可能耗時約 13.7 ms。累積數十層後，每一步可節省數百毫秒的時間。

激活檢查點是訓練大型模型的關鍵技術，因為它透過丟棄中間激活並在反向傳播時重新計算，節省 VRAM。當激活被卸載至釘住的 CPU 記憶體時，必須在反向計算前將其複製回 GPU。

在標準的單緩衝實作中，這個流程是串行的：

從 CPU 複製激活至 GPU $\rightarrow$ 2. 等待複製完成 $\rightarrow$ 3. 執行反向計算 $\rightarrow$ 4. 開始下一次複製。

Unsloth 引入了雙緩衝。當反向傳播在緩衝區 A 上計算時，複製串流會預先將下一個所需的激活載入緩衝區 B。計算完成後，兩者角色交換。這使系統能將複製延遲隱藏在有用的計算之後。

此最佳化對於反向計算量大的較大型密集模型特別有效。在 NVIDIA B200 Blackwell GPU 上的基準測試顯示，提升如下：

記憶體開銷保持在適度範圍，介於 0.23 GB 到 0.47 GB，成為性能提升的高效取捨。

混合專家（Mixture-of-Experts，MoE）模型需要路由機制將 token 分配給特定的專家。簡單的實作常在所有專家上使用 torch.where 迴圈。由於每個批次中每個專家的 token 數量會變化，這會產生資料依賴的輸出大小，進而觸發頻繁的 CPU‑GPU 同步。

Unsloth 改為一次性「分組」的方式，而非對每個專家逐一查詢執行時資訊：

此做法將開銷從與專家數量成正比（$\text{overhead} \propto \text{num_experts}$）轉變為幾乎恆定（$\text{overhead} \propto 1$）。

團隊驗證顯示，在 GPT-OSS 配置上可獲得 10‑15% 的加速，且特定路由路徑在前向傳播提升 +23%，在反向傳播提升 +13%。

這三項最佳化有共同的主題：它們針對「黏合程式碼」而非數學核心本身。隨著主要核心（如矩陣乘法與注意力）日益優化，先前不易察覺的剩餘開銷在總訓練時間中所佔比例逐漸增大。

此處的核心工程教訓是：當數學運算已被優化後，欲再取得加速需採取兩項策略：

透過聚焦於這些系統層面的瓶頸，Unsloth 與 NVIDIA 證明即使在高度優化的訓練堆疊中，仍有可能取得顯著的效能提升。