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

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

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

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

一、CodeGeeX模型簡述

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

為什么選擇CodeGeeX呢？因為：

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

下面我們就放出這兩張牛皮的架構圖：

模型架構

預訓練配置

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

二、預訓練代碼整體架構

2.1 預訓練代碼設計與使用規范

如下圖:

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

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

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

2.2 預訓練代碼整體設計

在pretrain_codegeex.py中，核心入口函數為pretrain，調用它則開啟預訓練過程：

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

如下圖，pretrain函數主要包含以下4個內容：

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

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

三、初始化Megatron

3.1 初始化在做一件什么事

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

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

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

• MP：模型并行組（Model Parallism）。假設一個完整的模型需要布在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）。對于一個模型的每一層，我們將其參數縱向切開，分別置于不同的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：數據并行組（Data Parallism）。經過上述切割，對維護有相同模型部分的GPU，我們就可以做數據并行，則圖中共有8個DP組。DP組為[[g0, g2], [g1, g3], [g4, g6], [g5, g7], [g8, g10], [g9, g11], [g12, g14], [g13, g15]]

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

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

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

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

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

總結來說，初始化Megatron做了如下事：

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

3.2 代碼整體解讀

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

回pretrain函數，它的第一行就通過initialize_megatron執行了分布式初始化：

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 函數位于megatron/initialize.py 文件中，我們直接來看它的核心函數_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:物理結點，1臺機器或者1個容器。圖中2個物理結點
 rank：進程在全局上的序號。圖中4個進程
 local_rank：進程在node上的序號。
 torch.cuda.device_count()：當前進程所在的node上可使用的GPU的數量
 device：GPU在某個node上的編號

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

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

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

#1.初始化進程，分配GPU，并設置進程大組（group）
ifdevice_count>0:
device=args.rank%device_count# 1塊進程1個GPU。device為GPU編號。例如圖例中的進程9，其所在機器上有8塊卡。因此進程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)#為當前進程分配GPU

#設置進程大組
init_method="tcp://"
master_ip=os.getenv("MASTER_ADDR","localhost")#獲取rank=0進程的ip
master_port=os.getenv("MASTER_PORT","6000")#獲取rank=0進程的端口
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分組策略，設置進程子組（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,
)

#設置DeepSpeedZeRO-R，對activation進行優化
ifargs.deepspeedandargs.deepspeed_activation_checkpointing:
setup_deepspeed_random_and_activation_checkpointing(args)

總體來說，這個代碼實現了3個目的：

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

我們來逐一講解。

3.3 代碼細節：torch.distributed，設置分布式環境

設置進程大組的目的是告知程序，從全局上看，有哪些進程共同組成了分布式訓練系統。我們先明確幾個術語：

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

 node:物理結點，1臺機器或者1個容器。圖中2個物理結點
 rank：進程在全局上的序號。圖中4個進程
 local_rank：進程在node上的序號。
 torch.cuda.device_count()：當前進程所在的node上可使用的GPU的數量
 device：GPU在某個node上的編號

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

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

init_method="tcp://"
master_ip=os.getenv("MASTER_ADDR","localhost")#獲取rank=0進程的ip
master_port=os.getenv("MASTER_PORT","6000")#獲取rank=0進程的端口
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，該函數實現了設置進程大組（group）的功能，它主要由以下幾個概念組成：

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

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

3.4 代碼細節：設置DP/TP/PP組

設置完進程大組（group）后，我們就可以進一步設置進程子組（subgroup）了，也即設置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,
)

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

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

你可能會問，為什么不設置DP相關的size？回想2.2中設計分布式的過程，我們根據TP+PP就可確認MP，進而推出DP。也就是定好了TP和PP，DP_size就能根據 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已經做了分布式初始化
world_size=torch.distributed.get_world_size()#得到全局進程的總數
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數，我們要保證前者能被后者整除
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//(#根據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的組數
num_pipeline_model_parallel_groups=world_size//pipeline_model_parallel_size#PP的組數
num_data_parallel_groups=world_size//data_parallel_size#DP的組數

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()#獲取當前進程的全局rank

#Buildthedata-parallelgroups.（設置DP組）
global_DATA_PARALLEL_GROUP#保存DP組，如[[0,2],[1,3]...]，數字表示進進程的全局序號
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)#設置DP組
ifrankinranks:
_DATA_PARALLEL_GROUP=group

#Buildthemodel-parallelgroups.（設置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)#設置MP組
ifrankinranks:
_MODEL_PARALLEL_GROUP=group

#Buildthetensormodel-parallelgroups.（設置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)#設置TP組
ifrankinranks:
_TENSOR_MODEL_PARALLEL_GROUP=group

#Buildthepipelinemodel-parallelgroupsandembeddinggroups
#(firstandlastrankineachpipelinemodel-parallelgroup).（設置PP組與embedding組）
global_PIPELINE_MODEL_PARALLEL_GROUP#設置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)#設置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)#設置embedding組
ifrankinembedding_ranks:
_EMBEDDING_GROUP=group

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

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

#這里展示和TP組相關的查詢操作。其余組也是類推。詳細代碼一樣都在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的架構圖）。因此，在計算完梯度，更新embedding權重前，輸入和輸出層需要進行通訊，保證word_embedding完全一致。也即PP組中的第一個和最后一個進程需要通訊。我們知道設置進程子組的目的就是進一步劃分通訊組，因此這里再添加一個embedding_group。

3.5 代碼細節：DeepSpeed ZeRO-R

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

但是，在實際應用中，通常采用DeepSpeed-Megatron的方式，借助微軟DeepSpeed庫，通過ZeRO技術，幫助我們更好節省顯存。例如codegeex就采用了ZeRO2 + Megatron的方式進行訓練。

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

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

那么ZeRO-R是怎么對顯存優化起作用的呢？

與ZeRO1，ZeRO2和ZeRO3是在DP組中做顯存優化不同，ZeRO-R是在TP組中特別針對activation做顯存優化。回想一下，在DP組里輸入數據X各不相同，對應的activation也不相同。這時對activation做切割是沒意義的。只有在輸入X相同的情況下，才有意義對activation進行不用時切割存儲，用時再gather回來的操作。

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

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