跟代码执行流程,读Megatron源码(四)megatron初始化脚本initialize.py之initialize_megatron()分布式环境初始化

  在前文中,我们讲述了pretrain函数的执行流程,其首要步骤是megatron分组的初始化与环境的配置。本文将深入initialize_megatron函数源码,剖析其初始化分布式训练环境的内部机制。

  注:在此假设读者具备3D并行相关知识

一. initialize_megatron函数的上下文调用关系(initialize.py)

  在Megatron-LM中,initialize.py文件中的initialize_megatron函数是分布式训练环境的初始化核心。该函数由trainning.py的pretrain函数调用,是整个pretrain流程中的第一个核心步骤,负责配置3D并行的环境和分组信息。

1. initialize_megatron函数源码

  initialize_megatron函数的核心代码段如下:

  需要注意的是,尽管initialize_megatron函数还涵盖了设置全局参数、分词器构建、自动恢复配置、TensorBoard日志记录、计时器设置以及依赖编译等辅助功能,但这些功能在初始化流程中虽具重要性,却非其核心职责所在。其核心功能聚焦于上述代码段所描述的分布式与模型并行初始化流程,而该流程主要是通过finish_mpu_init中调用的_initialize_distributed函数实现的,如下图。

2. _initialize_distributed

  _initialize_distributed函数主要有两个作用:

  a. 通过调用torch.distributed.init_process_group()初始化分布式环境,该函数设置了分布式训练所需的基本通信环境,包括进程间的通信后端、worldsize(参与训练的进程总数)、每个进程的rank号等。默认情况下,它会创建一个全局的进程组,这个进程组定义了哪些进程可以相互通信,也可以根据需要创建更多的进程组以支持更复杂的通信模式。

  在使用 init_process_group() 初始化分布式环境之后,我们可以使用PyTorch提供的分布式通信和同步 API 来实现跨进程的通信和数据同步。这包括使用dist.all_reduce() 来聚合梯度、dist.barrier() 同步所有进程的执行点等。

  请注意,在调用init_process_group()之前,需要确保已经正确设置了所有相关的环境变量(如MASTER_ADDR、MASTER_PORT),并且这些环境变量对于每张卡都是唯一的。

  b. 通过调用mpu.initialize_model_parallel()来初始化分布式训练环境中的数据并行(DP)、张量并行(TP)、和流水线并行(PP)的分组,如下图。

  mpu.initialize_model_parallel()的入参解释如下:

  • tensor_model_parallel_size:张量并行的大小。

  • pipeline_model_parallel_size:流水线并行的大小。

  • virtual_pipeline_model_parallel_size:虚拟流水线并行的大小,这是一个更高级的特性,允许在流水线阶段内部进一步分割模型。

  • pipeline_model_parallel_split_rank:该参数指定了流水线分割的起始rank,即决定了哪个rank的设备将开始处理流水线的第一个阶段,然后接下来的阶段按顺序分配给rank号递增的设备。

  • context_parallel_size、expert_model_parallel_size等参数用于特定的模型架构,如带有上下文并行或专家并行的Transformer模型。

  • distributed_timeout_minutes:分布式操作的超时时间(以分钟为单位)。

  • nccl_communicator_config_path:NCCL通信器的配置文件路径,NCCL是用于NVIDIA GPU的高效通信库。

  • order:指定并行策略的顺序,例如'tp-cp-ep-dp-pp'表示张量并行(Tensor Parallelism)、上下文并行(Context Parallelism)、专家并行(Expert Parallelism)、数据并行(Data Parallelism)和流水线并行(Pipeline Parallelism)的顺序。

  • encoder_pipeline_model_parallel_size:专门用于编码器的流水线并行大小。

  • get_embedding_ranks和get_position_embedding_ranks:用于获取特定用于embedding或position embedding的GPU rank号,以便为这些组件配置特定的并行策略。

