大模型分布式训练详解：从数据并行、模型并行到DeepSpeed实战

艺�

302人浏览 · 2026-11-10 09:38:53

艺� · 2026-11-10 09:38:53 发布

一、为什么大模型需要分布式训练？

随着大语言模型（LLM）参数规模不断增长，模型训练对 GPU 显存、计算能力和通信带宽提出了越来越高的要求。

早期的深度学习模型通常可以在单张 GPU 上完成训练，但是对于拥有数十亿、数百亿甚至数千亿参数的大模型来说，单卡训练已经很难满足需求。

例如，一个拥有 70B 参数的模型，如果使用 FP16 精度，仅模型参数就需要：

$[ 70\times10^9\times2=140GB ]$

这还没有计算梯度、优化器状态以及中间激活值所占用的显存。

如果使用 Adam 优化器，训练时还需要保存梯度和优化器状态，实际显存需求会进一步增加。

因此，大模型训练通常需要将计算任务分配到多张 GPU，甚至多个服务器上共同完成。

这就是分布式训练（Distributed Training）。

分布式训练的核心思想：将一个大模型的训练任务拆分成多个子任务，让多张 GPU 协同完成计算，从而提高训练效率、突破单卡资源限制。

二、大模型分布式训练整体架构

大模型分布式训练通常由数据集、训练任务调度、GPU 集群、并行计算策略、梯度同步和参数更新等模块组成。

整体训练流程如下：

从架构上看，分布式训练主要解决三个问题：

显存问题： 如何把模型参数、梯度和优化器状态分配到多张 GPU？
计算问题： 如何让多张 GPU 同时参与训练，提高计算效率？
通信问题： 如何在不同 GPU 之间高效交换梯度和参数信息？

围绕这三个问题，逐渐发展出了数据并行、张量并行、流水线并行以及 ZeRO 等技术。

三、大模型分布式训练的核心并行技术

3.1 数据并行（Data Parallelism）

数据并行是最基础、应用最广泛的分布式训练方式。

它的核心思想是：

每张 GPU 保存一份完整的模型副本，但处理不同的数据。

例如，使用 4 张 GPU 训练一个 Transformer 模型：

GPU 0：处理 Batch 0
GPU 1：处理 Batch 1
GPU 2：处理 Batch 2
GPU 3：处理 Batch 3

每张 GPU 独立完成前向传播和反向传播，得到自己的梯度。

随后通过 All-Reduce 操作对梯度进行聚合，使所有 GPU 得到一致的梯度结果，再同步更新模型参数。

训练流程：

                全局训练数据
                      |
              +-------+-------+
              |       |       |
              v       v       v
            GPU 0   GPU 1   GPU 2   GPU 3
              |       |       |       |
              v       v       v       v
            模型副本 模型副本 模型副本 模型副本
              |       |       |       |
              v       v       v       v
            前向传播 前向传播 前向传播 前向传播
              |       |       |       |
              v       v       v       v
            反向传播 反向传播 反向传播 反向传播
              |       |       |       |
              +-------+-------+-------+
                      |
                      v
                All-Reduce
                      |
                      v
                  梯度同步
                      |
                      v
                  参数更新

数据并行的优点：

实现简单，容易扩展。
可以提高整体训练吞吐量。
适合模型能够完整放入单张 GPU 的场景。

缺点：

每张 GPU 都需要保存完整模型。
模型参数、梯度和优化器状态会重复占用显存。
GPU 数量增加后，通信开销可能成为瓶颈。

在 PyTorch 中，常见的数据并行实现方式是 DistributedDataParallel（DDP）。

3.2 模型并行（Model Parallelism）

当模型本身太大，无法完整放入一张 GPU 时，就需要使用模型并行。

模型并行的核心思想是：

将一个模型拆分成多个部分，分别放在不同 GPU 上运行。

模型并行主要包括张量并行和流水线并行。

3.2.1 张量并行（Tensor Parallelism）

