企业官网建设流程全解析

Transformers 多 GPU 高效训练指南：数据并行、ZeRO、流水线与 3D 张量并行全解

【免费下载链接】transformers🤗 Transformers: the model-definition framework for state-of-the-art machine learning models in text, vision, audio, and multimodal models, for both inference and training.项目地址: https://gitcode.com/GitHub_Trending/tra/transformers

本文基于 Transformers 官方性能文档 Efficient Training on Multiple GPUs（英文版见 perf_train_gpu_many.md）编写，系统讲解当单卡训练太慢或模型装不进单卡显存时，如何在多 GPU 环境下选择并落地数据并行（DP/DDP）、ZeRO 分片、流水线并行（PP）、张量并行（TP）及其 2D/3D 组合方案。读完之后，你将掌握：各类并行的通信模式与适用边界、官方基准测试的完整复现命令，以及 Transformers 中通过TrainingArguments与 DeepSpeed 集成实际启用这些技术的方法。

核心概念总览

从单 GPU 切换到多 GPU，本质是为负载引入某种"并行度"。Transformers 文档将并行技术分为五种基本形态，也是后文所有策略讨论的术语基础：

1. DataParallel (DP，数据并行)—— 复制同一套模型副本，每份副本喂入数据切片，并行计算，并在每个训练步末尾同步。
2. TensorParallel (TP，张量并行)—— 把每个大张量切成若干分片，分片分布在不同 GPU 上并行计算，步末再同步。这是一种"水平"切分（按层内维度切，而非按层切）。
3. PipelineParallel (PP，流水线并行)—— 把模型沿"垂直"方向（按层）切到不同 GPU 上，每个 GPU 负责一个或多个阶段，并以小数据块（micro-batch）为单位流水推进。
4. ZeRO（Zero Redundancy Optimizer）—— 对参数、梯度、优化器状态做类似 TP 的分片，但前向/反向计算时会按需把完整张量重组回来，因此不需要修改模型代码；同时支持 CPU/NVMe 卸载来缓解显存压力。
5. Sharded DDP—— 与 ZeRO 同义，是各 ZeRO 实现中"分片数据并行"的别称（文献中也常见 Sharded、Partitioned 等称呼）。

可扩展性策略：先选路，再选技术

文档给出的核心决策框架按"单机/多机 × 模型能否装下"两个维度展开，可直接作为选路表：

单节点 / 多 GPU

• 模型能装进单卡：

  1. DDP（分布式数据并行）
  2. ZeRO（快慢取决于具体场景与配置）

• 模型装不进单卡：

  1. PP
  2. ZeRO
  3. TP

  若节点内互联极快（NVLink / NVSwitch），三者速度接近；否则 PP 通常快于 TP 和 ZeRO。TP 的切分度数也会影响结果——最稳妥的做法是针对自己的硬件做实验。实践中 TP 几乎总在单节点内使用，即 TP 尺寸 ≤ 单节点 GPU 数。

• 最大的一层都装不进单卡：

  • 不用 ZeRO：必须上 TP，因为仅靠 PP 装不下；
  • 用 ZeRO：回到"装进单卡"一栏的处理方式。

多节点 / 多 GPU

• 节点间有高速互联：
  1. ZeRO —— 基本无需改动模型
  2. PP+TP+DP —— 通信量更小，但需要对模型做大规模改造
• 节点间互联慢且显存仍不足：
  1. DP+PP+TP+ZeRO-1

一个提示：文档特意强调，单 GPU 篇 介绍的混合精度、梯度累积等通用策略在多卡场景同样适用，进入多卡调优前应先打好基础。

数据并行：DP vs DDP 的通信代价

多卡用户通常最先受益的是 PyTorch 内置的DataParallel（DP）与DistributedDataParallel（DDP）。Transformers 文档明确建议优先使用 DDP：DP 在部分模型上可能失败，PyTorch 官方文档本身也推荐 DDP。两者的通信模式差异是理解一切性能差异的起点：

DDP（多进程，规避 GIL）：

• 启动时，主进程把 GPU 0 上的模型复制到其余 GPU；
• 之后每个批次：
  1. 各 GPU 直接消费各自的 mini-batch；
  2. backward过程中，本地梯度就绪后即在所有进程间做平均（all-reduce）。

DP（Python 线程）：每个批次要走 5 轮数据交换——

