随着大模型、计算机视觉和自然语言处理任务的规模不断扩大，单张 GPU 已经很难满足模型训练需求。模型参数量从百万级增长到十亿级、百亿级，训练数据也从几 GB 增长到 TB 甚至 PB 级别。此时，仅仅提升单卡显存和计算能力已经无法解决问题，分布式训练逐渐成为深度学习工程中的核心技术。

分布式训练的本质，是将模型训练任务拆分到多张 GPU 或多台服务器上，让多个计算设备共同完成数据读取、前向计算、反向传播和参数更新。

本文将围绕以下内容展开：

• 分布式训练的基本原理
• 数据并行、模型并行与流水线并行
• PyTorch DistributedDataParallel 实战
• TensorFlow MultiWorkerMirroredStrategy 实战
• Horovod 分布式训练实战
• 通信机制、性能优化与常见问题
• 三种框架的选型建议

---

一、为什么需要分布式训练

假设一个模型训练过程如下：

读取数据
    ↓
前向传播
    ↓
计算损失
    ↓
反向传播
    ↓
更新参数

在单卡训练中，所有操作都由一张 GPU 完成。对于一个简单模型，单卡训练没有问题：

for images, labels in dataloader:
    outputs = model(images)
    loss = criterion(outputs, labels)

    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

但当模型和数据规模变大之后，通常会遇到三个问题。

1. 显存不足

模型参数、梯度、优化器状态和中间激活值都需要占用显存。

以 Adam 优化器为例，每一个参数通常需要保存：

• 模型参数
• 梯度
• 一阶动量
• 二阶动量

如果使用 FP32，每个参数实际占用的显存可能达到 12 到 16 字节甚至更多。

一个拥有 10 亿参数的模型，仅参数和优化器状态就可能占用几十 GB 显存，单张常见 GPU 无法容纳。

2. 训练时间过长

如果单卡训练一个模型需要 10 天，那么简单地提高 batch size 并不能有效解决问题。通过多张 GPU 并行处理数据，可以显著缩短训练时间。

3. 数据规模不断扩大

单机单卡的数据读取、预处理和计算吞吐有限。当数据量达到数百 GB 或 TB 时，单卡很容易成为整个系统的瓶颈。

---

二、分布式训练的基本概念

1. 进程、设备与节点

在分布式训练中，需要区分以下几个概念：

• Node：一台服务器
• GPU：服务器上的计算设备
• Process：执行训练代码的进程
• Rank：进程编号
• World Size：所有进程总数
• Local Rank：进程在当前节点上的编号

例如，使用两台服务器，每台服务器有 4 张 GPU：

节点 0:
    GPU 0 -> Rank 0
    GPU 1 -> Rank 1
    GPU 2 -> Rank 2
    GPU 3 -> Rank 3

节点 1:
    GPU 0 -> Rank 4
    GPU 1 -> Rank 5
    GPU 2 -> Rank 6
    GPU 3 -> Rank 7

此时：

World Size = 8

通常情况下，每个 GPU 对应一个训练进程。

2. 数据并行

数据并行是最常见的分布式训练方式。

假设全局 batch size 为 256，使用 8 张 GPU，则每张 GPU 处理：

Local Batch Size = 256 / 8 = 32

每个进程拥有一份完整模型副本，但每个进程读取不同的数据：

GPU 0: Batch 0
GPU 1: Batch 1
GPU 2: Batch 2
...
GPU 7: Batch 7

完成反向传播后，所有进程之间进行梯度同步，随后每个进程使用相同的梯度更新本地模型。

数据并行的核心流程如下：

每个进程:
    1. 获取不同的数据
    2. 使用本地模型进行前向计算
    3. 计算本地梯度
    4. 对所有进程的梯度进行 AllReduce
    5. 使用同步后的梯度更新参数

3. 模型并行

模型并行将一个模型拆分到多张 GPU 上。

例如：

GPU 0: Transformer 第 1～6 层
GPU 1: Transformer 第 7～12 层
GPU 2: Transformer 第 13～18 层
GPU 3: Transformer 第 19～24 层

模型并行适合单卡无法容纳整个模型的场景，但会产生较高的设备间通信开销。

4. 流水线并行

流水线并行将模型拆成多个阶段，并将一个 batch 拆分成多个 micro-batch：

