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

2024-02-03 12:52
文章标签 实现 源码 优化 系列 分布式 part2 distributed optimizer lm megatron

本文主要是介绍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分布式优化器实现Part1
  • NVIDIA/Megatron-LM

这篇关于Megatron-LM源码系列(七)：Distributed-Optimizer分布式优化器实现Part2的文章就介绍到这儿，希望我们推荐的文章对编程师们有所帮助！


http://www.chinasem.cn/article/674206。

相关文章

Vue3 的 shallowRef 和 shallowReactive：优化性能

Vue3 的 shallowRef 和 shallowReactive：优化性能

大家对 Vue3 的 ref 和 reactive 都很熟悉，那么对 shallowRef 和 shallowReactive 是否了解呢？ 在编程和数据结构中，“shallow”（浅层）通常指对数据结构的最外层进行操作，而不递归地处理其内部或嵌套的数据。这种处理方式关注的是数据结构的第一层属性或元素，而忽略更深层次的嵌套内容。 1. 浅层与深层的对比 1.1 浅层（Shallow） 定义
阅读更多...
Spring Security 从入门到进阶系列教程

Spring Security 从入门到进阶系列教程

Spring Security 入门系列 《保护 Web 应用的安全》 《Spring-Security-入门（一）：登录与退出》 《Spring-Security-入门（二）：基于数据库验证》 《Spring-Security-入门（三）：密码加密》 《Spring-Security-入门（四）：自定义-Filter》 《Spring-Security-入门（五）：在 Sprin
阅读更多...
HDFS—存储优化（纠删码）

HDFS—存储优化（纠删码）

纠删码原理 HDFS 默认情况下，一个文件有3个副本，这样提高了数据的可靠性，但也带来了2倍的冗余开销。 Hadoop3.x 引入了纠删码，采用计算的方式，可以节省约50％左右的存储空间。 此种方式节约了空间，但是会增加 cpu 的计算。 纠删码策略是给具体一个路径设置。所有往此路径下存储的文件，都会执行此策略。 默认只开启对 RS-6-3-1024k
阅读更多...
hdu1043（八数码问题，广搜 + hash（实现状态压缩） ）

hdu1043（八数码问题，广搜 + hash（实现状态压缩） ）

利用康拓展开将一个排列映射成一个自然数，然后就变成了普通的广搜题。 #include<iostream>#include<algorithm>#include<string>#include<stack>#include<queue>#include<map>#include<stdio.h>#include<stdlib.h>#include<ctype.h>#inclu
阅读更多...
使用opencv优化图片（画面变清晰）

使用opencv优化图片（画面变清晰）

文章目录 需求影响照片清晰度的因素 实现降噪测试代码 锐化空间锐化Unsharp Masking频率域锐化对比测试 对比度增强常用算法对比测试 需求 对图像进行优化，使其看起来更清晰，同时保持尺寸不变，通常涉及到图像处理技术如锐化、降噪、对比度增强等 影响照片清晰度的因素 影响照片清晰度的因素有很多，主要可以从以下几个方面来分析 1. 拍摄设备 相机传感器：相机传
阅读更多...
JAVA智听未来一站式有声阅读平台听书系统小程序源码

JAVA智听未来一站式有声阅读平台听书系统小程序源码

智听未来，一站式有声阅读平台听书系统 🌟&nbsp;开篇：遇见未来，从“智听”开始 在这个快节奏的时代，你是否渴望在忙碌的间隙，找到一片属于自己的宁静角落？是否梦想着能随时随地，沉浸在知识的海洋，或是故事的奇幻世界里？今天，就让我带你一起探索“智听未来”——这一站式有声阅读平台听书系统，它正悄悄改变着我们的阅读方式，让未来触手可及！ 📚&nbsp;第一站：海量资源，应有尽有 走进“智听
阅读更多...
【C++】_list常用方法解析及模拟实现

【C++】_list常用方法解析及模拟实现

相信自己的力量，只要对自己始终保持信心，尽自己最大努力去完成任何事，就算事情最终结果是失败了，努力了也不留遗憾。💓💓💓 目录   ✨说在前面 🍋知识点一：什么是list？ •🌰1.list的定义 •🌰2.list的基本特性 •🌰3.常用接口介绍 🍋知识点二：list常用接口 •🌰1.默认成员函数 🔥构造函数（⭐） 🔥析构函数 •🌰2.list对象
阅读更多...
【Prometheus】PromQL向量匹配实现不同标签的向量数据进行运算

【Prometheus】PromQL向量匹配实现不同标签的向量数据进行运算

✨✨ 欢迎大家来到景天科技苑✨✨ 🎈🎈 养成好习惯，先赞后看哦~🎈🎈 🏆 作者简介：景天科技苑 🏆《头衔》：大厂架构师，华为云开发者社区专家博主，阿里云开发者社区专家博主，CSDN全栈领域优质创作者，掘金优秀博主，51CTO博客专家等。 🏆《博客》：Python全栈，前后端开发，小程序开发，人工智能，js逆向，App逆向，网络系统安全，数据分析，Django，fastapi
阅读更多...
让树莓派智能语音助手实现定时提醒功能

让树莓派智能语音助手实现定时提醒功能

最初的时候是想直接在rasa 的chatbot上实现，因为rasa本身是带有remindschedule模块的。不过经过一番折腾后，忽然发现，chatbot上实现的定时，语音助手不一定会有响应。因为，我目前语音助手的代码设置了长时间无应答会结束对话，这样一来，chatbot定时提醒的触发就不会被语音助手获悉。那怎么让语音助手也具有定时提醒功能呢？ 我最后选择的方法是用threading.Time
阅读更多...
Android实现任意版本设置默认的锁屏壁纸和桌面壁纸（两张壁纸可不一致）

Android实现任意版本设置默认的锁屏壁纸和桌面壁纸（两张壁纸可不一致）

客户有些需求需要设置默认壁纸和锁屏壁纸  在默认情况下 这两个壁纸是相同的  如果需要默认的锁屏壁纸和桌面壁纸不一样 需要额外修改 Android13实现 替换默认桌面壁纸： 将图片文件替换frameworks/base/core/res/res/drawable-nodpi/default_wallpaper.*  （注意不能是bmp格式） 替换默认锁屏壁纸： 将图片资源放入vendo
阅读更多...

Copyright China编程 23002807@qq.com

沪ICP备2023033072号-2