跟代码执行流程，读Megatron源码（五）megatron训练脚本training.py之setup_model_and_optimizer()配置模型和优化器

上文详细讲解了megatron初始化的详细过程，尤其重点分析了megatron分组的逻辑，本文将继续跟踪代码执行流程，分析pretrain()的第二个重要步骤：配置模型和优化器。

一. setup_model_and_optimizer()的调用关系

　　setup_model_and_optimizer() 函数作为 training.py 脚本中 pretrain() 函数的一个核心组件被调用，其首要职责是实例化并配置模型（Model）与优化器（Optimizer）对象。这两个关键对象在随后的训练循环（training loop）中被频繁调用以执行模型的前向传播、计算损失、反向传播以及参数更新等核心步骤，如下图。

二. setup_model_and_optimizer()代码分析

　　setup_model_and_optimizer()函数定义位于training.py中，如下图。

　　主要包含以下几个关键步骤：

　　1. 初始化参数和工具

　　通过get_args(), get_timers(), 和 get_one_logger()函数分别获取配置参数、时间记录器和日志记录器。这些工具和参数将用于后续模型配置、训练和日志记录。

　　2. 构建模型

　　调用get_model(model_provider_func, model_type)函数，该函数依据输入的model_provider_func（模型提供函数）和model_type（模型类型）参数来实例化并构建模型。此过程不仅限于单纯获取模型实例，其核心功能在于对模型进行深度定制，特别是执行流水线并行（Pipeline Parallelism，PP）、数据并行（Data Parallelism, DP）的封装处理。这一封装步骤对于模型在分布式或多GPU环境下的并行训练至关重要，确保了模型能够在高效利用计算资源的同时，保持数据处理的同步性和一致性，从而优化训练过程的性能和效果，具体的模型实例化、和并行封装逻辑将在下一节详细讲解。

　　3. 配置优化器

　　通过dataclasses.fields(OptimizerConfig)和args（配置参数）构建优化器的配置字典kwargs。

　　使用这些配置信息实例化OptimizerConfig对象，并将其与模型、权重衰减条件（no_wd_decay_cond）、学习率缩放条件（scale_lr_cond）和学习率乘数（lr_mult）一起传递给get_megatron_optimizer函数，以获取优化器实例。

　　同时，通过get_optimizer_param_scheduler(optimizer)获取优化器参数调度器。

　　4. 加载检查点（checkpoint）

　　如果指定了加载检查点（通过args.load或args.pretrained_checkpoint），则记录加载开始时间，并调用load_checkpoint函数加载模型、优化器和优化器参数调度器的检查点。

　　最后，记录加载完成时间和整体耗时，并更新迭代次数（args.iteration）和已完成的浮点运算次数（args.num_floating_point_operations_so_far）。

　　5. 模型的特殊处理（具有init_state_dict_from_bert函数的模型）

　　如果当前迭代次数为0，且模型为单个未包装模型，并且该模型具有init_state_dict_from_bert函数，则调用此函数从预训练的BERT模型中初始化模型状态字典。

　　如果启用了FP16（半精度浮点数）训练，则调用优化器的reload_model_params函数重新加载模型参数，以确保它们以正确的精度格式加载。

　　6. 返回结果

　　最后，函数返回模型、优化器和优化器参数调度器的实例，供后续的训练或评估过程使用。

　　整个过程涵盖了从模型构建、优化器配置到可选的检查点加载和特定条件下的模型初始化，为后续的模型训练或评估工作做好了准备。

三. get_model()代码分析

　　get_model函数函数的核心职责在于根据配置构建模型实例，特别是在分布式并行计算环境中，其复杂性显著提升。

　　该函数不仅负责将基础模型实例化，更需依据流水线并行（Pipeline Parallelism）及虚拟流水线并行（Virtual Pipeline Parallelism）策略，对模型架构进行横向分割，确保各计算单元（如GPU卡）能够正确获取到分割后的模型片段。

　　此外，为实现资源的高效利用与训练速度的加速，get_model 还需集成数据并行（Data Parallelism）的封装机制，通过跨多个计算设备复制模型副本并分散处理数据子集，从而有效提升整体训练效率与吞吐量。

　　这一过程要求精确控制模型分割与数据分配，确保分布式训练的一致性与高效性。

　　代码主要流程解析如下：

　　1. 参数初始化

　　获取配置参数在之前的很多函数中都已经介绍过，此处和未来将不再赘述。

　　2. 虚拟流水线并行下的模型构建

　　这部分代码块考虑了虚拟流水线并行的特殊情况，其中模型被拆分成多个部分，在多个设备或进程中并行处理。

　　首先，判断流水线并行的world_size（即参与流水线并行的设备数）大于1，并且指定了虚拟流水线并行（为了进一步优化流水线并行中的模型划分）。

