技術サマリー：ShardTensor – 科学機械学習のためのドメイン並列化

1. 問題提起

科学機械学習（SciML）は、GPU メモリ制約により、高解像度の科学データに対するトレーニングおよび推論が不可能であるという重大なボトルネックに直面しています。パラメータ数によってメモリ使用量が支配されることが多い大規模言語モデル（LLM）とは異なり、SciML のワークロードは頻繁に中間アクティベーションによって制約を受けます。

本論文は、科学データにおける「次元の呪い」を特定しています。LLM においてテキスト長を倍増させることが入力トークンを約 2 倍にするのに対し、$N$ 次元科学データの解像度を倍増させることは、データサイズを $2^N$ 倍に増加させます。3 次元以上のデータ（流体力学、医療画像、気候モデリングなど）の場合、トレーニング中の逆伝播時に中間アクティベーションを保存するために必要なメモリは指数関数的に増大します。

既存の解決策は、スケーリングが失敗するか、精度が低下する傾向があります：

• 低精度化： メモリ節約はわずかな（2 倍）ものであり、科学用サロゲートモデルにおいて動的範囲の問題に陥る可能性があります。
• 空間的ダウンサンプリング： メモリを削減しますが、科学的な画期的発見に必要な高忠実度の解像度を犠牲にします。
• モデル縮小： 精度を低下させます。
• アクティベーションチェックポイント/CPU オフローディング： 効果的ですが、再計算またはホスト - デバイス転送による著しい計算オーバーヘッドを導入します。
• 標準的な並列化： データ並列化（バッチの分割）はバッチサイズ 1 で限界に達します。モデル並列化（重みの分割）は入力データではなくパラメータ向けに設計されており、畳み込みなどによってグローバル形状が変化する動的入力データ処理には当てはまらない静的なテンソル形状を前提としています。

したがって、デバイスあたりのバッチサイズ 1 未満の入力データを並列化するための汎用フレームワークが存在せず、SciML が科学者が必要とするネイティブ解像度に到達することを阻んでいます。

2. 手法：ShardTensor

著者らは、入力データの空間次元をハードウェア制約から解離させるドメイン並列化の新たなパラダイムであるShardTensorを提案します。これは、モデル重みではなく入力データを処理するように特別に設計された、PyTorch の DTensor（分散テンソル）抽象化の拡張です。

中核となる設計原則

• ドメイン並列化： バッチ次元を超えて、入力データに対するデバイス間の並列化として定義されます。バッチサイズが 1 に達すると、ShardTensor はさらに空間次元（高さ、幅、深さなど）を複数の GPU 間で細分化します。
• 動的形状の処理： モデル重みに適した静的な分散形状を前提とする DTensor とは異なり、ShardTensor は動的形状をサポートします。これには、メッシュ軸に沿った各テンソルのローカルチャンク形状を記述するメタデータの第 4 要素であるシャーディング形状が導入されます。これにより、非構造化または非一様なデータの任意のチャンキングが可能になります。
• 命令的ディスパッチ： このフレームワークは静的ではなく命令的です。PyTorch の __torch_dispatch__ および __torch_function__ メカニズムを利用して、操作を動的にルーティングします。これにより、事前にコンパイルされた計算グラフや各レイヤーへの独自のパッチを必要とすることなく、レイヤーごとのドメイン並列化を処理できます。
• 拡張性： ユーザーは、高レベルの関数インターフェース（@custom_op）または低レベルのディスパッチロジック（aten 操作）を介してシステムを拡張できます。これにより、進化しつつあるモデルアーキテクチャやカスタムカーネルとの互換性が保証されます。

実装の詳細

ShardTensor は、PyTorch の DTensor の拡張として構築されています。

• ディスパッチアーキテクチャ： Python レベルで操作をインターセプトします。特定の操作にドメイン並列ハンドラが存在する場合はそれが実行され、そうでない場合は DTensor のセマンティクスまたは標準的な PyTorch の動作にフォールバックします。
• 通信： このフレームワークは、畳み込みや正規化などの操作において数値的一貫性を維持するために必要な、集合操作（all-reduce、all-gather など）および「ハロー」交換（隣接デバイスからの隣接ピクセルの取得）を管理します。
• 統合： 既存の並列化戦略（データ並列化およびモデル並列化）と組み合わせ可能に設計されており、2 次元以上の並列化戦略（例：ある軸でバッチを分割し、別の軸で空間次元を分割する）を可能にします。

