🤖 machine learning

Model Parallelism With Subnetwork Data Parallelism

本論文では、モデルを構造化されたサブネットワークに分割し、アクティベーションの交換を行うことなくワーカー間で学習させることで、様々なアーキテクチャやスケールにおいて性能を維持または向上させつつ、大幅なメモリ削減（28%〜60%）を実現する分散学習フレームワークであるSubnetwork Data Parallelism（SDP）を提案する。

原著者： Vaibhav Singh, Zafir Khalid, Pietro Cagnasso, Edouard Oyallon, Eugene Belilovsky

公開日 2026-06-02

📖 1 分で読めます☕ さくっと読める

CC BY 4.0

原著者： Vaibhav Singh, Zafir Khalid, Pietro Cagnasso, Edouard Oyallon, Eugene Belilovsky

原論文は CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) でライセンスされています。 ✨ これは以下の論文のAI生成解説です。著者が執筆または承認したものではありません。技術的な正確性については原論文を参照してください。免責事項の全文を読む

あなたは、巨大で、極めて優秀だが、とてつもなく食いしん坊なロボット（大規模AIモデル）に、言葉を話したり画像を認識したりする方法を教えようとしていると想像してください。問題は、このロボットにはあまりにも巨大で高価なキッチンが必要であり、それを建設できるのはごく一部の人々だけだということです。AIの世界において、この「キッチン」とは、強力なグラフィックスカード（GPU）に搭載されたコンピュータメモリ（RAM）のことを指します。

この論文は、これら巨大なロボットを訓練するための新しい手法である**サブネットワーク・データ並列性（Subnetwork Data Parallelism: SDP）**を紹介しています。以下に、簡単な比喩を用いてその仕組みを説明します。

旧来の方法：「フル・レプリカ」の問題

従来、巨大なAIを訓練する場合、研究者は**データ並列性（DDP）**という手法を用いてきました。

比喩： 8人のシェフ（GPU）が、大規模な宴会の料理を作ろうとしている場面を想像してください。従来の方法では、すべてのシェフに、レシピ本一冊丸ごとと、すべての食材の完全なコピーが渡されます。彼らは全員、料理の全工程を作り、味見をし、そして次のラウンドに向けてレシピをどう改善すべきかについて、互いに意見を出し合います。
問題点： これは非常に無駄が多いです。各シェフは、自分自身のレシピと食材を保持するために、巨大なキッチンを必要とします。もしレシピが大きすぎると、キッチンが爆発（メモリ不足）してしまい、シェフたちは料理を作ることすらできなくなります。

新しい方法：「専門家チーム」（SDP）

著者らは、**サブネットワーク・データ並列性（SDP）**を提案しています。すべてのシェフにレシピ全体を与えるのではなく、レシピを特定のセクションに分割し、異なるセクションを異なるシェフに割り当てます。

比喩： 今度は、シェフAは前菜のみを担当し、シェフBはスープ、シェフCはメインディッシュを担当するとします。
- 料理全体を共有しない： 決定的なのは、彼らは料理全体をやり取りする必要はないということです。彼らは自分が担当している特定の材料についてのみ話し合います。
- キッチンが小さくなる： シェフAは前菜のレシピだけを持っていればよいため、キッチンを大幅に小さくできます。これにより、より巨大な宴会を、はるかに小さくて安価なキッチンで準備できるようになります。
- 結果： より大きなロボットを（あるいはより速く）訓練するために、超高価なキッチンを用意する必要がなくなります。

仕事を分割する2つの方法

論文では、これらの「サブ・レシピ」をシェフに割り当てる2つの異なる方法をテストしています。