1. GPU 0 读取数据批并切分发送到各 GPU；
2. 把最新模型权重从 GPU 0 复制到各 GPU；
3. 各 GPUforward，输出回传 GPU 0 计算损失；
4. 损失再广播回各 GPU 执行backward；
5. 各 GPU 梯度回传 GPU 0 做平均。

由此产生三条关键结论：

• 通信量：DDP 每批次只需同步梯度一次；DP 每批次 5 次数据交换。同步数据越多，慢互联对总耗时的拖累越大；
• 负载倾斜：DP 中 GPU 0 干的活远多于其他卡，GPU 利用率更低；DP 通过 Python 线程拷贝数据，DDP 通过torch.distributed；
• 跨机能力：DDP 可用于多机，DP 不能。

官方基准：NVLink 是否值得依赖

文档给出了一组 GPT-2 语言建模训练的实测（GPT-2 + wikitext，per_device_train_batch_size 4，max_steps 200，2× TITAN RTX 24GB + 2 条 NVLink，环境为当时版本的 PyTorch + CUDA 11.0）。完整可复现命令如下（NCCL_P2P_DISABLE=1用于在对应基准中关闭 NVLink）：

# DP rm -r /tmp/test-clm; CUDA_VISIBLE_DEVICES=0,1 \ python examples/pytorch/language-modeling/run_clm.py \ --model_name_or_path openai-community/gpt2 --dataset_name wikitext --dataset_config_name wikitext-2-raw-v1 \ --do_train --output_dir /tmp/test-clm --per_device_train_batch_size 4 --max_steps 200 {'train_runtime': 110.5948, 'train_samples_per_second': 1.808, 'epoch': 0.69} # DDP w/ NVlink rm -r /tmp/test-clm; CUDA_VISIBLE_DEVICES=0,1 \ torchrun --nproc_per_node 2 examples/pytorch/language-modeling/run_clm.py \ --model_name_or_path openai-community/gpt2 --dataset_name wikitext --dataset_config_name wikitext-2-raw-v1 \ --do_train --output_dir /tmp/test-clm --per_device_train_batch_size 4 --max_steps 200 {'train_runtime': 101.9003, 'train_samples_per_second': 1.963, 'epoch': 0.69} # DDP w/o NVlink rm -r /tmp/test-clm; NCCL_P2P_DISABLE=1 CUDA_VISIBLE_DEVICES=0,1 \ torchrun --nproc_per_node 2 examples/pytorch/language-modeling/run_clm.py \ --model_name_or_path openai-community/gpt2 --dataset_name wikitext --dataset_config_name wikitext-2-raw-v1 \ --do_train --output_dir /tmp/test-clm --per_device_train_batch_size 4 --max_steps 200 {'train_runtime': 131.4367, 'train_samples_per_second': 1.522, 'epoch': 0.69}

解读：带 NVLink 时 DDP 比 DP 快约 10%；而关闭 NVLink 后 DDP 反而比 DP 慢约 15%。这正说明并行方案的收益强依赖卡间互联带宽——梯度越需要频繁同步，慢链路对墙钟时间的放大越明显。上述命令中的 run_clm.py 在当前仓库中仍然存在，可直接按 examples/pytorch 的说明复制运行（注意结果会随 PyTorch/CUDA/驱动版本变化，数字仅作机制说明，不保证复现）。

在 Transformers 中，DDP 由 Trainer 经 Accelerate 自动接管，TrainingArguments提供了对应的控制项，如ddp_backend（指定torch.nn.DistributedDataParallel的后端）、ddp_find_unused_parameters（默认值取决于是否启用梯度检查点）等，定义见 training_args.py。

ZeRO 数据并行：把"每人背全套装备"变成"轮流共享"

ZeRO-DP（ZeRO 驱动的数据并行）可以理解为不复制完整副本的 DP：普通 DP 要在每张卡上保存完整的模型参数、梯度和优化器状态，而 ZeRO 让每张 GPU 只保存这些张量的一个分片；运行时某层需要完整参数做前向/反向计算时，所有 GPU 临时互补短板、重组出完整张量，用完即释放，并借助预取（prefetch）把同步成本藏进计算里。

用文档中的三层模型（每层 3 个参数）+ 3 张 GPU 的例子最直观：

La | Lb | Lc ---|----|--- a0 | b0 | c0 a1 | b1 | c1 a2 | b2 | c2

分片后：

