非均等テンソル並列性による大規模 LLM 学習におけるグッドプットの向上

大規模な LLM のトレーニングは、特に数千台の GPU にまたがり長期にわたって実行される場合、独自のインフラストラクチャ上の課題をもたらします。これらのジョブが長く実行されるほど、予期せぬ中断やリソースの変動に遭遇する可能性が高まります。たとえ頻度の低いデバイスの利用不可であっても、密接に相互接続されたクラスターにおいて過大な影響を及ぼし、特定のトレーニングランでスローダウンを引き起こすことがあります。

大規模なトレーニングジョブにおいては、利用可能な GPU の数に応じてジョブを弾力的に適応させることが、Goodput（実効処理量）を向上させる強力な手法です。AI トレーニングの文脈において、Goodput は単なるハードウェアのスループットではなく、有用で収束を促進する作業が完了したことを示す重要な指標を表します。

現在の弾力的スケーリングのための効果的な方法には、データレプリカの削除、高速なチェックポイントからの再起動の利用、またはホットスペアへの切り替えが含まれます。これらの手法により、LLM ジョブは GPU の利用状況の変化に適応しつつ、システム全体のバランスを維持することができます。しかし、トレーニングジョブが利用可能なリソースが減少した状態で実行されている期間中、これらにはある程度のスループットの損失とコストの上昇も伴います。

最近発表された Nonuniform Tensor Parallelism（NTP）に関する論文は、スループットオーバーヘッドを最小限に抑えながらこれらの手法を発展させる、先見的な実験的フレームワークを紹介しています。潜在的な動的電力増強と組み合わせることでパフォーマンス低下を相殺し、スループットを一定に保つことで、中断を管理可能かつ回復可能な事象へと変換します。

NTP の中核的な貢献は、一時的なデバイス障害が大規模で相互接続されたトレーニングジョブを停止させるのを防ぎながら、高い Goodput を維持する能力にあります。テンソル並列度の動的調整と必要なデータ再分割の知的な重なり合わせにより、NTP は失われる時間と計算リソースを最小化します。

これにより、クラスターは動作時間の最大限を有益な作業に費やすことになり、ハードウェア状況が変動してもトレーニング実行全体の整合性と効率性が維持されます。

大規模トレーニングにおける課題

AI モデルのトレーニングは、数千台の GPU にまたがる並列作業です。これらのワークロードを並列化するための一般的な手法としてテンソル並列性（Tensor Parallelism: TP）があります。これは、ニューラルネットワークの層を密結合された GPU グループに分割する技術です。このグループ内の GPU 数は、スケールアップドメインと一致しており、NVIDIA NVLink などの高速インターコネクトで相互接続されています。NVIDIA Blackwell および NVIDIA Blackwell Ultra システムでは、NVLink は最大 72 台の GPU を 1,800 GB/s で接続し、単一のホップ内でオール・トゥー・オール通信をサポートします。

典型的なシナリオでは、最先端の LLM はラックのクラスター全体でトレーニングされます。各ラックには複数のサーバーが収容されており、単一のラックのサーバー群がスケールアップドメインを形成し、1 つの TP グループを構成します。データ並列性（DP）は、モデルを複数のこのようなスケールアップドメインに複製し、それぞれが異なるデータバッチを処理します。スケールアップドメイン内の GPU ステータスの変化は、その TP グループの効率に影響を与える可能性があります。同じ TP グループ内の GPU は密結合された計算を共有するため、1 つのデバイスに問題が生じるとトレーニングパフォーマンスが低下したり、進捗を維持するために一時的な再バランスが必要になったりします。

データセンターアーキテクチャがより大規模なスケーリングドメイン（8 台から 72 台、あるいはそれ以上の GPU）をサポートするように進化するにつれ、すべての健全なデバイスの生産的な稼働時間を最大化することが、高いグッドプット（Goodput：実効スループット）を達成するための鍵となります。

論文では、局所的で一時的な中断が全体のトレーニングスループットを決定させるのではなく、システムはアクティブな GPU の绝大多数を継続的に処理し続けるように設計できると指摘しています。このハードウェアの変動に適応することで、クラスターは大規模モデルのトレーニング中にグッドプットを最適化するために必要な高効率なリソース利用を維持できます。

