Buckets:

hf-doc-build
/

doc

Files

xet

hf-doc-build/doc / diffusers /v0.43.0 /ja /optimization /tensor_parallelism.md

mishig

about 3 hours ago

preview code

download

raw

13.2 kB

# テンソル並列

[テンソル並列性](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)は両方のパックされたスタイルを使用します。

Xet Storage Details

Size:: 13.2 kB
Xet hash:: 36d0198e422e0766346a426c6d087b06e24585fef3e47393743c8d143b7b1b61

Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.

	# テンソル並列

	[テンソル並列性](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)は両方のパックされたスタイルを使用します。