GPU0: GPU1: GPU2: a0 | b0 | c0 a1 | b1 | c1 a2 | b2 | c2

前向过程：三张卡各自收到 mini-batch（x0→GPU0, x1→GPU1, x2→GPU2），处理层 La 时，GPU0 只有自己持有的 a0，需要从 GPU1 取 a1、从 GPU2 取 a2，拼出完整层参数后做前向；计算完立即丢弃临时数据。之后依次对 Lb、Lc 前向，反向按 Lc → Lb → La 重复同样的重组流程。

文档作者打过一个精彩的比方：三人露营，A 带帐篷、B 带炉子、C 带斧头，每晚共享装备、早上各归其位继续赶路——这就是 Sharded DDP / ZeRO-DP；而 PyTorch 的 DP/DDP 相当于每人背全套帐篷+炉子+斧头。从分片方式看，ZeRO 对权重的切法与张量并行神似（都是层内水平分片），与后述的"按层垂直切分"的 PP 形成对照。

在 Transformers 中启用 ZeRO

ZeRO 的完整实现来自 DeepSpeed（stage 1/2/3）。从源码看，Transformers 的接入点在 integrations/deepspeed.py：

• TrainingArguments.deepspeed接受DeepSpeed 配置 JSON 文件路径或 dict（见 training_args.py）：

  deepspeed: dict | str | None = field( default=None, metadata={"help": "Enable DeepSpeed integration. Value is a path to a JSON config file or a dict."}, )

• 创建TrainingArguments时会自动构造HfTrainerDeepSpeedConfig并用其中"auto"占位符同步TrainingArguments的取值（batch size、学习率等），保证两侧配置一致（见 HfTrainerDeepSpeedConfig）；

• 框架内部通过is_deepspeed_zero3_enabled()判断是否启用 ZeRO-3，模型加载路径会据此在from_pretrained时直接以分片形式落盘、避免临时占满显存（见 is_deepspeed_zero3_enabled）；

• ZeRO-3 下初始化权重也有专门路径initialize_weights_zero3：用deepspeed.zero.GatheredParameters在每个子模块上临时聚拢分片参数完成初始化，再交还分片（见 initialize_weights_zero3）。

因此启用 ZeRO 的实操是：安装 DeepSpeed 与 Accelerate，准备一个含zero_optimization（stage 1/2/3、offload 等）字段的 JSON 配置，然后以分布式方式启动脚本并在命令行传入--deepspeed ds_config.json，Trainer 会完成其余接线。更多细节可参考 Trainer 集成文档。

朴素模型并行（垂直 MP）与流水线并行（PP）

朴素模型并行（MP）的做法极简：用.to()把指定层搬到指定设备，数据流经哪层就自动跟到哪层的设备。文档称之为"垂直 MP"，因为相当于把模型图垂直切开：

=================== =================== || 0 | 1 | 2 | 3 | | 4 | 5 | 6 | 7 | =================== =================== gpu0 gpu1

数据从第 3 层跨到第 4 层时发生 GPU 间搬运：同节点内（同机）非常快，跨节点则开销陡增；最后一层算完后通常还要把数据送回第 0 层（或把标签送到最后一层）计算损失。其致命缺陷：

• 任一时刻除一张卡外其余卡全在等待——4 卡朴素 MP ≈ 单卡显存 ×4，计算资源被大量浪费。4 张 6GB 卡勉强容得下 1 张 24GB 卡能装的模型（还要扣掉数据拷贝开销），但训练速度远慢于单张大卡；
• 共享嵌入层可能需要在 GPU 间额外拷贝。

流水线并行（PP）与朴素 MP 结构相同，但通过把输入批切成 micro-batch 让各阶段"流水线"运转，消除卡空闲。两个核心概念：

• chunks（GAS）：同一阶段连续推送的数据块数，概念上等价于梯度累积步数（PyTorch 叫chunks，DeepSpeed 叫 GAS）。chunks=1时退化为低效的朴素 MP；chunks过大则 micro-batch 过小、利用率下降——需要实验找到让 GPU 利用最大化、"流水线气泡"（bubble，空闲时段）最小化的取值。
• micro-batch size (MBS)：DP 把全局批切成 mini-batch，PP 再把 mini-batch 切成 micro-batch。示例：DP 度 4、全局批 1024 → 4 个 256 的 mini-batch；chunks=32→ MBS = 256/32 = 8。PP 各阶段每次只处理一个 micro-batch。

