モデルシャーディングとは何か？ホームAIのメモリ制限にとってなぜ重要なのか？

モデルシャーディングは、必要なモデル状態を複数のデバイスに分散します。これにより、1台のアクセラレーターのメモリ容量を超えるモデルでも、ロードして実行できるようになります。

ホームAIサーバーで重要なのは、チェックポイントが複数のファイルとして保存されているかどうかではなく、実行時の重み、レイヤー、テンソル、その他の状態が実際に異なるデバイスへ配置されているかどうかです。シャーディングによって、1台のGPUでは不可能なメモリ要件を、複数デバイスによる実行で現実的な構成に変えられます。ただし、シャード間ではデータの交換やステージ間の処理引き継ぎが必要になるため、次の制約としてインターコネクト帯域幅、デバイス間の性能差、ランタイムの対応状況が問題になります。

ランタイムシャーディングはモデル状態をデバイス間に分散する

実行中のモデルには、各レイヤーの処理時に利用できなければならないテンソルが含まれています。シャーディングでは配置を変更し、モデル全体の状態をすべてのデバイスに複製するのではなく、異なるデバイスがそれぞれ異なる部分を管理します。

分散テンソルでは、シャーディングテンソルの配置を利用できます。これにより、各ランクで同一に保存するのではなく、デバイスメッシュ全体に分散次元を割り当てられます。

直接的なメモリ上のメリットは、各デバイスに常駐するデータ量が減ることです。一方、システム全体では、どのデバイスもすべての演算に必要なデータを単独で持たなくなるため、推論に調整処理が組み込まれます。

チェックポイントのシャードは、実行中のシャーディングモデルとは異なる

大規模モデルのリポジトリでは、ダウンロードや段階的なロードをしやすくするため、チェックポイントを多数のファイルに分割することがよくあります。しかし、このパッケージング方法だけで、ランタイムのロード完了後にテンソルがどこに配置されるかが決まるわけではありません。

ファイルレベルのシャーディングとランタイムでの配置は別の問題です。ローダーは、シャーディングされたチェックポイントを複数デバイスへのディスパッチと組み合わせられるためです。

そのため、ホームユーザーの環境では、ディスク上に数十個の`.safetensors`シャードがあっても、ランタイムが完全なモデルを1台のGPUに配置しようとすることがあります。逆に、ランタイムがロード時にチェックポイントを別の複数デバイス構成へ再分割することもあります。

容量計画では、リポジトリ内のファイル数から推論トポロジーを推測するのではなく、起動後の実際のデバイスマップと常駐割り当てを確認する必要があります。

シャーディングでは通信またはステージ間の処理引き継ぎが発生する

あるデバイスが、別のシャードで必要な値を生成した場合、そのデータはインターコネクト経由で送信するか、集合通信によって同期する必要があります。実際の通信量は、ランタイムがレイヤー内部のテンソルを分割するのか、異なるデバイスに異なるレイヤー範囲を配置するのか、または必要なときだけシャーディングされた状態を集約するのかによって変わります。

異なる複数デバイス推論戦略は、メモリ配置と通信パターンの異なる組み合わせを選択します。

そのため、合計VRAMが十分な2台のGPUでも、モデルを低速でしか提供できない場合があります。PCIeなどのリンクがアクセラレーターのローカルメモリより大幅に遅い場合、アクティベーションの転送や部分結果の同期がボトルネックになることがあります。

Featured

ZimaCube 2 パーソナルクラウドホームNAS

ZimaCube2

$799.00 - $2,499.00

Buy Now

デバイス性能の偏りによって、1つのシャードがボトルネックになることがある

ホームサーバーでは、メモリ容量、計算性能、リンク幅、世代の異なるGPUを組み合わせることがあります。数学的に均等な分割でも、最も遅いデバイスや最も小さいデバイスがリクエスト全体の処理速度を決めてしまうことがあります。

そのため、レイヤー配置やオフロードツールでは、対称的なハードウェアを前提にするのではなく、実際のデバイス容量を考慮する必要があります。明示的な分散モデル実行では、自動的な共有メモリ抽象化ではなく、並列処理の設定を使用します。

実用的な構成では、大容量のGPUにより多くのレイヤーを割り当てたり、遅延が重要なコンポーネントを最速の経路に配置したりします。目的はシャード数を均等にすることではなく、すべてのデバイスに収まり、クリティカルパスのバランスが取れた構成にすることです。

同じプロンプトを使い、デバイスごとのメモリ使用量、利用率、転送時間、アイドル時間を測定してください。あるシャードが常に待機したり、データをあふれさせてオフロードしたりするなら、性能を制限しているのは合計VRAMの容量ではなく、トポロジーである可能性があります。

モデルシャーディングは、まずメモリ上の実行可否を左右する手段として重要

シャーディングが最も役立つのは、シャーディングなしではモデル全体を1台のデバイスに収められない場合です。モデルをロードできるようになった後は、インターコネクトのコスト、バッチサイズ、KVキャッシュの配置、より小型または量子化されたモデルのほうが簡単ではないかといった点が最適化の中心になります。

具体的な計算方式の1つがテンソル並列推論です。モデルシャーディングはより広い配置の問題であり、必要な状態をどのデバイスに置くべきかを問うものです。

合計VRAMを、透過的に利用できる1つのプールとして扱わないでください。シャーディングによって個別のメモリを連携させることはできますが、すべてのランタイムには配置ルールと通信コストがあり、それらが実際に役立つデプロイになるかどうかを決定します。

よくある質問

シャーディングされたチェックポイントは、シャーディングされた実行中のモデルと同じですか？

いいえ。チェックポイントのシャードは、保存やロードのためにファイルを分割したものです。ランタイムシャーディングは、推論中にモデルの状態をどのデバイスに配置するかを決定します。

モデルシャーディングはテンソル並列処理と同じですか？

いいえ。テンソル並列処理は、テンソル演算を分割してシャーディングされたモデルを実行する方法の1つです。シャーディングには、レイヤー、ステージ、パラメーターなどを配置する他の戦略も含まれます。

12GBのGPUが2台あれば、自動的に利用可能な24GBのメモリプールになりますか？

いいえ。ランタイムがモデルを明示的に分割する必要があり、通信、重複する状態、KVキャッシュ、各デバイスに必要な余裕によって、合計容量のうち実際に利用できる量は少なくなります。