Stage 1 -> Stage 2 -> Stage 3 -> Stage 4

当 Stage 1 处理第 2 个 micro-batch 时，Stage 2 可以处理第 1 个 micro-batch，从而提高设备利用率。

在实际的大模型训练中，通常会组合使用：

数据并行 + 张量并行 + 流水线并行

---

三、PyTorch DistributedDataParallel 实战

PyTorch 提供了 torch.distributed 作为底层分布式通信接口，同时提供了更高级的 DistributedDataParallel，简称 DDP。

在单机多卡场景下，DDP 通常是最推荐的方案。

1. 为什么优先使用 DDP

PyTorch 早期提供过 DataParallel：

model = torch.nn.DataParallel(model)

它的工作方式是：

1. 主进程接收输入数据
2. 主进程将数据切分到多个 GPU
3. 多个 GPU 计算完成后，将结果返回主 GPU
4. 主 GPU 汇总并更新参数

这种方式存在明显问题：

• 主 GPU 负责数据分发和结果汇总
• 主 GPU 容易成为瓶颈
• 只有一个进程，无法充分利用多进程并行
• GPU 间通信效率不高

DDP 则为每张 GPU 启动一个独立进程，每个进程负责一张 GPU，通信效率更高。

---

2. 初始化分布式环境

一个典型的 DDP 初始化函数如下：

import os
import torch
import torch.distributed as dist


def setup_distributed():
    local_rank = int(os.environ["LOCAL_RANK"])

    torch.cuda.set_device(local_rank)

    dist.init_process_group(
        backend="nccl",
        init_method="env://"
    )

    return local_rank

这里有几个关键点。

LOCAL_RANK

LOCAL_RANK 表示当前进程在本机上的 GPU 编号。

例如：

LOCAL_RANK=0

表示当前进程使用第 0 张 GPU。

backend

常用通信后端包括：

• nccl：NVIDIA GPU 场景下的首选
• gloo：CPU 或部分调试场景
• mpi：依赖 MPI 环境

在 GPU 训练中，通常使用：

backend="nccl"

因为 NCCL 针对 NVIDIA GPU 通信进行了专门优化，能够利用 PCIe、NVLink 和 InfiniBand 等硬件。

---

3. 使用 DistributedSampler 切分数据

如果所有进程都读取完整数据集，那么每张 GPU 会重复处理相同样本，无法获得真正的并行效果。

因此，需要使用 DistributedSampler：

from torch.utils.data import DataLoader
from torch.utils.data.distributed import DistributedSampler


sampler = DistributedSampler(
    dataset,
    num_replicas=dist.get_world_size(),
    rank=dist.get_rank(),
    shuffle=True
)

dataloader = DataLoader(
    dataset,
    batch_size=32,
    sampler=sampler,
    num_workers=4,
    pin_memory=True
)

每个进程只会读取数据集中的一部分。

需要特别注意的是，每个 epoch 开始时必须调用：

sampler.set_epoch(epoch)

完整示例：

for epoch in range(num_epochs):
    sampler.set_epoch(epoch)

    for images, labels in dataloader:
        images = images.cuda(local_rank, non_blocking=True)
        labels = labels.cuda(local_rank, non_blocking=True)

        outputs = model(images)
        loss = criterion(outputs, labels)

        optimizer.zero_grad(set_to_none=True)
        loss.backward()
        optimizer.step()

如果不调用 set_epoch，不同 epoch 的 shuffle 顺序可能保持不变，从而影响训练效果。

---

4. 完整的 PyTorch DDP 示例

下面以图像分类任务为例，给出一个可运行的 DDP 训练框架。

import os
import random
import numpy as np

import torch
import torch.nn as nn
import torch.distributed as dist

from torch.utils.data import DataLoader
from torch.utils.data.distributed import DistributedSampler
from torch.nn.parallel import DistributedDataParallel as DDP
from torchvision import datasets, transforms, models


def seed_everything(seed):
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)


def setup():
    local_rank = int(os.environ["LOCAL_RANK"])

    torch.cuda.set_device(local_rank)

    dist.init_process_group(
        backend="nccl",
        init_method="env://"
    )

    return local_rank


def cleanup():
    dist.destroy_process_group()


