简介 || 什么是 DDP || 单节点多 GPU 训练 || 容错 || 多节点训练 || minGPT 训练

使用 DDP 训练“真实世界”模型

创建日期：2022 年 9 月 27 日 | 最后更新：2025 年 1 月 23 日 | 最后验证：2024 年 11 月 5 日

作者：Suraj Subramanian

你将学到什么

• 编写分布式训练脚本时的最佳实践

• 在云端保存/加载工件时提高灵活性

• DDP 不适用的情况

在 GitHub 上查看本教程中使用的代码

前提条件

• 熟悉多 GPU 训练和torchrun

• [可选] 熟悉多节点训练

• 2 台或更多可通过 TCP 访问的 GPU 机器（本教程使用 AWS p3.2xlarge 实例）

• 在所有机器上安装支持 CUDA 的 PyTorch

观看下方视频或在 youtube 上跟随学习。

在本视频中，我们将回顾多节点 DDP 中训练 GPT 模型的过程。我们首先克隆 minGPT 仓库，并重构 Trainer 使其类似于本系列中使用的结构。观看视频了解这些更改的详细信息。

我们使用 hydra 集中管理训练运行的所有配置。代码重构完成后，我们首先在带有 4 个 GPU 的单节点上运行它，然后在 slurm 集群上运行。

用于训练的文件

• trainer.py 包含 Trainer 类，该类使用提供的数据集在模型上运行分布式训练迭代。

• model.py 定义了模型架构。

• char_dataset.py 包含用于字符级别数据集的 Dataset 类。

• gpt2_train_cfg.yaml 包含数据、模型、优化器和训练运行的配置。

• main.py 是训练任务的入口点。它设置 DDP 进程组，读取所有配置并运行训练任务。

从云端保存和加载

在上面的视频中，我们将训练快照直接保存到云端。这使得我们可以灵活地从任何可以访问云存储桶的节点继续训练。

使用混合精度

为了加快速度，您可以使用混合精度来训练模型。在混合精度中，训练过程的某些部分以较低精度进行，而对精度下降更敏感的其他步骤则保持 FP32 精度。

何时 DDP 不够用？

典型的训练运行的内存占用包括模型权重、激活、梯度、输入批次和优化器状态。由于 DDP 在每个 GPU 上复制模型，因此只有当 GPU 有足够的容量容纳全部内存占用时才能工作。当模型变得更大时，更激进的技术可能会很有用

• 激活检查点：在正向传播期间，不保存中间激活，而是在反向传播期间重新计算激活。在这种方法中，我们运行更多的计算，但节省了内存占用。

• 全分片数据并行：在这里，模型不是复制的，而是在所有 GPU 上“分片”，计算与前向和后向传播中的通信重叠。阅读我们的博客，了解我们如何使用 FSDP 训练具有 1 万亿参数的模型。

延伸阅读

• 使用 DDP 进行多节点训练（本系列上一教程）

• 混合精度训练

• 全分片数据并行

• 使用 FSDP 训练 1 万亿参数模型

• FSDP 视频教程系列