圖解大模型系列之：Megatron源碼解讀1，分布式環(huán)境初始化

一晃過去大半個月，終于有時間來寫Megatron的源碼解讀篇了。

首先，請允許我介紹下封面。明明是Megatron，為什么放Bee??？還不是因為Megatron長得太丑了。翻遍了網(wǎng)絡(luò)全都是一坨灰加上兩只紅色眼睛，實在是有礙閱讀心情。。。放個明黃色舒爽下眼睛。

源碼閱讀類的文章很難寫。尤其是對Megatron這樣細(xì)節(jié)豐富，代碼結(jié)構(gòu)上又較為松散的項目而言。思考了一下，我決定依然用自己最擅長的【圖解】方式，和大家一起閱讀源碼。在這個系列里，我基本按以下3步驟來做解讀：

• 先通過【圖解】的方式，說明這塊代碼在做一件什么事
• 闡述代碼整體架構(gòu)，拆分邏輯
• 細(xì)節(jié)解讀

一、CodeGeeX模型簡述

使用Megatron來訓(xùn)練gpt類大模型的項目有很多。在這個系列里，我選擇了由THUDM開發(fā)的CodeGeeX項目，它是gpt在代碼生成方向上的應(yīng)用，對標(biāo)于openAI的CodeX。github地址在此。

為什么選擇CodeGeeX呢？因為：

• 完全開源。它開源了完整的預(yù)訓(xùn)練代碼。而很多號稱開源的項目，其實只公開了預(yù)訓(xùn)練參數(shù)。
• 簡潔精要的模型圖。在對應(yīng)論文中，用兩張圖就清晰描繪了整個預(yù)訓(xùn)練配置和模型架構(gòu)（精確到混合精度和矩陣維度）。極大提升了源碼閱讀的效率。

下面我們就放出這兩張牛皮的架構(gòu)圖：

模型架構(gòu)

預(yù)訓(xùn)練配置

在下一篇講解切割模型部分的源碼里，我們會配合模型架構(gòu)圖來讀。這一篇我們著重講分布式環(huán)境初始化。因此對gpt或codegeex模型架構(gòu)不熟悉的話，也不影響本文閱讀。特別說明的是，根據(jù)預(yù)訓(xùn)練配置，我們可知codegeex采用的是8頭TP，192頭DP，共1536塊GPU進(jìn)行訓(xùn)練，采用的訓(xùn)練框架為Megatron + DeepSpeed ZeRO2。

二、預(yù)訓(xùn)練代碼整體架構(gòu)

2.1 預(yù)訓(xùn)練代碼設(shè)計與使用規(guī)范

如下圖:

• 預(yù)訓(xùn)練入口函數(shù)在megatron/tools/pretrain_codegeex.py 這個路徑下
• 啟動腳本在pretrain_codegeex.sh這個文件中。

使用Megatron時，一般將預(yù)訓(xùn)練函數(shù)命名為pretrain_模型名.py的形式，例如pretrain_bert.py、pretrain_gpt.py等。在codegeex這個項目里，該代碼位于tools目錄下；在NVDIA提供的代碼中，則與tools目錄同級。放在哪里不重要，梳理出來只是方讀者查找閱讀。

在pretrain_codegeex.sh這個啟動腳本里，定義了模型訓(xùn)練的參數(shù)值，包括batch_size、hidden_size等；同時也定義了設(shè)置分布式環(huán)境的參數(shù)值，例如DP/TP/PP組的大小等。

2.2 預(yù)訓(xùn)練代碼整體設(shè)計

在pretrain_codegeex.py中，核心入口函數(shù)為pretrain，調(diào)用它則開啟預(yù)訓(xùn)練過程：

if__name__=="__main__":
pretrain(
train_valid_test_datasets_provider,
model_provider,
forward_step,
args_defaults={"tokenizer_type":"GPT2BPETokenizer"},
)

如下圖，pretrain函數(shù)主要包含以下4個內(nèi)容：

