Pytorch 多卡并行(2)—— 使用 torchrun 进行容错处理

  • 前文 Pytorch 多卡并行(1)—— 原理简介和 DDP 并行实践 介绍了使用 Pytorch 的 DDP 库进行单机多卡训练的方法,本文进一步说明如何用 torchrun 改写前文代码,以提高模型训练的效率和容错性
  • torchrun 是从 Pytorch 1.9.0 开始引入的一个命令,请保证您的 pytorch 版本符合要求
  • 完整代码下载:wxc971231/ddp-tutorial-series

文章目录

  • 1. torchrun
  • 2. 使用 torchrun 改写 DDP 代码
  • 3. 调试代码

1. torchrun

  • 在训练过程中,很容易遇到各种各样的错误,比如内存不足、网络故障、硬件故障等等。这些错误会导致训练过程中断或失败,从而浪费了训练时间和计算资源。 torchrun 允许我们在训练过程中按一定周期保存快照(snapshots),一旦某一并行进程出错退出,torchrun 会自动从最近 snapshots 重启所有进程。Snapshots 中要保存的参数由我们自行设定,它是模型 checkpoint 的超集,要包含恢复训练所需的全部参数,比如

    • 当前 epoch 值
    • 模型参数 model.state_dict()
    • 学习率调度器参数 lr_scheduler.state_dict()
    • 优化器参数 optimizer.state_dict()
    • 其他必要参数
  • 除了以上自动重启功能外,torchrun 还有其他一些功能

    1. torchrun 可以自动完成所有环境变量的设置,可以从环境变量中获取 rank 和 world size 等信息
      os.environ['RANK']          # 得到在所有node的所有进程中当前GPU进程的rank
      os.environ['LOCAL_RANK']    # 得到在当前node中当前GPU进程的rank
      os.environ['WORLD_SIZE']    # 得到GPU的数量
      
    2. torchrun 可以完成进程分配工作,不再需要使用 mp.spawn 手动分发进程,只需要设置一个通用的 main() 函数入口,然后用 torchrun 命令启动脚本即可
    3. 快照功能允许进行断点续训
  • 使用 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 保存 snapshotif __name__ == "__main__":# 加载参数args = parser.parse_args()	# 现在 torchrun 负责在各个 GPU 上生成进程并执行,不再需要 mp.spawn 了main(args)
    
  • 使用 torchrun 命令来启动程序

    torchrun --standalone --nproc_per_node=gpu XXX.py
    
    1. --standalone 代表单机运行
    2. --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 groupArgs:rank: Unique identifier of each processworld_size: Total number of processes"""# MASTER Node(运行 rank0 进程,多机多卡时的主机)用来协调各个 Node 的所有进程之间的通信os.environ["MASTER_ADDR"] = "localhost" # 由于这里是单机实验所以直接写 localhostos.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_idself.model = model.to(gpu_id)self.train_data = train_dataself.optimizer = optimizerself.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 = sizeself.data = [(torch.rand(20), torch.rand(1)) for _ in range(size)]def __len__(self):return self.sizedef __getitem__(self, index):return self.data[index]def load_train_objs():train_set = MyTrainDataset(2048)  # load your datasetmodel = torch.nn.Linear(20, 1)  # load your modeloptimizer = torch.optim.SGD(model.parameters(), lr=1e-3)return train_set, model, optimizerdef 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 argparseparser = 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 设置os.environ["MASTER_ADDR"] = "localhost" # 由于这里是单机实验所以直接写 localhostos.environ["MASTER_PORT"] = "12355"     # 任意空闲端口init_process_group(backend="nccl")torch.cuda.set_device(int(os.environ['LOCAL_RANK'])))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 会自动设置这个环境变量指出当前进程的 rankself.model = model.to(self.gpu_id)self.train_data = train_dataself.optimizer = optimizerself.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 = sizeself.data = [(torch.rand(20), torch.rand(1)) for _ in range(size)]def __len__(self):return self.sizedef __getitem__(self, index):return self.data[index]def load_train_objs():train_set = MyTrainDataset(2048)  # load your datasetmodel = torch.nn.Linear(20, 1)  # load your modeloptimizer = torch.optim.SGD(model.parameters(), lr=1e-3)return train_set, model, optimizerdef 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 argparseparser = 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的rankos.environ['LOCAL_RANK']    得到在当前node中当前GPU的rankos.environ['WORLD_SIZE']    得到GPU的数量通过 CUDA_VISIBLE_DEVICES 指定程序可见的GPU, 从而实现指定GPU运行:CUDA_VISIBLE_DEVICES=0,3 torchrun --standalone --nproc_per_node=gpu multi_gpu_torchrun.py
    '''
    

3. 调试代码

  • 如果使用 VScode 的话,可以如下编辑 launch.json 文件,然后像往常一样设置断点按 f5 调试即可
    {"version": "0.2.0","configurations": [{"name": "Python: torchrun","type": "python","request": "launch",// 设置 program 的路径为 torchrun 脚本对应的绝对路径"program": "/home/tim/anaconda3/envs/project/lib/python3.8/site-packages/torch/distributed/run.py",// 设置 torchrun 命令的参数"args":["--standalone","--nproc_per_node=gpu","multi_gpu_torchrun.py"],"console": "integratedTerminal","justMyCode": true}]
    }
    
    注意其中 “program” 是你的 torchrun 脚本路径,可使用 pip show torch 查看 torch 的安装路径进而找到它

本文来自互联网用户投稿,该文观点仅代表作者本人,不代表本站立场。本站仅提供信息存储空间服务,不拥有所有权,不承担相关法律责任。如若转载,请注明出处:http://www.mzph.cn/news/79437.shtml

如若内容造成侵权/违法违规/事实不符,请联系多彩编程网进行投诉反馈email:809451989@qq.com,一经查实,立即删除!

相关文章

[管理与领导-93]:IT基层管理者 - 扩展技能 - 5 - 职场丛林法则 -7- 复杂问题分析能力与复杂问题的解决能力:系统化思维

目录 前言: 一、系统化思维 VS 分解思维 1.1 系统化思维 1.2 分解思维 二、中医与西医思维模式的区别 三、正向闭环/正反馈 VS 负向闭环/负反馈 VS 开环 3.1 开环与管理 3.2 闭环与管理 3.3 生态系统是闭环系统 3.4 团队是一个闭环系统 3.5 正向闭环/正反…

华为云云服务器云耀L实例评测 | 在华为云耀L实例上搭建电商店铺管理系统:一次场景体验

🌷🍁 博主猫头虎(🐅🐾)带您 Go to New World✨🍁 🦄 博客首页——🐅🐾猫头虎的博客🎐 🐳 《面试题大全专栏》 🦕 文章图文…

最优化:建模、算法与理论(典型优化问题

第四章 典型优化问题 4.1 线性规划 4.1.1 基本形式和应用背景 再次说明一下,其实这本书很多的内容之前肯定大家都学过,但是我觉得这本书和我们之前学的东西的出发角度不一样,他更偏向数学,也多一个角度让我们去理解 线性规划问…

STM32-HAL库06-硬件IIC驱动FM24CL16B非易失存储器

STM32-HAL库06-IIC驱动FM24CL16B非易失存储器 一、所用材料: STM32VGT6自制控制板 STM32CUBEMX(HAL库软件) MDK5 二、所学内容: 通过HAL库的硬件IIC对FM24CL16B存储器进行写与读取操作。 三、CUBEMX配置: 第一步…

【Cocos Creator 3.5实现赛车游戏】10.实现汽车节点的运动逻辑

转载知识星球 | 深度连接铁杆粉丝,运营高品质社群,知识变现的工具 项目地址:赛车小游戏-基于Cocos Creator 3.5版本实现: 课程的源码,基于Cocos Creator 3.5版本实现 上一节的学习后,您已经完成了对汽车节点的控制逻…

简单记一下Vue router 路由中使用 vue-i18n 进行标题国际化

引入状态管理和国际化文件 import store from ../store import i18n from /configs/i18n使用状态管理设置路由当前国际化选项 // 使用状态管理 i18n.locale store.state.setStore.i18n??zh路由中使用i18n { path: /login, name: login, component: LoginPage, meta: { ti…

四川玖璨电子商务有限公司:短视频运营理解

短视频运营是一种通过策划、执行和管理短视频内容以达到品牌目标的一项工作。在如今面临信息爆炸的时代,短视频已经成为了吸引用户注意力的一种重要方式。作为一种新兴媒体形式,短视频拥有跨时空、生动有趣、内容丰富的特点,因此得到了越来越…

信道分类 | 信道

文章目录 高斯信道带限信道 高斯信道 定义:功率限制为P的高斯信道容量为 C m a x f ( x ) : E X 2 P I ( X ; Y ) Cmax_{f(x):EX^2P}I(X;Y) Cmaxf(x):EX2P​I(X;Y) I ( X ; Y ) ≤ 1 2 l o g ( 2 π e ( P N ) ) − 1 2 l o g ( 2 π e N ) I(X;Y)\le \frac{1}{…

华为云云耀云服务器L实例评测|云耀云服务器购买流程与功能介绍

目录 概述什么是云耀云服务器L实例功能概述 技术参数处理器和内存存储和网络GPU加速 强大的计算能力高性能计算集群大规模并行计算 灵活可扩展的存储本地SSD存储产品购买流程 概述 什么是云耀云服务器L实例 云耀云服务器L实例是新一代开箱即用、面向中小企业和开发者打造的全新…

CocosCreator3.8研究笔记(十)CocosCreator 图像资源的理解

一、图像资源导入 Cocos Creator 可使用图像文件格式,支持 JPG、PNG、BMP、TGA、HDR、WEBBP、PSD、TIFF 等。 将图像资源直接拖拽到 资源管理器 即可将其导入 二、图像资源的类型 在 属性检查器 面板中便可根据需要设置图像资源的使用类型:raw 、 textu…

flask bootstrap页面json格式化

html <!DOCTYPE html> <html lang"en"> <head><!-- 新 Bootstrap5 核心 CSS 文件 --> <link rel"stylesheet" href"static/bootstrap-5.0.0-beta1-dist/css/bootstrap.min.css"><!-- 最新的 Bootstrap5 核心 …

纯小白安卓刷机1

文章目录 常见的英文意思刷机是什么&#xff1f;为什么要刷机&#xff1f;什么是BL锁&#xff08;BootLoader锁&#xff09;&#xff1f;我的机能够刷机吗&#xff1f;什么是Boot镜像/分区&#xff1f;什么是Recovery镜像/分区&#xff08;缩写为rec&#xff09;&#xff1f;什…

数据分析与可视化 Numpy数组

1.创建数组 import numpy as np anp.array([1,2,3,4]) aarray([1, 2, 3, 4]) print(a)[1 2 3 4] b np.array((2,4,6,8)) print(b)[2 4 6 8] cnp.array([[1,2,3,4],[5,6,7,8]]) print(c)[[1 2 3 4] [5 6 7 8]] dnp.array([1,2,3,4],dtypefloat64) print(d)[1. 2. 3. 4.] …

重磅开赛!“山东工行杯”山东省第五届数据应用创新创业大赛报名火热进行中!

为进一步调动全社会参与数据价值挖掘和应用创新积极性&#xff0c;促进数据要素高效流通&#xff0c;释放数据价值&#xff0c;赋能经济社会高质量发展&#xff0c;探索公共数据资源与公共服务、社会生活、民生改善及经济建设的数字化结合方式&#xff0c;促进公共数据和企业数…

记录一次Flink安装记录

系统Ubuntu18 1.拉取镜像 #默认拉取最新版本 docker pull flink:1.14.4-scala_2.12-java8 #检查镜像 docker images | grep flink 2.编写flink的yaml文件 注意&#xff1a;云服务器需要设置安全策略放行8081端口&#xff0c;否则访问失败 注意2&#xff1a;docker-compose 版…

创建了一个名为nums_list的vector容器,其中存储了一系列的pair<int, int>

vector<pair<int, int>> nums_list;for (int i 0; i < nums.size(); i) {nums_list.emplace_back(i, nums[i]);}这段代码创建了一个名为nums_list的vector容器&#xff0c;其中存储了一系列的pair<int, int>。代码的逻辑如下&#xff1a;1. 创建一个空的…

Linux Debian12使用git将本地项目上传到码云(gitee)远程仓库

一、注册码云gitee账号 这个可以参考其他教程&#xff0c;本文不做介绍。 gitee官网&#xff1a;https://gitee.com/ 二、Linux Debian12安装git 如果Linux系统没有安装git&#xff0c;可以使用下面命令安装git sudo apt install git 三、gitee新建仓库 我这只做测试&…

自动化构建:提高开发流程效率与质量的关键工具

&#x1f482; 个人网站:【工具大全】【游戏大全】【神级源码资源网】&#x1f91f; 前端学习课程&#xff1a;&#x1f449;【28个案例趣学前端】【400个JS面试题】&#x1f485; 寻找学习交流、摸鱼划水的小伙伴&#xff0c;请点击【摸鱼学习交流群】 引言 自动化构建是现代…

便捷又炸街!Mate 60的智感支付,是如何做到快速又安全的?

扫码支付已成为线下消费的主流付款方式&#xff0c;平时出门&#xff0c;手机一带&#xff0c;钱包拜拜&#xff01; 以微信支付为例&#xff0c;正常线下支付&#xff0c;手机解锁状态下&#xff1a; 第一步&#xff1a;找到微信APP&#xff1b; 第二步&#xff1a;打开右上…

centos7 安装 rabbitmq

一、windows 上下载依赖文件以及安装文件 https://d28dx6y1hfq314.cloudfront.net/828/1033/el/7/package_files/1305009.rpm?t1694767188_8bd28780066a52b94cc2edfb2b7f574333d72318https://d28dx6y1hfq314.cloudfront.net/828/1039/el/7/package_files/1038460.rpm?t16947…