NVIDIA、生物学的基盤モデル向けMoE学習を効率化する技術を公開——Blackwell環境で最大2.21倍のスループットを達成

要点

• 米NVIDIAは2026年9月24日、生物学的基盤モデルに向けたMoE（Mixture-of-Experts）アーキテクチャの効率的な学習手法およびレシピを公開した。

• 「NVIDIA Transformer Engine」の「GroupedLinear」により、エキスパートごとの行列積をグループ化して一括実行し、カーネル起動オーバーヘッドを削減した。

• Blackwell GPUに最適化された「MXFP8」精度とカーネル融合技術により、メモリ消費の抑制と中間データ生成の回避を実現した。

• 8基のNVIDIA B200 Tensor コア GPUを用いた検証において、Hugging Faceの標準実装と比較して最大2.21倍の学習スループットを記録した。

• 米NVIDIAは2026年9月24日、生物学的基盤モデルの学習効率を高めるMoE（Mixture-of-Experts）向けの最適化手法と実装レシピを同社の技術ブログで公開した。同社の「NVIDIA Transformer Engine」や創薬・生物学向けフレームワーク「BioNeMo」を活用し、計算処理のグループ化やカーネル融合を組み合わせることでGPUの利用効率を大幅に引き上げたという。8基のNVIDIA B200 Tensor コア GPUを用いたベンチマークテストでは、一般的な実装と比較して最大2.21倍のスループットを達成したことが報告されている。

従来の密結合モデルが抱える課題とMoEの利点

生体分子やDNAなどの生物学的データを扱う基盤モデルの大規模化に伴い、すべての入力トークンが全層を通過する従来の密結合（Dense）なTransformerアーキテクチャでは、計算コストの増大が顕著になっている。これに対し、MoE（入力トークンごとに複数のサブネットワークの中から一部のみを選択して処理する手法）は、計算量を過度に増やさずにモデルの表現能力やパラメータ容量を効率的に拡張できる手法として関心を集めている。

しかし、MoEの性能向上効果は実装方法に大きく依存する。エキスパート計算が細分化されることでGPUの稼働効率が低下するほか、トークンを適切なエキスパートに割り振るルーティング処理に伴う通信負荷や、パラメータ増大によるメモリ圧迫が実用上の課題となっていた。

「GroupedLinear」によるカーネル起動オーバーヘッドの抑制

NVIDIAは計算の断片化を解消するため、「NVIDIA Transformer Engine（TE）」の「GroupedLinear」を採用した。

Hugging Faceなどの標準的なベースライン実装では、Pythonのループ文を用いてエキスパートを順次呼び出すため、エキスパートごとに個別のカーネル起動が発生し、スケジューリングの負荷が増大していた。これに対し、GroupedLinearは各エキスパートの重み行列を個別に保持しながら、複数の線形変換をグループ化されたGEMM（行列積演算）として1回の呼び出しでまとめてハードウェアに投入する。エキスパートごとに割り当てられるトークン数の偏り（split_sizes）を引数として受け取ることで一括処理を可能にし、カーネルの起動およびスケジューリングに伴うオーバーヘッドを大幅に抑制している。

MXFP8精度とカーネル融合によるメモリ・計算の最適化

メモリ使用量の削減とデータ転送の効率化に向けては、低精度フォーマット「MXFP8」とカーネル融合技術が導入された。

MXFP8（微小なブロック単位でスケーリングを行う8ビット浮動小数点形式）は、従来のBF16（16ビット脳浮動小数点数）と比較してメモリフットプリントを大幅に削減でき、NVIDIAの最新アーキテクチャであるBlackwell GPU上でハードウェアアクセラレーションが機能する。

さらに、Transformer EngineのSequential APIを用いて、GroupedLinear演算、活性化関数のScaledSwiGLU、そしてルーティング重みのスケーリング処理を単一のカーネル（ForwardGroupedMLP_CuTeGEMMSwiGLU_MXFP8）に統合した。これにより、演算途中で発生する中間データをメモリ上に実体化する必要がなくなり、メモリ帯域の負荷を抑えながら高速な処理を実現している。

B200 GPUでの性能検証と動作要件

8基のNVIDIA B200 Tensor コア GPUを用いた検証において、今回提供されたBioNeMoのレシピはHugging Faceのベースライン実装に対して最大2.21倍の学習スループットを記録したという。

本レシピを利用するための前提条件として、PythonやPyTorch、分散学習に関する環境のほか、CUDA対応の実行環境が必要となる。エキスパート並列（Expert Parallelism）を適用するには最低2基のGPUが求められ、融合MXFP8 GroupedMLPカーネルの動作にはNVIDIA Blackwell世代のGPUが必要とされている。

現在、GitHubのBioNeMoリポジトリにおいてMixtralベースの学習レシピが公開されているほか、Transformer Engineの公式ドキュメントで詳細な仕様が案内されている。

補足

生物学分野における基盤モデルはパラメータ数や配列長が拡大の一途をたどっており、限られた計算資源で効率的に学習を行う分散最適化技術の重要性が増している。