def build_model(num_classes):
    model = models.resnet18(weights=None)
    model.fc = nn.Linear(model.fc.in_features, num_classes)
    return model


def train():
    local_rank = setup()
    rank = dist.get_rank()

    seed_everything(42 + rank)

    transform = transforms.Compose([
        transforms.Resize((224, 224)),
        transforms.ToTensor(),
        transforms.Normalize(
            mean=[0.485, 0.456, 0.406],
            std=[0.229, 0.224, 0.225]
        )
    ])

    dataset = datasets.ImageFolder(
        root="./data/train",
        transform=transform
    )

    sampler = DistributedSampler(
        dataset=dataset,
        num_replicas=dist.get_world_size(),
        rank=rank,
        shuffle=True
    )

    dataloader = DataLoader(
        dataset,
        batch_size=64,
        sampler=sampler,
        num_workers=4,
        pin_memory=True,
        persistent_workers=True
    )

    model = build_model(num_classes=len(dataset.classes))
    model = model.cuda(local_rank)

    model = DDP(
        model,
        device_ids=[local_rank],
        output_device=local_rank,
        broadcast_buffers=False
    )

    criterion = nn.CrossEntropyLoss().cuda(local_rank)

    optimizer = torch.optim.AdamW(
        model.parameters(),
        lr=1e-3,
        weight_decay=1e-4
    )

    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(
        optimizer,
        T_max=50
    )

    num_epochs = 50

    for epoch in range(num_epochs):
        model.train()
        sampler.set_epoch(epoch)

        total_loss = 0.0
        total_samples = 0

        for images, labels in dataloader:
            images = images.cuda(local_rank, non_blocking=True)
            labels = labels.cuda(local_rank, non_blocking=True)

            optimizer.zero_grad(set_to_none=True)

            outputs = model(images)
            loss = criterion(outputs, labels)

            loss.backward()
            optimizer.step()

            total_loss += loss.item() * images.size(0)
            total_samples += images.size(0)

        scheduler.step()

        loss_tensor = torch.tensor(
            [total_loss, total_samples],
            dtype=torch.float64,
            device=local_rank
        )

        dist.all_reduce(
            loss_tensor,
            op=dist.ReduceOp.SUM
        )

        global_loss = loss_tensor[0] / loss_tensor[1]

        if rank == 0:
            print(
                f"Epoch [{epoch + 1}/{num_epochs}], "
                f"loss={global_loss.item():.6f}, "
                f"lr={scheduler.get_last_lr()[0]:.8f}"
            )

            torch.save(
                {
                    "epoch": epoch,
                    "model": model.module.state_dict(),
                    "optimizer": optimizer.state_dict()
                },
                "./checkpoint.pt"
            )

    cleanup()


if __name__ == "__main__":
    train()

这里有一个重要细节：

model.module.state_dict()

因为 DDP 包装后的模型结构是：

DDP
 └── module
      └── 原始模型

保存模型时使用 model.module.state_dict()，加载时可以直接加载到原始模型中。

---

5. 启动单机多卡训练

假设服务器上有 4 张 GPU，可以使用：

torchrun \
    --standalone \
    --nproc_per_node=4 \
    train_ddp.py

参数含义如下：

• --standalone：单机独立运行
• --nproc_per_node=4：启动 4 个进程
• 每个进程绑定一张 GPU

也可以使用：

CUDA_VISIBLE_DEVICES=0,1,2,3 torchrun \
    --standalone \
    --nproc_per_node=4 \
    train_ddp.py

---

6. 多机多卡训练

假设有两台服务器：

节点 0：192.168.1.10
节点 1：192.168.1.11

每台服务器拥有 4 张 GPU。

在节点 0 上执行：

torchrun \
    --nnodes=2 \
    --nproc_per_node=4 \
    --node_rank=0 \
    --master_addr=192.168.1.10 \
    --master_port=29500 \
    train_ddp.py

在节点 1 上执行：

torchrun \
    --nnodes=2 \
    --nproc_per_node=4 \
    --node_rank=1 \
    --master_addr=192.168.1.10 \
    --master_port=29500 \
    train_ddp.py

整个训练任务中：

World Size = 2 × 4 = 8

必须确保：

• 两台机器的代码版本一致
• 数据路径一致
• Python 环境一致
• 防火墙允许通信端口
• GPU 驱动和 CUDA 版本兼容
• 节点之间网络连通

