分布式训练指南

🔥 最新消息 🔥

• PyTorch 2025 大会幻灯片
• NeurIPS 2024 大会幻灯片

你是否曾想过如何在庞大的集群上训练一个大型神经网络？那就别再犹豫了！

本指南全面介绍了分布式训练的最佳实践、错误诊断方法，以及如何充分利用所有可用资源。它被组织成一系列循序渐进的章节，每个章节都包含一个 README.md 文件和一个 train_llm.py 脚本。README.md 不仅会讨论分布式训练的高层次概念，还会说明该章节引入的代码变更。

本指南完全使用最基础的 PyTorch 编写，并分别利用 transformers 和 datasets 库来处理模型和数据。分布式相关代码完全基于 PyTorch 实现，未使用任何其他第三方库。

1. 第1章 - 一个可在单个 GPU 上运行的标准因果语言模型训练脚本。
2. 第2章 - 将训练脚本升级为支持多 GPU 并使用 DDP。
3. 第3章 - 讨论如何在包含多个节点的集群上启动训练任务。
4. 第4章 - 将训练脚本升级为使用FSDP，以替代 DDP，从而实现更优的内存利用率。
5. 第5章 - 将训练脚本升级为用于训练 Llama-405B 模型。
6. 第6章 - 将我们的单 GPU 训练脚本升级为支持张量并行。
7. 第7章 - 将我们的张量并行训练脚本升级为使用二维并行（FSDP + TP）。
8. 替代框架 - 介绍几种底层均基于 PyTorch 的不同框架。
9. 错误诊断 - 在集群中快速诊断错误的最佳实践与操作指南。
10. 相关主题 - 进行分布式训练时需要了解的相关话题。

本指南解答的问题：

• 如何将单 GPU 的训练/微调脚本升级为在多 GPU 或多节点上运行？
• 如何诊断训练过程中出现的卡顿或错误？
• 我的模型/优化器太大，无法在单个 GPU 上运行——我该如何在我的集群上进行训练或微调？
• 如何在集群上调度和启动训练任务？
• 当工作节点数量增加时，我应该如何调整超参数？

此外，本指南还提供了记录 stdout/stderr 和 WandB 日志的最佳实践，因为在集群上进行训练时，日志记录对于诊断和调试至关重要。

每个训练脚本的目标都是训练一个因果语言模型（例如 GPT 或 Llama）。

环境搭建

克隆本仓库

git clone https://github.com/LambdaLabsML/distributed-training-guide.git

创建虚拟环境

cd distributed-training-guide
python3 -m venv venv
source venv/bin/activate
python -m pip install -U pip
pip install -U setuptools wheel
pip install -r requirements.txt
pip install flash-attn --no-build-isolation

配置 WandB

本教程使用 WandB 作为实验跟踪工具。

wandb login

🦄 Lambda 公司还有其他令人兴奋的机器学习项目：[ML Times](https://news.lambda.ai/news/today)、[Text2Video](https://lambdalabsml.github.io/Open-Sora/introduction/)、[GPU 基准测试](https://lambdalabs.com/gpu-benchmarks)。

distributed-training-guide 快速上手指南

本指南旨在帮助开发者快速掌握使用纯 PyTorch 进行大规模分布式训练的最佳实践，从单卡训练逐步进阶到多节点、FSDP 及张量并行训练。

环境准备

在开始之前，请确保您的开发环境满足以下要求：

• 操作系统: Linux (推荐 Ubuntu 20.04+)
• Python: Python 3.8 或更高版本
• GPU: 支持 CUDA 的 NVIDIA GPU（根据章节不同，可能需要单卡或多卡集群）
• 网络: 需要访问 Hugging Face 下载模型/数据集，以及访问 Weights & Biases (wandb) 进行实验追踪。
  • 国内用户提示: 若访问 Hugging Face 受阻，建议配置 HF_ENDPOINT 环境变量使用镜像源（如 https://hf-mirror.com）。

安装步骤

请按照以下顺序克隆仓库并配置虚拟环境。

1. 克隆项目仓库

   git clone https://github.com/LambdaLabsML/distributed-training-guide.git
   cd distributed-training-guide
   

2. 创建并激活虚拟环境

   python3 -m venv venv
   source venv/bin/activate
   

3. 升级基础工具并安装依赖

   python -m pip install -U pip
   pip install -U setuptools wheel
   
   # 安装核心依赖
   pip install -r requirements.txt
   
   # 安装 flash-attn (需确保已安装对应的 CUDA Toolkit)
   pip install flash-attn --no-build-isolation
   

   注意: 如果 pip install 速度较慢，可添加国内镜像源参数，例如：-i https://pypi.tuna.tsinghua.edu.cn/simple。

4. 配置实验追踪工具 (wandb) 本项目使用 wandb 记录训练日志，首次使用前需登录：

   wandb login
   

   按提示输入您的 API Key（可在 wandb.ai 官网获取）。

基本使用

本教程由一系列递进的章节组成，每个章节包含一个 README.md（理论讲解）和一个 train_llm.py（实战代码）。建议从第一章开始顺序学习。

1. 单卡训练入门 (Chapter 1)

这是最基础的起点，用于在单张 GPU 上训练因果语言模型（如 GPT/Llama）。

• 进入章节目录:

  cd 01-single-gpu/
  

• 运行训练脚本:

  python train_llm.py
  

  运行后，您将看到训练损失下降，并且数据会同步到您登录的 wandb 面板中。

2. 进阶至多卡分布式 (Chapter 2 & 3)

当您需要在多张 GPU 或多台机器上训练时，无需重写核心逻辑，只需切换目录并使用相应的启动器。

• 多卡数据并行 (DDP):

  cd ../02-distributed-data-parallel/
  # 假设使用 4 张卡
  torchrun --nproc_per_node=4 train_llm.py
  

• 多节点集群启动: 参考 03-job-launchers/ 中的说明，使用 torchrun 或集群调度系统（如 Slurm）指定主节点地址和端口来启动跨节点训练。

3. 大模型显存优化 (Chapter 4 & 5)

当模型过大无法放入单卡显存时，切换到 FSDP（Fully Sharded Data Parallel）模式。

• 启用 FSDP:

  cd ../04-fully-sharded-data-parallel/
  torchrun --nproc_per_node=8 train_llm.py
  

  此模式下，模型参数、梯度和优化器状态将在多卡间分片存储，显著降低显存占用，支持训练如 Llama-405b 等超大模型。

核心学习路径建议

1. 先运行 Chapter 1 确保环境无误且能跑通流程。
2. 若有多个 GPU，尝试 Chapter 2 观察多卡加速比。
3. 若需训练大模型，直接研读 Chapter 4 (FSDP) 和 Chapter 6 (Tensor Parallelism) 的代码实现。
4. 遇到训练挂起或报错时，查阅 Diagnosing Errors 章节获取调试技巧。