Pytorch 多卡并行（2）—— 使用 torchrun 进行容错处理
- 前文 Pytorch 多卡并行（1）—— 原理简介和 DDP 并行实践 介绍了使用 Pytorch 的 DDP 库进行单机多卡训练的方法，本文进一步说明如何用 TorchRun 改写前文代码，以提高模型训练的效率和容错性
- 完整代码下载：wxc971231/ddp-tutorial-series
文章目录
- 1. torchrun
- 2. 使用 torchrun 改写 DDP 代码
1. torchrun
-
在训练过程中，很容易遇到各种各样的错误，比如内存不足、网络故障、硬件故障等等。这些错误会导致训练过程中断或失败，从而浪费了训练时间和计算资源。 torchrun 允许我们在训练过程中按一定周期保存快照（snapshots），一旦某一并行进程出错退出，torchrun 会自动从最近 snapshots 重启所有进程。Snapshots 中要保存的参数由我们自行设定，它是模型 checkpoint 的超集，要包含恢复训练所需的全部参数，比如
- 当前 epoch 值
- 模型参数 model.state_dict()
- 学习率调度器参数 lr_scheduler.state_dict()
- 优化器参数 optimizer.state_dict()
- 其他必要参数
-
除了以上自动重启功能外，torchrun 还有其他一些功能
- torchrun 可以自动完成所有环境变量的设置，可以从环境变量中获取 rank 和 world size 等信息
os.environ['RANK'] # 得到在所有node的所有进程中当前GPU进程的rank os.environ['LOCAL_RANK'] # 得到在当前node中当前GPU进程的rank os.environ['WORLD_SIZE'] # 得到GPU的数量 - torchrun 可以完成进程分配工作，不再需要使用
mp.spawn手动分发进程，只需要设置一个通用的 main() 函数入口，然后用torchrun命令启动脚本即可 - 快照功能允许进行断点续训
- torchrun 可以自动完成所有环境变量的设置，可以从环境变量中获取 rank 和 world size 等信息
-
使用 torchrun 时，程序通常有以下结构
def main(args): ddp_setup() # 初始化进程池 load_train_objs(args) # 设置 dataset, model, optimizer, trainer 等组件，若存在 snapshot 则从中加载参数 trian(args) # 进行训练 destroy_process_group() # 销毁进程池 def train(args): for batch in iter(dataset): train_step(batch) if should_checkpoint: save_snapshot(snapshot_path) # 用 rank0 保存 snapshot if __name__ == "__main__": # 加载参数 args = parser.parse_args() # 现在 torchrun 负责在各个 GPU 上生成进程并执行，不再需要 mp.spawn 了 main(args) -
使用 torchrun 命令来启动程序
torchrun --standalone --nproc_per_node=gpu XXX.py--standalone代表单机运行--nproc_per_node=gpu代表使用所有可用GPU。等于号后也可写gpu数量n，这样会使用前n个GPU
如果想要进一步指定要运行的 GPU，可以通过 CUDA_VISIBLE_DEVICES 设置GPU可见性，比如
CUDA_VISIBLE_DEVICES=2,3 torchrun --standalone --nproc_per_node=gpu multi_gpu_torchrun.py这样会把本机上的 GPU2 和 GPU3 看做 GPU0 和 GPU1 运行
2. 使用 torchrun 改写 DDP 代码
- 使用 torchrun 改写以下 DDP 代码
# 使用 DistributedDataParallel 进行单机多卡训练 import torch import torch.nn.functional as F from torch.utils.data import Dataset, DataLoader import os # 对 python 多进程的一个 pytorch 包装 import torch.multiprocessing as mp # 这个 sampler 可以把采样的数据分散到各个 CPU 上 from torch.utils.data.distributed import DistributedSampler # 实现分布式数据并行的核心类 from torch.nn.parallel import DistributedDataParallel as DDP # DDP 在每个 GPU 上运行一个进程，其中都有一套完全相同的 Trainer 副本（包括model和optimizer） # 各个进程之间通过一个进程池进行通信，这两个方法来初始化和销毁进程池 from torch.distributed import init_process_group, destroy_process_group def ddp_setup(rank, world_size): """ setup the distribution process group Args: rank: Unique identifier of each process world_size: Total number of processes """ # MASTER Node（运行 rank0 进程，多机多卡时的主机）用来协调各个 Node 的所有进程之间的通信 os.environ["MASTER_ADDR"] = "localhost" # 由于这里是单机实验所以直接写 localhost os.environ["MASTER_PORT"] = "12355" # 任意空闲端口 init_process_group( backend="nccl", # Nvidia CUDA CPU 用这个 "nccl" rank=rank, world_size=world_size ) torch.cuda.set_device(rank) class Trainer: def __init__( self, model: torch.nn.Module, train_data: DataLoader, optimizer: torch.optim.Optimizer, gpu_id: int, save_every: int, ) -> None: self.gpu_id = gpu_id self.model = model.to(gpu_id) self.train_data = train_data self.optimizer = optimizer self.save_every = save_every # 指定保存 ckpt 的周期 self.model = DDP(model, device_ids=[gpu_id]) # model 要用 DDP 包装一下 def _run_batch(self, source, targets): self.optimizer.zero_grad() output = self.model(source) loss = F.cross_entropy(output, targets) loss.backward() self.optimizer.step() def _run_epoch(self, epoch): b_sz = len(next(iter(self.train_data))[0]) print(f"[GPU{self.gpu_id}] Epoch {epoch} | Batchsize: {b_sz} | Steps: {len(self.train_data)}") self.train_data.sampler.set_epoch(epoch) # 在各个 epoch 入口调用 DistributedSampler 的 set_epoch 方法是很重要的，这样才能打乱每个 epoch 的样本顺序 for source, targets in self.train_data: source = source.to(self.gpu_id) targets = targets.to(self.gpu_id) self._run_batch(source, targets) def _save_checkpoint(self, epoch): ckp = self.model.module.state_dict() # 由于多了一层 DDP 包装，通过 .module 获取原始参数 PATH = "checkpoint.pt" torch.save(ckp, PATH) print(f"Epoch {epoch} | Training checkpoint saved at {PATH}") def train(self, max_epochs: int): for epoch in range(max_epochs): self._run_epoch(epoch) # 各个 GPU 上都在跑一样的训练进程，这里指定 rank0 进程保存 ckpt 以免重复保存 if self.gpu_id == 0 and epoch % self.save_every == 0: self._save_checkpoint(epoch) class MyTrainDataset(Dataset): def __init__(self, size): self.size = size self.data = [(torch.rand(20), torch.rand(1)) for _ in range(size)] def __len__(self): return self.size def __getitem__(self, index): return self.data[index] def load_train_objs(): train_set = MyTrainDataset(2048) # load your dataset model = torch.nn.Linear(20, 1) # load your model optimizer = torch.optim.SGD(model.parameters(), lr=1e-3) return train_set, model, optimizer def prepare_dataloader(dataset: Dataset, batch_size: int): return DataLoader( dataset, batch_size=batch_size, pin_memory=True, shuffle=False, # 设置了新的 sampler，参数 shuffle 要设置为 False sampler=DistributedSampler(dataset) # 这个 sampler 自动将数据分块后送个各个 GPU，它能避免数据重叠 ) def main(rank: int, world_size: int, save_every: int, total_epochs: int, batch_size: int): # 初始化进程池 ddp_setup(rank, world_size) # 进行训练 dataset, model, optimizer = load_train_objs() train_data = prepare_dataloader(dataset, batch_size) trainer = Trainer(model, train_data, optimizer, rank, save_every) trainer.train(total_epochs) # 销毁进程池 destroy_process_group() if __name__ == "__main__": import argparse parser = argparse.ArgumentParser(description='simple distributed training job') parser.add_argument('--total-epochs', type=int, default=50, help='Total epochs to train the model') parser.add_argument('--save-every', type=int, default=10, help='How often to save a snapshot') parser.add_argument('--batch_size', default=32, type=int, help='Input batch size on each device (default: 32)') args = parser.parse_args() world_size = torch.cuda.device_count() # 利用 mp.spawn，在整个 distribution group 的 nprocs 个 GPU 上生成进程来执行 fn 方法，并能设置要传入 fn 的参数 args # 注意不需要 fn 的 rank 参数，它由 mp.spawn 自动分配 mp.spawn( fn=main, args=(world_size, args.save_every, args.total_epochs, args.batch_size), nprocs=world_size ) - 改写后的代码如下所示，请参考注释自行对比
# 使用 DistributedDataParallel 进行单机多卡训练的基础上，使用 torchrun 进行容错处理，增强程序稳定性 # torchrun 允许我们在训练过程中按一定保存 snapshots，其中应当包含当前 epoch、模型参数（ckpt）、优化器参数、lr调度器参数等恢复训练所需的全部参数 # 一旦程序出错退出，torchrun 会自动从最近 snapshots 重启所有进程 # 除了增强稳定性外，torchrun 还会自动完成所有环境变量设置和进程分配工作，所以不再需要手动设置 rank 或用 mp.spawn 生成并分配进程 import torch import torch.nn.functional as F from torch.utils.data import Dataset, DataLoader import os # 对 python 多进程的一个 pytorch 包装 import torch.multiprocessing as mp # 这个 sampler 可以把采样的数据分散到各个 CPU 上 from torch.utils.data.distributed import DistributedSampler # 实现分布式数据并行的核心类 from torch.nn.parallel import DistributedDataParallel as DDP # DDP 在每个 GPU 上运行一个进程，其中都有一套完全相同的 Trainer 副本（包括model和optimizer） # 各个进程之间通过一个进程池进行通信，这两个方法来初始化和销毁进程池 from torch.distributed import init_process_group, destroy_process_group def ddp_setup(): # torchrun 会处理环境变量以及 rank & world_size 设置 init_process_group(backend="nccl") class Trainer: def __init__( self, model: torch.nn.Module, train_data: DataLoader, optimizer: torch.optim.Optimizer, save_every: int, snapshot_path: str, # 保存 snapshots 的位置 ) -> None: self.gpu_id = int(os.environ['LOCAL_RANK']) # torchrun 会自动设置这个环境变量指出当前进程的 rank self.model = model.to(self.gpu_id) self.train_data = train_data self.optimizer = optimizer self.save_every = save_every # 指定保存 snapshots 的周期 self.epochs_run = 0 # 存储将要保存在 snapshots 中的 epoch num 信息 self.snapshot_path = snapshot_path # 若存在 snapshots 则加载，这样重复运行指令就能自动继续训练了 if os.path.exists(snapshot_path): print('loading snapshot') self._load_snapshot(snapshot_path) self.model = DDP(self.model, device_ids=[self.gpu_id]) # model 要用 DDP 包装一下 def _load_snapshot(self, snapshot_path): ''' 加载 snapshot 并重启训练 ''' loc = f"cuda:{self.gpu_id}" snapshot = torch.load(snapshot_path, map_location=loc) 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}") def _run_batch(self, source, targets): self.optimizer.zero_grad() output = self.model(source) loss = F.cross_entropy(output, targets) loss.backward() self.optimizer.step() def _run_epoch(self, epoch): b_sz = len(next(iter(self.train_data))[0]) print(f"[GPU{self.gpu_id}] Epoch {epoch} | Batchsize: {b_sz} | Steps: {len(self.train_data)}") self.train_data.sampler.set_epoch(epoch) for source, targets in self.train_data: source = source.to(self.gpu_id) targets = targets.to(self.gpu_id) self._run_batch(source, targets) def _save_snapshot(self, epoch): # 在 snapshot 中保存恢复训练所必须的参数 snapshot = { "MODEL_STATE": self.model.module.state_dict(), # 由于多了一层 DDP 包装，通过 .module 获取原始参数 "EPOCHS_RUN": epoch, } torch.save(snapshot, self.snapshot_path) print(f"Epoch {epoch} | Training snapshot saved at {self.snapshot_path}") def train(self, max_epochs: int): for epoch in range(self.epochs_run, max_epochs): # 现在从 self.epochs_run 开始训练，统一重启的情况 self._run_epoch(epoch) # 各个 GPU 上都在跑一样的训练进程，这里指定 rank0 进程保存 snapshot 以免重复保存 if self.gpu_id == 0 and epoch % self.save_every == 0: self._save_snapshot(epoch) class MyTrainDataset(Dataset): def __init__(self, size): self.size = size self.data = [(torch.rand(20), torch.rand(1)) for _ in range(size)] def __len__(self): return self.size def __getitem__(self, index): return self.data[index] def load_train_objs(): train_set = MyTrainDataset(2048) # load your dataset model = torch.nn.Linear(20, 1) # load your model optimizer = torch.optim.SGD(model.parameters(), lr=1e-3) return train_set, model, optimizer def prepare_dataloader(dataset: Dataset, batch_size: int): return DataLoader( dataset, batch_size=batch_size, pin_memory=True, shuffle=False, # 设置了新的 sampler，参数 shuffle 要设置为 False sampler=DistributedSampler(dataset) # 这个 sampler 自动将数据分块后送个各个 GPU，它能避免数据重叠 ) def main(save_every: int, total_epochs: int, batch_size: int, snapshot_path: str="snapshot.pt"): # 初始化进程池 ddp_setup() # 进行训练 dataset, model, optimizer = load_train_objs() train_data = prepare_dataloader(dataset, batch_size) trainer = Trainer(model, train_data, optimizer, save_every, snapshot_path) trainer.train(total_epochs) # 销毁进程池 destroy_process_group() if __name__ == "__main__": import argparse parser = argparse.ArgumentParser(description='simple distributed training job') parser.add_argument('--total-epochs', type=int, default=50, help='Total epochs to train the model') parser.add_argument('--save-every', type=int, default=10, help='How often to save a snapshot') parser.add_argument('--batch_size', default=32, type=int, help='Input batch size on each device (default: 32)') args = parser.parse_args() # 现在 torchrun 负责在各个 GPU 上生成进程并执行，不再需要 mp.spawn 了 main(args.save_every, args.total_epochs, args.batch_size) ''' 运行命令: torchrun --standalone --nproc_per_node=gpu multi_gpu_torchrun.py 参数说明： --standalone 代表单机运行 --nproc_per_node=gpu 代表使用所有可用GPU, 等于号后也可写gpu数量n, 这样会使用前n个GPU 运行后获取参数： os.environ['RANK'] 得到在所有机器所有进程中当前GPU的rank os.environ['LOCAL_RANK'] 得到在当前node中当前GPU的rank os.environ['WORLD_SIZE'] 得到GPU的数量 通过 CUDA_VISIBLE_DEVICES 指定程序可见的GPU, 从而实现指定GPU运行: CUDA_VISIBLE_DEVICES=0,3 torchrun --standalone --nproc_per_node=gpu multi_gpu_torchrun.py '''