---

四、TensorFlow 分布式训练

TensorFlow 主要通过 tf.distribute.Strategy 实现分布式训练。

常见策略包括：

• MirroredStrategy
• MultiWorkerMirroredStrategy
• TPUStrategy
• ParameterServerStrategy

1. MirroredStrategy

MirroredStrategy 适用于单机多 GPU。

它会在每张 GPU 上创建一个模型副本，并自动同步梯度。

import tensorflow as tf


strategy = tf.distribute.MirroredStrategy()

print("设备数量:", strategy.num_replicas_in_sync)

with strategy.scope():
    model = tf.keras.Sequential([
        tf.keras.layers.Input(shape=(224, 224, 3)),
        tf.keras.layers.Conv2D(32, 3, activation="relu"),
        tf.keras.layers.MaxPooling2D(),
        tf.keras.layers.Conv2D(64, 3, activation="relu"),
        tf.keras.layers.GlobalAveragePooling2D(),
        tf.keras.layers.Dense(10, activation="softmax")
    ])

    model.compile(
        optimizer=tf.keras.optimizers.Adam(learning_rate=1e-3),
        loss="sparse_categorical_crossentropy",
        metrics=["accuracy"]
    )

训练代码与普通 Keras 基本一致：

model.fit(
    train_dataset,
    epochs=20,
    validation_data=val_dataset
)

MirroredStrategy 会自动完成：

1. 数据分发
2. 多 GPU 前向计算
3. 梯度同步
4. 参数更新
5. 指标聚合

---

2. TensorFlow 数据集准备

使用 TensorFlow 分布式训练时，需要设置合适的全局 batch size：

num_replicas = strategy.num_replicas_in_sync
per_replica_batch_size = 32
global_batch_size = num_replicas * per_replica_batch_size

构建数据集：

train_dataset = tf.keras.utils.image_dataset_from_directory(
    "./data/train",
    image_size=(224, 224),
    batch_size=global_batch_size,
    shuffle=True
)

train_dataset = train_dataset.prefetch(
    tf.data.AUTOTUNE
)

与 PyTorch 不同，TensorFlow 的分布式策略通常会自动处理 batch 的分发，但开发者仍然需要理解全局 batch size 和单设备 batch size 的区别。

---

3. MultiWorkerMirroredStrategy

如果需要多机训练，可以使用：

strategy = tf.distribute.MultiWorkerMirroredStrategy()

TensorFlow 需要通过 TF_CONFIG 描述集群拓扑。

例如，节点 0 的配置：

{
  "cluster": {
    "worker": [
      "192.168.1.10:12345",
      "192.168.1.11:12345"
    ]
  },
  "task": {
    "type": "worker",
    "index": 0
  }
}

节点 1 的配置：

{
  "cluster": {
    "worker": [
      "192.168.1.10:12345",
      "192.168.1.11:12345"
    ]
  },
  "task": {
    "type": "worker",
    "index": 1
  }
}

在 Linux 环境中可以这样设置：

export TF_CONFIG='{
  "cluster": {
    "worker": [
      "192.168.1.10:12345",
      "192.168.1.11:12345"
    ]
  },
  "task": {
    "type": "worker",
    "index": 0
  }
}'

Python 代码：

import tensorflow as tf


strategy = tf.distribute.MultiWorkerMirroredStrategy()

with strategy.scope():
    model = create_model()

    model.compile(
        optimizer=tf.keras.optimizers.Adam(),
        loss=tf.keras.losses.SparseCategoricalCrossentropy(),
        metrics=["accuracy"]
    )

model.fit(
    train_dataset,
    epochs=20
)

---

4. TensorFlow 自定义训练循环

在复杂任务中，例如检测、分割、大模型训练，通常需要自定义训练循环。

import tensorflow as tf


strategy = tf.distribute.MirroredStrategy()

with strategy.scope():
    model = create_model()
    optimizer = tf.keras.optimizers.Adam()
    loss_object = tf.keras.losses.SparseCategoricalCrossentropy(
        from_logits=True,
        reduction=tf.keras.losses.Reduction.NONE
    )


def compute_loss(labels, predictions):
    per_example_loss = loss_object(labels, predictions)
    return tf.nn.compute_average_loss(
        per_example_loss,
        global_batch_size=GLOBAL_BATCH_SIZE
    )