フォワード・マスキング（「材料を切り分ける」アプローチ）：
- 仕組み： シェフは調理を開始する前に、使う予定のない材料を文字通り捨ててしまいます。彼らは割り当てられた料理の部分だけを調理します。
- メリット： 使わない材料を保持する必要さえないため、最もスペースを節約できます。これは、スープを作る際に、レシピ全体のリストを完全に無視して、必要なニンジンと玉ねぎだけを購入するようなものです。
- 難点： レシピの一部を無視しているため、全体像を確実に学習するためには、より多くのバッチをこなす必要があります。
バックワード・マスキング（「本全体を読み、自分のメモだけを書く」アプローチ）：
- 仕組み： シェフは文脈を理解するためにレシピ本（フルモデル）の全体を読みますが、料理をどう改善すべきかというメモを書くときは、自分が担当しているセクションについてのみ書き留めます。
- メリット： シェフは食事全体の文脈を理解した上で、書くメモの量だけを節約するため、数学的に見てより「誠実」です。これは、より安全で安定した学習方法となります。

研究の結果

研究者らは、これらを2つの全く異なるタイプのAIでテストしました。

言語モデル（LLM）： テキストを書く「LLaMA」のようなモデルです。5億および10億の「ニューロン（パラメータ）」を持つモデルの訓練を試みました。
- 結果： メモリ使用量を**28%から60%**削減しました。いくつかのケースでは、以前は巨大なキッチンを必要としていたモデルを、以前よりずっと小さなキッチンに収めることができ、かつAIの話し方の質も損なうことはありませんでした。
- ボーナス： SDPは、「アクティベーション・チェックポインティング」や「FSDP」といった他のメモリ節約テクニックとも完璧に併用できることが分かりました。これはレゴブロックを積み重ねるようなもので、SDPを他の手法と組み合わせることで、メモリ使用量を最大**85%**も削減できます。
画像分類器： 画像を認識するモデル（ResNetやSwin Transformerなど）です。
- 結果： 標準的な画像データセット（CIFAR）を用いてこれらのモデルを訓練しました。大幅に少ないメモリ（時には元の40%程度まで）を使用しても、このAIはフルサイズのバージョンと同等の画像認識能力を発揮しました。いくつかのケースでは、「専門家チーム」のアプローチが、一種の「正則化（過学習を防ぐ方法）」として機能し、AIの学習を実際に向上させました。

結論

この論文は、新しいタイプのAIを発明したり、病院や自動運転車でのAIの使い方を提案したりするものではありません。代わりに、ロジスティクス（物流・管理）の問題を解決しています。

これは、建設作業員の組織化に関する新しい方法だと考えてください。すべての作業員に超高層ビルの設計図一式を渡す（それだとポケットの中で場所を取りすぎる）代わりに、各作業員には彼らが担当する特定のフロアの設計図だけを渡すのです。彼らは依然として同じ超高層ビルを建設しますが、必要なスペースは少なくて済み、同じリソースでより高いビルを建てることができるようになります。

重要なポイント：

SDPは、巨大なAIモデルを小さなパーツに分割し、異なるコンピュータ間で分散させます。
すべてのコンピュータがモデルの「全体」を保持する必要をなくし、メモリを**28%〜60%**節約します。
既存のAIモデル（LLaMAやResNetなど）の使い道を変えることなく、そのまま利用できます。
他のメモリ節約テクニックと組み合わせることで、さらに多くのスペース（最大85%）を節約できます。
メモリ使用量が少ないにもかかわらず、AIのパフォーマンスは従来の方式と同等（あるいはそれ以上）です。

技術要約：サブネットワーク・データ並列性 (Subnetwork Data Parallelism: SDP)

問題提起

ディープニューラルネットワークの急速なスケーリングは、アクセラレータに対して前例のないメモリ需要を生み出しており、多くの場合、分散トレーニング戦略を必要とします。データ並列法 (DDP) は広く利用されていますが、モデルのフルセットを各GPUに複製するため、高いメモリオーバーヘッドと高価な勾配同期を招きます。対照的に、モデル並列法（パイプライン並列やテンソル並列など）は、より大きなアーキテクチャをデバイスに収めるためにモデルを分割しますが、活性化関数の通信コストが大きくなり、パイプラインバブルやロードインバランスが発生しやすくなります。

