简介 || 什么是 DDP || 单节点多 GPU 训练 || 容错 || 多节点训练 || minGPT 训练
使用 torchrun 进行容错分布式训练¶
请观看下面的视频或访问youtube。
在分布式训练中，单个进程的故障可能会中断整个训练作业。由于这里出现故障的可能性更高，因此使您的训练脚本具有鲁棒性在这里尤其重要。您可能也希望您的训练作业具有弹性，例如，计算资源可以在作业过程中动态加入和离开。
PyTorch 提供了一个名为torchrun的实用程序，它提供了容错和弹性训练。当发生故障时，torchrun会记录错误并尝试从训练作业上次保存的“快照”自动重新启动所有进程。
快照保存的内容不仅仅是模型状态；它可以包含有关已运行的 epoch 数量、优化器状态或训练作业的任何其他状态属性的详细信息，这些属性对于其连续性是必要的。
为什么要使用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都会终止并在可用设备上生成进程。拥有这种结构可以确保您的训练作业能够在无需人工干预的情况下继续进行。
进程组初始化¶
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)
...
恢复训练¶
训练可以从上次运行的 epoch 恢复，而不是从头开始。
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
进一步阅读¶
使用DDP进行多节点训练（本系列的下一个教程）
使用DDP进行多GPU训练（本系列的上一个教程）
