腾讯云个人网站备案,网站目标建设,阿里指数查询官网,做一个电商网站多少钱Megatron-LM默认支持GPT、T5、BERT等多个常见模型的预训练#xff0c;当下大模型流行#xff0c;故以pretrain_gpt.py为例做源码的走读。
一. 启动pretrain_gpt.py pretrain_gpt.py为GPT类模型的训练入口#xff0c;它通过命令行形式被调用#xff0c;其精确执行路径位于M… Megatron-LM默认支持GPT、T5、BERT等多个常见模型的预训练当下大模型流行故以pretrain_gpt.py为例做源码的走读。
一. 启动pretrain_gpt.py pretrain_gpt.py为GPT类模型的训练入口它通过命令行形式被调用其精确执行路径位于Megatron-LM框架的examples/gpt3目录下。具体而言启动过程依赖于train_gpt3_175b_distributed.sh这一脚本该脚本专为部署GPT-3模型在分布式环境下训练而设计当然也可以参照编写自定义的启动脚本。 在train_gpt3_175b_distributed.sh脚本内部核心操作是通过torchrun命令实现的该命令是PyTorch分布式训练的一部分用于在多个计算节点上高效并行地执行pretrain_gpt.py。此过程确保了模型训练任务能够充分利用集群资源加速训练过程代码如下图 二. torchrun简介 trochrun是PyTorch官方推荐用于替代torch.distributed.launch的分布式数据并行训练模块。它旨在提供一种更灵活、更健壮的方式来启动和管理分布式训练任务。 trochrun启动并行训练任务的原理如下
1. 初始化分布式环境 trochrun首先负责初始化分布式训练所需的环境。这包括设置通信后端如NCCL、GLOO等、分配工作进程的RANK和WORLD_SIZE即参与训练的总进程数以及处理其他与分布式训练相关的配置。
2. 分配工作进程 trochrun会根据指定的参数如--nnodes、--nproc-per-node等来分配工作进程。这些进程可以是同一台机器上的多个 GPU也可以是跨多台机器的GPU。每个进程都会加载相同的训练脚本如pretrain_gpt.py但会处理不同的数据子集以实现并行训练。
3. 同步与通信 在训练过程中torchrun 管理下的各个工作进程需要频繁地进行同步和通信。这包括梯度同步在反向传播后同步各GPU上的梯度、参数更新使用同步后的梯度更新模型参数等。PyTorch提供了丰富的API如torch.distributed.all_reduce、torch.distributed.barrier等来支持这些操作。
4. 优雅处理故障 trochrun相比torch.distributed.launch的一大改进是它能够更优雅地处理工作进程的故障。例如如果某个工作进程因为某种原因崩溃了torchrun可以尝试重新启动该进程以确保训练任务的连续性。此外torchrun还支持弹性训练elastic training即允许在训练过程中动态地增加或减少工作进程的数量。
5. 简化配置与启动 trochrun通过提供命令行接口和配置文件选项来简化分布式训练的配置和启动过程。用户只需指定少量的参数如节点数、每节点进程数等即可启动复杂的分布式训练任务。此外torchrun还支持从环境变量中读取配置信息这使得在不同环境中部署训练任务变得更加灵活。
6. 自动化资源分配 在某些情况下torchrun还可以与资源管理器如Kubernetes、Slurm等集成以自动化地分配和管理训练所需的计算资源。这包括GPU、CPU、内存和存储等资源。通过集成资源管理器torchrun可以进一步提高分布式训练的可扩展性和灵活性。 总之torchrun通过以上机制共同作用使得使用PyTorch进行分布式训练变得更加高效、可靠和易于管理。
三. 主要函数 pretrain_gpt.py脚本封装了多个核心功能组件具体包括model_provider()forward_step()train_valid_test_datasets_provider()以及pretrain()等主要函数。 其中model_provider()负责提供预训练所需的模型实例对象forward_step()定义了模型前向传播的具体步骤包括输入处理、模型计算等train_valid_test_datasets_provider()则负责准备训练、验证及测试所需的数据集确保数据的有效供给。 值得注意的是前三个函数model_provider()forward_step()train_valid_test_datasets_provider()更是作为pretrain()函数的入参共同构成了GPT模型训练入口。 这种设计确保了预训练过程的模块化、灵活性与可扩展性。下面会从model_provider()开始逐行解析源码。
四. 源码分析
1. model_provider def model_provider(pre_processTrue, post_processTrue) - 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 modelargs get_args()use_te args.transformer_impl transformer_engineprint_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_tokentypes0,parallel_outputTrue,pre_processpre_process,post_processpost_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(configconfig,transformer_layer_spectransformer_layer_spec,vocab_sizeargs.padded_vocab_size,max_sequence_lengthargs.max_position_embeddings,pre_processpre_process,post_processpost_process,fp16_lm_cross_entropyargs.fp16_lm_cross_entropy,parallel_outputTrue,share_embeddings_and_output_weightsnot args.untie_embeddings_and_output_weights,position_embedding_typeargs.position_embedding_type,rotary_percentargs.rotary_percent,)return modelmodel_provider函数是用于构建GPT生成预训练Transformer模型实例的函数它会以类似函数指针的形式作为pretrain()的入参传递到后续的训练代码中供训练过程调用。 该函数主要代码流程包含以下几个步骤 a. 获取参数和配置 通过 get_args() 函数获取命令行参数和配置文件中的参数。根据 args.transformer_impl 的值确定是否使用 Transformer Engine (use_te)。 b. 配置模型 如果指定了 YAML 配置文件 (args.yaml_cfg)则从 YAML 文件中加载模型结构。否则根据命令行参数 (args) 加载。 c. 选择模型类型 如果args.use_legacy_models为True则使用megatron.legacy.model.GPTModel构建模型。这通常用于向后兼容或测试旧版本的模型。 如果不使用旧版模型则直接运行构建GPTModel如图红框部分。 其中GPTModel的参数包括配置 (config)、词汇表大小 (vocab_size)、最大序列长度 (max_sequence_length)、是否进行前处理和后处理、是否使用 FP16 进行语言模型交叉熵计算、是否并行输出等。这些参数基本上都来源于配置文件关于配置文件的内容和解析将于下文详述。 注此处的model_provider是作为函数指针传递到pretrain()中函数指针只有在调用时才会真正执行故GPTModel的具体实现待到执行时再具体分析。
2. forward_step
def forward_step(data_iterator, model: GPTModel):Forward training step.Args:data_iterator : Input data iteratormodel (GPTModel): The GPT Modelargs get_args()timers get_timers()# Get the batch.timers(batch-generator, log_level2).start()global stimerwith stimer(bdataTrue):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,labelslabels)return output_tensor, partial(loss_func, loss_mask)forward_step顾名思义这个函数是GPT模型训练过程中前向处理函数负责处理一批输入数据并通过模型进行前向传播。 该函数的核心实现仍然是对modelforward_step函数的入参的forward的调用只是在调用之前封装了计时器计时逻辑以及批次数据获取的逻辑这部分逻辑会根据不同的业务场景变化而变化故不能直接封装到model的forward函数中而是应该在pretrain脚本中实现具体代码流程如下 a. 获取参数和计时器 通过get_args()和get_timers()函数分别获取训练参数和计时器对象用于控制训练过程和记录时间消耗。 b. 获取批次数据 使用timers对象记录获取批次数据的时间可选通过log_level2控制。 调用get_batch函数从data_iterator中获取一批数据包括tokens输入文本对应的token IDs、labels训练标签通常用于计算损失对于语言模型任务labels通常是tokens的右移一位版本、loss_mask损失掩码用于忽略某些位置的损失计算如填充位置、attention_mask注意力掩码用于指示哪些位置需要参与注意力计算和position_ids位置ID用于模型中的位置编码。 c. 模型前向传播 使用stimer可能是一个自定义的计时器记录模型前向传播的时间。 将获取到的数据tokens, position_ids, attention_mask, labels传递给模型model进行前向传播。这里labels是可选的用于计算损失但在前向传播阶段不一定需要。 模型输出output_tensor通常包含模型的预测结果如logits。 d. 返回输出和损失函数 返回output_tensor和partial函数该函数需要loss_mask作为参数来计算损失。这种方式允许延迟损失的计算直到所有相关的数据都已准备好。 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该函数接收一个参数train_val_test_num_samples这是一个列表包含了训练集、验证集和测试集的样本数量。函数的目的是根据提供的参数和配置构建GPT模型的训练、验证和测试数据集。主要代码流程如下 a. 获取参数和配置 使用get_args()函数获取训练过程中的全局参数并通过core_gpt_dataset_config_from_args根据这些参数生成数据集配置对象config。 b. 确定数据集类型 根据args.mock_data的值决定使用哪种数据集类型。如果mock_data为True则使用MockGPTDataset这是一种模拟数据集可能用于测试或快速原型开发。如果mock_data为False则使用GPTDataset这是实际的数据集类型包含真实的训练数据。 c. 构建数据集 使用BlendedMegatronDatasetBuilder类来构建数据集传递给BlendedMegatronDatasetBuilder的参数包括数据集类型dataset_type、训练/验证/测试集的样本数量train_val_test_num_samples、is_dataset_built_on_rank用于检查当前处理单元是否负责构建数据集以及配置对象config。 调用build()方法实际构建数据集该方法返回三个数据集对象训练集train_ds、验证集valid_ds和测试集test_ds。 其中BlendedMegatronDatasetBuilder来源于包“megatron.core.datasets.blended_megatron_dataset_builder”由于数据集构建逻辑比较简单故在此不做详述有兴趣的同学可以自行查看。 d. 返回值 函数返回三个数据集对象训练集train_ds、验证集valid_ds和测试集test_ds这些对象可以用于后续的训练、验证和测试过程。 4. pretrain pretrain函数是megatron/pretrain_gpt.py文件中的一个执行入口通常会将该函数写于文件的末尾。该函数被第一章中的启动脚本调用进而开启训练流程。 pretrain函数的入参如下 train_valid_test_datasets_provider这是第小节分析的函数指针负责提供训练、验证和测试数据集。 model_provider这是第小节分析的函数指针负责提供GPT模型的实例。它可能根据传入的配置或参数来初始化模型。 ModelType.encoder_or_decoder这个参数指定了模型的类型这里是编码器或解码器对于GPT模型它实际上是一个解码器。 forward_step这是第小节分析的函数指针定义了模型训练过程中的一个前向传播步骤包括数据的前向传递和损失的计算。 args_defaults这是一个字典包含了预训练过程中一些默认参数的键值对。在这个例子中它指定了默认的tokenizer_type为GPT2BPETokenizer这意味着在文本预处理时将使用基于BPEByte Pair Encoding的GPT-2分词器。 至此pretrain_gpt.py的源码基本解析完毕下一篇文章将以pretrain函数为入口跟随代码运行流程深入其内部实现详细解析。
