Megatron-LM源码系列(七)：Distributed-Optimizer分布式优化器实现Part2

1. 使用入口

  • DistributedOptimizer类定义在megatron/optimizer/distrib_optimizer.py文件中。创建的入口是在megatron/optimizer/__init__.py文件中的get_megatron_optimizer函数中。根据传入的args.use_distributed_optimizer参数来判断是用DistributedOptimizer还是Float16OptimizerWithFloat16Params。
def get_megatron_optimizer(model,no_weight_decay_cond=None,scale_lr_cond=None,lr_mult=1.0):...# Megatron optimizer.opt_ty = DistributedOptimizer \if args.use_distributed_optimizer else \Float16OptimizerWithFloat16Paramsreturn opt_ty(optimizer,args.clip_grad,args.log_num_zeros_in_grad,params_have_main_grad,args.use_contiguous_buffers_in_local_ddp,args.fp16,args.bf16,args.params_dtype,grad_scaler,model)
  • 相关的Optimizer的使用参考【Megatron-LM源码系列(六)：Distributed-Optimizer分布式优化器实现Part1】

2. 初始化init源码说明

在这里插入图片描述

  • 初始化的过程很大程度对应的上图grad buffer分片的实现，对应init函数如下：
    def __init__(self, optimizer, clip_grad, log_num_zeros_in_grad,params_have_main_grad, use_contiguous_buffers_in_local_ddp,fp16, bf16, params_dtype, grad_scaler, models):
  • init时会通过build_model_gbuf_range_map函数先创建grad buffer的范围映射，也就是对应图中的world_index/local_index/param_index三个。这里的self.models是一个list类型，对于使用了interleave流水线方式的训练来说，这里的self.models中会保存多份model, 其余情况list中只有一个元素。
        # Model grad buffer ranges.self.model_gbuf_ranges = []for model_index, model in enumerate(self.models):self.model_gbuf_ranges.append(self.build_model_gbuf_range_map(model))
  • build_model_gbuf_range_map会依次按grad buffer中类型来进行range的初始化build_model_gbuf_range。这里定义了一个单独的Range类。
@classmethoddef build_model_gbuf_range_map(cls, model):"""Create param-to-grad-buffer mappings, for grad buffer data typeswithin a specific virtual model."""return {dtype : cls.build_model_gbuf_range(model, dtype)for dtype in model._grad_buffers}class Range:"""A range represents a start and end points for indexing a shardfrom a full tensor."""def __init__(self, start, end):self.start = startself.end = endself.size = end - startdef normalize(self, start = 0):return Range(start, start + self.size)def __str__(self):return "%d,%d [%d]" % (self.start, self.end, self.size)def __len__(self):return self.end - self.start
  • build_model_gbuf_range初始化range的流程如下：
    • 获取DP的rank，计算单个Grad buffer切片的大小
    • 保存当前rank的world range和local range, 分别对应world index和local index
    • 计算param的range范围，对应param index
    • 返回当前rank的相关range范围
    @classmethoddef build_model_gbuf_range(cls, model, dtype):# 获取DP的rankdata_parallel_rank = mpu.get_data_parallel_rank()data_parallel_world_size = mpu.get_data_parallel_world_size()# 计算单个Grad buffer切片的大小grad_buffer = model._grad_buffers[dtype]gbuf_size = grad_buffer.numelmax_gbuf_range_size = int(math.ceil(gbuf_size / data_parallel_world_size))# 跟据DDP的rank总数，分别计算每个rank对应的全局rangegbuf_world_all_ranges = []for r in range(data_parallel_world_size):gbuf_world_start = r * max_gbuf_range_sizegbuf_world_end = min(gbuf_size, gbuf_world_start+max_gbuf_range_size)gbuf_world_range = Range(gbuf_world_start, gbuf_world_end)gbuf_world_all_ranges.append(gbuf_world_range)# 保存当前rank的world range和local range# Local DP's ranges.gbuf_world_range = gbuf_world_all_ranges[data_parallel_rank]gbuf_local_range = gbuf_world_range.normalize()# 计算param的range范围param_range_map = cls.build_model_gbuf_param_range_map(model,dtype,gbuf_world_range)# Group into dict.data = {"local" : gbuf_local_range,"world" : gbuf_world_range,"world_all" : gbuf_world_all_ranges,"param_map" : param_range_map,"max_range_size" : max_gbuf_range_size,}return data
  • 接着会根据当前rank相关的Range内容self.model_gbuf_ranges调用build_model_param_gbuf_map函数，主要作用是创建model_gbuf_ranges的逆映射，保存param->(modex_index, type)的映射。