DP+PP 下全局批的换算公式：global_batch = mbs × chunks × dp_degree（如8 × 32 × 4 = 1024）。

关于工程落地，文档区分了两类方案：

• 传统流水线 API：PyTorch、DeepSpeed、Megatron-LM。限制明显：必须把模型重写成nn.Sequential结构；接口只接受以 batch 为第一维的张量（或张量元组）输入输出；阶段内无法做条件控制流（如 T5 的编码器阶段需要特殊 hack）；还要把每层显式布置到设备以衔接输入输出。
• 更现代的方案：Varuna、SageMaker（AWS 专有）。据论文宣称，它们以远小的模型改动规避了上述问题；Varuna 还会用模拟搜索最优调度。DeepSpeed/Varuna/SageMaker 还采用交错式流水线（interleaved scheduling）——让反向路径优先，进一步压缩气泡。

文档同时给出了 Transformers 的立场（以文档写作时点为准）：当时核心库没有模型支持完整 PP，仅 GPT-2 与 T5 具备简单 MP 支持；主要障碍是模型无法机械转换为nn.Sequential、且所有输入必须是张量。社区项目 OSLO 则展示了基于 Transformers 的 PP/TP 实现（无需改写为nn.Sequential）。选择 PP 时应以目标版本中模型的device_map/ 并行支持现状为准，而不是照搬文档时点的结论。

张量并行（TP）：把一层切开，而不是把层搬走

这一节沿用 Megatron-LM 论文（《GPU 集群上高效的大语言模型训练》）的记法。Transformer 的关键构件可写成Y = GeLU(XA)：X输入、A权重矩阵、Y输出。由于矩阵乘法天然可按列拆分，把A按列切给 N 张 GPU 并行计算XA_1 … XA_N，得到 N 个输出分片Y_1 … Y_n，再各自独立过 GeLU 即可——整个 MLP 前段全程无需跨卡同步。Megatron 论文给出的分片处理图示进一步说明：只要把"需要全量输出"的合并点（AllReduce）放到必要的位置，任意深度的 MLP 都能保持分片状态。

多头自注意力更容易并行：它本来就由若干相互独立的 head 组成，本质已是并行的，只需把 head 分到不同 GPU 上。

两条工程红线：

• TP 需要极高带宽的互联，因此强烈建议不要跨节点做 TP。单节点 4 卡则 TP 度上限为 4；想要 TP 度 8，就需要一个至少 8 卡的节点；
• 术语对照：DeepSpeed 把同样的技术称为"tensor slicing"（张量切片）。

文档给出的 Transformers 当时状态：核心训练路径未内置 TP，推理侧可借助 parallelformers（当时仅推理）或 DeepSpeed-Inference 的 CUDA kernel 加速（支持 BERT、GPT-2、GPT-Neo）。需要说明，当前仓库版本的架构与文档写作时点已有差异，是否可用 TP 请以你使用的版本中模型与 Trainer 的实际能力为准。

组合并行：2D 与 3D

DP+PP：DeepSpeed 的经典 4 卡布局——DP rank 0 "看不见" GPU2，DP rank 1 "看不见" GPU3；对 DP 而言只有 GPU0/1 存在，而 GPU0 暗中把一部分负载经 PP 甩给 GPU2，GPU1 把 GPU3 拉入支援。每个维度至少 2 卡，故 2D 组合最少 4 卡。

DP+PP+TP（3D 并行）：在 DP×PP 的基础上，让每个 PP 阶段内部再用 TP 水平切分。每个维度仍需 ≥2 卡，故3D 并行最少 8 卡。实现方包括 DeepSpeed（内含 ZeRO-DP）、Megatron-LM、Varuna、SageMaker、OSLO；Microsoft 的《DeepSpeed 超大规模训练》博客是配套必读。

ZeRO + DP+PP+TP：ZeRO 通常单独使用即可，但与 PP/TP 叠加时有讲究——文档指出：

• ZeRO-DP 与 PP 组合时，通常只启用stage 1（优化器状态分片）；
• stage 2（梯度分片）不推荐与 PP 叠加：每个 micro-batch 都需要额外的 reduce-scatter 才能分片梯度，而 PP 恰恰依赖小 micro-batch 来压缩气泡，通信开销会直接侵蚀收益；且 PP 本身已把每卡梯度规模降到约 1/PP，再分片的边际收益有限；
• stage 3因引入更多节点间通信同样不适合；
• 附带红利：ZeRO-Offload可把 stage 1 优化器状态卸载到 CPU，进一步省显存。

