跟代码执行流程，读Megatron源码（一）从目录结构开始

最近看了很多篇写Megatron的文章，但大部分都是按其功能模块划分做源码的解析，很少有人会从程序的入口开始，跟随代码运行的流程，将Megatron从训练开始到结束，端到端分析完整。 我会用面向过程的方式，以代码debug的视角，分析Megatron源码，从入门到精通。本文假设读者已经具备了机器学习和深度学习的基础。

一. 为什么研究Megatron

　　2022年底~2023年初，随着ChatGPT的火爆，大模型在我司（也是在业界）成了热点，公司层面组织了专门的团队训练自己的大模型，而具体到每个部门，也纷纷投入大量的人力围绕大模型做起了文章，我们也不例外。

　　不可避免的，在大模型训练的场景下，Megatron成了必选。

　　Megatron是NVIDIA深度学习应用研究团队设计开发的AI模型分布式训练框架，它解决了训练超大规模模型时面临的计算和存储挑战（这些模型通常包含数十亿甚至数千亿个参数，单个GPU或计算节点无法容纳整个模型）。同时，国内也有基于Megatron思想的开源实现，如华为的AscendSpeed/ModelLink（由于AscendSpeed和ModelLink是迭代演进的关系，后文中统一称为ModelLink）。

　　不论是ModelLink还是Megatron，其实现思路和源码不能说完全一致，只能说几乎一模一样。不同的只是AscendSpeed/ModelLink中包含了大量适配了华为Atlas服务器（昇腾910系列）的模型源码，这对于使用昇腾产品的用户更为易用。故，文章会时不时的对两者做个比较分析，以便让大家同时了解国内外的两个开源框架。

　　本文会从代码入口开始，逐行debug程序运行流程，直至其运行结束，最终展示完整的Megatron代码流程图。期间涉及到的分布式训练相关概念，会在遇到时再行展开。

二. 源码目录分析

图一：Megatron-LM源码目录

图二：ModelLink源码目录

　　以上是英伟达开发的Megatron-LM并行训练框架（master分支）、以及华为开发的ModelLink并行训练框架（bk_origin_23分支）源码仓截图。其中ModelLink之所以没选择最新版本分支，是因为其代码几经修改，已经无法直观的跟Megatron-LM做对比了，故选择其原始版本，有助于我们对比学习。

　　在源码的根目录下会看到多个以pretrain_xxx.py命名的文件，这些文件主要用于执行不同的预训练任务，且每个pretrain_xxx.py文件通常对应一种或一类特定模型的预训练任务（如pretrain_bert.py用来执行bert的训练任务，pretrain_gpt.py用来执行gpt的训练任务）。由下图可以看出，ModelLink支持的模型种类明显少于Megatron-LM，但在主流大模型的支持上，两者能力差距不大。

　　下文中，我们将以pretrain_gpt.py为起点，逐步深入解读Megatron源码。

　　了解了预训练任务的入口，再看根目录下的其他文件夹（图三）。在Megatron-LM中，最主要的文件夹是megatron，它包含了模型的数据处理、模型定义、初始化、并行训练等所有关键模块，是megatron框架的核心实现。ModelLink中的ascendspeed文件夹的功能与megatron基本等同，在此不做赘述。

图三：根目录下文件夹

　　直接访问Megatron-LM与ModelLink的核心实现目录（即megatron与ascendspeed），鉴于版本发展的差异性（ascendspeed基于Megatron的早期版本，且未同步其后续更新），两者目录结构呈现出显著不同。尽管如此，它们之间仍保持着一定的映射关系，并非全然不相干，如图四所示。对于每个目录和文件的功能特性，我们将在追踪代码执行流程时逐一阐述，即遇即解。

图四：megatron和ascendspeed对应关系

　　至此，对Megatron-LM和ModelLink的源码目录的基本介绍已经完成，下面会从预训练入口（pretrain_gpt.py）开始，逐行拆解megatron的内核。