Tensor Parallelism

テンソル並列処理がGPU間でウェイト行列を分割し、大規模モデルをトレーニングする仕組みを学びます。Ultralyticsを用いたデータ並列処理との違いを探りましょう。

テンソル並列は、機械学習において大規模な個々の数理構造（テンソル）を、GPUやTPUなどの複数のハードウェアアクセラレータに分散させるために使用される高度な分散学習テクニックです。大規模なディープラーニングモデルのトレーニングでは、パラメータ数が1台のデバイスのメモリ容量を簡単に超えてしまうことがあります。ニューラルネットワークのレイヤー全体を1つのGPUに配置する代わりに、テンソル並列は重み行列をシャード化し、クラスター内の複数のデバイス間で（行列乗算などの）数学的演算を分割します。これにより、モデルはハードウェアセットアップ全体の結合されたメモリと演算能力を活用できるようになり、Single-Program Multiple-Data (SPMD)パラダイムで並列計算を実行しながら、NVIDIA NVLinkなどの高速インターコネクトを介して結果を同期します。

Tensor Parallelismの仕組み#

ニューラルネットワークの中核にあるのは行列乗算です。テンソル並列は、行列を行単位または列単位で分割することで、これらの演算を分散させます。たとえば、全結合層やTransformerのアテンション機構において、一方のGPUが行列の左半分を計算し、もう一方が右半分を計算することがあります。並列計算が終了した後、デバイスは通信を行い（多くの場合、高速なAll-Reduce集合通信オペレーションを使用）、部分的な結果を集約してから完全なテンソルを次のレイヤーに渡します。2025年の最近の学術的進歩により、大規模なコンピュートクラスターのボトルネックになりがちな通信オーバーヘッドを削減するため、部分的に同期されたアクティベーションを導入してこのプロセスがさらに最適化されています。

実社会での応用#

tensor parallelismは、特に大規模な計算規模を必要とする最先端アーキテクチャにおいて、現代のAIインフラストラクチャに不可欠です。

大規模言語モデル（LLM）のトレーニング：MetaのLlama 3やDeepSeek V3などの大規模な基盤モデルは、NVIDIA Megatron-LMなどのフレームワークを利用してテンソル並列を実装しています。これらのモデルの隠れ次元とアテンションヘッドは非常に大きいため、効率的にトレーニングを行い、リアルタイム推論中の低レイテンシを維持するには、8GPUノードにまたがって分割することが必須となります。
大規模ビジョンモデル（LVM）と3D生成：コンピュータビジョンが大規模なマルチモーダル推論システムに向けてスケールするにつれて、研究者はAWS SageMakerなどのサービス上でパイプライン並列と組み合わせたテンソル並列を使用して、巨大なビジョンTransformer（ViT）をトレーニングしています。このテクニックにより、膨大な連続メモリブロックを必要とする高解像度の画像処理や動画生成が可能になります。

PyTorchでのTensor Parallelismの実装#

歴史的に、エンジニアはテンソルをシャード化するために複雑なカスタム分散ロジックを書く必要がありました。最近、PyTorchは DTensor（Distributed Tensor）を導入し、このワークフローをネイティブに簡素化しました。公式のPyTorch Distributed Tensor APIを使用して行単位でシャード化されたテンソルを作成する例を以下に示します。

import torch
from torch.distributed.device_mesh import init_device_mesh
from torch.distributed.tensor import Shard, distribute_tensor

# Initialize a 1D device mesh across 2 GPUs
mesh = init_device_mesh("cuda", (2,))

# Create a standard PyTorch tensor (e.g., representing a layer's weights)
local_tensor = torch.randn(1024, 1024)

# Distribute the tensor across the GPUs by sharding along the first dimension (row-wise)
# Each GPU now holds a (512, 1024) chunk of the original tensor
distributed_tensor = distribute_tensor(local_tensor, mesh, [Shard(0)])

print(f"Global shape: {distributed_tensor.shape}, Local shape: {distributed_tensor.to_local().shape}")

エッジに最適化されたビジョンタスクや迅速なモデルデプロイのために、開発者は通常Ultralytics Platformに依存して、最適なハードウェア利用を自動的に処理します。数百億パラメータを持つ基盤モデルには手動のテンソル並列設定が必要ですが、シンプルなCLIコマンドをそのまま使用して、YOLO26などのモデルのトレーニングを効率的にスケールさせることができます。これにより、堅牢なモデルトレーニングのヒントと並行してネイティブのデータ並列テクニックをシームレスに活用し、最大のスループットが保証されます。

Explore solutions

ロボティクスにおけるコンピュータビジョン

Ultralytics YOLOモデルで、よりスマートなマシンを実現しましょう。ロボティクスにおけるビジョンAIは、自律航行、認識、物体追跡、リアルタイム制御を推進します。

Tensor Parallelism