• 初始化Megatron：設(shè)置分布式訓(xùn)練環(huán)境。主要目的是設(shè)置DP/TP/PP進(jìn)程組，并為每一個進(jìn)程分配GPU。
• 設(shè)置model，optimizer和lr schedule：在CPU上定義好模型，再將其按照第1步中定義好的分布式框架，把模型切割并搬運(yùn)到GPU上。
• 處理train/val/test數(shù)據(jù)集：按第1步定義好的分布式框架，對數(shù)據(jù)集進(jìn)行切分。
• 訓(xùn)練模型：在分布式環(huán)境中定義每個step的訓(xùn)練方式。

Megatron源碼解讀系列，也按上述邏輯分成4個部分。本篇將著重介紹第一部分：初始化Megatron。

三、初始化Megatron

3.1 初始化在做一件什么事

在閱讀代碼之前，我們先看初始化到底在完成一件什么事。

假設(shè)我們有2臺機(jī)器（node0和node1），每臺機(jī)器上有16塊GPU，GPU的編號為0~15。

我們使用這16塊GPU，做DP/TP/PP混合并行，如下圖：

• MP：模型并行組（Model Parallism）。假設(shè)一個完整的模型需要布在8塊GPU上，則如圖所示，我們共布了2個model replica（2個MP）。MP組為：[[g0, g1, g4, g5, g8, g9, g12, g13], [g2, g3, g6, g7, g10, g11, g14, g15]]
• TP：張量并行組（Tensor Parallism）。對于一個模型的每一層，我們將其參數(shù)縱向切開，分別置于不同的GPU上，則圖中一共有8個TP組。TP組為：[[g0, g1], [g4, g5],[g8, g9], [g12, g13], [g2, g3], [g6, g7], [g10, g11], [g14, g15]]
• PP：流水線并行組（Pipeline Parallism）。對于一個模型，我們將其每一層都放置于不同的GPU上，則圖中一共有4個PP組。PP組為：[[g0, g4, g8, g12], [g1, g5, g9, g13], [g2, g6, g10, g14], [g3, g7, g11, g15]]
• DP：數(shù)據(jù)并行組（Data Parallism）。經(jīng)過上述切割，對維護(hù)有相同模型部分的GPU，我們就可以做數(shù)據(jù)并行，則圖中共有8個DP組。DP組為[[g0, g2], [g1, g3], [g4, g6], [g5, g7], [g8, g10], [g9, g11], [g12, g14], [g13, g15]]

明確了分組設(shè)計，我們再來看下面幾個問題。

（1）分組的原則是什么？

• MP設(shè)定原則：MP其實由TP+PP共同決定。在開始訓(xùn)練前，需要我們根據(jù)實際模型，預(yù)估訓(xùn)練時顯存消耗（特別注意峰值顯存），來為模型安排GPU資源。
• TP、DP和PP設(shè)定原則：在這三種并行模式的原理篇中，我們分析過三者的通訊量。一般而言，TP>DP>PP。通訊量大的盡量放入一臺機(jī)器內(nèi)，因為機(jī)器內(nèi)帶寬高。所以在圖例中，TP和DP不跨機(jī)，PP跨機(jī)。再提一點(diǎn)，在使用Megatron時，很多項目是不用PP，僅用TP+DP的，此時一般將TP放入一臺機(jī)器內(nèi)，令DP跨機(jī)（比如codegeex）

（2）分組的目的是什么？

• 分配進(jìn)程：
  • 確認(rèn)分組方案后，在每塊GPU上啟動一個進(jìn)程（process），每個進(jìn)程獨(dú)立執(zhí)行自己所維護(hù)的那部分模型的計算，實現(xiàn)并行訓(xùn)練。
  • 進(jìn)程0~15，為一個進(jìn)程大組（group），其下的每一個DP/MP/PP組，為一個進(jìn)程子組（subgroup）
• 組間通訊：確認(rèn)好DP/TP/PP組，并分配好進(jìn)程后，我們就能進(jìn)一步設(shè)置不同進(jìn)程間的通訊方案。例如屬于一個DP組的g0和g2需要進(jìn)行梯度通訊，屬于一個PP組的g4和g8需要進(jìn)行層間輸出結(jié)果的通訊。

總結(jié)來說，初始化Megatron做了如下事：