代表实现：Megatron-DeepSpeed（含 BigScience 的 fork）、OSLO；关键论文是 Megatron-Turing NLG 530B 的训练报告（DeepSpeed + Megatron 组合）。

FlexFlow：让算法替你搜索并行方案

FlexFlow（论文《超越数据与模型并行》）走了一条不同的路——4D 并行：

1. Sample：样本维度（即数据并行），如 10×512 的批切成 5×2×512 分到两个设备；
2. Operator：把单个算子拆成子操作并行，如 LayerNorm 的 std 与 mean 同时算（先把输入复制到两台设备）；
3. Attribute：序列长度维度，如 10×512 切成 10×2×256；
4. Parameter：参数维度，即张量/层模型并行（水平或垂直皆可）。

框架的价值在于：把 GPU/TPU/CPU 对、RAM/DRAM 对、快内网/慢外网等资源全部作为变量，由算法自动决定在哪用哪种并行。代价与前提也要清楚：它面向静态、固定形态的 workload，对迭代行为动态变化的模型未必适配。工作流是在目标集群上跑约 30 分钟模拟，得到针对该硬件的最优并行规划；拓扑变更后重新规划，然后开训——不同部署各有其最优解。

何时用哪个策略：完整决策清单

文档末尾给出了覆盖面更全的落地清单（每条列表的靠前项通常更快）：

单 GPU

• 模型装得下：正常训练即可。
• 装不下：
  1. ZeRO + CPU offload（可选 NVMe offload）；
  2. 若最大单层也装不进单卡：在上述基础上启用Memory Centric Tiling（DeepSpeed ZeRO-3 文档中的特性，用于把超大单层切成块计算）。
• 最大单层装不进、又不走 ZeRO：必须启用 TP（PP 单独装不下）。

单节点 / 多 GPU

• 装得下：1) DDP 2) ZeRO（视配置而定）
• 装不下：1) PP 2) ZeRO 3) TP —— 节点内若有 NVLink/NVSwitch 三者接近，否则 PP 更快；TP 度数也有影响，实验定胜负；TP 基本限于节点内（TP 尺寸 ≤ 单节点卡数）
• 最大单层装不进：不走 ZeRO 则必须 TP；走 ZeRO 则回看"单 GPU"条目。

多节点 / 多 GPU

• 高速节点间互联：1) ZeRO（几乎不改模型）2) PP+TP+DP（通信少，但要大改模型）
• 慢互联且显存紧张：DP+PP+TP+ZeRO-1

从概念到命令行：在 Transformers 中落地

把上述策略落到 Transformers 工程上，核心抓手有三个：

1. DDP：什么都不用改，用torchrun --nproc_per_node N（多机则加--nnodes/--node_rank/--master_addr）启动任意run_*.py脚本，Trainer 自动包装 DDP。可用--ddp_backend指定后端、按 training_args.py 的字段说明调整ddp_find_unused_parameters与 bucket 上限。
2. ZeRO：pip install deepspeed accelerate后，编写 DeepSpeed JSON 配置，脚本加--deepspeed ds_config.json启动；"auto"占位符由HfTrainerDeepSpeedConfig自动对齐到TrainingArguments，避免两处配置打架。
3. 模型切分（垂直 MP 的近似替代）：对"模型太大但暂时不想上 PP/TP"的场景，from_pretrained(..., device_map="auto")按层把模型铺满多卡，配合max_memory控制每卡配额；训练场景下再结合梯度累积、混合精度等 单 GPU 篇 的通用手段。

最后回到文档的核心提醒：没有放之四海皆准的并行配方——同一组 GPU，NVLink 的有无、TP/PP 的度数、chunks与 micro-batch 的取值都会显著改变结果（前文 DDP 基准从 101s 到 131s 的摆动就是最直观的例子）。先按决策清单锁定候选方案，再在自己的硬件上跑小规模实验，是 Transformers 多卡训练最高效的收敛路径。

更多文档与入口：性能篇总览、DeepSpeed 集成说明、Trainer 主类文档、语言建模示例。

【免费下载链接】transformers🤗 Transformers: the model-definition framework for state-of-the-art machine learning models in text, vision, audio, and multimodal models, for both inference and training.项目地址: https://gitcode.com/GitHub_Trending/tra/transformers