业余做网站,天津哪家制作网站好,泰安可信的网站建设,移动ui界面设计课程总结1. 使用入口
DistributedOptimizer类定义在megatron/optimizer/distrib_optimizer.py文件中。创建的入口是在megatron/optimizer/__init__.py文件中的get_megatron_optimizer函数中。根据传入的args.use_distributed_optimizer参数来判断是用DistributedOptimizer还是Float16O…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_condNone,scale_lr_condNone,lr_mult1.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.startbuild_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_startmax_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 DPs 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 rangeOptimizer原本有param_groups包括多个参数组这里build_optimizer_group_ranges为了创建param参数到group_index的map映射也就是model_parameter:group_indexself.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_indexworld_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分布式优化器实现Part2Megatron-LM源码系列(六)Distributed-Optimizer分布式优化器实现Part1NVIDIA/Megatron-LM
