model_forward_1f1b_overlap.py — 1F1B 流水线重叠前向

文件路径

verl/models/mcore/model_forward_1f1b_overlap.py

文件概述

为 Megatron-Core 的 1F1B (One Forward One Backward) 流水线调度提供前向函数。1F1B 是流水线并行的核心调度策略，通过重叠前向和反向计算来减少流水线气泡（bubble）。

关键代码讲解

def gptmodel_forward_1f1b_overlap(model, input_ids, position_ids, attention_mask,
                                   labels=None, labels_mask=None, temperature=1.0, ...):
    # 构建调度计划
    schedule_plan = model.build_schedule_plan(
        input_ids=input_ids_rmpad,
        attention_mask=attention_mask,
        labels=labels,
        position_ids=position_ids,
        packed_seq_params=packed_seq_params,
    )

    # 设置后处理回调
    if post_process:
        def postprocess_fn(output):
            # 使用 linear_cross_entropy 融合计算
            log_probs, entropy = linear_cross_entropy(...)
            return postprocess_packed_seqs_for_dict_output(...)
        schedule_plan.set_postprocess_fn(postprocess_fn)

    return schedule_plan  # 返回调度计划，由调度器执行

关键区别：这个函数不直接执行前向计算，而是构建一个 TransformerModelChunkSchedulePlan，交给上层的 1F1B 调度器执行。调度器会负责在合适的时间点交替执行前向和反向。

核心函数列表

函数名	作用
gptmodel_forward_1f1b_overlap()	构建 1F1B 调度计划

与其他模块的关系

• 使用 mcore 的 TransformerModelChunkSchedulePlan
• 使用 util.py 的打包函数
• 使用 linear_cross_entropy 融合 kernel

小结

1F1B 重叠前向是流水线并行的高级优化。通过返回调度计划而非直接执行，允许调度器在流水线的不同阶段灵活安排前向和反向计算的重叠执行。