• 定義模型的切割框架
• 在此框架上，初始化進(jìn)程，分配GPU，設(shè)置進(jìn)程組（DP/TP/PP）

3.2 代碼整體解讀

明確了初始化代碼要做的事情，現(xiàn)在可以來看代碼實現(xiàn)了。

回pretrain函數(shù)，它的第一行就通過initialize_megatron執(zhí)行了分布式初始化：

defpretrain(
train_valid_test_dataset_provider,
model_provider,
forward_step_func,
valid_forward_step_func=None,
extra_args_provider=None,
args_defaults={},
):
initialize_megatron(
extra_args_provider=extra_args_provider,args_defaults=args_defaults
)
...

initialize_megatron 函數(shù)位于megatron/initialize.py 文件中，我們直接來看它的核心函數(shù)_initialize_distributed。代碼如下：

def_initialize_distributed():
"""Initializetorch.distributedandmpu.
|Node1|Node2|
____________|p1|p2|p3|p4|
local_rank|0|1|0|1|
rank|0|1|2|3|

 node:物理結(jié)點(diǎn)，1臺機(jī)器或者1個容器。圖中2個物理結(jié)點(diǎn)
 rank：進(jìn)程在全局上的序號。圖中4個進(jìn)程
 local_rank：進(jìn)程在node上的序號。
 torch.cuda.device_count()：當(dāng)前進(jìn)程所在的node上可使用的GPU的數(shù)量
 device：GPU在某個node上的編號

該函數(shù)作用：
 1、設(shè)置分布式環(huán)境：初始化進(jìn)程，分配GPU，并設(shè)置進(jìn)程大組（group）
2、制定DP/TP/PP分組策略，設(shè)置進(jìn)程子組（subgroup）
3、設(shè)置DeepSpeedZeRO-R，對activation進(jìn)行優(yōu)化
"""
args=get_args()

device_count=torch.cuda.device_count()#當(dāng)前進(jìn)程所在的node上可使用的GPU的數(shù)量
iftorch.distributed.is_initialized():#如果已創(chuàng)建好分布式環(huán)境
ifargs.rank==0:#在0號進(jìn)程上打印出“創(chuàng)建完畢”的日志
print(
"torchdistributedisalreadyinitialized,"
"skippinginitialization...",
flush=True,
)
args.rank=torch.distributed.get_rank()#取得當(dāng)前進(jìn)程的全局序號
args.world_size=torch.distributed.get_world_size()#取得全局進(jìn)程的個數(shù)

else:#如果未創(chuàng)建好分布式環(huán)境
ifargs.rank==0:
print(">initializingtorchdistributed...",flush=True)

#1.初始化進(jìn)程，分配GPU，并設(shè)置進(jìn)程大組（group）
ifdevice_count>0:
device=args.rank%device_count# 1塊進(jìn)程1個GPU。device為GPU編號。例如圖例中的進(jìn)程9，其所在機(jī)器上有8塊卡。因此進(jìn)程9使用的gpu編號為8%9=1
ifargs.local_rankisnotNone:
assert(
args.local_rank==device
),"expectedlocal-ranktobethesameasrank%device-count."
else:
args.local_rank=device

ifargs.force_deviceisnotNone:
print(
f">forcefullysetthedeviceto{args.force_device},originally{device}"
)
device=args.force_device
torch.cuda.set_device(device)#為當(dāng)前進(jìn)程分配GPU

#設(shè)置進(jìn)程大組
init_method="tcp://"
master_ip=os.getenv("MASTER_ADDR","localhost")#獲取rank=0進(jìn)程的ip
master_port=os.getenv("MASTER_PORT","6000")#獲取rank=0進(jìn)程的端口
init_method+=master_ip+":"+master_port
print(
f">(rank={args.rank})initializingprocessgroup:"
f"world_size={args.world_size}"
f"backend={args.distributed_backend}"
f"init_method={init_method}",
flush=True,
)
timeout=datetime.timedelta(minutes=args.dist_timeout)
torch.distributed.init_process_group(
backend=args.distributed_backend,
world_size=args.world_size,
rank=args.rank,
init_method=init_method,
timeout=timeout
)
print(f">(rank={args.rank})processgroupinitialized")

