pipeline_parallel.py — 流水线并行

文件路径: verl/utils/megatron/pipeline_parallel.py

文件概述

流水线并行 (Pipeline Parallelism) 的工具函数，处理 micro-batch 的输入形状计算和虚拟流水线并行 (VPP) 的批次生成。

核心函数

1. 计算输入形状

def compute_transformers_input_shapes(batches, meta_info):
    """预计算每个 micro-batch 在每个 PP stage 的输入形状"""
    for model_inputs in batches:
        input_ids_rmpad = unpad_input(input_ids, attention_mask)
        if meta_info["sequence_parallel"]:
            input_ids_rmpad = pad_to_sequence_parallel(input_ids_rmpad)
            # 形状: [total_nnz // tp_size, 1, hidden_size]
        else:
            # 形状: [total_nnz, 1, hidden_size]

PP 需要提前知道每个 stage 的输入形状，以便分配缓冲区。

2. 批次生成器

def make_batch_generator(batches, vpp_size):
    """创建适合 Megatron PP 的批次迭代器"""
    if vpp_size > 1:
        # VPP: 每个虚拟 stage 复制一份迭代器
        batch_generator = [iter(batches)] * vpp_size
    else:
        batch_generator = iter(batches)

与其他模块的关系

• 依赖 sequence_parallel.py 进行 padding
• 被 Megatron 训练 step 的 PP 调度使用

小结

流水线并行的批次管理工具。