如何使用 NVIDIA Apex（FusedAdam、FusedLayerNorm）和原生 torch.amp 加速 Transformer 訓練

本教程詳細實現了 NVIDIA Apex，重點關注現代 GPU 訓練工作流程中仍然重要的元件。我們首先檢查 CUDA 執行環境，從原始碼構建 Apex（包含 CUDA 和 C++ 擴充套件），並檢測環境中實際可用的融合核心。這對於確保高效能核心可用至關重要，因為純 Python 安裝可能看似成功但實際上缺少關鍵核心。

構建完成後，我們基準測試 FusedAdam 與 PyTorch AdamW，比較 FusedLayerNorm 和 FusedRMSNorm 與標準歸一化層，並執行已棄用的 apex.amp 和現代 torch.amp 示例。最後，在一個小型 Transformer 訓練實驗中整合所有內容，比較純 FP32 PyTorch 路徑與融合 Apex+AMP 路徑，評估對吞吐量的實際影響。

實驗設定包括：檢測 CUDA GPU 可用性、安裝必要依賴、克隆並構建 Apex 原始碼。透過檢查 amp_C 和 fused_layer_norm_cuda 模組判斷核心可用性。

基準測試結果：

• FusedAdam 在最佳化器步驟上比 AdamW 快約 2.67 倍。
• FusedLayerNorm 的前向+反向傳播比標準 LayerNorm 快約 1.87 倍。
• FusedRMSNorm 與 FusedLayerNorm 類似，適用於 LLaMA 風格模型。

混合精度部分演示了棄用的 apex.amp 和推薦的 torch.amp 用法。端到端 Transformer 訓練（詞彙量 2000，隱藏維度 256，4 層）顯示，使用 FusedAdam、FusedLayerNorm 和 torch.amp（FP16）相比純 FP32 基線，吞吐量提升約 1.48 倍，且損失收斂相似。

結論：FusedAdam、FusedLayerNorm 和 FusedRMSNorm 是 Apex 中仍然相關的部分；apex.amp 已棄用，應使用 torch.amp；融合核心與原生 torch.amp 配合良好；對於實際工作負載，建議使用更大模型和 bf16 自動轉換。