#2、制定DP/TP/PP分組策略，設(shè)置進(jìn)程子組（subgroup）
ifdevice_count>0:
ifmpu.model_parallel_is_initialized():
print("modelparallelisalreadyinitialized")
else:
mpu.initialize_model_parallel(#megatron/mpu/initialize.py
args.tensor_model_parallel_size,
args.pipeline_model_parallel_size,
args.virtual_pipeline_model_parallel_size,
)

#設(shè)置DeepSpeedZeRO-R，對activation進(jìn)行優(yōu)化
ifargs.deepspeedandargs.deepspeed_activation_checkpointing:
setup_deepspeed_random_and_activation_checkpointing(args)

總體來說，這個代碼實現(xiàn)了3個目的：

• 設(shè)置分布式環(huán)境：初始化進(jìn)程，分配GPU，并設(shè)置進(jìn)程大組（group）。也即例子中的0~15號進(jìn)程同屬一個分布式進(jìn)程大組
• 制定DP/TP/PP分組策略，設(shè)置進(jìn)程子組（subgroup）
• 設(shè)置DeepSpeed ZeRO-R，對activation進(jìn)行優(yōu)化

我們來逐一講解。

3.3 代碼細(xì)節(jié)：torch.distributed，設(shè)置分布式環(huán)境

設(shè)置進(jìn)程大組的目的是告知程序，從全局上看，有哪些進(jìn)程共同組成了分布式訓(xùn)練系統(tǒng)。我們先明確幾個術(shù)語：

|Node1|Node2|
____________|p1|p2|p3|p4|
local_rank|0|1|0|1|
rank|0|1|2|3|

 node:物理結(jié)點(diǎn)，1臺機(jī)器或者1個容器。圖中2個物理結(jié)點(diǎn)
 rank：進(jìn)程在全局上的序號。圖中4個進(jìn)程
 local_rank：進(jìn)程在node上的序號。
 torch.cuda.device_count()：當(dāng)前進(jìn)程所在的node上可使用的GPU的數(shù)量
 device：GPU在某個node上的編號

特別說明，在2.2的圖例中，我們用g0~g15表示GPU編號，但更準(zhǔn)確地應(yīng)理解為進(jìn)程編號。GPU的編號與local_rank一樣，是相對于node而言的，即0~8，0~8。

我們借助torch.distributed 來實現(xiàn)這一步，它是pytorch用于設(shè)置分布式訓(xùn)練環(huán)境的偏底層API（distributed communication package）。如果你看過pytorch的文檔，可能會發(fā)現(xiàn)對于該API的闡述比較抽象。所以我把它單獨(dú)拎出來做說明。

init_method="tcp://"
master_ip=os.getenv("MASTER_ADDR","localhost")#獲取rank=0進(jìn)程的ip
master_port=os.getenv("MASTER_PORT","6000")#獲取rank=0進(jìn)程的端口
init_method+=master_ip+":"+master_port
print(
f">(rank={args.rank})initializingprocessgroup:"
f"world_size={args.world_size}"
f"backend={args.distributed_backend}"
f"init_method={init_method}",
flush=True,
)
timeout=datetime.timedelta(minutes=args.dist_timeout)
torch.distributed.init_process_group(
backend=args.distributed_backend,
world_size=args.world_size,
rank=args.rank,
init_method=init_method,
timeout=timeout
)
print(f">(rank={args.rank})processgroupinitialized")

我們聚焦于torch.distributed.init_process_group，該函數(shù)實現(xiàn)了設(shè)置進(jìn)程大組（group）的功能，它主要由以下幾個概念組成：

