distributed-training-guide
用纯 PyTorch 原生 API 从单 GPU 讲到 405B 分布式训练的实战指南
加载项目详情…
本应用为开源项目，仅供学习研究，请遵守其开源协议。
微信扫码咨询客服
POMOAI — 让开源 AI 即刻运行
用纯 PyTorch 原生 API 从单 GPU 讲到 405B 分布式训练的实战指南
加载项目详情…
本应用为开源项目，仅供学习研究，请遵守其开源协议。
图1：LambdaLabsML 组织头像
想象一下：你本地电脑上跑通了一个 GPT-2 微调脚本，效果不错。但当你想把它迁移到公司集群——8 张 H100、跨 2 个节点——突然所有代码都要重写。分布式训练的门槛，不在于理解原理，而在于知道该改哪几行。这就是 Lambda Labs 团队写这份指南的原因。
Lambda Labs 是一家专注于 GPU 云服务的公司，他们每天都要帮客户在超大规模集群上跑训练任务。2023-2025 年间，他们将内部最佳实践整理成这份完全开源的指南，发布在 GitHub 并在 NeurIPS 2024 和 PyTorch Conference 2025 上公开演讲。指南的核心理念：用最少的外部依赖，讲清楚分布式训练每一步发生了什么——不依赖 DeepSpeed 不依赖 Megatron，只用纯 PyTorch 原生 API。
这不是一个传统意义上的项目——它更像是一本活的技术手册。仓库结构按章节组织，从最简单的单 GPU 脚本出发，逐步升级到多 GPU、多节点、FSDP、张量并行、2D 并行，最终覆盖 Llama 3.1 405B 的完整训练流程。每章都有对应的 train_llm.py 脚本，README 详细解释为什么要这样改代码。
十大章节覆盖的内容：
| 章节 | 主题 | 关键技能 |
|---|---|---|
| Ch1 | 单 GPU 基线 | GPT-2 因果语言模型训练 |
| Ch2 | DDP 多 GPU | torchrun 多进程启动 |
| Ch3 | 集群任务调度 | SLURM / MPI / DeepSpeed launcher |
| Ch4 | FSDP 全分片数据并行 | optimizer sharding |
| Ch5 | Llama-3.1 405B | 64 x H100 全参数微调 |
| Ch6 | 张量并行 | ColumnLinear / RowLinear 切分 |
| Ch7 | 2D 并行 FSDP+TP | 6GPU 混合并行实战 |
| Alt | DeepSpeed 替代框架 | 与原生 PyTorch 对比 |
| Err | 集群错误诊断 | hanging / OOM / NCCL 超时 |
| Rel | 相关主题 | batch size scaling / wandb 配置 |
整个指南只用两个外部库：transformers（加载模型）和 datasets（加载数据）。分布式训练本身完全用 PyTorch 原生的 torchrun、DistributedDataParallel（DDP）、FullyShardedDataParallel（FSDP）实现。这意味着你学到的技能可以无缝迁移到任何 PyTorch 项目，而不是被某个框架绑定。
依赖清单（requirements.txt）：
torch==2.8.0
tqdm
datasets==4.1.1
transformers==4.57.0
版本锁定策略非常保守——torch 2.8.0 是 2025 年初的稳定版，保证和 FSDP2 / Flash Attention 的兼容性。
Chapter 1 的 train_llm.py 是整份指南的起点。它包含标准的参数解析、数据集加载、模型初始化、训练循环和 wandb 日志。没有任何分布式相关的代码，纯粹是一个能跑的基线。
Chapter 2 引入了 DDP（Distributed Data Parallel）。关键改动是：
代码只有 5-6 行变化，但背后的原理涉及 all-reduce 梯度同步、进程组初始化等核心概念。README 都有图文解释。
Chapter 4 的 FSDP（全分片数据并行）是整个指南最有价值的部分。当模型 + 优化器太大，单个 GPU 放不下时，DDP 就不够用了。FSDP 通过分片存储 optimizer state、gradients、model parameters 来解决这个问题。PyTorch 原生 FSDP2 相比 DDP 可以节省约 60-70% 的显存。
Chapter 5 则展示了用纯 PyTorch + FSDP 训练 Llama-3.1 405B 的完整流程——需要 64 张 H100，约 8 个节点。这章重点讲：
Chapter 6-7 讲解张量并行（Tensor Parallelism）和 2D 并行（FSDP + TP 混合）。这是训练超大模型的必备技能。Ch7 展示了 6 GPU 场景下，2-way FSDP x 3-way TP 如何将模型切分到每个设备上。
图2：PyTorch 是本指南的核心技术栈，所有分布式原语均基于 PyTorch 原生 API
这是一份纯 CLI 的指南，没有 Web UI。所有操作都在终端完成。上手流程非常直接：
# 1. 克隆
git clone https://github.com/LambdaLabsML/distributed-training-guide.git
cd distributed-training-guide
# 2. 创建虚拟环境
python3 -m venv venv
source venv/bin/activate
pip install -U pip wheel setuptools
pip install -r requirements.txt
# 3. 单 GPU 训练（GPT-2，alpaca 数据集）
cd 01-single-gpu
python train_llm.py -d tatsu-lab/alpaca -m openai-community/gpt2
# 4. 多 GPU 训练（需要 torchrun）
cd ../02-distributed-data-parallel
export TORCHELASTIC_ERROR_FILE=../error.json
export OMP_NUM_THREADS=1
torchrun --nproc-per-node=gpu train_llm.py -d tatsu-lab/alpaca -m openai-community/gpt2
wandb 日志集成是开箱即用的——只需要一个 wandb login 就可将所有训练指标上传到云端。这在调试集群训练时非常重要，因为多节点日志是分布式训练中最容易出问题的环节。
这份指南专注于单技术栈（PyTorch）和同步训练（bulk synchronous parallel）。几个它没有覆盖的领域：
Pipeline Parallelism（流水线并行）：虽然 Ch6 提到了 Llama 论文的 4D parallelism 包含 pipeline parallel，但指南没有独立章节讲 PP。PP 是训练超大模型的千亿参数必备技能，需要配合气泡（bubble）填充策略（如 GPipe / PipeDream）。
异步训练：指南只覆盖同步 SGD 风格的训练，没有涉及异步 Parameter Server 等架构。
生产级调度：Chapter 3 覆盖了 SLURM / MPI，但 Kubernetes 上的训练调度（Kubeflow / KServe）没有展开。
在 Lambda Labs 发布这份指南之前，市面上关于分布式训练的教程要么过于理论（学术论文式推导），要么过于依赖特定框架（DeepSpeed / Megatron 文档）。Lambda Labs 的指南开创了一种新的范式：用最小可行代码，解释最核心的原理。
从社区数据来看，该项目在 2024-2025 年间增长显著，621 颗星背后是大量 ML 工程师、研究者的实际使用。NeurIPS 2024 和 PyTorch Conference 2025 的演讲进一步扩大了影响力。这份指南代表了 GPU 云服务商在教育培训方向的投入——帮助更多人学会在集群上训练模型，从而更愿意使用他们的云服务。
对于国内开发者而言，这份指南也有特殊价值：它详细覆盖了 NCCL（NVIDIA Collective Communications Library）的使用和调优，而 NCCL 正是国内多卡服务器（尤其是国产 GPU 集群）分布式训练的核心通信库。诊断错误章节中对 NCCL 超时和 Bus error 的处理方法，对国内 GPU 运维团队极具参考价值。
Lambda Labs 的这份分布式训练指南是目前 GitHub 上最系统的 PyTorch 分布式训练实战教程。它用 10 个章节、10 个脚本、零额外框架依赖的方式，从单 GPU 讲到了 405B 参数的跨节点训练。对于 AI 研究者和工程师来说，这是一份值得收藏的分布式训练工具箱——当你需要在集群上跑训练任务时，这套代码可以直接改写复用。
推荐阅读顺序：从头到尾通读一遍（2-3 小时），然后根据自己当前的任务选择对应章节作为参考手册。重点推荐 Chapter 2（DDP）和 Chapter 4（FSDP），这是大多数人的必经之路。