张量并行是将模型中的大型张量运算拆分到多张 GPU 上执行。

例如，一个大型线性层：

$[ Y=XW ]$

可以将权重矩阵 (W) 拆分成多个子矩阵：

$[ W=[W_1,W_2,W_3,W_4] ]$

每张 GPU 负责计算其中一部分。

最后通过通信操作，将各 GPU 的计算结果组合起来。

张量并行通常应用于 Transformer 中的：

Attention 线性投影层。
Feed Forward Network（FFN）。
大规模矩阵乘法。
Embedding 层。

例如，一个 Transformer 的 Attention 模块可以按照注意力头或者矩阵维度进行切分。

优点：

能够将大型矩阵运算分配到多张 GPU。
降低单卡参数存储压力。
适合超大规模 Transformer 模型。

缺点：

GPU 之间需要频繁通信。
对 GPU 互联带宽要求较高。
模型切分和通信调度相对复杂。

3.2.2 流水线并行（Pipeline Parallelism）

流水线并行的核心思想是：

将模型不同层划分到不同 GPU，让不同 GPU 像生产流水线一样协同工作。

例如，一个拥有 24 层 Transformer 的模型，可以拆分为 4 个阶段：

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

数据依次通过不同 GPU，完成前向传播。

为了减少 GPU 空闲时间，可以将一个 Batch 拆分成多个 Micro-Batch，让不同阶段交错执行。

优点：

可以将模型不同层分配到不同 GPU。
降低单张 GPU 的模型存储压力。
适合层数较多的大型神经网络。

缺点：

存在流水线气泡（Pipeline Bubble）。
不同阶段的计算负载需要合理平衡。
需要处理阶段之间的激活值传输。

3.3 ZeRO 优化技术

ZeRO（Zero Redundancy Optimizer）是大模型分布式训练中非常重要的显存优化技术，由 Microsoft DeepSpeed 推广应用。

传统数据并行中，每张 GPU 都保存完整的：

模型参数（Parameters）
梯度（Gradients）
优化器状态（Optimizer States）

这会产生大量冗余显存开销。

ZeRO 的核心思想是：

将模型训练状态切分到不同 GPU，减少数据并行中的重复存储。

ZeRO 主要分为三个阶段。

ZeRO Stage 1：优化器状态切分

将优化器状态分配到不同 GPU。

例如，Adam 优化器中的一阶动量和二阶动量不再在每张 GPU 上完整保存，而是由不同 GPU 分别负责。

这样可以显著减少优化器状态占用的显存。

ZeRO Stage 2：梯度切分

在 Stage 1 的基础上，进一步对梯度进行切分。

每张 GPU 只保留自己负责的梯度分片，降低梯度存储开销。

ZeRO Stage 3：模型参数切分

在 Stage 2 的基础上，进一步切分模型参数。

每张 GPU 只保存部分模型参数，需要计算时再通过通信收集所需参数。

三种阶段对比如下：

技术	参数切分	梯度切分	优化器状态切分
ZeRO Stage 1	否	否	是
ZeRO Stage 2	否	是	是
ZeRO Stage 3	是	是	是

需要注意：ZeRO Stage 3 并不是简单地把模型永久拆成互不关联的几部分，而是在训练过程中根据计算需要进行参数收集、使用和释放。

四、四种核心并行技术对比

技术	核心思想	主要解决问题	典型应用
数据并行	多 GPU 处理不同数据	提升吞吐量	DDP
张量并行	拆分模型张量	单层计算和显存压力	Megatron-LM
流水线并行	拆分模型层	超深模型训练	Pipeline
ZeRO	切分训练状态	降低显存冗余	DeepSpeed

实际训练中，这些技术并不是互斥的。

对于超大规模模型，通常会组合使用：

数据并行 + 张量并行 + 流水线并行 + ZeRO 优化。

这种组合也被称为 3D 并行等混合并行训练方案。

五、PyTorch 分布式训练实战：使用 DDP 训练模型

