Megatron-LM默認(rèn)支持GPT、T5、BERT等多個(gè)常見(jiàn)模型的預(yù)訓(xùn)練，當(dāng)下大模型流行，故以pretrain_gpt.py為例做源碼的走讀。

一. 啟動(dòng)pretrain_gpt.py

　　pretrain_gpt.py為GPT類模型的訓(xùn)練入口，它通過(guò)命令行形式被調(diào)用，其精確執(zhí)行路徑位于Megatron-LM框架的examples/gpt3目錄下。具體而言，啟動(dòng)過(guò)程依賴于train_gpt3_175b_distributed.sh這一腳本，該腳本專為部署GPT-3模型在分布式環(huán)境下訓(xùn)練而設(shè)計(jì)（當(dāng)然，也可以參照編寫自定義的啟動(dòng)腳本）。

　　在train_gpt3_175b_distributed.sh腳本內(nèi)部，核心操作是通過(guò)torchrun命令實(shí)現(xiàn)的，該命令是PyTorch分布式訓(xùn)練的一部分，用于在多個(gè)計(jì)算節(jié)點(diǎn)上高效并行地執(zhí)行pretrain_gpt.py。此過(guò)程確保了模型訓(xùn)練任務(wù)能夠充分利用集群資源，加速訓(xùn)練過(guò)程，代碼如下圖：

二. torchrun簡(jiǎn)介

　　trochrun是PyTorch官方推薦用于替代torch.distributed.launch的分布式數(shù)據(jù)并行訓(xùn)練模塊。它旨在提供一種更靈活、更健壯的方式來(lái)啟動(dòng)和管理分布式訓(xùn)練任務(wù)。

　　trochrun啟動(dòng)并行訓(xùn)練任務(wù)的原理如下：

1. 初始化分布式環(huán)境

　　trochrun首先負(fù)責(zé)初始化分布式訓(xùn)練所需的環(huán)境。這包括設(shè)置通信后端（如NCCL、GLOO等）、分配工作進(jìn)程的RANK和WORLD_SIZE（即參與訓(xùn)練的總進(jìn)程數(shù)），以及處理其他與分布式訓(xùn)練相關(guān)的配置。

2. 分配工作進(jìn)程

　　trochrun會(huì)根據(jù)指定的參數(shù)（如--nnodes、--nproc-per-node等）來(lái)分配工作進(jìn)程。這些進(jìn)程可以是同一臺(tái)機(jī)器上的多個(gè) GPU，也可以是跨多臺(tái)機(jī)器的GPU。每個(gè)進(jìn)程都會(huì)加載相同的訓(xùn)練腳本（如pretrain_gpt.py），但會(huì)處理不同的數(shù)據(jù)子集，以實(shí)現(xiàn)并行訓(xùn)練。

3. 同步與通信

　　在訓(xùn)練過(guò)程中，torchrun 管理下的各個(gè)工作進(jìn)程需要頻繁地進(jìn)行同步和通信。這包括梯度同步（在反向傳播后同步各GPU上的梯度）、參數(shù)更新（使用同步后的梯度更新模型參數(shù)）等。PyTorch提供了豐富的API（如torch.distributed.all_reduce、torch.distributed.barrier等）來(lái)支持這些操作。

4. 優(yōu)雅處理故障

　　trochrun相比torch.distributed.launch的一大改進(jìn)是它能夠更優(yōu)雅地處理工作進(jìn)程的故障。例如，如果某個(gè)工作進(jìn)程因?yàn)槟撤N原因崩潰了，torchrun可以嘗試重新啟動(dòng)該進(jìn)程，以確保訓(xùn)練任務(wù)的連續(xù)性。此外，torchrun還支持彈性訓(xùn)練（elastic training），即允許在訓(xùn)練過(guò)程中動(dòng)態(tài)地增加或減少工作進(jìn)程的數(shù)量。

5. 簡(jiǎn)化配置與啟動(dòng)

　　trochrun通過(guò)提供命令行接口和配置文件選項(xiàng)來(lái)簡(jiǎn)化分布式訓(xùn)練的配置和啟動(dòng)過(guò)程。用戶只需指定少量的參數(shù)（如節(jié)點(diǎn)數(shù)、每節(jié)點(diǎn)進(jìn)程數(shù)等），即可啟動(dòng)復(fù)雜的分布式訓(xùn)練任務(wù)。此外，torchrun還支持從環(huán)境變量中讀取配置信息，這使得在不同環(huán)境中部署訓(xùn)練任務(wù)變得更加靈活。

6. 自動(dòng)化資源分配