@tf.function
def distributed_train_step(dataset_inputs):
    def train_step(inputs):
        images, labels = inputs

        with tf.GradientTape() as tape:
            predictions = model(images, training=True)
            loss = compute_loss(labels, predictions)

        gradients = tape.gradient(
            loss,
            model.trainable_variables
        )

        optimizer.apply_gradients(
            zip(gradients, model.trainable_variables)
        )

        return loss

    per_replica_losses = strategy.run(
        train_step,
        args=(dataset_inputs,)
    )

    return strategy.reduce(
        tf.distribute.ReduceOp.SUM,
        per_replica_losses,
        axis=None
    )

这里最重要的是：

strategy.run(...)

它负责在每个副本上运行训练函数。

strategy.reduce(...)

它负责对多个设备上的结果进行聚合。

---

五、Horovod 分布式训练

Horovod 是 Uber 开源的分布式深度学习框架，支持：

• TensorFlow
• PyTorch
• Keras
• MXNet

Horovod 的核心理念是将 MPI 的通信能力应用到深度学习训练中。它通过 AllReduce 同步各个进程的梯度。

1. Horovod 的基本工作原理

训练流程如下：

进程 0: 计算梯度 g0
进程 1: 计算梯度 g1
进程 2: 计算梯度 g2
进程 3: 计算梯度 g3

AllReduce:
g = (g0 + g1 + g2 + g3) / 4

所有进程使用 g 更新模型

这样可以保证每个进程上的模型参数保持一致。

---

2. Horovod PyTorch 示例

首先导入 Horovod：

import horovod.torch as hvd

初始化：

hvd.init()

local_rank = hvd.local_rank()
torch.cuda.set_device(local_rank)
device = torch.device("cuda", local_rank)

构建模型：

model = create_model().to(device)

optimizer = torch.optim.Adam(
    model.parameters(),
    lr=1e-3 * hvd.size()
)

这里对学习率乘以进程数量，是因为全局 batch size 增大之后，通常需要采用线性学习率缩放策略。

使用 Horovod 包装优化器：

optimizer = hvd.DistributedOptimizer(
    optimizer,
    named_parameters=model.named_parameters()
)

广播初始参数：

hvd.broadcast_parameters(
    model.state_dict(),
    root_rank=0
)

hvd.broadcast_optimizer_state(
    optimizer,
    root_rank=0
)

完整训练代码：

import torch
import torch.nn as nn
import horovod.torch as hvd

from torch.utils.data import DataLoader
from torch.utils.data.distributed import DistributedSampler


def train():
    hvd.init()

    local_rank = hvd.local_rank()
    rank = hvd.rank()
    world_size = hvd.size()

    torch.cuda.set_device(local_rank)
    device = torch.device("cuda", local_rank)

    dataset = MyDataset("./data/train")

    sampler = DistributedSampler(
        dataset,
        num_replicas=world_size,
        rank=rank,
        shuffle=True
    )

    dataloader = DataLoader(
        dataset,
        batch_size=32,
        sampler=sampler,
        num_workers=4,
        pin_memory=True
    )

    model = MyModel().to(device)

    base_lr = 1e-3
    optimizer = torch.optim.Adam(
        model.parameters(),
        lr=base_lr * world_size
    )

    optimizer = hvd.DistributedOptimizer(
        optimizer,
        named_parameters=model.named_parameters()
    )

    criterion = nn.CrossEntropyLoss().to(device)

    hvd.broadcast_parameters(
        model.state_dict(),
        root_rank=0
    )

    hvd.broadcast_optimizer_state(
        optimizer,
        root_rank=0
    )

    for epoch in range(20):
        model.train()
        sampler.set_epoch(epoch)

        for images, labels in dataloader:
            images = images.to(device, non_blocking=True)
            labels = labels.to(device, non_blocking=True)

            optimizer.zero_grad()

            outputs = model(images)
            loss = criterion(outputs, labels)

            loss.backward()
            optimizer.step()

        if rank == 0:
            print(f"epoch={epoch}, loss={loss.item():.6f}")


if __name__ == "__main__":
    train()

启动命令：

horovodrun -np 4 -H localhost:4 python train_horovod.py

也可以使用 MPI：