完全シャード・データ並列法 (FSDP) や ZeRO といった現在のソリューションは、パラメータ、勾配、およびオプティマイザの状態をシャードすることでデバイスあたりのメモリを削減しますが、特に勾配同期において依然として相当な通信オーバーヘッドが発生します。活性化関数の帯域幅とノードあたりのメモリ使用量を、収束の質を損なったり高帯域幅のインターコネクトを必要としたりすることなく削減するという課題が残っています。

手法：サブネットワーク・データ並列性 (SDP)

著者らは、モデルを構造化されたサブネットワークに分割し、ワーカー間で活性化関数を交換することなく学習を行う分散トレーニングフレームワークである Subnetwork Data Parallelism (SDP) を提案しています。計算を逐次的に分割するパイプライン型のアプローチとは異なり、SDPは各ワーカーに、入力から損失への完全なパスを保持する構造的に完全な部分（サブネットワーク）を割り当て、独立した勾配計算を可能にします。

コアメカニズム

サブネットワークの構築:
- ブロックレベル (B-SDP): アーキテクチャのブロック全体（例：TransformerブロックやResNetの基本ブロック）を削除します。
- ニューロン/チャネルレベル (N-SDP): 全結合層における特定のニューロンや、畳み込み層における特定のチャネルを削除しますが、隣接する層との一貫性を確保します。
- ワーカーは、ステップごとの平均化を通じて、重複するパラメータを同期します。
マスキング・レジーム:
本論文では、学習中にサブネットワークがどのように利用されるかを定義するために、2つの補完的なマスキング戦略を調査しています。
- フォワード・マスキング (Forward-Masking): モデルはマスクされたパラメータ ( $m_{fwd} = m$ ) を用いて評価されます。これにより、マスクされたコンポーネントのパラメータ、活性化関数、および勾配が除去され、大幅なメモリと計算の節約が得られます。ただし、勾配はマスクされたフォワードパスを反映します。
- バックワード・マスキング (Backward-Masking): フォワードパスではフルモデル ( $m_{fwd} = m_{uni}$ ) を使用しますが、バックプロパゲーションと集計の際にのみスパース性が適用されます ( $m_{bwd} = m$ )。これは、理論的に裏付けられた偏りのない勾配推定を維持し、原理的なベースラインを提供しますが、メモリ節約は主に勾配とオプティマイザの状態に対して行われ、活性化関数に対しては行われません。
理論的根拠:
著者らは、 $L$ -平滑な設定におけるバックワード・マスキング・レジームの収束解析を提供しています。彼らは、収束率が「マスキング誤差」（ユニフォームなマスキングと特定のマスク分布との距離）によって支配されることを示しています。この解析は、収束の質をマスクグラフのスペクトルギャップに関連付け、周期的な平均化を伴う重複するパラメータ割り当てが、部分的な同期を維持することを示しています。

主な貢献

新しいパラダイム: サブネットワークの学習をノード間で分散させるメモリ効率の高い分散トレーニング・パラダイムとして、SDPを導入。これにより、パフォーマンスを維持または向上させつつ、デバイスあたりのメモリを28%〜60%削減します。
理論的保証: バックワード・マスキングされたSDPの収束をマスクグラフのスペクトルギャップ条件に結びつけ、特定のマスキング条件下で偏りのない勾配が維持されることを証明しました。
構成可能性 (Composability): SDPが既存のメモリ節約技術と直交することを示しました。SDPはモデルのコンポーネントを物理的に除去することで、パラメータ、勾配、オプティマイザの状態、および活性化関数を同時に縮小します。これにより、SDPは FSDP (ZeRO-3) や 活性化関数チェックポインティング (AC) とクリーンに組み合わせることが可能です。
汎用性: 大規模言語モデル (LLaMA) や画像分類モデル (ResNet, Swin Transformer) を含む多様なアーキテクチャにおける検証を行い、SDPが標準的なデータ並列法の「ドロップイン」代替案であることを示しました。

