生物基盤モデルのMoE学習を効率化する
NVIDIAは生物情報科学分野における大規模モデルの効率化を目指し、BioNeMoレシピとTransformer Engineを活用したMoEアーキテクチャの最適化手法を開示した。従来の密集型トランスフォーマーではパラメータ増に伴い計算コストが直線的に増大する一方、MoEは各トークンに対して限られた専門サブネットワークのみを活性化させることでスケーラビリティを向上させる。しかし実装には計算の断片化やルーティング通信オーバーヘッド、大規模パラメータに伴うメモリ制約といった課題が存在した。 同レシピはTransformer Engineの専用プリミティブを駆使し、これらのボトルネックを解消する。Expert計算の分散実行によるオーバーヘッドを軽減するため、GroupedLinearにより複数のExpert処理を単一のGEMM呼び出しに統合し、Pythonループ起因のカーネル起動頻度を大幅に削減した。また、ゲノムワークロードで必須の長文シーケンス処理に伴うアクティベーションメモリの圧迫に対し、BF16に替わり8ビット表現のFP8およびブロック単位のスケールファクターを持つMXFP8量子化を採用。NVIDIA Blackwell世代GPUでは専用Tensor Core命令によりハードウェア支援が実現され、メモリ利用効率が向上する。さらに低精度訓練における量子化のオーバヘッドを解消するため、Sequential APIを用いて活性化関数とルーティングスケーリングを一貫したグループ化MLPカーネルに融合。中間結果の一時保存を回避し、フレームワークの管理負荷を最小化した。 8基のNVIDIA B200 Tensor Core GPUを用いたベンチマークでは、標準的なオープンソース実装と比較し最大2.21倍の訓練スループットを実現した。本レシピは並列設定を柔軟に調整可能であり、生体分子・ゲノミクス分野における次世代ファウンデーションモデルの開発基盤として広く適用が見込まれる。詳細な実装コードと実行手順はBioNeMo公式レシピおよびTransformer Engineドキュメントにて公開されている。