• backend：直譯為后端。但本質(zhì)上是在定義IPC通信機(jī)制（對數(shù)據(jù)實現(xiàn)reduce, gather, broadcase等通信操作）。取值有gloo，nccl 等。粗暴來說，使用CPU時，用gloo；使用GPU時，用nccl。
• world_size：全局進(jìn)程數(shù)。例如圖例中的world_size = 16。
• rank：當(dāng)前進(jìn)程在全局上的序號。例如圖例中進(jìn)程序號的取值范圍0~15，我們需要對每個進(jìn)程操作init_process_group，將其納入進(jìn)程大組中。
• init_method：這個概念較難理解，官方文檔也寫得比較抽象。通俗來說，這個參數(shù)指明了一個地址，進(jìn)程組內(nèi)的進(jìn)程通過該地址中存放的信息進(jìn)行交流。這些信息包括：哪些進(jìn)程間應(yīng)該相互通訊；各進(jìn)程的計算進(jìn)度如何等。還是以圖例說明，g1和g3屬于一個DP組，當(dāng)它們把各自梯度算完后，需要對梯度做AllReduce。g1算完自己的梯度后，它就會去這個地址下，聲明自己已算完，并去查找自己應(yīng)該和誰通訊，通訊方是否已算完等信息。借助這個地址中存儲的信息，進(jìn)程組內(nèi)的進(jìn)程就能相互知道彼此狀態(tài)，并聯(lián)系彼此。一般來說，為避免冗余，這個信息只存一份，存在rank 0 進(jìn)程上（rank 0進(jìn)程又稱為master進(jìn)程）。
• store：默認(rèn)值為None。它的作用和init_method一樣，只不過init_method指定的是一個地址，指定后在該地址下創(chuàng)建存儲交流信息的數(shù)據(jù)對象，這個數(shù)據(jù)對象就是store。也就是說，store顯式地指明了交流信息的內(nèi)容。因此store和init_method是互斥的，即store非空時，會忽略init_method。
• timeout：設(shè)置每個進(jìn)程等待的時間。進(jìn)程間的計算速度不一樣，還是以DP組的g1和g3為例，可能g1都算完梯度了，g3還在執(zhí)行forward。在等待g3算梯度的過程中，g1可能會timeout。因此這個參數(shù)就用于設(shè)置每個進(jìn)程的最大等待時間。

現(xiàn)在回頭再看這個代碼片段，是不是好理解很多~torch.distributed.init_process_group 非常重要，它貫穿了Megatron，也是使用pytorch做分布式訓(xùn)練不可略過的一環(huán)。關(guān)于torch.distributed的更多信息，推薦大家閱讀官方文檔，以及這篇blog。

3.4 代碼細(xì)節(jié)：設(shè)置DP/TP/PP組

設(shè)置完進(jìn)程大組（group）后，我們就可以進(jìn)一步設(shè)置進(jìn)程子組（subgroup）了，也即設(shè)置DP/TP/PP組。

mpu.initialize_model_parallel(#megatron/mpu/initialize.py
args.tensor_model_parallel_size,
args.pipeline_model_parallel_size,
args.virtual_pipeline_model_parallel_size,
)

核心函數(shù)initialize_model_parallel 在megatron/mpu/initialize.py 下。mpu的含義是model parallisim utils，也就是和模型并行設(shè)置相關(guān)的函數(shù)，都放在這個目錄下，它接收3個參數(shù)：

• tensor_model_parallel_size：每個TP組的進(jìn)程數(shù)量。例如圖例中是2
• pipeline_model_parallel_size：每個PP組的進(jìn)程數(shù)量。例如圖例中是4
• virtual_pipeline_model_parallel_size：每個virtual PP組的進(jìn)程數(shù)量。這是NVIDIA對Megatron做后續(xù)迭代時提出的一種優(yōu)化方法。我們之后會單獨(dú)開一篇文章來講解。這里可暫時忽略（不是必須參數(shù)，可以傳None值）。

你可能會問，為什么不設(shè)置DP相關(guān)的size？回想2.2中設(shè)計分布式的過程，我們根據(jù)TP+PP就可確認(rèn)MP，進(jìn)而推出DP。也就是定好了TP和PP，DP_size就能根據(jù) world_size // (TP_size * PP_size)計算得出。因此不用定義。

我們來看具體代碼：