二. mpu.initialize_model_parallel函数的调用关系(_init_.py)

  mpu.initialize_model_parallel()函数的调用关系在import中表明,如下图:

  其中mpu定义在megatron/core/_init_.py中:

  如上图,mpu指向parallel_state,因此,对于mpu.initialize_model_parallel()的调用既是对parallel_state.initialize_model_parallel()的调用,该函数实现在parallel_state.py中。

三. 分组逻辑的具体实现(parallel_state.py)

1. parallel_state.initialize_model_parallel的调用关系

  parallel_state.initialize_model_parallel函数在分布式训练架构中扮演着关键但非终结性的角色,其核心功能是启动模型并行所需的基础设置。该函数专注于预配置模型并行相关的进程群组(process groups)与全局状态变量,确保并行执行环境的基础通信框架得以确立。

  下面以dp为例,展示该函数的代码执行逻辑。

  首先,该函数创建RankGenerator对象实体,该对象的作用是根据用户输入的tp/dp/pp大小,以及总卡数(world_size),确定最终每张卡的分组,如下图。

  其次,它首先调用RankGenerator组件动态生成高效的rank分配方案,这一步骤是优化资源利用与通信效率的关键。随后,通过调用后端(backend)接口,根据RankGenerator产出的分组策略,构建起实际用于数据交换的通信群组(communication groups)。这些通信群组覆盖了数据并行(dp)、张量并行(tp)、以及流水线并行(pp)等维度(这里只以dp为例),确保模型训练过程中的数据流通与参数同步能够高效且有序地进行,如下图。

  随着所有必需通信群组的成功建立,以及全局分组变量的初始化,模型并行的核心通信网络得以全面搭建完成。这一网络的构建不仅标志着模型并行化训练环境的初步就绪,更为后续的高性能计算任务奠定了坚实的基础,确保了分布式训练过程中数据一致性与效率的最优化实现。

2. RankGenerator.get_ranks

  RankGenerator,顾名思义,其主要职责为生成rank分组,它管理着用户的tp/pp/dp/ep/cp数值,以及全局卡数(world_size)等分组相关的配置项,并在get_ranks函数中调用generate_masked_orthogonal_rank_groups函数获取用户需要的最终分组信息,如下图:

3. generate_masked_orthogonal_rank_groups

  generate_masked_orthogonal_rank_groups函数是rank分组的最终实现,其代码逻辑如下:

  a. 筛选并行性尺寸:

  masked_shape:从parallel_size和mask中筛选出被掩码(即True)的并行尺寸,这些尺寸将用于生成组内的rank。

  unmasked_shape:同样从parallel_size和mask中筛选出未被掩码的并行尺寸,这些尺寸将用于在更广泛的并行环境中(跨组)定位每个组。

  b. 计算步长:

  global_stride:通过prefix_product(parallel_size)计算得到全局步长。

  masked_stride和unmasked_stride:分别根据mask从global_stride中筛选出被掩码和未被掩码的步长。这些步长用于计算全局rank。

  c. 确定组大小和组数:

  group_size:通过prefix_product(masked_shape)[-1]计算得到每个组的大小。

  num_of_group:通过world_size // group_size计算得到组的数量,即全局大小除以每个组的大小。

  d. 生成rank:

  遍历每个组(group_index从0到num_of_group-1)。

  使用decompose(group_index, unmasked_shape)根据未被掩码的并行尺寸分解组索引,得到该组在全局并行环境中的位置(decomposed_group_idx)。

  在每个组内,遍历每个rank(rank_in_group从0到group_size-1)。

  使用decompose(rank_in_group, masked_shape)根据被掩码的并行尺寸分解组内rank,得到该rank在组内的位置(decomposed_rank_idx)。

  计算每个rank的全局索引,通过将被掩码和未被掩码的索引分别与其对应的步长进行内积(inner_product),然后将两个内积相加得到。

  最后,将计算得到的每个组内的rank添加到ranks列表中。

  e. 返回rank列表:

  函数最终返回ranks列表,其中包含了每个组内的所有rank,这些rank在全局并行环境中是唯一的。

  至此分布式训练分组的全部逻辑均介绍完毕,后续文章将继续解析分组完成后的训练逻辑。

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

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

