テンソル並列

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

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

Pass a TensorParallelConfig to the parallel_config argument of the model’s 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.

別のモデルにプランを追加するには、モデルにテンソル並列性を追加するを参照してください。

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

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で起動します。

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 (transformer のみ、32B パラメータ、bf16) について得られた結果を示しています。

制約事項

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

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

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

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

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

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

Colwise と rowwise

"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を使用してください。これにより、各パックされたブロックが独立してシャードされます。

# 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] を使用します。

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

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

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

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

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() と enable_parallelism() は、重みの読み込みやシャーディングを行う前に、両方をチェックします。

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

類似アーキテクチャの既存のプランから開始します。

• transformer_qwenimage.py には融合された投影が含まれていないため、すべてのエントリは単なる "colwise" または "rowwise" です。
• transformer_flux.pyは、1つのパックされた行方向の射影を追加します。
• transformer_flux2.py は、両方のパックスタイルを使用しています。