definitialize_model_parallel(
tensor_model_parallel_size_=1,
pipeline_model_parallel_size_=1,
virtual_pipeline_model_parallel_size_=None,
):
"""
Initializemodeldataparallelgroups.

Arguments:
tensor_model_parallel_size:numberofGPUsusedtoparallelizemodeltensor.
pipeline_model_parallel_size:numberofGPUsusedtoparallelizemodelpipeline.

Let'ssaywehaveatotalof16GPUsdenotedbyg0...g15andwe
use2GPUstoparallelizethemodeltensor,and4GPUstoparallelize
themodelpipeline.Thepresentfunctionwill
create8tensormodel-parallelgroups,4pipelinemodel-parallelgroups
and8data-parallelgroupsas:
8data_parallelgroups:
[g0,g2],[g1,g3],[g4,g6],[g5,g7],[g8,g10],[g9,g11],[g12,g14],[g13,g15]
8tensormodel-parallelgroups:
[g0,g1],[g2,g3],[g4,g5],[g6,g7],[g8,g9],[g10,g11],[g12,g13],[g14,g15]
4pipelinemodel-parallelgroups:
[g0,g4,g8,g12],[g1,g5,g9,g13],[g2,g6,g10,g14],[g3,g7,g11,g15]
2model-parallelgroup:
[g0,g1,g4,g5,g8,g9,g12,g13],[g2,g3,g6,g7,g10,g8,g14,g15]

Notethatforefficiency,thecallershouldmakesureadjacentranks
areonthesameDGXbox.Forexampleifweareusing2DGX-1boxes
withatotalof16GPUs,rank0to7belongtothefirstboxand
ranks8to15belongtothesecondbox.
"""
iftorch.distributed.get_rank()==0:
print(
">initializingtensormodelparallelwithsize{}".format(
tensor_model_parallel_size_
)
)
print(
">initializingpipelinemodelparallelwithsize{}".format(
pipeline_model_parallel_size_
)
)
#Getworldsizeandrank.Ensuresomeconsistencies.
asserttorch.distributed.is_initialized()#確保torch已經(jīng)做了分布式初始化
world_size=torch.distributed.get_world_size()#得到全局進(jìn)程的總數(shù)
tensor_model_parallel_size=min(tensor_model_parallel_size_,world_size)
pipeline_model_parallel_size=min(pipeline_model_parallel_size_,world_size)

ensure_divisibility(#后者表示一個完整模型所占的gpu數(shù)，我們要保證前者能被后者整除
world_size,tensor_model_parallel_size*pipeline_model_parallel_size
)
#在codegeex中，TP_size=8,PP_size=1，world_size=1536，因此DP_size是1536/(8*1)=192
data_parallel_size=world_size//(#根據(jù)TP_size和PP_size，求出DP_size
tensor_model_parallel_size*pipeline_model_parallel_size
)

num_tensor_model_parallel_groups=world_size//tensor_model_parallel_size#TP的組數(shù)
num_pipeline_model_parallel_groups=world_size//pipeline_model_parallel_size#PP的組數(shù)
num_data_parallel_groups=world_size//data_parallel_size#DP的組數(shù)

ifvirtual_pipeline_model_parallel_size_isnotNone:
global_VIRTUAL_PIPELINE_MODEL_PARALLEL_RANK
global_VIRTUAL_PIPELINE_MODEL_PARALLEL_WORLD_SIZE
_VIRTUAL_PIPELINE_MODEL_PARALLEL_RANK=0
_VIRTUAL_PIPELINE_MODEL_PARALLEL_WORLD_SIZE=(
virtual_pipeline_model_parallel_size_
)

rank=torch.distributed.get_rank()#獲取當(dāng)前進(jìn)程的全局rank

#Buildthedata-parallelgroups.（設(shè)置DP組）
global_DATA_PARALLEL_GROUP#保存DP組，如[[0,2],[1,3]...]，數(shù)字表示進(jìn)進(jìn)程的全局序號
assert_DATA_PARALLEL_GROUPisNone,"dataparallelgroupisalreadyinitialized"
all_data_parallel_group_ranks=[]
foriinrange(pipeline_model_parallel_size):
start_rank=i*num_pipeline_model_parallel_groups
end_rank=(i+1)*num_pipeline_model_parallel_groups
forjinrange(tensor_model_parallel_size):
ranks=range(start_rank+j,end_rank,tensor_model_parallel_size)
all_data_parallel_group_ranks.append(list(ranks))
group=torch.distributed.new_group(ranks)#設(shè)置DP組
ifrankinranks:
_DATA_PARALLEL_GROUP=group