　　符合以上判断条件，则按虚拟流水线并行的策略构建模型，如下图：

　　在此循环迭代过程中，针对虚拟流水线并行的每一阶段，均独立构建了相应的模型片段实例。通过执行mpu.set_virtual_pipeline_model_parallel_rank(i)函数，我们明确了当前处理的模型片段在整体流水线中的唯一标识符（rank），这一编号是区分流水线中不同模型片段的关键。而在模型的初始化过程中，也会根据该编号来初始化不同的layer，如下图，图中的Transformer layer的初始化就是根据virtual_pipeline_model_parallel_rank来确定的（只做举例，后续章节会详述模型初始化的逻辑）：

　　pre_process与post_process布尔变量精确指示了当前模型片段在流水线中的角色：作为首阶段时需执行输入数据的预处理任务，而作为末阶段时则负责处理模型的最终输出，可能包括必要的后处理逻辑。

　　随后，利用model_provider_func函数定制化地实例化模型片段，该函数智能地根据是否为首阶段或末阶段调整模型的具体配置与结构。具体而言，首阶段实例可能会集成额外的输入处理层以增强数据处理能力，而末阶段实例则可能增设输出处理层或激活函数以优化输出结果的格式与质量。这一过程确保了模型片段能够无缝集成于虚拟流水线并行架构中，高效协同工作。

　　3. 流水线并行下的模型构建

　　以上是虚拟流水线并行场景的处理，对于非虚拟流水线并行场景的处理如下：

　　首先，程序设置了pre_process和post_process标志，与上一小节中一样，这些标志用于判断当前阶段是否是流水线的第一个或最后一个阶段。

　　然后，根据模型类型（model_type）来进一步配置模型。

　　如果模型类型是encoder_and_decoder，则代码会检查是否启用了流水线模型并行（mpu.get_pipeline_model_parallel_world_size() > 1）。如果是，则根据当前流水线阶段的编号（rank）和编码器流水线模型并行大小（args.encoder_pipeline_model_parallel_size）来动态调整预处理、后处理以及是否添加编码器和解码器的逻辑。最终会根据这些配置调用model_provider_func来构建模型实例。

　　如果模型类型不是encoder_and_decoder（即只是编码器或解码器），则直接调用model_provider_func，仅传入预处理和后处理标志作为参数。

　　4. 设置张量并行属性

　　这段代码遍历模型（或模型列表中的每个模型）的所有参数，并调用tensor_parallel.set_defaults_if_not_set_tensor_model_parallel_attributes函数来确保每个参数都设置了张量模型并行的默认属性。这些属性对于优化器来说是必要的，因为它们可能需要根据这些属性来调整参数的更新方式。

　　5. 打印参数数量

　　这段代码仅在数据并行编号为0的进程中执行（为了避免在多个进程中重复打印信息）。它计算并打印当前（张量、流水线）并行编号下的参数总数（通过遍历模型并累加所有参数的数量来实现）。

　　6. GPU分配和半精度转换

　　首先，这段代码将模型（或模型列表中的每个模型）移动到当前CUDA设备上。torch.cuda.current_device()返回当前CUDA设备的索引，通常与代码中先前设置的设备相对应。

　　其次，如果命令行参数指定了使用FP16（半精度）或BF16（混合精度），则这段代码将模型（或模型列表中的每个模型）转换为相应的半精度或混合精度模型。这是通过创建一个Float16Module实例（一个自定义的包装器类，用于将模型转换为半精度）并将原始模型作为参数传递给它的构造函数来实现的。这个新创建的半精度模型实例将被添加到新的模型列表中，该列表随后替换原始的模型列表。

　　7. 封装模型为DDP

　　下图的代码主要处理了在分布式训练环境中，将模型封装进DistributedDataParallel（DDP）的过程，以及一些与DDP配置和数据并行性相关的额外步骤。

　　首先，检查是否需要将模型封装进DDP（wrap_with_ddp），并使用get_model_config函数获取模型配置。

　　创建一个DistributedDataParallelConfig实例，用于配置DDP的各种参数，如梯度聚合是否使用FP32、是否重叠梯度聚合和计算、是否使用分布式优化器等。

　　然后，遍历模型列表为每个模型片段创建一个DDP实例。对于第二个及以后的模型片段，禁用bucketing（一种优化通信效率的技术），因为这些块的通信可以与计算重叠。

　　最后，调用每个DDP模型实例的broadcast_params()方法，该方法将从数据并行的0号卡广播参数到组内的其他所有卡，以确保它们具有相同的初始参数值。

　　至此setup_model_and_optimizer()函数的主要代码流程均分析完毕，下一篇文章将进入训练循环（training loop），看一看3D并行场景下模型的详细训练流程。