相关文章

【MARL】MADDPG + attention 实现(+论文解读)

文章目录 前言注意力机制论文里的attention回顾知识-MADDPG讲解1.Q的定义2.Q的恒等式3.论文里的attention4.好处 实现 和 修改结果展示原论文代码 翻改版修改后原maddpg代码 前言 导师让在MADDPG上加一个注意力机制,试了很多种,下面的参考的论文的效果最…

C++——保持原有库头文件不变的情况下,成功编译运行工程

问:想要保持原来库方式,应该怎么操作呢? 答:如果想保持原来的方式,则只需要将 库所在路径 tracker/detector/rknn_model_zoo/utils 加入到 工程库包含中即可。

基于jeecgboot-vue3的Flowable流程-自定义业务表单流程历史信息显示

因为这个项目license问题无法开源&#xff0c;更多技术支持与服务请加入我的知识星球。 1、对于自定义业务表单的流程历史记录信息做了调整&#xff0c;增加显示自定义业务表单 <el-tab-pane label"表单信息" name"form"><div v-if"customF…

ESP32开发进阶:OLED屏幕显示旋转的3D模型

一、硬件接线 我选择的是最常见的一块板子&#xff1a;ESP-WROOM-32&#xff0c;硬件接线如下&#xff1a; 21 - SDA 22 - SCL 二、Arduino端代码 我们使用Arduino和Adafruit SSD1306库在OLED显示屏上绘制和旋转一个3D立方体。 首先&#xff0c;定义立方体顶点和…

国际金融入门:国际收支与平衡表解析

在全球化的经济体系中&#xff0c;国际金融已成为我们日常生活不可或缺的一部分。了解国际金融的基础知识&#xff0c;可以帮助我们更好地理解世界经济的动态和趋势。今天&#xff0c;我们将深入探讨国际收支及其平衡表&#xff0c;以及它们是如何影响国家经济。 国际收支&…

2024.07-电视版免费影视App推荐和猫影视catvod、TVBox源(最新接口地址)

文章目录 电视版免费影视App推荐精选列表&#xff08;2024.07可用筛选列表&#xff09;&#xff1a;2024.07可用筛选列表&#xff0c;盲盒资源打包合集下载安装说明真的是盲盒&#xff1f; 猫影视catvod、TVBoxTVBox源推荐可用列表目前不可用列表&#xff08;前缀为错误状态码&…

探索 Odoo 前端源码

今天我们要深入探索一段 Odoo 前端源码。这段代码看似简单,却隐藏着 Odoo 应用高效运行的秘密。我们一起来解读这段源码,了解它在 Odoo 中的作用,并教你如何正确使用它。准备好了吗?让我们开始这段奇妙之旅吧! /** @odoo-module **/import { registry } from "./cor…

CSS(七)——CSS 列表和CSS Table(表格)

目录 CSS 列表 列表 作为列表项标记的图像 列表 - 简写属性 移除默认设置 所有的CSS列表属性 CSS 表格 表格边框 折叠边框&#xff08;border-collapse&#xff09; 表格宽度和高度 表格文字对齐 表格填充 表格颜色 CSS 列表 CSS 列表属性作用如下&#xff1a; 设…

mysql 主从复制 读写分离 MHA

mysql 的主从复制和读写分离&#xff1a; 读写分离和MHA高可用的前提 主从复制 主从复制的模式&#xff1a; 1.mysql的默认模式 异步模式&#xff1a;主库在更新完事务之后会立即把结果返回给从服务器&#xff0c;不关心从库是否接收到&#xff0c;是否处理成功 网络问题可能…

大模型学习应用 1:用 itrex 创新高效实现 LLM 的部署和微调

用 itrex 创新高效实现 LLM 的部署和微调 - 项目作业 目录 准备工作Task 1 完成在线环境的工具包安装&#xff0c;包含 基础环境包、Extension for Transformers 包、加速计算包Task 2 利用 Intel Extension for Transformers 部署通义千问 Qwen-7B Chat&#xff0c;并根据 pr…

zynqmpsoc On-chip Memory

