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

使用`torchrun`进行容错分布式训练¶

创建于：2022年9月27日 | 最后更新：2024年11月12日 | 最后验证：2024年11月5日

作者: Suraj Subramanian

What you will learn

使用torchrun启动多GPU训练任务
保存和加载训练任务的快照
为优雅重启构建您的训练脚本

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

Prerequisites

DDP的概述
熟悉DDP代码
一台配备多个GPU的机器（本教程使用AWS p3.8xlarge实例）
PyTorch installed 已安装 CUDA

跟随下面的视频或在youtube上观看。

在分布式训练中，单个进程的失败可能会中断整个训练任务。由于在这里失败的易感性可能更高，因此使您的训练脚本健壮尤为重要。您可能还希望您的训练任务具有弹性，例如，计算资源可以在任务过程中动态加入和离开。

PyTorch 提供了一个名为 torchrun 的工具，用于提供容错和弹性训练。当发生故障时，torchrun 会记录错误并尝试从训练作业的最后保存的“快照”自动重启所有进程。

快照保存的不仅仅是模型状态；它还可以包括有关运行的周期数、优化器状态或训练作业的任何其他有状态属性的详细信息，这些信息对于其连续性来说是必要的。

为什么使用 `torchrun`¶

torchrun 处理分布式训练的细节，这样你就不需要了。例如，

你不需要设置环境变量或显式传递rank和world_size；torchrun会分配这些以及其他几个环境变量。
无需在脚本中调用mp.spawn；你只需要一个通用的main()入口点，并使用torchrun启动脚本。这样，相同的脚本可以在非分布式以及单节点和多节点设置中运行。
从最后保存的训练快照优雅地重新开始训练。

优雅重启¶

为了实现优雅的重启，您应该按照以下方式构建您的训练脚本：

def main():
  load_snapshot(snapshot_path)
  initialize()
  train()

def train():
  for batch in iter(dataset):
    train_step(batch)

    if should_checkpoint:
      save_snapshot(snapshot_path)

如果发生故障，torchrun 将终止所有进程并重新启动它们。每个进程入口点首先加载并初始化最后保存的快照，并从那里继续训练。因此，在任何故障发生时，您只会丢失自上次保存快照以来的训练进度。

在弹性训练中，每当有任何成员资格变更（添加或删除节点）时，torchrun 将终止并在可用设备上生成进程。这种结构确保您的训练作业可以在无需人工干预的情况下继续。

差异对比 multigpu.py 与 multigpu_torchrun.py

进程组初始化¶

torchrun 自动分配 RANK 和 WORLD_SIZE，以及其他环境变量

- def ddp_setup(rank, world_size):
+ def ddp_setup():
-     """
-     Args:
-         rank: Unique identifier of each process
-         world_size: Total number of processes
-     """
-     os.environ["MASTER_ADDR"] = "localhost"
-     os.environ["MASTER_PORT"] = "12355"
-     init_process_group(backend="nccl", rank=rank, world_size=world_size)
+     init_process_group(backend="nccl")
     torch.cuda.set_device(int(os.environ["LOCAL_RANK"]))

使用torchrun提供的环境变量¶

- self.gpu_id = gpu_id
+ self.gpu_id = int(os.environ["LOCAL_RANK"])

保存和加载快照¶

定期将所有相关信息存储在快照中，使我们的训练任务在中断后能够无缝恢复。

+ def _save_snapshot(self, epoch):
+     snapshot = {}
+     snapshot["MODEL_STATE"] = self.model.module.state_dict()
+     snapshot["EPOCHS_RUN"] = epoch
+     torch.save(snapshot, "snapshot.pt")
+     print(f"Epoch {epoch} | Training snapshot saved at snapshot.pt")

+ def _load_snapshot(self, snapshot_path):
+     snapshot = torch.load(snapshot_path)
+     self.model.load_state_dict(snapshot["MODEL_STATE"])
+     self.epochs_run = snapshot["EPOCHS_RUN"]
+     print(f"Resuming training from snapshot at Epoch {self.epochs_run}")

在Trainer构造函数中加载快照¶

当重新启动一个中断的训练任务时，您的脚本将首先尝试加载一个快照以从中恢复训练。

class Trainer:
   def __init__(self, snapshot_path, ...):
   ...
+  if os.path.exists(snapshot_path):
+     self._load_snapshot(snapshot_path)
   ...

恢复训练¶

训练可以从最后一个运行的时期恢复，而不是从头开始。

def train(self, max_epochs: int):
-  for epoch in range(max_epochs):
+  for epoch in range(self.epochs_run, max_epochs):
      self._run_epoch(epoch)

运行脚本¶

只需像非多进程脚本一样调用你的入口点函数；torchrun会自动生成进程。

if __name__ == "__main__":
   import sys
   total_epochs = int(sys.argv[1])
   save_every = int(sys.argv[2])
-  world_size = torch.cuda.device_count()
-  mp.spawn(main, args=(world_size, total_epochs, save_every,), nprocs=world_size)
+  main(save_every, total_epochs)

- python multigpu.py 50 10
+ torchrun --standalone --nproc_per_node=4 multigpu_torchrun.py 50 10

使用torchrun进行容错分布式训练¶

为什么使用 torchrun¶