理解原理之后，我们通过 PyTorch 实现一个简单的分布式训练程序。

这里使用 DistributedDataParallel（DDP）实现多 GPU 数据并行。

5.1 环境准备

安装 PyTorch：

pip install torch torchvision

确保已经正确安装 NVIDIA 驱动、CUDA 环境以及支持当前 GPU 的 PyTorch 版本。

查看 GPU：

nvidia-smi

5.2 完整训练代码

创建文件：

train_ddp.py

代码如下：

import os
import torch
import torch.nn as nn
import torch.optim as optim

import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP
from torch.utils.data import DataLoader
from torch.utils.data import TensorDataset
from torch.utils.data.distributed import DistributedSampler


# 1. 初始化分布式环境
def setup():

    dist.init_process_group(
        backend="nccl"
    )

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

    torch.cuda.set_device(local_rank)

    return local_rank


# 2. 构建简单模型
class SimpleModel(nn.Module):

    def __init__(self):

        super().__init__()

        self.net = nn.Sequential(
            nn.Linear(100, 256),
            nn.ReLU(),
            nn.Linear(256, 128),
            nn.ReLU(),
            nn.Linear(128, 10)
        )

    def forward(self, x):

        return self.net(x)


# 3. 创建数据集
def create_dataset():

    x = torch.randn(10000, 100)

    y = torch.randint(
        0, 10, (10000,)
    )

    dataset = TensorDataset(x, y)

    return dataset


# 4. 训练函数
def train(local_rank):

    device = torch.device(
        f"cuda:{local_rank}"
    )

    dataset = create_dataset()

    sampler = DistributedSampler(
        dataset,
        shuffle=True
    )

    loader = DataLoader(
        dataset,
        batch_size=64,
        sampler=sampler
    )

    model = SimpleModel().to(device)

    model = DDP(
        model,
        device_ids=[local_rank]
    )

    criterion = nn.CrossEntropyLoss()

    optimizer = optim.Adam(
        model.parameters(),
        lr=0.001
    )

    epochs = 5

    for epoch in range(epochs):

        sampler.set_epoch(epoch)

        model.train()

        total_loss = 0

        for x, y in loader:

            x = x.to(device)
            y = y.to(device)

            optimizer.zero_grad()

            output = model(x)

            loss = criterion(output, y)

            loss.backward()

            optimizer.step()

            total_loss += loss.item()

        if dist.get_rank() == 0:

            print(
                f"Epoch {epoch + 1}, "
                f"Loss: {total_loss / len(loader):.4f}"
            )

    if dist.get_rank() == 0:

        torch.save(
            model.module.state_dict(),
            "model.pth"
        )

    dist.destroy_process_group()


if __name__ == "__main__":

    local_rank = setup()

    train(local_rank)

5.3 启动分布式训练

假设机器上有 4 张 GPU：

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

其中：

--standalone：使用单机独立运行模式。
--nnodes=1：指定训练节点数量为 1。
--nproc_per_node=4：每台机器启动 4 个训练进程。

每个进程负责一张 GPU，通过 NCCL 完成 GPU 间通信。

注意：这里使用的是随机生成的分类数据，目的是演示 DDP 的完整训练流程，并非真正的大模型训练任务。

六、DeepSpeed 实战：使用 ZeRO 降低显存占用

DeepSpeed 是微软开源的深度学习优化库，广泛用于大模型训练和推理。

它提供了：

ZeRO 优化器。
分布式训练。
混合精度训练。
梯度累积。
CPU/NVMe Offload。
大规模模型训练支持。

6.1 安装 DeepSpeed

pip install deepspeed

6.2 配置 ZeRO Stage 2

创建配置文件：

{
  "train_micro_batch_size_per_gpu": 4,
  "gradient_accumulation_steps": 4,
  "fp16": {
    "enabled": true
  },
  "zero_optimization": {
    "stage": 2,
    "allgather_partitions": true,
    "reduce_scatter": true,
    "contiguous_gradients": true
  },
  "steps_per_print": 100,
  "wall_clock_breakdown": false
}

