推論のためのテンソル並列

テンソル並列では、モデルのレイヤーを分割し、複数のハードウェアアクセラレーターが同時に処理を行います。これにより、単一 GPU のメモリ容量を超えるモデルの実行や、スループットの向上が可能になります。各レイヤーで GPU 同士が部分的な結果をやり取りするため、ノード内での高速な通信が必要です。

モデルがテンソル並列をサポートしているかどうかは、その設定にbase_model_tp_planが定義されているかどうかで判断できます。読み込まれたモデルについては、supports_tp_planプロパティで確認できます。

from transformers import AutoModelForCausalLM

model = AutoModelForCausalLM.from_pretrained("meta-llama/Meta-Llama-3-8B-Instruct")
print(model.supports_tp_plan)

このガイドでは、Transformers でテンソル並列を有効にする方法と、利用可能なパーティショニング戦略について説明します。

モデルのパーティショニング

DistributedConfig の tp_size で、テンソル並列デバイス数を設定します。

DistributedConfig(tp_size=N) を設定して、モデルの事前定義されたプランを使用します。
手動でtp_planを定義し、tp_sizeとともにDistributedConfigに渡します。

tp_plan="auto" を設定して、事前定義されたプランを明示的に指定することもできます。tp_size が省略され、tp_plan が設定されている場合、tp_size は WORLD_SIZE を他の並列サイズで割った値から導出されます。tp_plan を from_pretrained() に直接渡すことは非推奨となり、v5.18 で削除されます。

自動プラン

手動プラン

パーティショニング戦略

ALL_PARALLEL_STYLES は distributed/parallel_utils.py で定義された ParallelInterface であり、プランで使用できる各戦略名を、設定済みの戦略インスタンスにマッピングします。モデルをシャーディングする際に直接操作することはありませんが、利用可能な名の公式なリストです。distributed/tensor_parallel.py がテンソル並列戦略をこれに登録します。

ALL_PARALLEL_STYLES.register("embedding_rowwise", EmbeddingParallel(shard_dim=0))
ALL_PARALLEL_STYLES.register("embedding_colwise", EmbeddingParallel(shard_dim=1))
ALL_PARALLEL_STYLES.register("colwise_gather_output", ColwiseParallel(input_layouts=Replicate(), output_layouts=Replicate()))
ALL_PARALLEL_STYLES.register("colwise_rep", ColwiseParallel(input_layouts=Replicate(), output_layouts=Replicate()))
ALL_PARALLEL_STYLES.register("colwise", ColwiseParallel(input_layouts=Replicate(), output_layouts=Shard(-1)))
ALL_PARALLEL_STYLES.register("rowwise", RowwiseParallel(input_layouts=Shard(-1), output_layouts=Replicate()))
ALL_PARALLEL_STYLES.register("rowwise_split_input", RowwiseParallel(input_layouts=Replicate(), output_layouts=Replicate()))
ALL_PARALLEL_STYLES.register("rowwise_rep", RowwiseParallel(input_layouts=Replicate(), output_layouts=Replicate()))
ALL_PARALLEL_STYLES.register("packed_colwise", PackedColwiseParallel())
ALL_PARALLEL_STYLES.register("packed_rowwise", PackedRowwiseParallel())
ALL_PARALLEL_STYLES.register("sequence_parallel", SequenceParallel(use_local_output=True))
ALL_PARALLEL_STYLES.register("replicated_with_grad_allreduce", ReplicatedWithGradAllReduce())
ALL_PARALLEL_STYLES.register("mla_kv_a_proj", MlaKvAProjParallel())
ALL_PARALLEL_STYLES.register("all_reduce", AllReduceParallel())
ALL_PARALLEL_STYLES.register("moe_tp_experts", MoeExpertsParallel())
ALL_PARALLEL_STYLES.register("moe_identity_expert", MoeIdentityParallel())
ALL_PARALLEL_STYLES.register("local_params", LocalParamsParallel())

すべてのテンソル並列戦略は、そのファイル内のTensorParallelLayerのサブクラスです。ep_dispatch_expertsなどのエキスパート並列戦略はdistributed/expert_parallel.pyに定義されており、同じフックを持つExpertParallelLayerをサブクラス化し、同じALL_PARALLEL_STYLESに登録されます。最もよく使用するColwiseParallelとRowwiseParallelは、input_layoutsとoutput_layoutsの配置を受け取り、これにより1つのクラスが複数のプラン名をカバーします。colwiseは出力を最後の次元でシャードしたままにして、後続のrowwiseレイヤーが消費できるようにし、colwise_gather_outputはそれをオールギャザーして完全なテンソルに戻します。

