「大規模モデルの学習:GPU/TPUによる大規模AI学習を再考するための実践者ガイド」

featured-image

金融、銀行、電気通信といったハイリスクな世界では、洗練されたAIモデルを導入するための競争が絶え間なく続いています。競争優位性は、新しいモデルをいかに迅速かつ費用対効果高く学習・展開できるかにかかっていることがよくあります。BnK Solutionでの私たちの仕事は、銀行向けのLLMから保険向けのオーダーメイドのビジョン言語モデルに至るまで、あらゆる目的のために膨大なドメイン特化型モデルを学習させることを含んでいます。この経験から、私たちは重要な洞察を得ました。それは、従来のAIスケーリング手法が、新しいソフトウェアのパラダイムによってではなく、それを支えるハードウェアの力ずくな進化によって書き換えられつつあるということです。

巨大なモデルを学習させるための従来の手法は、ますます複雑になる落とし穴へと私たちを導いてきました。しかし、ハードウェアの相互接続能力における根本的な変化が、その複雑さから抜け出し、よりエレガントで効率的なシンプルさへと戻る道を開いています。これは単なる理論的な好奇心ではなく、AIを大規模に扱う方法を変えつつある実践的な優位性です。

従来のスケーリング手法:並列性の詳細

現在の変化を理解するためには、まず私たちをここまで導いた道筋を評価する必要があります。数十億パラメータのモデルの学習を分散させることは、3つの基本的な技術にわたる複雑なオーケストレーションです。

  • データ並列性(DP): これは分散学習の中核となるものです。モデル全体を各ワーカー(例:GPU)に複製し、各ワーカーがデータの異なるシャードを処理します。バックプロパゲーションを通じてローカルな勾配が計算された後、通常はAll-Reduceと呼ばれる集合通信操作が、これらの勾配をすべてのワーカー間で平均化します。更新された重みはその後分散し直され、すべてのモデルレプリカが同期された状態を保ちます。概念的にはシンプルですが、DPのスケーラビリティは、このDP通信ステップのコストに極めて敏感です。遅延の大きいネットワークを持つクラスターでは、All-Reduce操作が主要なボトルネックとなり、システム全体が最も遅いワーカーが通信を完了するのを待つことになります。

  • テンソル並列性(TP): モデルの単一のレイヤーでさえ、1つのデバイスのメモリには大きすぎる場合、TPが必要になります。これは、巨大な重み行列をいくつかのデバイスに分割するもので、通常はNVIDIAのNVLinkのような超高速相互接続によって密接に結合された単一のノード内で行われます。フォワードパスおよびバックワードパスの間、これらのデバイスは極めて低い遅延で中間結果(活性化と勾配)を交換する必要があります。TPは単一の論理的な計算ユニットを「スケールアップ」するための強力なツールですが、より広範で遅いノード間のネットワークをまたいで効果的にスケールすることはできません。

  • パイプライン並列性(PP): これは、兆パラメータ規模のモデルを可能にしたブレークスルーでした。モデルは、連続したステージに垂直にスライスされます(例:レイヤー1〜8がステージ1、レイヤー9〜16がステージ2など)。データのバッチは、より小さな「マイクロバッチ」に分割され、最初のステージに供給されます。ステージ1がマイクロバッチの処理を終えるとすぐに、その出力活性化をステージ2に渡し、次のマイクロバッチの処理を直ちに開始します。これにより、理論的にはすべてのステージを忙しくさせる計算上の「パイプライン」が作成されます。PPは、TPでスケーリングされたノードでも大きすぎるモデルにとって唯一の実行可能な解決策であったため、コミュニティはこれを機能させるために複雑なスケジューラに多大な投資をしました。

隠されたコスト:パイプライン化の欠点が露呈する場所

PPはAIの新たなスケールを解き放ちましたが、私たちの実践的な経験は、それが効率と複雑さの両方で大きな負担を伴い、その負担がしばしば過小評価されていることを示しています。

最も明白なコストは「パイプラインバブル」です。物理的な組み立てラインと同様に、パイプラインは開始時に満たされる時間と、終了時に空になる時間が必要で、この間、高価なアクセラレータチップの多くが完全にアイドル状態になります。これは、トレーニング効率の重要な指標であるモデルFLOPs利用率(MFU)を直接低下させます。大規模なパイプライン化されたモデルでのMFUが45%であれば、許容範囲と見なされることが多いですが、これは、計算予算の半分以上を事実上無駄にしていることを意味します。

