NVIDIA 最佳化生物基礎模型訓練：TE 與 B200 執行 MoE 架構
為何重要
這項技術更新對於致力於生命科學領域 DeepTech 的 AI 團隊至關重要，它提供了一套經證實的實作路徑來最大化新一代 Blackwell 架構的效能。對於 NVIDIA 的生態圈而言，這展示了軟體堆疊（TE）對於硬體利用率（2.21x 提升）的直接價值，有利於推動 Bio-LLM 的開發與商業化落地。
在語言模型與生物醫療演算所需的計算度持續攀升下，專家混合（MoE）架構被視為比稠密 Transformer 更高效的擴容方式。NVIDIA 在 Technical Blog 中揭露，透過結合 Transformer Engine (TE) 與 BioNeMo 管道，能解決 MoE 訓練中碎片化 Kernel 啟動與記憶體瓶頸的問題。
- NVIDIA 的 GroupedLinear 將分散的專家運算打包成一次 grouped GEMM 指令，透過接受每個專家的 token 數量作為引數（split_sizes），顯著降低 Python loop 迴圈帶來的開銷。
- 採用 MXFP8 (block-scaled 8-bit) 混合精度與 TE Sequential API，將量化、SwiGLU 以及路由權重縮放層融合為單一 ForwardGroupedMLP_CuTeGEMMSwiGLU_MXFP8 kernel，並利用 NVIDIA Blackwell GPU 硬體加速，大幅節省中間資料暫存與記憶體需求。
- 在 8 部 NVIDIA B200 Tensor Core GPU 的基準測試中，BioNeMo 食譜的輸送量達到 Hugging Face 基準的 2.21 倍。