下表では、各戦略について説明しています。

戦略	説明
`ColwiseParallel`	線形重みを出力特徴次元でシャード化し（2次元の`nn.Linear`重みでは`Shard(0)`）、1次元のバイアスもシャード化します。入力を`Replicate()`に再分散し、出力を`output_layouts`に配置します。
`RowwiseParallel`	入力特徴量次元(`Shard(-1)`)で線形重みをシャードし、リダクション後に加算されるバイアスは複製します。各ランクのフォワードは、和の自分担を含む`Partial()`出力を生成し、`transform_output_post_forward`がそれを`output_layouts`に再分配します。`Replicate()`へリデュースするとオールリデュースが発行され、`Shard(1)`へリデュースするとリデューススキャッターが発行されます。
`EmbeddingParallel`	`nn.Embedding` を、トークン ID を複製した状態でシャードします。`shard_dim=0`（`embedding_rowwise`、通常はこれを選択）では、語彙をシャードし、各ランクの部分的なルックアップを all-reduce して完全な出力を生成します。`shard_dim=1`（`embedding_colwise`）では、隠れ次元をシャードし、各ランクの出力のスライスを all-gather します。
`PackedColwiseParallel`	融合された重み用の`ColwiseParallel`のバリアントです。例えば、`up_proj`と`gate_proj`が`gate_up_proj`にパックされている場合に使います。重みが2以外の数のレイヤーをパックしている場合は、`split_factor`を使用してください。
`PackedRowwiseParallel`	最終次元に沿ってパックされた重み用の、行方向の対応する処理です。1次元のパラメータを複製します。
`SequenceParallel`	モジュールのパラメータを複製し、入力を`sequence_dim`(デフォルトは`1`)に沿ってシャードします。`LayerNorm`や`RMSNorm`のように、トークンごとに処理を行う正規化に使用されます。
`ReplicatedWithGradAllReduce`	パラメータを複製し、その勾配をall-reduceします。カラム方向レイヤーと行方向レイヤーの間に位置し、シャードされた軸に沿って正規化を行うノルムに必要です。この場合、各ランクは自身のヘッドのみを確認できます。
`AllReduceParallel`	モジュールの`Partial()`のフォワード出力を`Replicate()`にオールリデュースします。計算が部分和で終わるモジュールの同期ポイントとして使用します。
`MlaKvAProjParallel`	DeepSeek-V2 形式の MLA 注意機構における `kv_a_proj_with_mqa` の出力を分割し、RoPE 半分の勾配を all-reduce します。これにより `kv_b_proj` をバイパスするため、それ以外の場合に生じる部分勾配の保持を回避できます。モデル設定に `qk_rope_head_dim` が含まれている必要があります。
`MoeExpertsParallel`	Tensor parallel MoE experts. 前方伝播でエキスパートの出力を all-reduce し、隠れ状態とルーティング重みに対する後方伝播の all-reduce を加算します。
`MoeIdentityParallel`	`moe_tp_experts`が下流で適用するall-reduceを相殺するため、ゼロまたは恒等変換のエキスパートの入力をメッシュサイズで事前除算します。
`LocalParamsParallel`	モジュールの forward を、自身のシャード化されたパラメータをローカルテンソルとして使用して実行します。自身のプランエントリによってシャード化される、アテンションシンクのような素の `nn.Parameter` の所有者に対して使用します。
`MoEParamShard`	指定された配置でMoEエキスパートの重みをシャードします。`grouped_gemm`名に対応しており、`shards_expert_dim=True`の場合、`module.num_experts`もランクごとのエキスパート数に書き換えられます。

パック戦略

Weight packing は複数の線形層を単一の大きな層に結合します。PackedColwiseParallel と PackedRowwiseParallel 戦略は、パックされた重みを正しくシャードします。基本的な ColwiseParallel または RowwiseParallel 戦略では、パックされた重みのシャードが正しく行われません。

以下の例では、up_proj と gate_proj を単一の gate_up_proj モジュールにパックし、gate_up_proj をシャーディングするには packed_rowwise 戦略が必要です。

class Llama4TextExperts(nn.Module):
    ...
    self.gate_up_proj = nn.Parameter(torch.zeros(self.num_experts, self.hidden_size, 2 * self.expert_dim))