3. 主要な貢献

1. ShardTensor フレームワーク： バッチサイズ 1 未満の入力データに対するドメイン並列化を可能にする、一般化されたオープンソースフレームワーク（NVIDIA PhysicsNeMo を通じて利用可能）。
2. 新たな抽象化： DTensor などの重み中心の並列化ツールとは区別される、動的入力テンソル分散を処理するための「シャーディング形状」の導入。
3. 非侵襲的プログラミングモデル： 並列化のためにレイヤーやモデルを書き換える必要なく、最小限のコード変更でモデル入力をラップできる設計。
4. 構成可能性： 直交するメッシュ軸上でドメイン並列化をデータ並列化およびモデル並列化と組み合わせる能力の実証。

4. 結果とベンチマーク

本論文は、NVIDIA Blackwell (GB200) および H100 GPU 上での合成ベンチマークと実世界の産業応用を通じて ShardTensor を検証しました。

• リングアテンション： 非常に小さなシーケンスサイズでは、並列化オーバーヘッドが支配的です。しかし、非常に大きなシーケンスサイズでは、トレーニングおよび推論の両方において、GPU 数に対してほぼ線形の強スケーリングを達成します。
• ビジョントランスフォーマー（ViT）：
  • 強スケーリング： 高解像度（例：$4096^2$）において、推論は単一 GPU と比較して 16 GPU で15 倍高速です。$2048^2$ におけるトレーニングは 8 GPU で5 倍高速です。
  • メモリ効率： メモリ使用量は 2 次元データでは二次関数的に、3 次元データでは三次関数的にスケーリングし、中間アクティベーションが支配的であることを確認しました。ShardTensor による強スケーリングは、デバイスあたりのメモリを比例して削減し、標準ハードウェア上で10 億を超える入力点を持つ 3 次元モデルのトレーニングを可能にします。
• Transolver（自動車空力）：
  • DrivaerML データセットで、ドメインあたり最大120 万点でトレーニングされました。
  • ドメインサイズが 2 倍、4 倍、8 倍に増加しても、単一 GPU ベースラインと比較して数値的安定性と競争力のある精度（圧力および速度の L2 誤差）を維持しました。
  • 前処理パイプライン全体も並列化され、エンドツーエンドのスケーラビリティが確保されました。
• StormScope（気象予報）：
  • 大陸規模のモデル（CONUS）で、3 km 解像度（$1024 \times 1792$ グリッド）を扱います。
  • ピークメモリ使用量（114 GB）は、単一 H100 GPU の容量（80 GB）を超えていました。
  • ShardTensor は、16 のデータ並列グループ内の 2 GPU にアクティベーションを分割することで安定したトレーニングを可能にしました。モデルは安定して収束し、6 km 解像度の単一 GPU ベースラインの損失軌跡を追跡し、ドメイン範囲を犠牲にすることなく以前はアクセス不可能だった個々の対流性嵐の解像を可能にしました。

5. 意義と主張

本論文は、ShardTensor が SciML の根本的なボトルネック、すなわち科学者が必要とするネイティブ解像度でのデータに対するトレーニングおよび推論の不可能性を解決すると主張しています。

• 解像度制限ワークロードの打開： メモリ上限により以前は不可能だった計算を実行可能なものに変換し、体積医療画像、計算流体力学、気候モデリングなどの分野での高忠実度トレーニングを可能にします。
• 構成可能性： データ量、モデルサイズ、解像度の間のトレードオフを排除します。科学者は現在、複数の並列化軸を同時にスケーリングできます。
• アクセシビリティ： 非侵襲的プログラミングモデルを提供することで、分散システム専門家ではない実践者の参入障壁を下げ、以前のアーキテクチャ固有の並列化努力とは対照的です。

認められた限界

著者らは謙虚に、命令的でレイヤーごとのディスパッチモデルが Python レベルのオーバーヘッドを導入し、小規模な操作や低解像度データでは効率が劣ることを指摘しています。性能は相互接続帯域幅にも依存します。さらに、現在すべての PyTorch 操作にドメイン並列ディスパッチパスが存在するわけではなく、静的グラフ最適化のための torch.compile との完全な統合は進行中です。

結論として、ShardTensor は、今日のデータ並列化と同様に、科学機械学習において高ドメイン並列化を一般的なものにするための一歩であり、科学基盤モデルの時代を促進します。