Transformers documentation
用于训练的张量并行
用于训练的张量并行
张量并行（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 支持针对其 config 中定义了 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，请创建一个 DistributedConfig，在 tp_size 中指定设备数量，并将其传递给 from_pretrained()。Transformers 会使用模型预定义的规划，初始化设备网格，并为你分片受支持的层。
您也可以在 DistributedConfig 中设置 tp_plan="auto"。当省略 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。
在单节点上使用 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。