mpirun -np 4 \
    -H localhost:4 \
    -bind-to none \
    -map-by slot \
    -x NCCL_DEBUG=INFO \
    python train_horovod.py

---

3. Horovod TensorFlow 示例

import tensorflow as tf
import horovod.tensorflow as hvd


hvd.init()

gpus = tf.config.list_physical_devices("GPU")

if gpus:
    tf.config.set_visible_devices(
        gpus[hvd.local_rank()],
        "GPU"
    )

strategy = tf.distribute.MirroredStrategy()

with strategy.scope():
    model = create_model()

    optimizer = tf.keras.optimizers.Adam(
        learning_rate=0.001 * hvd.size()
    )

    optimizer = hvd.DistributedOptimizer(
        optimizer
    )

    model.compile(
        optimizer=optimizer,
        loss="sparse_categorical_crossentropy",
        metrics=["accuracy"]
    )

callbacks = [
    hvd.callbacks.BroadcastGlobalVariablesCallback(
        root_rank=0
    ),
    hvd.callbacks.MetricAverageCallback()
]

if hvd.rank() == 0:
    callbacks.append(
        tf.keras.callbacks.ModelCheckpoint(
            "checkpoint.h5",
            save_best_only=True
        )
    )

model.fit(
    train_dataset,
    epochs=20,
    callbacks=callbacks
)

Horovod 中需要特别注意回调函数：

hvd.callbacks.BroadcastGlobalVariablesCallback(
    root_rank=0
)

该回调会将第 0 个进程的初始变量广播给其他进程，保证所有进程从相同参数开始训练。

---

六、梯度同步与通信机制

分布式训练的计算部分通常不是最大瓶颈，真正影响训练速度的往往是通信。

1. AllReduce

AllReduce 是分布式训练中最关键的通信操作之一。

它的功能是：

1. 收集所有进程的数据
2. 执行求和或求平均
3. 将结果发送回所有进程

以 4 个进程为例：

Rank 0: g0
Rank 1: g1
Rank 2: g2
Rank 3: g3

AllReduce 后:

Rank 0: (g0 + g1 + g2 + g3) / 4
Rank 1: (g0 + g1 + g2 + g3) / 4
Rank 2: (g0 + g1 + g2 + g3) / 4
Rank 3: (g0 + g1 + g2 + g3) / 4

DDP 通常会将梯度划分为多个 bucket，并在反向传播过程中逐步执行通信，从而实现计算与通信重叠。

2. Ring AllReduce

Ring AllReduce 将所有进程组织成环：

Rank 0 -> Rank 1 -> Rank 2 -> Rank 3 -> Rank 0

整个操作分为两个阶段：

• Reduce-Scatter
• AllGather

它不需要一个中心节点，因此可以降低单点通信压力。

---

七、分布式训练中的关键性能优化

1. 正确设置 batch size

设：

单卡 batch size = 32
GPU 数量 = 8

那么全局 batch size 为：

Global Batch Size = 32 × 8 = 256

如果训练代码中将学习率从 0.001 调整为：

lr = 0.001 * world_size

则需要观察训练稳定性。

学习率线性缩放并不是绝对规则。在以下情况下可能需要谨慎：

• 模型使用 LayerNorm
• 训练 batch 较小
• 使用 AdamW
• 使用大规模语言模型
• 训练初期梯度不稳定

通常可以配合 warmup：

def adjust_learning_rate(
    optimizer,
    base_lr,
    epoch,
    warmup_epochs,
    world_size
):
    if epoch < warmup_epochs:
        lr = (
            base_lr
            * world_size
            * (epoch + 1)
            / warmup_epochs
        )
    else:
        lr = base_lr * world_size

    for param_group in optimizer.param_groups:
        param_group["lr"] = lr

2. 使用混合精度

混合精度可以使用 FP16 或 BF16 完成部分计算，从而降低显存和提升吞吐。

PyTorch 示例：

scaler = torch.cuda.amp.GradScaler()

for images, labels in dataloader:
    optimizer.zero_grad(set_to_none=True)

    with torch.cuda.amp.autocast():
        outputs = model(images)
        loss = criterion(outputs, labels)

    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()

新版本 PyTorch 也可以使用：