　　在某些情況下，torchrun還可以與資源管理器（如Kubernetes、Slurm等）集成，以自動(dòng)化地分配和管理訓(xùn)練所需的計(jì)算資源。這包括GPU、CPU、內(nèi)存和存儲(chǔ)等資源。通過(guò)集成資源管理器，torchrun可以進(jìn)一步提高分布式訓(xùn)練的可擴(kuò)展性和靈活性。

　　總之，torchrun通過(guò)以上機(jī)制共同作用，使得使用PyTorch進(jìn)行分布式訓(xùn)練變得更加高效、可靠和易于管理。

三. 主要函數(shù)

　　pretrain_gpt.py腳本封裝了多個(gè)核心功能組件，具體包括model_provider()，forward_step()，train_valid_test_datasets_provider()，以及pretrain()等主要函數(shù)。

　　其中，model_provider()負(fù)責(zé)提供預(yù)訓(xùn)練所需的模型實(shí)例對(duì)象；forward_step()定義了模型前向傳播的具體步驟，包括輸入處理、模型計(jì)算等；train_valid_test_datasets_provider()則負(fù)責(zé)準(zhǔn)備訓(xùn)練、驗(yàn)證及測(cè)試所需的數(shù)據(jù)集，確保數(shù)據(jù)的有效供給。

　　值得注意的是，前三個(gè)函數(shù)model_provider()，forward_step()，train_valid_test_datasets_provider()更是作為pretrain()函數(shù)的入?yún)?#xff0c;共同構(gòu)成了GPT模型訓(xùn)練入口。

　　這種設(shè)計(jì)確保了預(yù)訓(xùn)練過(guò)程的模塊化、靈活性與可擴(kuò)展性。下面會(huì)從model_provider()開(kāi)始逐行解析源碼。

四. 源碼分析

1. model_provider

　　def model_provider(pre_process=True, post_process=True) -> Union[GPTModel, megatron.legacy.model.GPTModel]:"""Builds the model.If you set the use_legacy_models to True, it will return the legacy GPT model and if not the mcore GPT model.Args:pre_process (bool, optional): Set to true if you need to compute embedings. Defaults to True.post_process (bool, optional): Set to true if you need to want to compute output logits/loss. Defaults to True.Returns:Union[GPTModel, megatron.legacy.model.GPTModel]: The returned model"""args = get_args()use_te = args.transformer_impl == "transformer_engine"print_rank_0('building GPT model ...')# Experimental loading arguments from yamlif args.yaml_cfg is not None:config = core_transformer_config_from_yaml(args, "language_model")else:config = core_transformer_config_from_args(args)if args.use_legacy_models:model = megatron.legacy.model.GPTModel(config,num_tokentypes=0,parallel_output=True,pre_process=pre_process,post_process=post_process,)else: # using core modelsif args.spec is not None:transformer_layer_spec = import_module(args.spec)else:if use_te:transformer_layer_spec = get_gpt_layer_with_transformer_engine_spec(args.num_experts, args.moe_grouped_gemm, args.qk_layernorm)else:transformer_layer_spec = get_gpt_layer_local_spec(args.num_experts, args.moe_grouped_gemm, args.qk_layernorm)model = GPTModel(config=config,transformer_layer_spec=transformer_layer_spec,vocab_size=args.padded_vocab_size,max_sequence_length=args.max_position_embeddings,pre_process=pre_process,post_process=post_process,fp16_lm_cross_entropy=args.fp16_lm_cross_entropy,parallel_output=True,share_embeddings_and_output_weights=not args.untie_embeddings_and_output_weights,position_embedding_type=args.position_embedding_type,rotary_percent=args.rotary_percent,)return model

　　model_provider函數(shù)是用于構(gòu)建GPT（生成預(yù)訓(xùn)練Transformer）模型實(shí)例的函數(shù)，它會(huì)以類似函數(shù)指針的形式，作為pretrain()的入?yún)鬟f到后續(xù)的訓(xùn)練代碼中，供訓(xùn)練過(guò)程調(diào)用。

　　該函數(shù)主要代碼流程包含以下幾個(gè)步驟：

　　a. 獲取參數(shù)和配置：

　　通過(guò) get_args() 函數(shù)獲取命令行參數(shù)和配置文件中的參數(shù)。根據(jù) args.transformer_impl 的值確定是否使用 Transformer Engine (use_te)。

　　b. 配置模型：

　　如果指定了 YAML 配置文件 (args.yaml_cfg)，則從 YAML 文件中加載模型結(jié)構(gòu)。否則，根據(jù)命令行參數(shù) (args) 加載。

　　c. 選擇模型類型：

　　如果args.use_legacy_models為True，則使用megatron.legacy.model.GPTModel構(gòu)建模型。這通常用于向后兼容或測(cè)試舊版本的模型。

　　如果不使用舊版模型，則直接運(yùn)行構(gòu)建GPTModel，如圖紅框部分。