forward パスでバッチ行列積を用いて、gate_up_proj モジュールの出力を計算します。

def forward(self, hidden_states):
    ...
    gate_up = torch.bmm(hidden_states, self.gate_up_proj) # Compute the output of the gate_up_proj module
    gate, up = gate_up.chunk(2, dim=-1) # Split the output into gate and up

プレーンなShardは、その次元を連続するブロックに分割するため、各ランクはgateの末尾とupの先頭を受け取り、それぞれのスライスを受け取る代わりに、そのようになります。パック戦略では_StridedShardにsplit_factor(デフォルトは2)を指定して分割をインターリーブし、各ランクが両方の半分の対応するスライスを受け取れるようにすることで、シャーディング後もchunkが正しく揃うようにします。

カスタムパーティショニング戦略

TensorParallelLayer を継承して、distributed/tensor_parallel.py にカスタムのパーティショニング戦略を作成します。すべてのフックには何もしないデフォルト実装が用意されているため、戦略に必要なフックのみをオーバーライドしてください。

フック	目的
`validate_param`	この戦略ではシャード化できないパラメータを、重みの読み込み前に拒否します。
`shard_param`	`DTensor`のプレースホルダーに1つのパラメータを置き換えることで、ローダーがこのランクに属するシャードを特定できるようにします。
`transform_inputs_pre_forward`	モジュールの入力を、その forward が期待するレイアウトに再分配します。
`context_around_forward`	forward をコンテキストマネージャでラップし、カーネルにローカルテンソルを公開するなどの処理を行います。
`transform_output_post_forward`	モジュールの出力を再分配または削減します。
`should_use_local_tensors`	このモジュールの forward は、`DTensor` ではなくプレーンなテンソルを必要とする旨を報告します。
`install_forward`	`module.forward`をそのまま置き換えてください。上記のフックでは不十分な場合のみ、`ReplicatedWithGradAllReduce`がバックワードフックを登録するようにオーバーライドしてください。

以下の例では、ColwiseParallelの簡略化されたバージョンを順を追って説明します。

TensorParallelLayer を継承し、ストラテジーが扱うプレースメントを保存します。ベースクラスは __init__ を定義していないため、super() を呼び出す必要はありません。

class ColwiseParallel(TensorParallelLayer):
    def __init__(self, *, input_layouts=None, output_layouts=None, use_local_output: bool = True):
        self.input_layouts = input_layouts or Replicate()
        self.output_layouts = output_layouts if output_layouts is not None else Shard(-1)
        self.use_local_output = use_local_output

shard_param を実装し、1つのパラメータをDTensorプレースホルダーとしてラップします。これはメタテンソル上で実行されるため、distribute_tensorはメタデータのみを構築し、データを移動しません。まだ分散元の完全なテンソルが存在しないため、src_data_rank=Noneを渡します。

def shard_param(self, module, param, mesh):
    meta = module._parameters.get(param)
    if meta is None:
        return
    # Output features live on dim 0 for a 2D weight and dim -1 for a 1D bias
    placement = Shard(1) if isinstance(module, torch.nn.Embedding) else Shard(meta.ndim - 2)
    module._parameters[param] = torch.nn.Parameter(
        distribute_tensor(meta, mesh, [placement], src_data_rank=None),
        requires_grad=meta.requires_grad,
    )

入力と出力の変換を実装します。install_forward はモジュールの元の forward の前後でこれらを呼び出すため、テンソルをレイアウト間で移動するだけで構いません。カラム方向のパーティショニングでは、レプリケートされた入力を想定し、最後の次元でシャードされた出力を生成します。
```
def transform_inputs_pre_forward(self, module, args, kwargs, mesh):
    x = args[0]
    if not isinstance(x, DTensor):
        x = DTensor.from_local(x, mesh, [self.input_layouts], run_check=False)
    if x.placements != (Replicate(),):
        x = x.redistribute(placements=[Replicate()])
    return (x,) + args[1:], kwargs

def transform_output_post_forward(self, module, output, mesh):
    if not isinstance(output, DTensor):
        output = DTensor.from_local(output, mesh, [Shard(-1)], run_check=False)
    if output.placements != (self.output_layouts,):
        output = output.redistribute(placements=[self.output_layouts])
    return output.to_local() if self.use_local_output else output
```
同梱されているColwiseParallelは、この仕組みの上に高速パスを追加し、通常のnn.Linearの推論や、ローカルテンソルを必要とする量子化モジュールにおいて、DTensorの往復をスキップします。同じ処理を戦略に必要とする場合は、コピーする前にソースコードを確認してください。