with torch.autocast(
    device_type="cuda",
    dtype=torch.bfloat16
):
    outputs = model(images)
    loss = criterion(outputs, labels)

混合精度的优势包括：

• 减少显存占用
• 提升 Tensor Core 利用率
• 提高矩阵乘法吞吐
• 支持更大的 batch size

但需要注意数值稳定性，尤其是：

• Softmax
• Loss 计算
• 梯度归一化
• 指数运算
• 归一化层

3. 使用梯度累积

当显存无法容纳较大的 batch 时，可以使用梯度累积：

accumulation_steps = 4

optimizer.zero_grad(set_to_none=True)

for step, (images, labels) in enumerate(dataloader):
    images = images.cuda()
    labels = labels.cuda()

    with torch.autocast(
        device_type="cuda",
        dtype=torch.float16
    ):
        outputs = model(images)
        loss = criterion(outputs, labels)
        loss = loss / accumulation_steps

    scaler.scale(loss).backward()

    if (step + 1) % accumulation_steps == 0:
        scaler.step(optimizer)
        scaler.update()
        optimizer.zero_grad(set_to_none=True)

有效 batch size 近似为：

Effective Batch Size
= 单卡 Batch Size
× GPU 数量
× 梯度累积步数

4. 优化数据读取

如果 GPU 利用率很低，问题可能不在模型，而在数据加载。

建议：

DataLoader(
    dataset,
    batch_size=64,
    num_workers=8,
    pin_memory=True,
    persistent_workers=True,
    prefetch_factor=2
)

将数据移动到 GPU 时：

images = images.cuda(
    non_blocking=True
)

同时可以使用：

• LMDB
• WebDataset
• TFRecord
• Parquet
• Arrow
• 内存映射文件
• 本地 NVMe SSD

对于大规模训练，随机读取网络存储通常会显著降低吞吐。

5. 只让主进程保存文件

分布式训练中，每个进程都会执行同一份代码。如果所有进程都保存模型，会导致文件覆盖和 I/O 冲突。

PyTorch 中：

if dist.get_rank() == 0:
    torch.save(
        model.module.state_dict(),
        "model.pt"
    )

TensorFlow 中：

if hvd.rank() == 0:
    model.save("model.h5")

日志、评估和可视化也通常只在主进程执行。

---

八、常见问题与排查方法

1. 程序启动后一直卡住

常见原因：

• 不同进程的 world_size 不一致
• 多机之间端口不通
• MASTER_ADDR 设置错误
• MASTER_PORT 被占用
• 某个进程提前退出
• NCCL 网络配置错误

可以增加 NCCL 日志：

set NCCL_DEBUG=INFO

Linux 环境：

export NCCL_DEBUG=INFO

也可以测试：

python -c "import torch; print(torch.cuda.device_count())"

2. 训练速度没有线性提升

如果使用 8 张 GPU，理论上训练速度可能接近单卡的 8 倍，但实际往往只有 5～7 倍。

原因包括：

• GPU 之间通信成本
• 数据读取速度不足
• batch size 过小
• 模型计算量太少
• CPU 预处理成为瓶颈
• PCIe 带宽不足
• 网络带宽不足
• 梯度同步等待时间过长

可以通过 Nsight Systems、PyTorch Profiler 或 TensorBoard Profiler 定位瓶颈。

PyTorch Profiler 示例：

from torch.profiler import profile, record_function
from torch.profiler import ProfilerActivity


with profile(
    activities=[
        ProfilerActivity.CPU,
        ProfilerActivity.CUDA
    ],
    record_shapes=True,
    profile_memory=True
) as prof:
    with record_function("train_step"):
        outputs = model(images)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()

print(
    prof.key_averages().table(
        sort_by="cuda_time_total",
        row_limit=20
    )
)

3. 不同进程的 loss 不一致

本地 loss 不一致并不一定是错误，因为每个进程处理的数据不同。

如果希望得到全局平均 loss，需要进行聚合：

loss_tensor = torch.tensor(
    [loss_sum, sample_count],
    device=device,
    dtype=torch.float64
)

dist.all_reduce(
    loss_tensor,
    op=dist.ReduceOp.SUM
)

global_loss = loss_tensor[0] / loss_tensor[1]

不能简单地直接平均进程平均值，因为不同进程处理的样本数量可能不同。

4. 验证集被重复计算