zynqmpsoc 开发手册 On-chip Memory 章节 Introduction The on-chip memory (OCM) contains 256 KB of RAM. It supports a 128-bit AXI slave interface port. The OCM has eight exclusive access monitors that can simultaneously monitor up to eight exclusive access tr…

前端系列-7 Vue3响应式数据

1.响应式数据和实现 响应式数据的核心是当数据模型发生变化时&#xff0c;与之相关(使用该数据模型)的视图或组件可以自动更新&#xff0c;以反映最新的数据状态。 实现原理是数据劫持、依赖收集和分发更新&#xff0c;数据劫持是一种AOP策略。 在vue2中通过Object.defineProp…

(四)js前端开发中设计模式之简单工厂模式

简单工厂模式,又叫静态工厂方法&#xff0c;由一个工厂对象决定创建出哪一种产品类的实例&#xff0c;主要用来创建同一类对象 let LoginAlert function (msg) {this.content msg } LoginAlert.prototype {show() {const div document.createElement(div)div.style.cssText…

C#开发的全屏图片切换效果应用 - 开源研究系列文章 - 个人小作品

这天无聊&#xff0c;想到上次开发的图片显示软件《 PhotoNet看图软件 》&#xff0c;然后想到开发一个全屏图片切换效果的应用&#xff0c;类似于屏幕保护程序&#xff0c;于是就写了此博文。这个应用比较简单&#xff0c;主要是全屏切换换图片效果的问题。 1、 项目目录&…

设计模式实战:文件系统的设计与实现

问题描述 设计一个文件系统模拟器,用户可以创建文件和文件夹,对文件进行复制、删除等操作,并支持克隆文件和文件夹。系统需要实现组合模式来处理文件和文件夹的层次结构,使用命令模式来封装操作,并利用原型模式实现文件和文件夹的克隆。 设计分析 组合模式 组合模式用…

自动驾驶-机器人-slam-定位面经和面试知识系列04之高频面试题(02)

这个博客系列会分为C STL-面经、常考公式推导和SLAM面经面试题等三个系列进行更新&#xff0c;基本涵盖了自己秋招历程被问过的面试内容&#xff08;除了实习和学校项目相关的具体细节&#xff09;。在知乎和牛客&#xff08;牛客上某些文章上会附上内推码&#xff09;也会同步…

【Vue3】watch 监视 ref 定义的数据

【Vue3】watch 监视 ref 定义的数据 背景简介开发环境开发步骤及源码参数说明 背景 随着年龄的增长&#xff0c;很多曾经烂熟于心的技术原理已被岁月摩擦得愈发模糊起来&#xff0c;技术出身的人总是很难放下一些执念&#xff0c;遂将这些知识整理成文&#xff0c;以纪念曾经努…

【C++进阶学习】第八弹——红黑树的原理与实现——探讨树形结构存储的最优解

二叉搜索树&#xff1a;【C进阶学习】第五弹——二叉搜索树——二叉树进阶及set和map的铺垫-CSDN博客 AVL树&#xff1a; ​​​​​​【C进阶学习】第七弹——AVL树——树形结构存储数据的经典模块-CSDN博客 前言&#xff1a; 在前面&#xff0c;我们已经学习了二叉搜索树和…

PCIe 6.0为什么需要14-bit tag

1.TLP中的tag是什么 在PCIe TLP&#xff08;Transaction Layer Packet&#xff09;中&#xff0c;tag是分配给特定Non-Posted Request的编号&#xff0c;协议要求CPL/CPLD中的tag 与对应non-post request TLP中的tag保持一致&#xff0c;因此Requester可以使用tag来识别CPL…

Windows通过命令查看mac : getmac

要查看本机网卡mac&#xff0c;可以通过ipconfig /all 显示&#xff0c;但输出内容过多 可以通过getmac命令查看 示例 C:\Users\Desktop> getmac物理地址 传输名称暂缺 没有硬件 1C-1B-B5-04-E2-7D \Device\Tcpip_{80096E40-D51D-490C-9AF7-…