tp_plan から参照できるように、戦略を登録します。登録には戦略のインスタンスを渡すため、__init__ に渡した placements がその名前とともに保存されます。以下の () に注意してください。

import torch

from transformers import AutoModelForCausalLM, DistributedConfig
from transformers.distributed.parallel_utils import ALL_PARALLEL_STYLES

ALL_PARALLEL_STYLES.register("colwise_custom", ColwiseParallel())
tp_plan = {
    "model.layers.*.self_attn.q_proj": "colwise_custom",
    ...
}
distributed_config = DistributedConfig(tp_size=4, tp_plan=tp_plan)
model = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Meta-Llama-3-8B-Instruct",
    dtype=torch.bfloat16,
    distributed_config=distributed_config,
)

ベンチマーク

Tensor parallelismは、特に大規模なバッチサイズや長いシーケンスにおいて、推論を大幅に高速化します。

このグラフは、シーケンス長が512のLlamaにおける1回のフォワードパスの期待される高速化を示しています。

実装の設計

Transformers は、フレームワークに依存しない方法でテンソル並列を実装しています。torch.distributed の DeviceMesh と DTensor に依存することで、シンプルで拡張可能なインターフェースを提供します。

DeviceMesh

DeviceMeshは、相互に通信するデバイスの多次元グリッドを作成します。並列化戦略によって必要な通信パターンが異なります。これらのパターンを処理するために、複数のサブメッシュを持つDeviceMeshを作成します。

import torch
from torch.distributed.device_mesh import init_device_mesh

# Create a 1D mesh of 4 accelerators
device_type = torch.accelerator.current_accelerator().type
device_mesh = init_device_mesh(device_type, (4,), mesh_dim_names=["tp"])

torch.distributedのほとんどの並列化戦略は、メッシュ自体またはそのサブメッシュに適用されます。メッシュは通信パターンを自動的に処理します。

DTensor

DTensor（Distributed Tensor）は、通常のテンソル演算の上に分散ロジックを処理します。テンソル並列化下でのモデルの重みはDTensorとして保存され、これにより、戦略が明示的な集合通信ではなく、レイアウトの変更として通信を記述できるようになります。

placements 属性は、DeviceMesh 内のデバイス間でテンソルがどのように配置されるかを PyTorch に伝えます。この属性には、次の値を指定できます。

Shard(dimension) は、構築時に使用された DeviceMesh 上で、指定された次元に沿って DTensor を分割します。列方向のパーティションでは、重みの出力特徴次元とバイアスの唯一の次元がシャード化されます。
```
weight = DTensor.from_local(weight, device_mesh["tp"], placements=[Shard(0)]) # Shard the output features
bias = DTensor.from_local(bias, device_mesh["tp"], placements=[Shard(-1)]) # Shard the ONLY dimension
```
Row-wise partitioning は、代わりに入力特徴量次元をシャードし、バイアスを複製します。これは、バイアスが各ランクで加算されるのではなく、集約後に一度だけ加算されるためです。
```
weight = DTensor.from_local(weight, device_mesh["tp"], placements=[Shard(-1)]) # Shard the input features
bias = DTensor.from_local(bias, device_mesh["tp"], placements=[Replicate()]) # Replicate bias across all GPUs
```
Replicate() は、DTensor を DeviceMesh にわたって複製し、各デバイス上にテンソルの完全なコピーを作成します。
Partial() は、テンソルがリダクション待ちであることを示します。行方向レイヤーの forward 出力は Partial() であり、これを Replicate() に再分散させることで all-reduce が発行されます。

リソース

テンソル並列性に関する詳細は、Ultra-Scale Playbookセクションを参照してください。
エキスパート並列化のガイドを確認してください。Mixture-of-Experts (MoE) モデルを使用している場合、これらのモデルはテンソル並列化とエキスパート並列化をサポートしています。
テンソル並列の概要を素早く把握し、カラム並列と行並列の構成の違いを学ぶには、Tensor Parallelism (TP) in Transformers: 5 Minutes to Understand というブログ記事をお読みください。
学習環境での使用方法については、Tensor parallelismの学習ガイドを参照してください。

Update on GitHub