#Buildthemodel-parallelgroups.（設(shè)置MP組）
global_MODEL_PARALLEL_GROUP#保存MP組
assert_MODEL_PARALLEL_GROUPisNone,"modelparallelgroupisalreadyinitialized"
foriinrange(data_parallel_size):
ranks=[
data_parallel_group_ranks[i]
fordata_parallel_group_ranksinall_data_parallel_group_ranks
]
group=torch.distributed.new_group(ranks)#設(shè)置MP組
ifrankinranks:
_MODEL_PARALLEL_GROUP=group

#Buildthetensormodel-parallelgroups.（設(shè)置TP組）
global_TENSOR_MODEL_PARALLEL_GROUP#保存TP組
assert(
_TENSOR_MODEL_PARALLEL_GROUPisNone
),"tensormodelparallelgroupisalreadyinitialized"
foriinrange(num_tensor_model_parallel_groups):
ranks=range(
i*tensor_model_parallel_size,(i+1)*tensor_model_parallel_size
)
group=torch.distributed.new_group(ranks)#設(shè)置TP組
ifrankinranks:
_TENSOR_MODEL_PARALLEL_GROUP=group

#Buildthepipelinemodel-parallelgroupsandembeddinggroups
#(firstandlastrankineachpipelinemodel-parallelgroup).（設(shè)置PP組與embedding組）
global_PIPELINE_MODEL_PARALLEL_GROUP#設(shè)置PP組
global_PIPELINE_GLOBAL_RANKS
assert(
_PIPELINE_MODEL_PARALLEL_GROUPisNone
),"pipelinemodelparallelgroupisalreadyinitialized"
global_EMBEDDING_GROUP
assert_EMBEDDING_GROUPisNone,"embeddinggroupisalreadyinitialized"
foriinrange(num_pipeline_model_parallel_groups):
ranks=range(i,world_size,num_pipeline_model_parallel_groups)
group=torch.distributed.new_group(ranks)#設(shè)置PP組
ifrankinranks:
_PIPELINE_MODEL_PARALLEL_GROUP=group
_PIPELINE_GLOBAL_RANKS=ranks
#Setupembeddinggroup(toexchangegradientsbetween
#firstandlaststages).
iflen(ranks)>1:
embedding_ranks=[ranks[0],ranks[-1]]
else:
embedding_ranks=ranks
group=torch.distributed.new_group(embedding_ranks)#設(shè)置embedding組
ifrankinembedding_ranks:
_EMBEDDING_GROUP=group

總結(jié)來說，我們采用torch.distributed.new_group(ranks)在進(jìn)程大組下設(shè)置子組。ranks是list of list，表示對進(jìn)程序號的劃分，例如設(shè)置DP組，則ranks為[[0,2], [1,3]...]，以此類推。我們將劃分結(jié)果存在全局變量中（例如_DATA_PARALLEL_GROUP），方便我們在后續(xù)切割模型時使用。

同時，我們定義以下函數(shù)，使得對于任意一個進(jìn)程，我們都能查到它在DP/TP/PP組中的局部序號（local_rank），以及它對應(yīng)的DP/TP/PP組的world_size。這也是為后續(xù)切割模型使用:

#這里展示和TP組相關(guān)的查詢操作。其余組也是類推。詳細(xì)代碼一樣都在megatron/mpu/initialize.py中
defget_tensor_model_parallel_group():
"""Getthetensormodelparallelgroupthecallerrankbelongsto."""
assert(
_TENSOR_MODEL_PARALLEL_GROUPisnotNone
),"intra_layer_modelparallelgroupisnotinitialized"
return_TENSOR_MODEL_PARALLEL_GROUP


