meta-pytorch/torchft
GitHub: meta-pytorch/torchft
torchft 为 PyTorch 分布式训练提供以训练步长为粒度的容错能力，在 worker 节点故障时无需中断全局训练即可自动恢复并继续运行。
Stars: 519 | Forks: 70
为 PyTorch 提供简单的每步容错能力
| Documentation
| Poster
| Design Doc
|
本代码仓库实现了执行每步容错的技术，
因此当发生错误时，您可以继续训练，而无需中断整个训练
任务。
[这基于 PyTorch
Conference 2024 上提出的大规模训练技术。](./media/fault_tolerance_poster.pdf)
## 概述
torchft 旨在提供所需的基础组件，以在任意应用程序/训练脚本中实现容错，并提供所需的基础组件来
实现自定义容错策略。
开箱即用，torchft 提供以下算法：
* 容错 DDP
* 容错 HSDP：在复制维度上提供容错能力，并支持在其他维度上任意组合 FSDP/TP 等。
* LocalSGD
* DiLoCo
为了实现这些功能，torchft 提供了一些关键的可重用组件：
1. 协调组件，可以通过
每步心跳检测来确定哪些 worker 是健康的
2. 容错 ProcessGroup 实现，能够合理报告错误并
平滑地重新初始化。
3. 检查点传输组件，可用于在执行扩容操作时从健康的
对等节点进行实时恢复。
以下组件图展示了高层组件及其相互
关联方式：
有关更多详细信息，请参阅 [torchft 文档](https://pytorch.org/torchft)。
## 示例
### torchtitan（容错 HSDP）
torchtitan 提供了一个开箱即用的、基于
torchft 构建的容错 HSDP 训练循环，可用于训练诸如 Llama 3 70B 之类的模型。
它也是一个很好的示例，展示了如何将 torchft 集成到您自己的训练脚本中以用于 HSDP。
有关端到端的用法，请参阅 [torchtitan 的文档](https://github.com/pytorch/torchtitan/blob/main/docs/torchft.md)。
### 容错 DDP
我们提供了一个最简的 DDP 训练循环，突出了 torchft 中所有关键组件。
有关更多信息，请参阅 [train_ddp.py](./train_ddp.py)。
### DiLoCo
LocalSGD 和 DiLoCo 目前处于实验阶段。
有关如何将这些算法集成到您的训练循环中的示例，请参阅
[diloco_train_loop/local_sgd_train_loop 测试](./torchft/local_sgd_integ_test.py)。
## 设计
torchft 旨在使用 DDP 或 HSDP（带有 DDP 的 FSDP）等复制权重进行训练时提供容错能力。
有关最详细的解释，请参阅[设计文档](https://docs.google.com/document/d/1OZsOsz34gRDSxYXiKkj4WqcD9x0lP9TcsfBeu_SsOY4/edit)。
### Lighthouse
torchft 实现了一个 lighthouse server，用于在不同的
副本组之间进行协调，随后是一个针对单个副本组的 manager 和容错库，
它们可以在标准的 PyTorch 训练循环中使用。
这允许以训练步长为粒度进行成员变更，从而
通过避免在发生错误时停止全局训练，大幅提高效率。
### 容错 HSDP 算法
torchft 提供了容错 HSDP/DDP 算法的实现。
以下图表展示了在训练循环中为确保在恢复操作期间一切保持一致而需要执行的
高层操作。
有关更多详细信息，请参阅上面链接的设计文档。
## 从 PyPI 安装
我们在 https://pypi.org/project/torchft-nightly/ 提供了每日构建版本
要安装具有最简依赖项的 torchft，您可以运行：
```
pip install torchft-nightly
```
如果您想要所有开发依赖项，可以安装：
```
pip install torchft-nightly[dev]
```
## 从源码安装
### 前置条件
在继续之前，请确保您已安装以下内容：
- Rust（包含必要的依赖项）
- `protobuf-compiler` 以及相应的 Protobuf 开发包。
- PyTorch 2.7 RC+ 或 Nightly 版本
请注意，许多 conda 环境中可用的 Rust 版本可能已过时。要安装最新版本的 Rust，我们建议直接从官方网站下载，如下面的命令所示：
```
curl --proto '=https' --tlsv1.2 https://sh.rustup.rs -sSf | sh
```
要在基于 Debian 的系统（如 Ubuntu）上使用 apt 安装所需的软件包，请运行：
```
sudo apt install protobuf-compiler libprotobuf-dev
```
或者在基于 Red Hat 的系统上，运行：
```
sudo dnf install protobuf-compiler protobuf-devel
```
### 安装说明
```
pip install .
```
这使用 pyo3+maturin 来构建包，您需要安装 maturin。
如果安装命令由于无法获取清单而导致无法调用 `cargo update`，这可能是由于您的 .`gitconfig` 文件中的 `proxy`、`proxySSLCert` 和 `proxySSLKey` 设置影响了 `cargo` 命令。要解决此问题，请在运行安装命令之前尝试从您的 `.gitconfig` 中临时移除这些字段。
要在安装 Rust 扩展和开发依赖项的同时以可编辑模式安装，您可以使用标准的 pip install 命令：
```
pip install -e '.[dev]'
```
## 用法
### Lighthouse
Lighthouse 用于在使用同步训练时，跨复制的 worker (DDP/FSDP)
提供容错能力。
您可以通过运行以下命令来启动 lighthouse server：
```
RUST_BACKTRACE=1 torchft_lighthouse --min_replicas 1 --quorum_tick_ms 100 --join_timeout_ms 10000
```
### 训练循环示例 (DDP)
有关完整示例，请参阅 [train_ddp.py](./train_ddp.py)。
使用以下命令调用：
```
TORCHFT_LIGHTHOUSE=http://localhost:29510 torchrun --master_port 29501 --nnodes 1 --nproc_per_node 1 train_ddp.py
```
train.py:
```
from torchft import Manager, DistributedDataParallel, Optimizer, ProcessGroupGloo
manager = Manager(
pg=ProcessGroupGloo(),
load_state_dict=...,
state_dict=...,
)
m = nn.Linear(2, 3)
m = DistributedDataParallel(manager, m)
optimizer = Optimizer(manager, optim.AdamW(m.parameters()))
for i in range(1000):
batch = torch.rand(2, 2, device=device)
optimizer.zero_grad()
out = m(batch)
loss = out.sum()
loss.backward()
optimizer.step()
```
### 运行 DDP
在通过运行以下命令启动 lighthouse server 之后：
```
RUST_BACKTRACE=1 torchft_lighthouse --min_replicas 1 --quorum_tick_ms 100 --join_timeout_ms 10000
```
可以使用 torchX 启动一个测试 DDP 脚本：
```
torchx run
```
或者使用以下命令启动 DiLoCo：
```
USE_STREAMING=True torchx run ./torchft/torchx.py:hsdp --script='train_diloco.py'
```
请参阅 [.torchxconfig](.torchxconfig)、[torchx.py](./torchft/torchx.py) 以及 [torchX 文档](https://pytorch.org/torchx/latest/) 以了解 DDP 是如何运行的。
如果训练脚本支持，当 `workers_per_replica` 设置为 > 1 时，`torchx.py` 也可以启动 HSDP 任务。有关启用了 torchFT 的 HSDP 训练实现示例，请参阅 [torchtitan](https://github.com/pytorch/torchtitan)。
或者，要在具有两个 GPU 的节点上进行测试，您可以通过以下方式启动两个副本组来运行 [train_ddp.py](./train_ddp.py)：
在 shell 1 中（一个副本组开始初始训练）：
```
export REPLICA_GROUP_ID=0
export NUM_REPLICA_GROUPS=2
CUDA_VISIBLE_DEVICES=0 TORCHFT_LIGHTHOUSE=http://localhost:29510 torchrun --master_port=29600 --nnodes=1 --nproc_per_node=1 -- train_ddp.py
```
在 shell 2 中（第二个副本组加入）：
```
export REPLICA_GROUP_ID=1
export NUM_REPLICA_GROUPS=2
CUDA_VISIBLE_DEVICES=1 TORCHFT_LIGHTHOUSE=http://localhost:29510 torchrun --master_port=29601 --nnodes=1 --nproc_per_node=1 -- train_ddp.py
```
通过观察两个 shell 的输出，您应该能看到 process group 的重新配置和实时的检查点恢复。
### Parameter Server 示例
torchft 具有一个基于其可重新配置的 ProcessGroups 构建的容错 parameter server 实现。这不需要/使用 Lighthouse server。
有关示例，请参阅 [parameter_server_test.py](./torchft/parameter_server_test.py)。
## 许可证
torchft 采用 BSD 3-Clause 许可证。有关更多详细信息，请参阅 [LICENSE](./LICENSE)。
标签：AI基础设施, Apex, PyTorch, 凭据扫描, 分布式训练, 可视化界面, 容错机制, 机器学习, 逆向工具