class DistributedOptimizer(MixedPrecisionOptimizer):def __init__(...):...self.model_param_gbuf_map = \self.build_model_param_gbuf_map(self.model_gbuf_ranges)...def build_model_param_gbuf_map(cls, model_gbuf_ranges):"""Create a reverse of the model_gbuf_ranges, for referencing inopposite direction."""param_gbuf_map = {}for model_index, model_gbuf_range_map in enumerate(model_gbuf_ranges):for dtype, gbuf_range_map in model_gbuf_range_map.items():for param, param_range_map in gbuf_range_map["param_map"].items():param_gbuf_map[param] = (model_index, dtype)return param_gbuf_map
  • 在self.build_model_param_gbuf_map之后是初始化Optimizer对应的local group range，Optimizer原本有param_groups包括多个参数组，这里build_optimizer_group_ranges为了创建param参数到group_index的map映射，也就是<model_parameter:group_index>；self.build_model_param_gbuf_map最后对每个group_range中增加新的orig_group和orig_group_idx两个key，原来group_range初始化的时候只有params一个key
class DistributedOptimizer(MixedPrecisionOptimizer):def __init__(...):...# Optimizer ranges.self.model_param_group_index_map, self.opt_group_ranges = \self.build_optimizer_group_ranges(self.optimizer.param_groups,self.model_gbuf_ranges)...def build_optimizer_group_ranges(cls, param_groups, model_gbuf_ranges):# 获取param_groups中组的个数num_groups = len(param_groups)# 创建全局的参数到group_index的map映射，也就是<model_parameter:group_index>world_param_group_map = {}for group_index, group in enumerate(param_groups):for param in group["params"]:assert param.requires_gradworld_param_group_map[param] = group_index# 创建当前rank的local_param_group_map, local_param_group_map是param与(group_index, group_params_len)的映射， local_param_group_map虽然返回了但后面没用local_param_group_map = {}group_ranges = [ {"params": []} for _ in param_groups ]for model_gbuf_range_map in model_gbuf_ranges:for dtype, gbuf_range_map in model_gbuf_range_map.items():for param in gbuf_range_map["param_map"]:group_index = world_param_group_map[param]group_range = group_ranges[group_index]group_range["params"].append(param)local_param_group_map[param] = \(group_index, len(group_range["params"]) - 1)# Squeeze zero-size group ranges.for group_index, group_range in enumerate(group_ranges):group_range["orig_group"] = param_groups[group_index]group_range["orig_group_idx"] = param_groups[group_index]return local_param_group_map, group_ranges
  • 在初始化Optimizer之后，是通过创建self.build_model_and_main_param_groups创建optimizer step要用到的main parameter groups, 这里的group一方面是要进行reduce和gather通信操作，另一方面是被优化器用于梯度的更新操作。
