跟代码执行流程，读Megatron源码（二）训练入口pretrain

　　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_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 yaml
    if 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 models
        if 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函数是用于构建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 iterator
        model (GPTModel): The GPT Model
    """
    args = get_args()
    timers = get_timers()

    # Get the batch.
    timers('batch-generator', log_level=2).start()
    global stimer
    with 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顾名思义，这个函数是GPT模型训练过程中前向处理函数，负责处理一批输入数据并通过模型进行前向传播。

　　该函数的核心实现，仍然是对model（forward_step函数的入参）的forward的调用，只是在调用之前封装了计时器计时逻辑以及批次数据获取的逻辑（这部分逻辑会根据不同的业务场景变化而变化，故不能直接封装到model的forward函数中，而是应该在pretrain脚本中实现），具体代码流程如下：

　　a. 获取参数和计时器：

　　通过get_args()和get_timers()函数分别获取训练参数和计时器对象，用于控制训练过程和记录时间消耗。

　　b. 获取批次数据：

　　使用timers对象记录获取批次数据的时间（可选，通过log_level=2控制）。

　　调用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 = MockGPTDataset
    else:
        dataset_type = GPTDataset

    print_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。