問題は、より洗練されたモデルアーキテクチャで深まります。金融モデルに関する私たちの仕事では、非常に深いネットワークの学習を安定させるために補助的な損失が必要です。これらは、バックプロパゲーション中に付加的な勾配信号を提供し、恐ろしい勾配消失問題を回避するために中間レイヤーにアタッチされる追加の損失関数です。パイプライン化された設定では、これは悪夢です。補助的な損失を計算するためには、パイプラインがその時点まで完全にフラッシュされなければなりません。これは、「すべてフォワード、すべてバックワード」(all f all b)と呼ばれる同期パターンです。これにより、損失が計算され、バックワードパスが開始されるまで、進行中のすべてのマイクロバッチがフォワードパスを完了することを強制されます。これは、トレーニングの途中でバブルを再生成するだけでなく、バックワードパスがトリガーされるまで、すべての単一マイクロバッチからの活性化を保存しなければならないため、ピークメモリ使用量に巨大なスパイクを引き起こします。

さらに、負荷分散の課題もあります。レイヤーを単純に分割すると、あるステージが他のステージよりも著しく多くの作業を持つ可能性があります。このパイプライン内の「長い棒」が全体のペースを決定し、他のステージをアイドル状態にさせます。PPVP(可変パーティショニングによるパイプライン並列性)のような高度な技術は、計算グラフを精巧に分析することでこれを解決しようとしますが、これはまた別の複雑さの層を追加します。DP、TP、PPVPの組み合わせを使用するモデルのデバッグは非常に困難であり、研究開発を遅らせます。これは、電気通信や金融といった動きの速い業界における重要なボトルネックです。

パラダイムシフト:帯域(Bandwidth)が解になるとき

PPの複雑さを受容してきた根本原因は、「DPのAll-Reduceは遅すぎて拡張できない」という前提でした。しかし、現代のインフラ(GoogleのTPU pod、NVIDIA NVL72をベースにした企業クラスのシステムなど)は、インターコネクト帯域が桁違いに向上しており、この前提を覆しています。

例えるなら物流です。コンテナ輸送に1ヶ月かかるなら複雑な倉庫とバッチ処理(パイプライン)が必要になりますが、コンテナを瞬時に転送できるなら倉庫は不要で直送(データ並列)で十分です。高速インターコネクトはデータの“テレポート”技術に相当します。

BnKではこの変化が大きな意味を持ちます。銀行顧客データを用いたagentic AIアプリ向けのLLMトレーニングでは、ハードウェアトポロジーが戦略を決定します。高帯域クラスター上では、我々はほとんどPPを使わずに数百億〜数百兆パラメータ級のモデルを学習させることに成功しています。実務的には、TPを用いて高速ノード内にモデルレプリカを収め、その構成を大規模なDPドメインで多数ノードへスケールアウトする方針を取ります。All-Reduceは依然通信集約的ですが、設計の良いネットワーク上では、パイプラインのバブルやスタールによる確実な損失よりも速く効率的です。

結果は明確です。MFUが60%超(場合によってはさらに高い)を安定して達成でき、これはクライアントにとって直接的な競争優位性、すなわちモデル反復の高速化、新機能の市場投入スピード向上、そしてトレーニング総コストの大幅低減に繋がります。

結論:ハードウェアを意識したAIの新時代

AIのスケーリングの旅は、データ並列性のシンプルさから、パイプライン並列性という必要な悪へと私たちを導きました。今、ハードウェアの進歩が、より堅牢で、効率的で、エレガントなシンプルさへと私たちを回帰させています。エリートAIエンジニアリングの未来は、最も複雑なスケジューリングアルゴリズムを習得することではなく、モデル、ソフトウェアスタック、そして物理的なハードウェアの間の相互作用を深く理解することにあります。

金融や電気通信という厳しい分野での私たちの経験は、私たちの核となる信念を鍛え上げました。AIにおける真のイノベーションは、システムのこの深い、ハードウェアを意識した共同設計から生まれるということです。それは、時として最も強力なアルゴリズムが、新しい最適化関数ではなく、プロセッサ間を静かに流れる1秒あたりテラビットの帯域幅であることを認識することです。この力を活用することで、私たちは次世代のAIをよりスマートに、より速く、より効率的に構築し、生の計算能力を目に見えるビジネス価値に変えることができます。

 

How can we help you? Contacts Us