実験結果

大規模言語モデル (LLaMA)

実験は、FineWebを用いて134M、500M、および1BパラメータのLLaMAモデルに対して行われました。

メモリ効率: 1Bスケールにおいて、69%の活性コンポーネント ( $C=0.69$ ) を持つ B-SDP は、DDPと比較してピークGPUメモリを 28% 削減し、同等またはそれ以下の検証損失を達成しました。500Mスケールでは、50%の活性コンポーネントを持つ B-SDP がメモリを 40% 削減しました。
ダウンストリーム性能: 5つのダウンストリーム・ベンチマーク (ARC-E, BoolQ, HellaSwag, OBQA, SciQ) において、すべてのSDPバリアントはマクロ平均でDDPの性能に約1ポイント以内で一致し、一部のバリアント (B-SDP) では特定のタスクでわずかな改善が見られました。
構成可能性: B-SDPをFSDPおよび活性化関数チェックポインティングと組み合わせることで、1Bモデルにおいてデバイスあたりのピークメモリ使用量は 9.7 GB となり、DDP (70.9 GB) に対して 85%の削減 を達成し、FSDP+AC単独 (11.0 GB) よりも厳密に低くなりました。
通信: SDPは、ステップあたりの通信量を（DDPの3,143 MBに対し2,253 MBへと）約28%削減しました。インターノード設定において、B-SDPとFSDPの組み合わせは、DDPと比較してウォールクロック・トレーニング時間を 19.9% 短縮しました。

画像分類

実験は、CIFAR-10およびCIFAR-100を用い、ResNet-18、WideResNet-18、およびSwin-Tinyに対して行われました。

ResNet-18: SDP構成は、DDPのメモリ予算のわずか 40% を使用して、DDPと同等または高い精度を達成しました。N-SDP (ニューロンレベル) は強力な正則化効果を示し、メモリ使用量64%においてDDPを上回りました。
Swin Transformer: B-SDP は、CIFAR-10において68%のメモリでDDPの精度に一致し、CIFAR-100においては32%のメモリ削減で精度を2%向上させました。
スパース性の限界: フォワード・マスキング (B-SDP/N-SDP) は中程度のスパース性では良好に機能しましたが、バックワード・マスキング (Bb-SDP) は極端なスパース性レベル (例: $C=3/8$ ) において優位であり、他のバリアントがモデル崩壊を起こす場面でもこれを防ぎました。

意義と主張

本論文は、SDPがデータ並列法とモデル並列法を補完する、実用的な「第3の軸」のパラダイムを提供すると主張しています。その主な意義は以下の通りです。

メモリと帯域幅のデカップリング: 活性関数をワーカー間で交換する必要を排除することで、SDPはパイプライン並列やテンソル並列を制限することが多い帯域幅のボトルネックを緩和します。
FLOP一致のスケーリングの実現: SDPは、同じハードウェア予算の下で、より大きなモデルやより長いシーケンスの学習を可能にします。1イテレーションあたりの収束率のわずかな低下は、減少したステップあたりの計算量と通信量により、イテレーション数を増やすことができる（FLOP一致）能力によって相殺されます。
ドロップイン・ソリューションの提供: SDPは、マスキング論理を除いて分散インフラストラクチャに変更を必要とせず、標準的なデータ並列法のドロップイン代替として機能します。

著者らは、SDPが特にメモリと帯域幅が主要な制約となるシナリオにおいて価値があり、複雑なパイプラインスケジューリングの複雑さやフルシャードの通信オーバーヘッドなしに、メモリ制限のあるアクセラレータ上で大規模モデルを学習するためのスケーラブルな道筋を提供すると強調しています。