class DistributedOptimizer(MixedPrecisionOptimizer):def __init__(...):...# Allocate main param shards.(self.model_float16_groups,self.model_fp32_groups,self.shard_float16_groups,self.shard_fp32_groups,self.shard_fp32_from_float16_groups,) = self.build_model_and_main_param_groups(self.model_gbuf_ranges,self.model_param_gbuf_map,self.opt_group_ranges)...
  • self.build_model_and_main_param_groups的实现主要是关于fp32/fp16/bf16三种类型训练时优化器内的显存分配。
    @classmethoddef build_model_and_main_param_groups(cls,model_gbuf_ranges,param_gbuf_map,opt_group_ranges):...# 保存原本fp16类型parammodel_float16_groups = []# 保存原本fp32类型parammodel_fp32_groups = []# 保存原本fp16类型param的切片shard_float16_groups = []# 保存原本fp32类型param的切片shard_fp32_groups = []# 保存原本fp16类型param的fp32类型param的副本shard_fp32_from_float16_groups = []# 分配每个group的param参数切片for group_index, group_range in enumerate(opt_group_ranges):for model_param in group_range["params"]:if model_param.type() in ['torch.cuda.HalfTensor','torch.cuda.BFloat16Tensor']:# 如果是fp16/bf16类型参数，clone为fp32类型的切片.shard_model_param = model_param.detach().view(-1) \[param_range.start:param_range.end]shard_main_param = shard_model_param.clone().float()...# 添加到group中model_float16_params_this_group.append(model_param)shard_float16_params_this_group.append(shard_model_param)shard_fp32_from_float16_params_this_group.append(shard_main_param)elif model_param.type() == 'torch.cuda.FloatTensor':# 如果是fp32类型参数，不进行clone，直接引用shard_model_param = model_param.view(-1) \[param_range.start:param_range.end]model_fp32_params_this_group.append(model_param)shard_fp32_params_this_group.append(shard_model_param)...# 更新优化器的参数group_range["orig_group"]["params"] = [*shard_fp32_params_this_group,*shard_fp32_from_float16_params_this_group,]return (model_float16_groups,model_fp32_groups,shard_float16_groups,shard_fp32_groups,shard_fp32_from_float16_groups,)
  • 在Optimizer init中，接下来是初始化self.param_buffers，这里的self.param_buffers是DDP模型的grad buffer的view示图，跟grad buffer共享存储，但是用自己的数据类型；最后更新优化器的param_groups。
class DistributedOptimizer(MixedPrecisionOptimizer):def __init__(...):...# 初始化self.param_buffersself.param_buffers = []for model_index, model in enumerate(self.models):current_param_buffers = {}for dtype, grad_buffer in model._grad_buffers.items():# 获取存储，这里是兼容的写法.try:storage = grad_buffer.data.storage()._untyped()except:storage = grad_buffer.data.storage().untyped()# 基于grad_buffer的storage创建param_buffer类型，这里的params_dtype是参数类型; 这里的torch.tensor没有autograd的历史。param_buffer = torch.tensor(storage,dtype = params_dtype,device = grad_buffer.data.device)param_buffer = param_buffer[:grad_buffer.numel_padded]# 这里的dtype是grad_buffer的类型current_param_buffers[dtype] = param_bufferself.param_buffers.append(current_param_buffers)# 最后更新优化器的param_groupsself.optimizer.param_groups = \[ g["orig_group"] for g in self.opt_group_ranges ]self.optimizer.load_state_dict(self.optimizer.state_dict())

3. 参考

  • Megatron-LM源码系列(七)：Distributed-Optimizer分布式优化器实现Part2
  • Megatron-LM源码系列(六)：Distributed-Optimizer分布式优化器实现Part1
  • NVIDIA/Megatron-LM

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

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

相关文章

C#，雅各布斯塔尔—卢卡斯（Jacobsthal Lucas Number）的算法与源代码

C#，雅各布斯塔尔—卢卡斯（Jacobsthal Lucas Number）的算法与源代码

1 雅各布斯塔尔序列 雅各布斯塔尔序列是一个与斐波那契序列类似的加法序列&#xff0c;由递归关系JnJn-12Jn-2定义&#xff0c;初始项J00&#xff0c;J11。序列中的一个数字称为雅可布沙尔数。它们是卢卡斯序列Un&#xff08;P&#xff0c;Q&#xff09;的一种特殊类型&#x…
阅读更多...
股票交易

股票交易

这里尝试利用单调队列优化&#xff0c;这里不好直接用单调队列的原因是因为(以买为例)\(-ap[i]*k_1\)不是只与下标有关的 所以解决方案一&#xff1a;我们将下标变成一个整体&#xff0c;再把后面的代价换掉然后将与下标无关的直接提出去 解决方案二&#xff1a;利用“蚯蚓”那…
阅读更多...
＜设计模式＞单例模式懒汉和饿汉

＜设计模式＞单例模式懒汉和饿汉

目录 一、单例模式概述 二、懒汉模式和饿汉模式 1.饿汉模式 1.1代码实现 1.2实现细节 1.3模式优劣 2.懒汉模式 2.1代码实现 2.2实现细节 2.3模式优劣 三、多线程下的线程安全问题 1.懒汉和饿汉线程安全问题分析 1.1安全的饿汉模式 1.2不安全的懒汉模式 2.懒汉线程…
阅读更多...
YOLOv5算法进阶改进（15）— 引入密集连接卷积网络DenseNet