参数说明：

参数	作用
train_micro_batch_size_per_gpu	每张 GPU 的微批次大小
gradient_accumulation_steps	梯度累积次数
fp16.enabled	是否启用 FP16
zero_optimization.stage	ZeRO 优化阶段
reduce_scatter	使用 Reduce-Scatter 通信
contiguous_gradients	优化梯度内存布局

6.3 DeepSpeed 训练代码

import torch
import torch.nn as nn
import deepspeed


class SimpleModel(nn.Module):

    def __init__(self):

        super().__init__()

        self.net = nn.Sequential(
            nn.Linear(100, 256),
            nn.ReLU(),
            nn.Linear(256, 10)
        )

    def forward(self, x):

        return self.net(x)


def main():

    model = SimpleModel()

    engine, optimizer, _, _ = (
        deepspeed.initialize(
            model=model,
            model_parameters=model.parameters(),
            config="ds_config.json"
        )
    )

    device = engine.device

    for step in range(100):

        x = torch.randn(
            4, 100,
            device=device
        )

        y = torch.randint(
            0, 10, (4,),
            device=device
        )

        output = engine(x)

        loss = nn.CrossEntropyLoss()(
            output, y
        )

        engine.backward(loss)

        engine.step()

        if step % 10 == 0:
            print(
                f"Step {step}, "
                f"Loss: {loss.item():.4f}"
            )


if __name__ == "__main__":

    main()

启动：

deepspeed --num_gpus=4 train.py

这段代码演示了 DeepSpeed 的基本训练接口。实际使用时，应根据数据集、模型规模和显卡环境调整 Batch Size、精度以及 ZeRO 配置。

七、大模型分布式训练中的关键优化技术

除了并行策略之外，还有一些重要技术可以进一步提高训练效率。

7.1 混合精度训练（Mixed Precision Training）

混合精度训练通过使用 FP16 或 BF16 降低计算和存储开销。

常见数据类型：

数据类型	位宽	特点
FP32	32 bit	精度较高，显存占用较大
FP16	16 bit	显存占用较低，但动态范围较小
BF16	16 bit	动态范围较大，适合大模型训练
FP8	8 bit	进一步降低计算和存储开销，需硬件与软件支持

其中，BF16 因为具有与 FP32 接近的指数范围，在现代大模型训练中得到广泛应用。

7.2 梯度累积（Gradient Accumulation）

当显存不足以容纳较大的 Batch Size 时，可以使用梯度累积。

例如：

每次训练处理 4 个样本。
累积 8 次梯度。
再进行一次参数更新。

有效 Batch Size 为：

$B_{\text{micro}}\times K\times N ]$

其中：

$(B_{\text{micro}})$ ：每张 GPU 的微批次大小。
(K)：梯度累积次数。
(N)：数据并行 GPU 数量。

梯度累积能够在不显著增加单步显存需求的情况下，扩大有效 Batch Size。

7.3 梯度检查点（Gradient Checkpointing）

梯度检查点的核心思想是：

前向传播时不保存所有中间激活值，在反向传播时重新计算部分激活值。

这样可以减少显存占用，但会增加一定的计算开销。

适合训练层数较多、激活值占用显存较大的 Transformer 模型。

7.4 通信优化

分布式训练不仅需要计算，还需要频繁进行 GPU 间通信。

常见通信操作包括：

All-Reduce：聚合多个 GPU 的数据。
All-Gather：收集不同 GPU 上的数据。
Reduce-Scatter：聚合并分发数据。
Broadcast：将数据广播到多个 GPU。

在多机训练中，通信性能直接影响训练效率。

常见技术包括 NCCL、InfiniBand、RoCE 和 NVLink。

八、大模型分布式训练框架对比

目前比较常见的分布式训练框架包括 PyTorch DDP、DeepSpeed、Megatron-LM 和 FSDP。

