Pipeline Parallelism 调度机制

pretrain_gpt.py  __main__
  pretrain(model_provider=partial(model_provider_fn, gpt_builder), ...)   ← A
    setup_model_and_optimizer(model_provider_func=partial(...), ...)       ← B
      get_model(model_provider_func=partial(...), ...)                     ← C
        for vp_stage in [0, 1]:          ← VPP 循环，每次建一个 chunk
          model_provider_func(pre_process=?, post_process=?, vp_stage=?)  ← D
            model_provider_fn(gpt_builder, pre_process, post_process, vp_stage)  ← E
              gpt_builder(args, pre_process, post_process, vp_stage)       ← F
                GPTModel(config, ..., pre_process, post_process, vp_stage) ← G
                  TransformerBlock(config, ..., vp_stage)                  ← H
                    get_num_layers_to_build(config, vp_stage)  → 2        ← I
                    _build_layers()                                        ← J
                      get_transformer_layer_offset(config, vp_stage, pp_rank) → offset
                      for n in [1, 2]:  global_n = n + offset
                        TransformerLayer(config, layer_number=global_n)
        ← 返回 GPTModel 对象，append 进 model list
      ← 返回 model = [GPTModel(vp=0), GPTModel(vp=1)]

from model_provider import model_provider as model_provider_fn   # E 步的函数

pretrain(
    train_valid_test_datasets_provider,
    partial(model_provider_fn, gpt_builder),  # ← 打包成 model_provider_func
    ...                                        #   以后调用它只需传 pre/post_process, vp_stage
)

model, optimizer, opt_param_scheduler = setup_model_and_optimizer(
    model_provider,   # 就是 partial(model_provider_fn, gpt_builder)
    ...
)

model = get_model(model_provider_func, model_type)
# model_provider_func 还是 partial(model_provider_fn, gpt_builder)

def build_model():
    if pp_world_size > 1 and args.virtual_pipeline_model_parallel_size is not None:
        # VPP 路径
        model = []
        for i in range(vp_size):   # i = 0, 1
            # 问：这个 chunk 需要 Embedding 吗？需要 Output Layer 吗？
            pre_process  = mpu.is_pipeline_first_stage(ignore_virtual=False, vp_stage=i)
            post_process = mpu.is_pipeline_last_stage(ignore_virtual=False, vp_stage=i)
            # ↑ 两个函数的返回值见下面

            chunk = model_provider_func(pre_process=pre_process,
                                        post_process=post_process,
                                        vp_stage=i)
            # 调用 partial(model_provider_fn, gpt_builder)(pre_process, post_process, vp_stage=i)
            # 等价于  model_provider_fn(gpt_builder, pre_process, post_process, vp_stage=i)

            model.append(chunk)
    else:
        # 普通 PP 路径：只建一个模型对象
        pre_process  = mpu.is_pipeline_first_stage()
        post_process = mpu.is_pipeline_last_stage()
        model = model_provider_func(pre_process=pre_process, post_process=post_process)
    return model

def model_provider(model_builder, pre_process=True, post_process=True, vp_stage=None):
    # model_builder = gpt_builder（由 partial 在 A 步绑定好的）
    # pre_process / post_process / vp_stage 由 D 步的循环传入

    args = get_args()
    # ... 一些内存追踪、ModelOpt 判断 ...
    return model_builder(args, pre_process, post_process, vp_stage)
    # 调用 gpt_builder(args, pre_process, post_process, vp_stage)

def gpt_builder(args, pre_process, post_process, vp_stage=None, config=None):
    if config is None:
        config = core_transformer_config_from_args(args)
        # config.num_layers = 8
        # config.pipeline_model_parallel_size = 2
        # config.virtual_pipeline_model_parallel_size = 2

    transformer_layer_spec = get_gpt_layer_local_spec(...)
    # transformer_layer_spec 描述每一层用什么模块实现（Attention/MLP 用哪个类）

    return GPTModel(
        config=config,
        transformer_layer_spec=transformer_layer_spec,
        vocab_size=args.padded_vocab_size,
        pre_process=pre_process,    # ← 从 D 步传来
        post_process=post_process,  # ← 从 D 步传来
        vp_stage=vp_stage,          # ← 从 D 步传来（0 或 1）
    )

def __init__(self, config, transformer_layer_spec, vocab_size,
             pre_process=True, post_process=True, vp_stage=None, ...):
    self.pre_process = pre_process
    self.post_process = post_process
    self.vp_stage = vp_stage

    # ── 只有 pre_process=True 的 chunk 才建 Embedding ──
    if self.pre_process:
        self.embedding = Embedding(vocab_size=vocab_size, hidden_size=config.hidden_size, ...)
        # 仅 GPU0/chunk0 有这个属性

    # ── 每个 chunk 都建 TransformerBlock，但层数和全局编号不同（见 H）──
    self.decoder = TransformerBlock(
        config=config,
        spec=transformer_layer_spec,
        pre_process=pre_process,
        post_process=post_process,
        vp_stage=vp_stage,          # ← 传下去，决定层偏移
    )

    # ── 只有 post_process=True 的 chunk 才建 Output Layer ──
    if self.post_process:
        self.output_layer = Linear(config.hidden_size, vocab_size, bias=False)
        # 仅 GPU1/chunk1 有这个属性