defset_tensor_model_parallel_world_size(world_size):
"""Setthetensormodelparallelsize"""
global_MPU_TENSOR_MODEL_PARALLEL_WORLD_SIZE
_MPU_TENSOR_MODEL_PARALLEL_WORLD_SIZE=world_size

defget_tensor_model_parallel_rank():
"""Returnmyrankforthetensormodelparallelgroup.
 my_rank指的就是local_rank，例如[g2, g3]這一個TP組，rank為2，3；local_rank為0，1
"""
global_MPU_TENSOR_MODEL_PARALLEL_RANK
if_MPU_TENSOR_MODEL_PARALLEL_RANKisnotNone:
return_MPU_TENSOR_MODEL_PARALLEL_RANK
returntorch.distributed.get_rank(group=get_tensor_model_parallel_group())

最后，你可能想問，為什么還有一個embedding_group?

在GPT類模型中，輸入層和輸出層共享一個word_embedding（可參見第一部分中codegeex的架構(gòu)圖）。因此，在計算完梯度，更新embedding權(quán)重前，輸入和輸出層需要進(jìn)行通訊，保證word_embedding完全一致。也即PP組中的第一個和最后一個進(jìn)程需要通訊。我們知道設(shè)置進(jìn)程子組的目的就是進(jìn)一步劃分通訊組，因此這里再添加一個embedding_group。

3.5 代碼細(xì)節(jié)：DeepSpeed ZeRO-R

到目前為止，對于初始化，我們設(shè)置了全局的分布式，完成了進(jìn)程大組的設(shè)置；同時根據(jù)DP/TP/PP設(shè)計劃分了進(jìn)程子組。NVIDIA Megatron初始化部分的代碼，其實已經(jīng)結(jié)束了。

但是，在實際應(yīng)用中，通常采用DeepSpeed-Megatron的方式，借助微軟DeepSpeed庫，通過ZeRO技術(shù)，幫助我們更好節(jié)省顯存。例如codegeex就采用了ZeRO2 + Megatron的方式進(jìn)行訓(xùn)練。

總結(jié)來說，在Megatron中使用ZeRO的方法很簡單，按照這篇官方教程，秉持著萬物皆可wrap的原則，在原始代碼特定的幾個位置，把DeepSpeed提供的API包進(jìn)去，就能幫我們在訓(xùn)練中管理顯存了。使用ZeRO-R，對activation做顯存管理，是一個可選項。當(dāng)activation大小成為顯存瓶頸時，可以按照教程指導(dǎo)，在初始化Megatron的代碼里引入這部分優(yōu)化：

#設(shè)置ZeRO-R
ifargs.deepspeedandargs.deepspeed_activation_checkpointing:
setup_deepspeed_random_and_activation_checkpointing(args)

那么ZeRO-R是怎么對顯存優(yōu)化起作用的呢？

與ZeRO1，ZeRO2和ZeRO3是在DP組中做顯存優(yōu)化不同，ZeRO-R是在TP組中特別針對activation做顯存優(yōu)化?；叵胍幌?，在DP組里輸入數(shù)據(jù)X各不相同，對應(yīng)的activation也不相同。這時對activation做切割是沒意義的。只有在輸入X相同的情況下，才有意義對activation進(jìn)行不用時切割存儲，用時再gather回來的操作。

回顧Megatron每一層的計算，在TP組中，各GPU上的模型部分計算完畢后，需要經(jīng)過一次AllReduce將聚合后的結(jié)果取回，然后才能進(jìn)行下一層計算。此時，不同的GPU都擁有了同一個輸入X，也意味著在后續(xù)計算中會產(chǎn)生相同的activation，這時我們就能通過ZeRO-R來避免冗余了。如下圖，提供了TP下transfomer MLP層的計算：

關(guān)于初始化Megatron，就講解到這了，本文列舉了核心代碼，各位讀者可去官方github上，閱讀更多細(xì)節(jié)。在下一篇里，我們將進(jìn)入預(yù)訓(xùn)練的第二部分：模型切割，這也是整個Megatron的核心。這部分代碼細(xì)節(jié)較多，代碼架構(gòu)上也比較分散，我依然會通過圖解+細(xì)節(jié)解讀的模式，和大家一起閱讀～