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

[テンソル並列](./perf_train_gpu_many#tensor-parallelism)は、モデルのレイヤーを分割し、複数のハードウェアアクセラレータが同時に処理を行うようにします。これにより、単一のGPUのメモリ容量を超えるモデルの実行が可能になり、スループットの向上も実現できます。各レイヤーでGPUが部分的な結果を交換するため、ノード内での高速な通信が必要です。

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

```py
from transformers import AutoModelForCausalLM

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

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

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

[DistributedConfig](/docs/transformers/main/ja/expert_parallelism#transformers.DistributedConfig) の `tp_size` で、テンソル並列デバイス数を設定します。

- モデルの事前定義された計画を使用するには、`DistributedConfig(tp_size=N)`を設定します。
- 手動で`tp_plan`を定義し、`tp_size`とともに[DistributedConfig](/docs/transformers/main/ja/expert_parallelism#transformers.DistributedConfig)に渡します。

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

```py
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, DistributedConfig

# model_id = "meta-llama/Llama-4-Scout-17B-16E-Instruct" # better to visualize all the possible strategies
distributed_config = DistributedConfig(tp_size=4)
model = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Meta-Llama-3-8B-Instruct",
    dtype=torch.bfloat16,
    distributed_config=distributed_config,
)
print(model.tp_plan)

tokenizer = AutoTokenizer.from_pretrained("meta-llama/Meta-Llama-3-8B-Instruct")
prompt = "Can I help"
inputs = tokenizer(prompt, return_tensors="pt").input_ids.to(model.device)

# distributed run
outputs = model(inputs)
```

[torchrun](https://pytorch.org/docs/stable/elastic/run.html)で推論スクリプトを起動します。GPUごとに1つのプロセスを使用してください。

```bash
torchrun --nproc-per-node 4 demo.py
```

`tp_plan` 内の各レイヤーに対してテンソル並列計画を定義し、[DistributedConfig](/docs/transformers/main/ja/expert_parallelism#transformers.DistributedConfig) に渡します。以下の例では、カラム分割と行分割を使用しています。サポートされている他の戦略については、[分割戦略](#partitioning-strategies) セクションを参照してください。

手動パーティショニングには、モデルアーキテクチャと戦略間の相互作用に関する深い理解が必要です。不適切なパーティショニングの選択は、失敗したり誤った結果を生成したりする遅いモデルを生み出します。[Ultra-Scale Playbook](https://huggingface.co/spaces/nanotron/ultrascale-playbook?section=tensor_parallelism)では、パーティショニング戦略について詳しく説明しています。

キーはモジュール名またはパラメータ名で、`*` はレイヤーのインデックスを表します。認識されない戦略名が指定されると、サポートされている名前の一覧を含む `ValueError` が発生し、Transformers は何もマッチしなかったプランのルールや、どのルールにもカバーされなかったパラメータに対して警告をログに記録します。

```py
from transformers import AutoModelForCausalLM, DistributedConfig

tp_plan = {
    "model.layers.*.self_attn.q_proj": "colwise",
    "model.layers.*.self_attn.k_proj": "colwise",
    "model.layers.*.self_attn.v_proj": "colwise",
    "model.layers.*.self_attn.o_proj": "rowwise",
    ...
}