class TransformerBlock(MegatronModule):
    def __init__(self, config, spec, pre_process, post_process, vp_stage=None, ...):
        self.vp_stage = vp_stage

        # I: 这个 chunk 要建几层？
        self.submodules = _get_block_submodules(config, spec, vp_stage, pp_rank)
        # _get_block_submodules 内部调 get_num_layers_to_build(config, vp_stage, pp_rank)
        # → 8 // 2 // 2 = 2 层（所有 chunk 都一样）

        self._build_layers()   # J
        self.num_layers_per_pipeline_rank = len(self.layers)   # = 2

    def _build_layers(self):
        def build_layer(layer_spec, layer_number):   # layer_number = 1 或 2（本地编号）
            # J: 算全局编号
            global_layer_number = layer_number + get_transformer_layer_offset(
                self.config, self.vp_stage, pp_rank
            )
            # get_transformer_layer_offset 公式：
            #   offset = vp_stage * (num_layers // vp_size) + pp_rank * (num_layers // pp_size // vp_size)
            #          = vp_stage * 4                       + pp_rank * 2

            return build_module(layer_spec, config=config, layer_number=layer_number, ...)
            # TransformerLayer 内部会用 global_layer_number 初始化 RoPE 等

        self.layers = torch.nn.ModuleList([
            build_layer(spec, i + 1)          # i+1 = 1, 2
            for i, spec in enumerate(self.submodules.layer_specs)
        ])

GPU 0:  model = [
    # model[0]  gpt_builder(pre_process=True,  post_process=False, vp_stage=0) 的返回值
    GPTModel {
        embedding:         Embedding(50257→1024)    ← pre_process=True，有
        decoder.layers[0]: TransformerLayer(global_layer_number=1)
        decoder.layers[1]: TransformerLayer(global_layer_number=2)
        output_layer:      (无)                     ← post_process=False
    },
    # model[1]  gpt_builder(pre_process=False, post_process=False, vp_stage=1) 的返回值
    GPTModel {
        embedding:         (无)
        decoder.layers[0]: TransformerLayer(global_layer_number=5)
        decoder.layers[1]: TransformerLayer(global_layer_number=6)
        output_layer:      (无)
    },
]

GPU 1:  model = [
    # model[0]  gpt_builder(pre_process=False, post_process=False, vp_stage=0)
    GPTModel {
        embedding:         (无)
        decoder.layers[0]: TransformerLayer(global_layer_number=3)
        decoder.layers[1]: TransformerLayer(global_layer_number=4)
        output_layer:      (无)
    },
    # model[1]  gpt_builder(pre_process=False, post_process=True,  vp_stage=1)
    GPTModel {
        embedding:         (无)
        decoder.layers[0]: TransformerLayer(global_layer_number=7)
        decoder.layers[1]: TransformerLayer(global_layer_number=8)
        output_layer:      Linear(1024→50257)       ← post_process=True，有
    },
]

def forward_step_helper(virtual_microbatch_id, ...):
    chunk_id = get_model_chunk_id(virtual_microbatch_id, forward=True)
    # chunk_id = 0 或 1，对应 vp_stage

    output = forward_step(
        forward_step_func,          # pretrain_gpt.py 里的 forward_step()
        data_iterator[chunk_id],    # 每个 chunk 有自己的数据迭代器
        model[chunk_id],            # ← 拿对应 chunk 的 GPTModel
        ...
    )
    # output = hidden_states（[seq_len, batch, hidden]），
    # 如果是最后一层则是 loss scalar

# forward_backward_pipelining_with_interleaving
num_warmup_microbatches = (pipeline_parallel_size - pipeline_parallel_rank - 1) * 2
num_warmup_microbatches += (num_model_chunks - 1) * microbatch_group_size_per_vp_stage

def get_forward_backward_func(pp_size=None, vp_size=None):
    if pp_size is None and vp_size is None:
        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:
            forward_backward_func = forward_backward_pipelining_with_interleaving
        else:
            forward_backward_func = forward_backward_pipelining_without_interleaving
    else:
        forward_backward_func = forward_backward_no_pipelining
    return forward_backward_func

if group.rank() % 2 == 0:
    # 偶数 rank：先发送，再接收
    if tensor_send_next is not None:
        reqs["send_next"] = isend(tensor_send_next, dst=next_pipeline_rank, ...)
    if tensor_recv_prev is not None:
        reqs["recv_prev"] = irecv(tensor_recv_prev, src=prev_pipeline_rank, ...)
    if tensor_send_prev is not None:
        reqs["send_prev"] = isend(tensor_send_prev, dst=prev_pipeline_rank, ...)
    if tensor_recv_next is not None:
        reqs["recv_next"] = irecv(tensor_recv_next, src=next_pipeline_rank, ...)
else:
    # 奇数 rank：先接收，再发送
    if tensor_recv_prev is not None:
        reqs["recv_prev"] = irecv(tensor_recv_prev, src=prev_pipeline_rank, ...)
    if tensor_send_next is not None:
        reqs["send_next"] = isend(tensor_send_next, dst=next_pipeline_rank, ...)
    if tensor_recv_next is not None:
        reqs["recv_next"] = irecv(tensor_recv_next, src=next_pipeline_rank, ...)
    if tensor_send_prev is not None:
        reqs["send_prev"] = isend(tensor_send_prev, dst=prev_pipeline_rank, ...)

# max_outstanding_backprops = warmup 数 + 1（当前正在跑的 forward）
max_outstanding_backprops = None
if config.num_microbatches_with_partial_activation_checkpoints is not None:
    max_outstanding_backprops = num_warmup_microbatches + 1

# Warmup 阶段决定是否对当前 micro-batch 做 full recompute
for i in range(num_warmup_microbatches):
    if max_outstanding_backprops is not None:
        checkpoint_activations_microbatch = (
            i % max_outstanding_backprops
            >= config.num_microbatches_with_partial_activation_checkpoints
        )
    else:
        checkpoint_activations_microbatch = None
    # forward_step 内部检查 checkpoint_activations_microbatch
    # 不为 None 时传入 forward_step_func 以触发 gradient checkpointing