学習のためのテンソル並列

Tensor parallelism (TP) は、重み行列を GPU 間で列方向または行方向に分割します。各 GPU はシャードを保持し、部分結果を計算したうえで、all-reduce によって同期して完全な出力を生成します。

TPはGPU間の頻繁な通信に依存します。NVLinkのような高速なノード内リンクを備えたハードウェアで最も効果的に動作します。

    ┌─────────────────────────────┐
    │       X  (replicated)       │
    └────┬──────────┬─────────┬───┘
         │          │         │
    ┌────▼───┐ ┌────▼───┐ ┌───▼────┐
    │ ▓▓▓ W₀ │ │ ░░░ W₁ │ │ ███ W₂ │
    │  X@W₀  │ │  X@W₁  │ │  X@W₂  │
    └────┬───┘ └────┬───┘ └───┬────┘
         └──────────┼─────────┘
               Y₀+Y₁+Y₂
    ┌────────────────────────────┐
    │          Y (full)          │
    └────────────────────────────┘

Transformers は、base_model_tp_plan が定義されているアーキテクチャに対して TP をサポートしています。モデルがネイティブ TP をサポートしているかどうかを確認するには、まずこのフィールドを確認してください。

from transformers import AutoConfig

config = AutoConfig.from_pretrained("Qwen/Qwen3-0.6B")
print(config.base_model_tp_plan is not None)
print(config.base_model_tp_plan)

モデルがTPをサポートしている場合、tp_sizeにデバイス数を指定してDistributedConfigを作成し、from_pretrained()に渡します。Transformersは、モデルの事前定義されたプランを使用してデバイスメッシュを初期化し、サポートされているレイヤーをシャードします。

tp_plan="auto"をDistributedConfigに設定することもできます。tp_sizeが省略された場合、WORLD_SIZEから推論されます。tp_planを直接from_pretrained()に渡すことは非推奨となり、v5.18で削除されます。

device_map と distributed_config を併用しないでください。両者は重みの読み込みレベルで競合します。device_map はモジュール全体を特定の GPU に配置するのに対し、テンソル並列では同じパラメータをすべての GPU にシャードします。

import torch

from transformers import AutoModelForCausalLM, DistributedConfig

distributed_config = DistributedConfig(tp_size=4)

model = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen3-0.6B",
    dtype=torch.bfloat16,
    distributed_config=distributed_config,
)

Trainer はテンソル並列の計画を検出し、モデルから tp_size を読み取って、ParallelismConfig を自動的に作成します。

1ノードの4つのGPUでトレーニングを開始します。

torchrun --nproc-per-node 4 train_tp.py

ParallelismConfig

FSDP などの他の並列化技術と TP を組み合わせる場合は、ParallelismConfig を明示的に渡してください。

import torch

from accelerate import ParallelismConfig
from transformers import AutoModelForCausalLM, DistributedConfig, TrainingArguments

distributed_config = DistributedConfig(tp_size=4)

model = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen3-0.6B",
    dtype=torch.bfloat16,
    distributed_config=distributed_config,
)

parallelism_config = ParallelismConfig(tp_size=4)

args = TrainingArguments(
    ...,
    parallelism_config=parallelism_config,
)

PEFT アダプター

テンソル並列ベースモデルにPEFTアダプターを読み込むことができます（peft >= 0.21.1 が必要です）。PEFTはアダプター重みのシャーディングを処理します。アダプターの追加、読み込み、トレーニングの方法（分散トレーニングに関する注記を含む）については、PEFTを参照してください。

次のステップ

Read the Tensor Parallelism chapter from The Ultra-Scale Playbook for more details about how it works.
テンソル並列推論ガイドを読み、パーティショニング戦略、手動TPプラン、実装の詳細について詳しく学んでください。

Update on GitHub