　　其中，GPTModel的參數(shù)包括配置 (config)、詞匯表大小 (vocab_size)、最大序列長(zhǎng)度 (max_sequence_length)、是否進(jìn)行前處理和后處理、是否使用 FP16 進(jìn)行語(yǔ)言模型交叉熵計(jì)算、是否并行輸出等。這些參數(shù)基本上都來(lái)源于配置文件，關(guān)于配置文件的內(nèi)容和解析將于下文詳述。

　　注：此處的model_provider是作為函數(shù)指針傳遞到pretrain()中，函數(shù)指針只有在調(diào)用時(shí)才會(huì)真正執(zhí)行，故，GPTModel的具體實(shí)現(xiàn)待到執(zhí)行時(shí)再具體分析。

2. forward_step

def forward_step(data_iterator, model: GPTModel):"""Forward training step.Args:data_iterator : Input data iteratormodel (GPTModel): The GPT Model"""args = get_args()timers = get_timers()# Get the batch.timers('batch-generator', log_level=2).start()global stimerwith stimer(bdata=True):tokens, labels, loss_mask, attention_mask, position_ids = get_batch(data_iterator)timers('batch-generator').stop()with stimer:output_tensor = model(tokens, position_ids, attention_mask,labels=labels)return output_tensor, partial(loss_func, loss_mask)

　　forward_step顧名思義，這個(gè)函數(shù)是GPT模型訓(xùn)練過(guò)程中前向處理函數(shù)，負(fù)責(zé)處理一批輸入數(shù)據(jù)并通過(guò)模型進(jìn)行前向傳播。

　　該函數(shù)的核心實(shí)現(xiàn)，仍然是對(duì)model（forward_step函數(shù)的入?yún)?#xff09;的forward的調(diào)用，只是在調(diào)用之前封裝了計(jì)時(shí)器計(jì)時(shí)邏輯以及批次數(shù)據(jù)獲取的邏輯（這部分邏輯會(huì)根據(jù)不同的業(yè)務(wù)場(chǎng)景變化而變化，故不能直接封裝到model的forward函數(shù)中，而是應(yīng)該在pretrain腳本中實(shí)現(xiàn)），具體代碼流程如下：

　　a. 獲取參數(shù)和計(jì)時(shí)器：

　　通過(guò)get_args()和get_timers()函數(shù)分別獲取訓(xùn)練參數(shù)和計(jì)時(shí)器對(duì)象，用于控制訓(xùn)練過(guò)程和記錄時(shí)間消耗。

　　b. 獲取批次數(shù)據(jù)：

　　使用timers對(duì)象記錄獲取批次數(shù)據(jù)的時(shí)間（可選，通過(guò)log_level=2控制）。

　　調(diào)用get_batch函數(shù)從data_iterator中獲取一批數(shù)據(jù)，包括tokens（輸入文本對(duì)應(yīng)的token IDs）、labels（訓(xùn)練標(biāo)簽，通常用于計(jì)算損失，對(duì)于語(yǔ)言模型任務(wù)，labels通常是tokens的右移一位版本）、loss_mask（損失掩碼，用于忽略某些位置的損失計(jì)算，如填充位置）、attention_mask（注意力掩碼，用于指示哪些位置需要參與注意力計(jì)算）和position_ids（位置ID，用于模型中的位置編碼）。

　　c. 模型前向傳播：

　　使用stimer（可能是一個(gè)自定義的計(jì)時(shí)器）記錄模型前向傳播的時(shí)間。

　　將獲取到的數(shù)據(jù)（tokens, position_ids, attention_mask, labels）傳遞給模型model進(jìn)行前向傳播。這里labels是可選的，用于計(jì)算損失，但在前向傳播階段不一定需要。

　　模型輸出output_tensor，通常包含模型的預(yù)測(cè)結(jié)果（如logits）。

　　d. 返回輸出和損失函數(shù)：

　　返回output_tensor和partial函數(shù)，該函數(shù)需要loss_mask作為參數(shù)來(lái)計(jì)算損失。這種方式允許延遲損失的計(jì)算，直到所有相關(guān)的數(shù)據(jù)都已準(zhǔn)備好。

3. train_valid_test_datasets_provider

def train_valid_test_datasets_provider(train_val_test_num_samples):"""Build the train test and validation datasets.Args:train_val_test_num_samples : A list containing the number of samples in train test and validation."""args = get_args()config = core_gpt_dataset_config_from_args(args)if args.mock_data:dataset_type = MockGPTDatasetelse:dataset_type = GPTDatasetprint_rank_0("> building train, validation, and test datasets for GPT ...")train_ds, valid_ds, test_ds = BlendedMegatronDatasetBuilder(dataset_type,train_val_test_num_samples,is_dataset_built_on_rank,config).build()print_rank_0("> finished creating GPT datasets ...")return train_ds, valid_ds, test_ds