在 DDP 中，如果每个进程都完整遍历验证集，会造成重复计算。

可以使用 DistributedSampler：

val_sampler = DistributedSampler(
    val_dataset,
    num_replicas=dist.get_world_size(),
    rank=dist.get_rank(),
    shuffle=False
)

但要注意最后一个 batch 的补齐策略，必要时需要自定义验证集切分逻辑，避免样本重复统计。

---

九、PyTorch、TensorFlow 与 Horovod 如何选择

PyTorch DDP

适合：

• PyTorch 项目
• 自定义训练循环
• 目标检测、图像分割、NLP
• 大模型训练
• 对底层控制要求较高的任务

优势：

• 性能稳定
• API 清晰
• 调试相对方便
• 社区生态丰富
• 与 PyTorch 模型兼容性好

TensorFlow Strategy

适合：

• Keras 项目
• 生产级 TensorFlow 工作流
• TPU 训练
• 需要高度封装的工程
• 使用 tf.data 数据管道的项目

优势：

• API 集成度高
• Keras 使用方便
• 多种硬件策略统一
• 适合快速搭建标准训练流程

Horovod

适合：

• 已有 MPI 集群
• 需要统一支持 PyTorch 和 TensorFlow
• 现有项目已经采用 Horovod
• 多机通信基础设施比较完善

优势：

• 框架适配范围广
• 通信逻辑清晰
• 迁移单卡代码相对容易
• 适合跨框架集群环境

不过，在新建 PyTorch 项目时，通常优先考虑原生 DDP；在新建 TensorFlow 项目时，优先考虑官方 Strategy。

---

十、一个完整的分布式训练检查清单

在正式启动大规模训练前，可以按照以下步骤检查。

环境检查

python --version
nvidia-smi
python -c "import torch; print(torch.__version__)"
python -c "import torch; print(torch.cuda.device_count())"

GPU 检查

确认：

• 每张 GPU 都可以正常使用
• GPU 显存大小一致
• CUDA 版本兼容
• 驱动版本满足要求
• GPU 之间通信正常

数据检查

确认：

• 不同进程不会读取完全相同的数据
• DistributedSampler 配置正确
• 每个 epoch 调用 set_epoch
• 训练集和验证集划分正确
• 数据加载不会成为瓶颈

模型检查

确认：

• 模型已经移动到正确 GPU
• DDP 包装发生在模型移动之后
• 保存模型时使用 model.module
• BatchNorm 在多卡环境下是否需要 SyncBatchNorm

如果使用同步 BatchNorm：

model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(
    model
)

训练检查

确认：

• 全局 batch size 设置合理
• 学习率是否根据 GPU 数量调整
• 是否需要 warmup
• 是否启用混合精度
• 是否只让 Rank 0 保存模型
• 是否只让 Rank 0 输出日志

---

总结

分布式训练并不是简单地将一张 GPU 扩展到多张 GPU，而是一个涉及计算、通信、数据和工程管理的系统性问题。

从技术原理看：

• 数据并行解决了数据规模和训练速度问题
• 模型并行解决了单卡显存不足问题
• 流水线并行提高了大模型多阶段计算效率
• AllReduce 保证了多个进程之间的梯度一致性
• NCCL 和 MPI 决定了设备之间的通信效率

从框架选择看：

• PyTorch 项目优先使用 DistributedDataParallel
• TensorFlow 项目优先使用 tf.distribute
• 跨框架或 MPI 环境可以考虑 Horovod

一个高性能的分布式训练系统，需要同时关注以下几个指标：

GPU 利用率
数据读取吞吐
显存占用
通信时间
单步训练时间
扩展效率
最终模型效果

理想情况下，增加 GPU 数量后，训练时间应明显下降，同时模型收敛效果保持稳定。但在实际工程中，GPU 数量越多，通信和数据处理的影响越明显。因此，分布式训练的核心目标并不是“使用更多 GPU”，而是在计算效率、通信开销、显存消耗和训练效果之间取得平衡。

对于大多数 PyTorch 项目而言，建议从单机多卡 DDP 开始，先保证模型、数据和指标计算正确，再逐步扩展到多机训练。在此基础上，结合混合精度、梯度累积、数据预取和通信优化，才能真正构建稳定、高效、可扩展的分布式训练系统。