YOLOv5算法进阶改进（15）— 引入密集连接卷积网络DenseNet

前言：Hello大家好，我是小哥谈。DenseNet（密集连接卷积网络）是一种深度学习神经网络架构，它在2017年由Gao Huang等人提出。DenseNet的核心思想是通过密集连接（dense connection）来促进信息的流动和共享。在传统的卷积神经网络中，每个层的输入只来自于前一层的输出。而在…
阅读更多...
Linux下编译EtherCAT主站SOEM-1.4.1

Linux下编译EtherCAT主站SOEM-1.4.1

目录 1、SOEM下载 2、CMake安装​​​​​​ 3、编译 环境&#xff1a;Ubuntu1604. 1、SOEM下载 最新版为SOEM-v1.4.0&#xff0c;可以从github下载地址&#xff1a; https://github.com/OpenEtherCATsociety/SOEM 2、CMake安装​​​​​​ 3、编译 解压文件&#xff0c…
阅读更多...
WebSocket学习笔记以及用户与客服聊天案例简单实现（springboot+vue）

WebSocket学习笔记以及用户与客服聊天案例简单实现（springboot+vue）

一&#xff1a;介绍&#xff1a; 二&#xff1a;http协议与websocket对比&#xff1a; 三&#xff1a;websocket协议&#xff1a; 四&#xff1a;实现&#xff1a; 4.1客户端&#xff1a; 4.2服务端&#xff1a; 五&#xff1a;案例&#xff1a; 环境&#xff1a;做一个书店…
阅读更多...
分布式任务调度框架XXL-JOB详解

分布式任务调度框架XXL-JOB详解

分布式任务调度 概述 场景: 如12306网站根据不同车次设置放票时间点&#xff0c;商品成功发货后向客户发送短信提醒等任务,某财务系统需要在每天上午10天前统计前一天的账单数据 任务的调度是指系统为了完成特定业务&#xff0c;基于给定的时间点&#xff0c;时间间隔&#…
阅读更多...
onnx转换为rknn置信度大于1，图像出现乱框问题解决

onnx转换为rknn置信度大于1，图像出现乱框问题解决

前言 环境介绍&#xff1a; 1.编译环境 Ubuntu 18.04.5 LTS 2.RKNN版本 py3.8-rknn2-1.4.0 3.单板 迅为itop-3568开发板 一、现象 采用yolov5训练并将pt转换为onnx&#xff0c;再将onnx采用py3.8-rknn2-1.4.0推理转换为rknn出现置信度大于1&#xff0c;并且图像乱框问题…
阅读更多...
【服务器】RAID（独立磁盘冗余阵列）

【服务器】RAID（独立磁盘冗余阵列）

RAID&#xff08;独立磁盘冗余阵列&#xff09; 一、RAID的介绍二、RAID的分类#2-1 RAID 02-2 RAID 1#2-3 RAID 32-4 RAID 52-5 RAID 62-6 RAID 10(先做镜像&#xff0c;再做条带化)2-7 RAID 01&#xff08;先做条带&#xff0c;再做镜像&#xff09;2-8 RAID比较 三、磁盘阵列…
阅读更多...
代码随想录刷题第24天

代码随想录刷题第24天

今天正式进入回溯。看了看文章介绍&#xff0c;回溯并不是很高效的算法&#xff0c;本质上是穷举操作。代码形式较为固定。 第一题为组合问题&#xff0c;用树形结构模拟&#xff0c;利用回溯算法三部曲&#xff0c;确定终止条件与单层逻辑&#xff0c;写出如下代码。 不难发现…
阅读更多...
负载均衡下webshell连接

负载均衡下webshell连接

目录 一、什么是负载均衡 分类 负载均衡算法 分类介绍 分类 均衡技术 主要应用 安装docker-compose 2.1上传的文件丢失 2.2 命令执行时的漂移 2.3 大工具投放失败 2.4 内网穿透工具失效 3.一些解决方案 总结 一、什么是负载均衡 负载均衡&#xff08;Load Balanc…
阅读更多...
网络安全挑战：威胁建模的应对策略与实践

网络安全挑战：威胁建模的应对策略与实践