トレーニング通常は、各ステージが先行するステップのタイムリーな完了に依存するパイプライン形式で行われます。ある TP グループ（Tensor Parallelism Group：テンソル並列処理グループ）内で 1 つの GPU に遅延が生じると、そのグループ全体での同期や処理が遅れ、システムが回復するか負荷を再分配するまで、一時的な停止やスループットの低下を引き起こす可能性があります。

NTP がトレーニングを維持する方法

NTP（Nonuniform Tensor Parallelism：非均等テンソル並列）の根本的な考え方は、一時的なハードウェア中断中であっても DP レプリカ（Data Parallel replica：データ並列レプリカ）が生産性を保ち続けることを保証することです。最新のチェックポイントから再開すると、モデルは利用可能なハードウェアに合わせて自動的に設定を適応させ、部分的な機能を維持しながらパイプラインが完全に停止することなく稼働し続け、そのレプリカの貢献を失わないようにします。

以下では、NTP がこのレジリエンス（耐障害性）を実現する方法について詳しく解説します。

動的 TP 度合い適応

スケールアップドメイン内の GPU が中断に遭遇した場合、システムは影響を受けたグループを特定し、テンソル並列性を自動的に再構成して、残りの機能的な GPU のみを活用します。例えば、8 個の GPU で構成される TP グループで 1 つがダウンした場合、TP デグリーを動的に 7 に切り替えることができます。このモデルシャードは計算を継続し、その貢献分の完全な喪失を防ぎます。グループ内の残りの GPU は個別の負荷を増加させ、リソースの一部に問題が発生している場合でも、トレーニングジョブが高いグッドプットと可用性を維持できるようにします。

パフォーマンス補償のためのパワーブースト

TP デグリーを削減するだけでは、グローバルスループットを維持することはできません。GPU 数が少ない DP レプリカは本質的に低速で動作するため、最も遅いレプリカの完了を待って、全体の DP システムが停止してしまいます。これを防ぐため、本研究では電気的・熱的な能力を向上させたラック設計を提案しています。この設計により、可用性が低下したスケールアップドメインに対してパワーブーストが可能になります。影響を受けたドメイン内のアクティブな GPU に供給される電力を動的に増加させることで、クロック周波数と計算スループットを一時的に引き上げることができます。

TP デグリーを低減した DP レプリカは、他の完全に機能するレプリカと効果的に追いつき、ペースを維持することで、グローバル同期のボトルネックを防ぎ、局所的なハードウェアの変動があってもクラスターの Goodput が高度に最適化された状態を保証します。

効率的な再シャード

TP デグリーの動的調整には、残りの GPU 間でモデルのテンソル シャードを再分配するための効率的なメカニズムが必要です。NTP は、他の計算ステージと重畳させる巧妙な 再シャード (resharding) 技術を採用しています。

この再シャードを逆伝播計算およびパラメータ同期フェーズ中に実行することで、健全なレプリカに導入されるオーバーヘッドは最小化され、通常は 1% を下回る程度になります。この慎重なスケジューリングにより計算効率が最大化され、適応機構自体がパフォーマンスのボトルネックとなることなく、最適な Goodput がシームレスに維持されます。

NTP は、スケーリングされた AI 学習へのより堅牢な道筋を構築する

本稿は、大規模 AI 学習に内在する課題に対処するために、ハードウェアとソフトウェアを協調設計することの決定的な重要性を強調しています。動的電力ブーストに必要な電気的・熱的な余裕を提供する先進的なラック設計との NTP（Nonuniform Tensor Parallelism）の統合は、思慮深いハードウェア革新が洗練されたソフトウェアソリューションをいかに深く補完し得るかを示す好例です。

このハードウェアとソフトウェアの相乗関係は、次世代 AI システムにおけるより高いパフォーマンス、効率性、安定したグッドプット（Goodput）、そして回復力を引き出すために不可欠です。先駆的な実験的機能として、NTP は回復力が並列化戦略に直接組み込まれた場合に何が実現可能かを示しています。

この基盤の上に、研究はすでに、標準的なテンソル並列化が最適ではない混合専門家（Mixture-of-Experts: MoE）モデルにおける回復力を最適化する非均質専門家並列化（Nonuniform Expert Parallelism: NEP）への概念の拡張に向けて進められています。

開発者ブランチへの最近の追加を含む詳細については、NVIDIA Resiliency Extension (NVRx) または Nonuniform Tensor Parallelism Readme で利用可能な本番環境対応の障害耐性および回復力機能をご覧ください。これらは、NVIDIA Megatron Core の開発者ブランチへの最近の追加に関する情報を得るためのものです。