Megatron-LM training.py 源码精读 · 训练循环篇

# 1. 将 optimizer 的 loss scaling 函数注入 config
#    PP 调度器需要用它来 scale 梯度
config.grad_scale_func = optimizer.scale_loss

# 2. 设置 no_sync_func：控制何时执行梯度 AllReduce
#    overlap_grad_reduce 开启时，DDP 自行管理通信时机
if isinstance(model[0], DDP) and args.overlap_grad_reduce:
    config.no_sync_func = [model_chunk.no_sync for model_chunk in model]
    if args.align_grad_reduce:
        config.grad_sync_func = [model_chunk.start_grad_sync ...]

# 3. 设置 param_sync_func：控制何时执行参数 AllGather
#    overlap_param_gather 开启时，参数会在前向传播之前异步收集
if args.overlap_param_gather and args.align_param_gather:
    config.param_sync_func = [model_chunk.start_param_sync ...]

# 4. 设置梯度归一化函数
config.finalize_model_grads_func = finalize_model_grads

# 5. 选择 PP 调度函数
forward_backward_func = get_forward_backward_func()

rerun_state_machine = get_rerun_state_machine()

# 注意：这不是简单的 "执行一次"
# rerun 状态机可能要求重新执行前向反向（容错重试）
while rerun_state_machine.should_run_forward_backward(data_iterator):
    # 清零梯度
    for model_chunk in model:
        model_chunk.zero_grad_buffer()
    optimizer.zero_grad()

    # 执行前向 + 反向
    losses_reduced = forward_backward_func(
        forward_step_func=forward_step_func,
        data_iterator=data_iterator,
        model=model,
        num_microbatches=get_num_microbatches(),
        seq_length=args.seq_length,
        micro_batch_size=args.micro_batch_size,
        forward_only=False,                        # 训练模式
    )

# 重试结束后，检查是否需要存盘退出
should_checkpoint, should_exit, exit_code = \
    rerun_state_machine.should_checkpoint_and_exit()

if mpu.is_pipeline_last_stage(ignore_virtual=True):
    loss_reduced = {}
    for key in losses_reduced[0].keys():
        val = [x[key].view(-1) for x in losses_reduced]

        if val[0].numel() == 2:
            # 新模式：per-token 归一化
            # 每个 microbatch 返回 [loss_sum, num_tokens]
            # 先对所有 microbatch 求和，再 AllReduce 跨 DP
            val = torch.vstack(val).sum(dim=0)   # [total_loss, total_tokens]
            torch.distributed.all_reduce(
                val,
                group=mpu.get_data_parallel_group(with_context_parallel=True)
            )
            loss_reduced[key] = val[0] / val[1]  # total_loss / total_tokens

        elif val[0].numel() == 1:
            # 旧模式：per-microbatch 平均
            val = torch.cat(val).mean()
            loss_reduced[key] = val

def get_forward_backward_func():
    pp_size = parallel_state.get_pipeline_model_parallel_world_size()
    vp_size = parallel_state.get_virtual_pipeline_model_parallel_world_size()

    if pp_size > 1:
        if vp_size is not None:
            # VPP 模式：交错 1F1B
            return forward_backward_pipelining_with_interleaving
        else:
            # 标准 PP：1F1B
            return forward_backward_pipelining_without_interleaving
    else:
        # 无 PP：简单遍历所有 microbatch
        return forward_backward_no_pipelining

# ========== 循环开始前：禁用 pre-hook ==========
if should_disable_forward_pre_hook(args):
    disable_forward_pre_hook(model, param_sync=False)
    param_sync_func = config.param_sync_func
    config.param_sync_func = None      # 也暂时禁用 param_sync
    pre_hook_enabled = False

# ========== 首次成功训练后：启用 pre-hook ==========
if iteration == start_iteration:
    if skipped_iter:
        # FP16 loss scaling 还没稳定，继续等待
        start_iteration = iteration + 1
    else:
        # 首次成功！启用 pre-hook
        enable_forward_pre_hook(model)
        config.param_sync_func = param_sync_func
        pre_hook_enabled = True

# 每 eval_interval 个 iteration 做一次评估
if args.eval_interval and iteration % args.eval_interval == 0 and args.do_valid:
    # 1. 暂停训练计时器（不计入训练时间）
    timers('interval-time').stop()

    # 2. 禁用 forward pre-hook（评估不需要梯度通信）
    disable_forward_pre_hook(model)

    # 3. 执行评估
    evaluate_and_print_results(
        prefix=f'iteration {iteration}',
        forward_step_func=forward_step_func,
        data_iterator=valid_data_iterator,
        model=model,
        iteration=iteration,
        config=config,
    )

    # 4. 重新启用 pre-hook + 恢复计时器
    enable_forward_pre_hook(model)
    timers('interval-time', log_level=0).start(barrier=True)