# テンソル並列

[テンソル並列性](https://huggingface.co/spaces/nanotron/ultrascale-playbook?section=tensor_parallelism)は、モデルの重み行列をデバイス間で分割します。各デバイスは、各レイヤーの列方向（`"colwise"`）または行方向（`"rowwise"`）のスライスを保持し、部分結果を計算します。レイヤーの境界での `AllReduce` または `AllGather` により、部分結果が結合されて完全な出力が生成されます。[コンテキスト並列性](./context_parallelism) とは異なり、デバイスごとの*重み*メモリを削減するため、モデルが1つのデバイスに収まらない場合に有効です。他の戦略との比較については、[分散推論](../training/distributed_inference)を参照してください。

## シャード化されたモデルを読み込む

Pass a [TensorParallelConfig](/docs/diffusers/v0.43.0/ja/api/parallel#diffusers.TensorParallelConfig) to the `parallel_config` argument of the model's [from_pretrained()](/docs/diffusers/v0.43.0/ja/api/models/overview#diffusers.ModelMixin.from_pretrained). `tp_degree` is the number of devices to shard across and must divide the model's number of attention heads. The model's class must define a `_tp_plan`, which maps module names to a sharding style. To check whether a model supports tensor parallelism, look for this attribute on its class, such as [Flux2Transformer2DModel](/docs/diffusers/v0.43.0/ja/api/models/flux2_transformer#diffusers.Flux2Transformer2DModel).

別のモデルにプランを追加するには、[モデルにテンソル並列性を追加する](#add-tensor-parallelism-to-a-model)を参照してください。

この方法で読み込むと、チェックポイントが読み込まれる際にシャード化されます。各ランクは各重みの自分のスライスのみを読み取り、それを直接その GPU に配置するため、どのランクもフルサイズの重みを保持することはありません。ランクあたりのメモリは、`tp_degree` が増加するにつれて減少します。

```py
import os
import torch
from torch import distributed as dist
from diffusers import DiffusionPipeline, Flux2Transformer2DModel, TensorParallelConfig

def main():
    dist.init_process_group(backend="nccl")
    rank, world_size = dist.get_rank(), dist.get_world_size()
    # torchrun sets LOCAL_RANK to this process's GPU index on its node
    local_rank = int(os.environ["LOCAL_RANK"])
    device = torch.device(f"cuda:{local_rank}")
    torch.cuda.set_device(device)

    # Each rank reads only its own shard of every planned weight, straight onto its own GPU.
    transformer = Flux2Transformer2DModel.from_pretrained(
        "black-forest-labs/FLUX.2-dev",
        subfolder="transformer",
        dtype=torch.bfloat16,
        parallel_config=TensorParallelConfig(tp_degree=world_size),
    )

    pipeline = DiffusionPipeline.from_pretrained(
        "black-forest-labs/FLUX.2-dev", transformer=transformer, dtype=torch.bfloat16
    )
    # The transformer is already on its device; move the remaining components individually. Do not call
    # `pipeline.to(device)` — that would move every rank's shards onto the same device.
    pipeline.text_encoder.to(device)
    pipeline.vae.to(device)

    generator = torch.Generator().manual_seed(42)
    image = pipeline(prompt="a cat holding a sign that says hello", generator=generator).images[0]
    if rank == 0:
        image.save("output.png")
    dist.destroy_process_group()

if __name__ == "__main__":
    main()
```

スクリプトを`tensor_parallel_flux2.py`として保存し、`torchrun`で起動します。

```shell
torchrun --nproc-per-node 4 tensor_parallel_flux2.py
```

スクリプトは`tp_degree`を`world_size`に設定するため、`--nproc-per-node 4`によりtransformerが4つのGPUにシャードされます。

トランスフォーマーのみがシャーディングされます。各ランクはテキストエンコーダーの完全なコピーを保持するため、各GPUにはテキストエンコーダーとトランスフォーマーの割り当て分を収めるのに十分なメモリが必要です。

## 読み込みパフォーマンス

`parallel_config` を使用して読み込むと、フルモデルを読み込んだ後に `enable_parallelism()` を呼び出す場合よりも、ランクあたりの CPU メモリ使用量が少なく、読み込みが高速になります。下表は、4x A10G (23GB) 上で [black-forest-labs/FLUX.2-dev](https://huggingface.co/black-forest-labs/FLUX.2-dev) (transformer のみ、32B パラメータ、bf16) について得られた結果を示しています。

| tp_degree | メソッド | 読み込み時間 | ピークCPU/ランク |
|---|---|---|---|
| 4 | `from_pretrained(parallel_config=...)` | 12.5s | 6.8GB |
| 4 | `from_pretrained` + `enable_parallelism` | 30.4s | 64.1GB |

## 制約事項

Tensor parallelismは一部の読み込みオプションや機能と互換性がなく、シャード化されたモデルを保存することはできません。

- `parallel_config` を使用したテンソル並列での読み込みは、`device_map`、`quantization_config`、`low_cpu_mem_usage=False`、`use_flashpack=True`、または非safetensors形式の重みとは互換性がありません。これらを使用すると、完全なチェックポイントを静かに読み込む代わりにエラーが発生します。
- テンソル並列性を量子化、オフロード、またはLoRAアダプターと組み合わせると、モデルがどのようにシャーディングされていても、エラーが発生します。
- メモリに既に読み込まれているモデルをシャーディングするには、同じ設定で`enable_parallelism()`を呼び出します。この場合、まずモデル全体を読み込んでからシャーディングを行うため、すべてのランクで完全なチェックポイントのメモリが必要です。`parallel_config`で読み込むことで、このコストを回避できます。
- テンソル並列モデルを保存すると、[save_pretrained()](/docs/diffusers/v0.43.0/ja/api/models/overview#diffusers.ModelMixin.save_pretrained)でエラーが発生します。シャーディングする前にモデルを保存してください。

## モデルにテンソル並列性を追加する

このセクションは、Diffusersモデルにテンソル並列性のサポートを追加したいコントリビューター向けです。

Tensor parallelism は、モジュール名とシャーディングスタイルを対応付けるクラス属性である `_tp_plan` を定義しているモデルでのみ機能します。この定義を作成する際は、主に、隠れ次元を *拡張* する射影と、それを元に戻して *縮小* する射影をペアにする必要があります。

各キーには`*`を最大1つまで含めることができ、`*`の前の部分は[`nn.ModuleList`](https://pytorch.org/docs/stable/generated/torch.nn.ModuleList.html)に解決されなければなりません。これにより、1つのエントリがすべてのブロックをカバーします。`*`を含まないキーはモデル自体に適用されます。パスはモデルを基準とした相対パスです。

### Colwise と rowwise

| スタイル | シャード | 各ランク | 用途 |
|---|---|---|---|
| `"colwise"` | 出力特徴(`weight`の次元0) | 出力のスライスを計算します | `to_q`、`to_k`、`to_v`、FFN の入力射影 |
| `"rowwise"` | 入力特徴(`weight` 次元 1) | 部分和を計算します | `to_out.0`、FFN 出力投影 |
| `"rowwise_split_input"` | 入力特徴(`weight` 次元 1) | レプリケートされた入力のスライスから部分和を計算します | 入力値が複製された独立した投影。例えば、タイムステップ埋め込みを読み込む AdaLN 変調投影など |

`"colwise"` と `"rowwise"` は、必ずこの順序でペアにしてください。`"colwise"` 投影は出力をシャードしたまま残し、続く `"rowwise"` 投影はそのシャードを直接消費し、ブロック境界での単一の `AllReduce` によって結果が再構築されます。このペアを他の方法でシャードすると、中間で gather が発生し、通信量が大幅に増加します。

Attention において、列方向の Q、K、V 投影により、各ランクにヘッドのサブセットが割り当てられるため、`tp_degree` はヘッド数で割り切れる必要があります。エンコーダストリームの複製（`add_q_proj`、`to_add_out`、`ff_context`）は、画像ストリームの対応する部分と同じパターンに従います。

### 融合された投影

1つの`Linear`がシャード対象の次元に沿って複数の論理テンソルをパックしている場合、単純な`"colwise"`/`"rowwise"`スライスでは連結を横断してそのままスライスされてしまい、各パーツがずれてしまいます。代わりに`PackedColwiseParallel`/`PackedRowwiseParallel`を使用してください。これにより、各パックされたブロックが独立してシャードされます。

```py
# in src/diffusers/models/transformers/your_model.py
from ...hooks.tensor_parallel import PackedColwiseParallel, PackedRowwiseParallel
```

`blocks` は、パックされた次元を割り切れる整数のリストを指定します。例えば、等しい半分を持つ SwiGLU の gate と up projection には `[1, 1]` を、`mlp_ratio=3` を持つ融合された Q、K、V、gate、up projection には `[1, 1, 1, 3, 3]` を使用します。

```py
"transformer_blocks.*.ff.linear_in": PackedColwiseParallel([1, 1]),
```

ブロックサイズが設定ファイルからのみ判明する場合、`blocks`を省略し、代わりに`__init__`の時点で`Linear`に絶対サイズを`_tp_packed_col_blocks`または`_tp_packed_row_blocks`として保存してください。

```py
# in the attention module's __init__
self.to_out._tp_packed_row_blocks = [self.inner_dim, self.mlp_hidden_dim]
```

```py
# in the model's _tp_plan
"single_transformer_blocks.*.attn.to_out": PackedRowwiseParallel(),
```

[`transformer_flux2.py`](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/transformers/transformer_flux2.py) の計画では、両方の形式を使用しています。ダブルストリームの SwiGLU ペアは `blocks` を直接渡します。融合されたシングルストリームの QKV および MLP 投影は、`mlp_ratio` および `mlp_mult_factor` に依存するため、そのサイズを `Flux2ParallelSelfAttention.__init__` 内の `Linear` に保存します。

### 除外する項目

計画に含まれていないモジュールは、すべてのランクに複製されます。正規化レイヤー、パッチおよびテキストの埋め込み、そして最終的な`norm_out`と`proj_out`は除外してください。これらはサイズが小さいため、シャーディングしてもメモリ削減の効果は限定的であり、通信コストが増加します。AdaLNのモジュレーション（`img_mod`と`txt_mod`）も、通常は小さいため除外して問題ありません。MiniMax-H3の`adaln_proj.linear`のようにサイズが大きい場合は、`"rowwise_split_input"`でシャーディングしてください。

### 制約と検証

- `tp_degree`は`config.num_attention_heads`で割り切れる必要があります。
- すべてのパックされたブロックは、合計値だけでなく、*それぞれ個別に*`tp_degree`で割り切れる必要があります。

[from_pretrained()](/docs/diffusers/v0.43.0/ja/api/models/overview#diffusers.ModelMixin.from_pretrained) と `enable_parallelism()` は、重みの読み込みやシャーディングを行う前に、両方をチェックします。

目視ではなく数値的に新しいプランを検証します。固定のシードで単一デバイス上で生成し、次にテンソル並列環境下で再度生成して、出力を比較します。`"colwise"`/`"rowwise"`の配置が誤っていても、通常は実行され、妥当に見えますが誤った画像が生成されます。

> [!TIP]
> 類似アーキテクチャの既存のプランから開始します。
>
> - [`transformer_qwenimage.py`](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/transformers/transformer_qwenimage.py)には融合された投影がなく、すべてのエントリは`"colwise"`または`"rowwise"`です。 - [`transformer_flux.py`](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/transformers/transformer_flux.py)は1つのパックされた行方向の投影を追加します。 - [`transformer_flux2.py`](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/transformers/transformer_flux2.py)は両方のパックされたスタイルを使用します。