藉助NVIDIA Transformer Engine、融合核心、BF16、FP8與GPU基準測試加速Transformer訓練

本教程展示瞭如何利用NVIDIA Transformer Engine（TE）將融合GPU核心、BF16計算與硬體感知的FP8執行結合起來，以加速Transformer工作負載。首先，需要安裝transformer_engine[pytorch]，並檢測當前GPU的計算能力：只有Ampere或更高架構（計算能力>=8.0）才支援TE融合核心；而支援sm_89及以上（如L4、H100、Ada、Blackwell）的GPU還能啟用FP8張量核心。若GPU過舊或匯入失敗，程式碼會自動回退到純PyTorch路徑，使用BF16/FP32。隨後，教程介紹了te.Linear、te.LayerNorm、te.LayerNormLinear、te.LayerNormMLP和te.TransformerLayer等核心融合模組，並配置了延遲縮放FP8配方，使用HYBRID格式（前向E4M3、反向E5M2），amax_history_len=16，amax_compute_algo=max。

在模型構建部分，教程定義了一個緊湊的GPT風格因果語言模型MiniGPT_TE，每個Transformer塊都是單個融合的te.TransformerLayer；同時實現了一個等價純PyTorch版本MiniGPT_PT，包含多頭注意力、層歸一化、殘差連線和前饋網路。程式碼會根據GPU支援情況自動選擇模型，並報告引數數量（約數百萬）和架構維度。為了訓練，教程生成了確定性的算術模式序列（phase + stride * steps）作為資料，使用AdamW最佳化器，並在支援FP8時用te.fp8_autocast包裹前向傳播。模型訓練60步後，損失降至遠低於隨機猜測基線，說明其學會了序列中的恆定步長規律。

在基準測試環節，作者以batch size 32、序列長度256，對比了高精度（BF16/FP32）與FP8模式下的前向+反向+最佳化器更新。測量結果顯示，FP8模式在訓練速度和峰值視訊記憶體上均優於高精度模式，且速度提升會隨模型規模增大而更加明顯（例如D_MODEL=2048、N_LAYERS=12）。對於不支援FP8的GPU，教程會跳過FP8基準測試。此外，程式碼還檢查了第一個Transformer塊內部的FP8狀態，包括scaling_fwd的縮放因子和amax_history，以便理解延遲縮放如何穩定FP8張量。

最後，教程實現了貪心自迴歸生成：將最新上下文反覆輸入模型，並驗證連續生成的token是否保持了訓練資料中的恆定算術步長。作者總結了實際擴充套件方向：增大模型維度、嘗試E4M3與HYBRID等不同FP8格式、延長amax歷史長度、在自定義架構中使用te.LayerNormMLP/te.LayerNormLinear，以及使用fp8_model_init()在推理時直接以FP8儲存權重。總體而言，將Transformer Engine整合到端到端訓練流程中，能在保持Colab不同GPU環境相容性的同時，透過融合核心和合理的低精度策略，明顯降低核心啟動開銷和記憶體流量，從而提升訓練效率。