　　該函數(shù)接收一個(gè)參數(shù)train_val_test_num_samples，這是一個(gè)列表，包含了訓(xùn)練集、驗(yàn)證集和測(cè)試集的樣本數(shù)量。函數(shù)的目的是根據(jù)提供的參數(shù)和配置，構(gòu)建GPT模型的訓(xùn)練、驗(yàn)證和測(cè)試數(shù)據(jù)集。主要代碼流程如下：

　　a. 獲取參數(shù)和配置：

　　使用get_args()函數(shù)獲取訓(xùn)練過(guò)程中的全局參數(shù)，并通過(guò)core_gpt_dataset_config_from_args根據(jù)這些參數(shù)生成數(shù)據(jù)集配置對(duì)象config。

　　b. 確定數(shù)據(jù)集類型：

　　根據(jù)args.mock_data的值決定使用哪種數(shù)據(jù)集類型。如果mock_data為True，則使用MockGPTDataset，這是一種模擬數(shù)據(jù)集，可能用于測(cè)試或快速原型開(kāi)發(fā)。如果mock_data為False，則使用GPTDataset，這是實(shí)際的數(shù)據(jù)集類型，包含真實(shí)的訓(xùn)練數(shù)據(jù)。

　　c. 構(gòu)建數(shù)據(jù)集：

　　使用BlendedMegatronDatasetBuilder類來(lái)構(gòu)建數(shù)據(jù)集，傳遞給BlendedMegatronDatasetBuilder的參數(shù)包括數(shù)據(jù)集類型dataset_type、訓(xùn)練/驗(yàn)證/測(cè)試集的樣本數(shù)量train_val_test_num_samples、is_dataset_built_on_rank（用于檢查當(dāng)前處理單元是否負(fù)責(zé)構(gòu)建數(shù)據(jù)集），以及配置對(duì)象config。

　　調(diào)用build()方法實(shí)際構(gòu)建數(shù)據(jù)集，該方法返回三個(gè)數(shù)據(jù)集對(duì)象：訓(xùn)練集train_ds、驗(yàn)證集valid_ds和測(cè)試集test_ds。

　　其中BlendedMegatronDatasetBuilder來(lái)源于包“megatron.core.datasets.blended_megatron_dataset_builder”，由于數(shù)據(jù)集構(gòu)建邏輯比較簡(jiǎn)單，故，在此不做詳述，有興趣的同學(xué)可以自行查看。

　　d. 返回值

　　函數(shù)返回三個(gè)數(shù)據(jù)集對(duì)象：訓(xùn)練集train_ds、驗(yàn)證集valid_ds和測(cè)試集test_ds，這些對(duì)象可以用于后續(xù)的訓(xùn)練、驗(yàn)證和測(cè)試過(guò)程。

4. pretrain

　　pretrain函數(shù)是megatron/pretrain_gpt.py文件中的一個(gè)執(zhí)行入口，通常會(huì)將該函數(shù)寫于文件的末尾。該函數(shù)被第一章中的啟動(dòng)腳本調(diào)用，進(jìn)而開(kāi)啟訓(xùn)練流程。

　　pretrain函數(shù)的入?yún)⑷缦?#xff1a;

　　train_valid_test_datasets_provider：這是第３小節(jié)分析的函數(shù)指針，負(fù)責(zé)提供訓(xùn)練、驗(yàn)證和測(cè)試數(shù)據(jù)集。

　　model_provider：這是第１小節(jié)分析的函數(shù)指針，負(fù)責(zé)提供GPT模型的實(shí)例。它可能根據(jù)傳入的配置或參數(shù)來(lái)初始化模型。

　　ModelType.encoder_or_decoder：這個(gè)參數(shù)指定了模型的類型，這里是編碼器或解碼器（對(duì)于GPT模型，它實(shí)際上是一個(gè)解碼器）。

　　forward_step：這是第２小節(jié)分析的函數(shù)指針，定義了模型訓(xùn)練過(guò)程中的一個(gè)前向傳播步驟，包括數(shù)據(jù)的前向傳遞和損失的計(jì)算。

　　args_defaults：這是一個(gè)字典，包含了預(yù)訓(xùn)練過(guò)程中一些默認(rèn)參數(shù)的鍵值對(duì)。在這個(gè)例子中，它指定了默認(rèn)的tokenizer_type為GPT2BPETokenizer，這意味著在文本預(yù)處理時(shí)將使用基于BPE（Byte Pair Encoding）的GPT-2分詞器。

　　至此，pretrain_gpt.py的源碼基本解析完畢，下一篇文章將以pretrain函數(shù)為入口，跟隨代碼運(yùn)行流程，深入其內(nèi)部實(shí)現(xiàn)，詳細(xì)解析。