在数字威胁不断演变的时代&#xff0c;了解和降低网络安全风险对各种规模的组织都至关重要。威胁建模作为安全领域的一个关键流程&#xff0c;提供了一种识别、评估和应对潜在安全威胁的结构化方法。本文将深入探讨威胁建模的复杂性&#xff0c;探索其机制、方法、实际应用、优…
阅读更多...
python爬虫5

python爬虫5

1.selenium交互 无页面浏览器速度更快 #配置好的自己不用管 from selenium import webdriverfrom selenium.webdriver.chrome.options import Optionschrome_options Options()chrome_options.add_argument(‐‐headless)chrome_options.add_argument(‐‐disable‐gpu)# path…
阅读更多...
109.乐理基础-五线谱-五线谱的附点、休止符、连线、延音线

109.乐理基础-五线谱-五线谱的附点、休止符、连线、延音线

内容参考于&#xff1a;三分钟音乐社 上一个内容&#xff1a;五线谱的拍号、音符与写法-CSDN博客 上一个内容里练习的答案&#xff1a; 附点&#xff1a;写在符头的右方&#xff0c;附点的作用与简谱一样&#xff0c;延长前面音符本身时值的一半&#xff08;附点&#xff09;…
阅读更多...
Hadoop3.x基础（3）- Yarn

Hadoop3.x基础（3）- Yarn

来源&#xff1a;B站尚硅谷 目录 Yarn资源调度器Yarn基础架构Yarn工作机制作业提交全过程Yarn调度器和调度算法先进先出调度器&#xff08;FIFO&#xff09;容量调度器&#xff08;Capacity Scheduler&#xff09;公平调度器&#xff08;Fair Scheduler&#xff09; Yarn常用命…
阅读更多...
回归预测 | Matlab实现POA-CNN-LSTM-Attention鹈鹕算法优化卷积长短期记忆网络注意力多变量回归预测（SE注意力机制）

回归预测 | Matlab实现POA-CNN-LSTM-Attention鹈鹕算法优化卷积长短期记忆网络注意力多变量回归预测（SE注意力机制）

回归预测 | Matlab实现POA-CNN-LSTM-Attention鹈鹕算法优化卷积长短期记忆网络注意力多变量回归预测&#xff08;SE注意力机制&#xff09; 目录 回归预测 | Matlab实现POA-CNN-LSTM-Attention鹈鹕算法优化卷积长短期记忆网络注意力多变量回归预测&#xff08;SE注意力机制&…
阅读更多...
RocketMQ—RocketMQ发送同步、异步、单向、延迟、批量、顺序、批量消息、带标签消息

RocketMQ—RocketMQ发送同步、异步、单向、延迟、批量、顺序、批量消息、带标签消息

RocketMQ—RocketMQ发送同步、异步、单向、延迟、批量、顺序、批量消息、带标签消息 发送同步消息 生产者发送消息&#xff0c;mq进行确认&#xff0c;然后返回给生产者状态。这就是同步消息。 前文demo程序就是发送的同步消息。 发送异步消息 异步消息通常用在对响应时间敏…
阅读更多...
gorm day1

gorm day1

gorm day1 gorm简介gorm声明模型 代码样例基本来自官方文档 Gorm简介 什么是ORM&#xff1f; 对象关系映射(Objection Relational Mapping,简称ORM)模式是一种为了解决面向对象与关系数据库(如mysql数据库&#xff09;存在的互不匹配现象的计数。简单来说&#xff0c;ORM是通…
阅读更多...
计算机毕设医院挂号预约系统ssm

计算机毕设医院挂号预约系统ssm

项目运行 环境配置&#xff1a; Jdk1.8 Tomcat7.0 Mysql HBuilderX&#xff08;Webstorm也行&#xff09; Eclispe&#xff08;IntelliJ IDEA,Eclispe,MyEclispe,Sts都支持&#xff09;。 项目技术&#xff1a; vue mybatis Maven mysql5.7或8.0等等组成&#xff0c;B…
阅读更多...
【Redis】整理

【Redis】整理

对于现代大型系统而言&#xff0c;缓存是一个绕不开的技术话题&#xff0c;一提到缓存我们很容易想到Redis。 Redis整理&#xff0c;供回顾参考
阅读更多...
推荐文章
最新文章

Copyright @ 2022~2023