框架	核心能力	典型应用
PyTorch DDP	数据并行	多 GPU 常规模型训练
DeepSpeed	ZeRO、Offload、混合精度	大模型显存优化
Megatron-LM	张量并行、流水线并行	超大规模 Transformer
PyTorch FSDP	参数分片	大模型分布式训练

如何选择？

普通模型、多 GPU 训练：优先学习 PyTorch DDP。
模型参数和优化器状态占用显存过高：学习 DeepSpeed ZeRO。
超大规模 Transformer 训练：学习 Megatron-LM。
希望使用 PyTorch 原生生态进行参数分片：学习 FSDP。

需要注意，这些框架的能力存在交叉，实际选型还要考虑硬件、模型结构、通信拓扑和训练规模。

九、大模型分布式训练面临的挑战

虽然分布式训练能够突破单卡限制，但也会带来新的工程问题。

9.1 GPU 显存不均衡

不同 GPU 上的模型分片、激活值和通信缓存可能存在差异。

如果某张 GPU 显存率先达到上限，就可能导致整个训练任务失败。

9.2 通信开销过高

当 GPU 计算速度很快，但 GPU 之间通信速度较慢时，训练效率可能明显下降。

因此，需要合理规划：

GPU 数量。
节点之间的网络带宽。
模型切分策略。
通信与计算的重叠。

9.3 训练任务容错

大模型训练通常持续数天甚至数周。

如果某个节点出现故障，可能导致整个训练任务中断。

因此，需要建立完善的 Checkpoint 机制，定期保存：

模型参数。
优化器状态。
学习率调度器状态。
训练步数和随机数状态。

这样才能在故障后尽可能恢复训练进度。

9.4 负载不均衡

在流水线并行中，如果不同阶段的计算量差异过大，就会出现部分 GPU 等待其他 GPU 的情况。

因此，需要合理划分模型层数，并尽量平衡不同阶段的计算负载。

十、总结：如何系统学习大模型分布式训练？

大模型分布式训练并不是单纯地增加 GPU 数量，而是需要同时理解模型计算、显存管理、并行策略和通信机制。

建议按照以下路线学习：

第一阶段：掌握深度学习训练基础

PyTorch 基础。
Tensor 和自动求导。
前向传播与反向传播。
Optimizer 和 Loss。
GPU 显存管理。

第二阶段：掌握分布式训练基础

多 GPU 训练。
PyTorch DDP。
NCCL 通信。
All-Reduce。
梯度同步。

第三阶段：掌握大模型并行技术

Data Parallelism。
Tensor Parallelism。
Pipeline Parallelism。
ZeRO Stage 1/2/3。
混合精度训练。

第四阶段：掌握主流训练框架

DeepSpeed。
Megatron-LM。
PyTorch FSDP。
Hugging Face Transformers。

第五阶段：完成实战项目

建议从小到大完成以下项目：

使用 PyTorch DDP 训练 ResNet。
使用 DDP 训练小型 Transformer。
使用 DeepSpeed ZeRO Stage 2 训练语言模型。
使用 ZeRO Stage 3 进行模型参数切分。
使用张量并行训练 Transformer。
结合数据并行和张量并行完成混合并行训练。

最后需要记住：

大模型训练的核心不只是模型架构，更是计算、显存、通信和工程系统的综合优化。

理解数据并行、张量并行、流水线并行和 ZeRO，才能真正掌握大模型分布式训练的底层技术。

北京朝阳AI社区

更多推荐

防拍照系统的工程分水岭

北京朝阳AI社区

从一笔投资到一个生态：浦东创投把AI创业者带进了“朋友圈”

从2025年人工智能产业生态推介会上种子基金亮相，到2026年训练营两期结营、第三期面向全球招募，浦东创投持续把资源投向AI创业的早期阶段。这背后，是浦东对创新链条前端的重视：大模型、具身智能、AI for Science等方向需要产业厚度，也需要能接住科学家和青年创业者的资本机制、产业资源与成长空间。