distributed_config = DistributedConfig(tp_size=4, tp_plan=tp_plan)
model = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Meta-Llama-3-8B-Instruct",
    dtype="auto",
    distributed_config=distributed_config,
)
print(model.tp_plan)
```

## パーティショニング戦略

`ALL_PARALLEL_STYLES` は [distributed/parallel_utils.py](https://github.com/huggingface/transformers/blob/main/src/transformers/distributed/parallel_utils.py) で定義された `ParallelInterface` であり、プランで使用できる各戦略名を、設定済みの戦略インスタンスにマッピングします。モデルのシャーディングを直接行う際にこれとやり取りすることはありませんが、利用可能な名の公式なリストです。[distributed/tensor_parallel.py](https://github.com/huggingface/transformers/blob/main/src/transformers/distributed/tensor_parallel.py) がテンソル並列戦略をこれに登録します:

```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](https://github.com/huggingface/transformers/blob/main/src/transformers/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`、通常はこれを選択）では、語彙をシャードし、各ランクの部分的なルックアップを全減算して完全な出力を生成します。`shard_dim=1`（`embedding_colwise`）では、隠れ次元をシャードし、各ランクの出力のスライスを全収集します。 |
| `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` | モジュールのフォワードを、そのモジュール自身のシャード化されたパラメータをローカルテンソルとして使用して実行します。自身のプランエントリによってシャード化される、アテンションシンクのような素の`nn.Parameter`の所有者に対して使用します。 |
| `MoEParamShard` | 指定された配置でMoEエキスパートの重みをシャードします。`grouped_gemm`名に対応しており、`shards_expert_dim=True`の場合、`module.num_experts`もランクごとのエキスパート数に書き換えられます。 |
| `EpRouterParallel` | 非ローカルなエキスパートのルータースコアをマスクし、グローバルなエキスパート ID をローカルな ID に再マッピングすることで、各ランクが所有するエキスパートのみを実行するようにします。`num_experts` がメッシュサイズで割り切れる必要があります。 |
| `RouterParallelMegaMoe` | DeepGEMM Mega MoE のルーターバリアントです。このルーターは、カーネル内でエキスパートをディスパッチし、ルーターの出力がそのまま（未加工のまま）であることを要求します。 |
| `MoeTensorParalellMegaMoeExperts` | DeepGEMM Mega MoE 用の推論専用エキスパート層です。勾配の同期をスキップし、プロセスグループをモジュールに渡すことで、カーネルが最初のフォワード時に共有バッファを設定できるようにします。 |

### パック戦略

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

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

```python
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` モジュールの出力を計算します。

```python
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](https://github.com/huggingface/transformers/blob/main/src/transformers/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`の簡略化されたバージョンを順を追って説明します。

1. `TensorParallelLayer` を継承し、戦略が扱う配置を保存します。ベースクラスは `__init__` を定義していないため、`super()` を呼び出す必要はありません。

    ```python
    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
    ```

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

    ```python
    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,
        )
    ```

3. 入力と出力の変換を実装します。`install_forward` はモジュールの元の forward の前後でこれらを呼び出すため、レイアウト間のテンソルの移動のみを行う必要があります。カラム方向のパーティショニングでは、レプリケートされた入力を想定し、最後の次元でシャードされた出力を生成します。

    ```python
    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`の往復をスキップします。同じ処理を戦略に必要とする場合は、コピーする前にソースコードを確認してください。

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

    ```python
    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](./model_doc/llama)における1回のフォワードパスの期待される高速化を示しています。

    

## 実装の設計

Transformers は、フレームワークに依存しない方法でテンソル並列を実装しています。[torch.distributed](https://docs.pytorch.org/tutorials/beginner/dist_overview.html) の [DeviceMesh](https://docs.pytorch.org/tutorials/recipes/distributed_device_mesh.html) と [DTensor](https://docs.pytorch.org/docs/stable/distributed.tensor.html) に依存することで、シンプルで拡張可能なインターフェースを提供します。

### DeviceMesh

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

```python
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` を分割します。列方向のパーティションでは、重みの出力特徴次元とバイアスの唯一の次元がシャード化されます。

    ```python
    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 は、代わりに入力特徴量次元をシャードし、バイアスを複製します。これは、バイアスが各ランクで加算されるのではなく、集約後に一度だけ加算されるためです。

    ```python
    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](https://huggingface.co/spaces/nanotron/ultrascale-playbook?section=tensor_parallelism)セクションを参照してください。

- [エキスパート並列化](./expert_parallelism)のガイドを確認してください。Mixture-of-Experts (MoE) モデルを使用している場合、これらのモデルはテンソル並列化とエキスパート並列化をサポートしています。

- テンソル並列の概要を素早く把握し、カラム並列と行並列の構成の違いを学ぶには、[Tensor Parallelism (TP) in Transformers: 5 Minutes to Understand](https://huggingface.co/blog/qgallouedec/tp) というブログ記事をお読みください。

- 学習環境での使用方法については、[Tensor parallelism](./tensor_parallelism)の学習ガイドを参照してください。