藉助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環境兼容性的同時，通過融合內核和合理的低精度策略，明顯降低內核啓動開銷和內存流量，從而提升訓練效率。