Interleaved Pipeline Schedule 深度源码精读

def forward_step_helper_preprocess(chunk_id):
    # input_tensors[chunk_id] 中，offset 指向下一个待消费的 tensor
    input_tensor = input_tensors[chunk_id][input_tensors_offset[chunk_id]]
    input_tensors_offset[chunk_id] += 1
    # 用索引而非 pop(0) 的原因：
    #   pop(0) 对 list 是 O(n)，每次消费都要移动所有后续元素
    #   索引 + counter 是 O(1)，避免 warmup 阶段的性能退化
    return input_tensor

# ── 调用时机：每次 forward / backward 计算"完成之后"，决定本次 P2P 是否 recv ──
# 函数从不在"计算之前"调用，它回答的是：
#   "我刚算完 vid=X，对端（邻居 stage）现在有没有东西要发给我？
#    如果有，应该放进 input_tensors[哪个 chunk] 的队列？"

# warmup 阶段（对每个 k 调用一次，决定 k→k+1 的 recv）
for k in range(num_warmup_microbatches):
    output_tensor = forward_step(...)                                # 先算
    recv_prev, next_chunk_id = recv_tensor_from_previous_stage(k, forward=True)  # 再问
    input_tensor = send_forward_recv_forward(output_tensor, recv_prev)
    if recv_prev:
        input_tensors[next_chunk_id].append(input_tensor)           # 存进下一步的 chunk 队列

# steady 阶段（每轮两次：forward 后问一次，backward 后问一次）
for k in range(num_microbatches_remaining):
    fid, bid = k + num_warmup, k
    output_tensor    = forward_step(fid)                            # forward 算完
    input_tensor_grad = backward_step(bid)                          # backward 算完
    recv_prev, next_fwd_chunk = recv_tensor_from_previous_stage(fid, forward=True)
    recv_next, next_bwd_chunk = recv_tensor_from_previous_stage(bid, forward=False)
    (input_tensor, output_tensor_grad) = send_fwd_bwd_recv_fwd_bwd(
        output_tensor, input_tensor_grad, recv_prev, recv_next)
    if recv_prev: input_tensors[next_fwd_chunk].append(input_tensor)
    if recv_next: output_tensor_grads[next_bwd_chunk].append(output_tensor_grad)

def recv_tensor_from_previous_stage(vid, forward):
    # 前向时：leading = GPU0（第一个处理数据的）
    # 反向时：leading = GPU1（第一个做 backward 的，因为它是 last stage）
    is_leading = is_pp_first_stage() if forward else is_pp_last_stage()

    # 前向时 last_chunk = c1（最后一个 chunk 的激活值不发给任何人，只产生 loss）
    # 反向时 last_chunk = c0（最后一个反向 chunk 不 recv 梯度，只 recv None）
    last_chunk = (num_model_chunks - 1) if forward else 0

    if is_leading:   # ── GPU0 视角（pp=2, V=2） ────────────────────────────────
        # 关键约束：GPU0 比 GPU1 早 (pp_size-1)=1 步
        # 当 GPU0 做完 vid=X 时，GPU1 刚做完 vid = X - 1
        # GPU0 想 recv 的是 GPU1 的 c0 输出（循环绕回给 GPU0 的 c1）
        # 但如果 X < pp_size-1=1，即 X=0，GPU1 还没开始，没有任何输出
        if vid < (pipeline_parallel_size - 1):   # X=0 时
            recv = False
            # next_chunk_id 设为 vid+1 对应的 chunk（仅用于边界判断，recv=False 后不实际使用）
            next_chunk_id = get_model_chunk_id(vid + 1, forward)
        else:
            # GPU1 刚做完的是 vid = X-(pp_size-1) = X-1，对应 chunk：
            next_chunk_id = get_model_chunk_id(vid - (pipeline_parallel_size - 1), forward)
            # 例：vid=1 → GPU1 刚做完 vid=0 (c0,m0) → next_chunk_id = c0
            #     vid=2 → GPU1 刚做完 vid=1 (c0,m1) → next_chunk_id = c0
            #     vid=3 → GPU1 刚做完 vid=2 (c1,m0) → next_chunk_id = c1

        # 如果 GPU1 刚做完的是 last_chunk (c1)，说明这是最后一个前向 chunk，
        # 它只产生 loss，不向 GPU0 发激活值 → 不 recv
        if next_chunk_id == last_chunk:   # vid=3 时 next_chunk_id=c1=last_chunk → recv=False
            recv = False

        # GPU1 把 c0 的输出发给 GPU0 的 c1（即下一个 chunk）
        # 所以存储目标队列是 next_chunk_id + 1
        # 前向时 chunk_id 单调递增；反向时单调递减
        next_chunk_id = next_chunk_id + 1 if forward else next_chunk_id - 1

    else:   # ── GPU1 视角（非 leading）────────────────────────────────────────
        # 极简：vid=X 做完，下一步是 vid=X+1，recv 它对应 chunk 的 input
        next_chunk_id = get_model_chunk_id(vid + 1, forward)

    return recv, next_chunk_id

# GPU0 完成 fid=X 后调用 recv_tensor_from_previous_stage(X, forward=True)
#
# fid=0 (c0,m0 做完)：X=0 < pp_size-1=1 → recv=False
#   GPU1 此时才刚开始处理 vid=0，还没任何输出
#
# fid=1 (c0,m1 做完)：X=1 >= 1 → 看 GPU1 刚做完的 vid=1-1=0 → schedule[0]=(m0,c0) → c0
#   c0 != last_chunk(c1) → recv=True
#   next_chunk_id = c0 + 1 = c1
#   → recv GPU1 发来的 c0,m0 激活值，存入 input_tensors[c1]  ← 这是 fid=2(c1,m0) 的 input
#
# fid=2 (c1,m0 做完)：X=2 → GPU1 刚做完 vid=1 → schedule[1]=(m1,c0) → c0
#   recv=True, next=c1
#   → recv GPU1 发来的 c0,m1 激活值，存入 input_tensors[c1]  ← 这是 fid=3(c1,m1) 的 input
#
# fid=3 (c1,m1 做完)：X=3 → GPU1 刚做完 vid=2 → schedule[2]=(m0,c1) → c1
#   c1 == last_chunk → recv=False
#   GPU1 的 c1 chunk 是最后一层，产生 loss，不发激活值给 GPU0
#
# ── 结论：warmup 中 GPU0 只在 fid=1 和 fid=2 之后 recv，共 recv 2 次（c1,m0 和 c1,m1 的 input）

###############################################################################
# ① 前向 preprocess：对 c0（pp_first_stage + vp_first_stage）APPEND None 占位，
#    然后 INDEX 读取（不 pop）
###############################################################################
# schedules.py:1271
if _is_vp_first_stage(chunk_id) and is_pp_first_stage():
    # c0 永远从 data_iterator 读，不走 P2P recv。
    # 但 input_tensors[c0] 里仍需留 None 占位，保持下面 index 运算的一致性。
    # 触发条件：len(input) == len(output)，即"上一次 c0 forward 已经配对了一个 output"
    # → 要再做一次 c0 forward，需要再加一个 None 槽位
    if len(input_tensors[chunk_id]) == len(output_tensors[chunk_id]):
        input_tensors[chunk_id].append(None)   # ← 占位 None

# offset = 截止到本 fid，已经 pop 掉的 input_tensors[chunk_id] 条目数
# （等于"对应 chunk 已完成的反向次数"）
# 公式（非 forward_only 模式）：
#   if fid < num_warmup: offset = 0         ← warmup 阶段没有任何反向
#   else:
#     bid_so_far = fid - num_warmup          ← 到目前为止已完成的反向步数
#     offset = model_chunk_id_table[:bid_so_far].count(V-1-chunk_id)
#             ← 数前 bid_so_far 个反向步里，有几个处理了对应 chunk（反向时 chunk 翻转）
offset = num_released_microbatches(fid, chunk_id)

mb_id = microbatch_id_table[fid]   # 本 fid 处理的是第几个 microbatch
input_tensor = input_tensors[chunk_id][mb_id - offset]   # ← INDEX，不 pop！
# 为什么 mb_id - offset：
#   input_tensors[c] 原本按 mb_id 0,1,2,... 顺序 append
#   每次反向 pop(0) 会把队列头删掉，使列表缩短 1
#   缩短 offset 次后，原来下标 mb_id 的元素现在在下标 mb_id-offset

###############################################################################
# ② 前向 postprocess：APPEND output（激活值入队等反向消费）
###############################################################################
# schedules.py:1288
output_tensors[chunk_id].append(output_tensor)

###############################################################################
# ③ P2P recv 后：APPEND 到下一步前向对应的 chunk 队列（next_chunk_id，非当前）
###############################################################################
# schedules.py:1590
recv, next_chunk_id = recv_tensor_from_previous_stage(fid, forward=True)
if recv:
    input_tensor = send_forward_recv_forward(output_tensor, ...)
    input_tensors[next_chunk_id].append(input_tensor)

###############################################################################
# ④ 反向 preprocess：POP × 3（唯一真正释放内存的地方）
###############################################################################
# schedules.py:1348
if _is_vp_last_stage(chunk_id) and is_pp_last_stage():
    if len(output_tensor_grads[chunk_id]) == 0:
        output_tensor_grads[chunk_id].append(None)  # last stage 梯度来自 loss

input_tensor       = input_tensors[chunk_id].pop(0)        # ← 消费前向 input
output_tensor      = output_tensors[chunk_id].pop(0)        # ← 消费前向 output
output_tensor_grad = output_tensor_grads[chunk_id].pop(0)   # ← 消费下游梯度

# ─── 启动前（主循环开始前的单次 recv）────────────────────────────────────────
# p2p_communicator.recv_forward(is_pp_first=True) → 对 pp_first_stage 返回 None
input_tensors[0].append(None)
# 状态 ▶ input[0]=[None]  input[1]=[]  output[0]=[]  output[1]=[]

# ─── Warmup k=0 : fid=0 → (c0, m0) ─────────────────────────────────────────
# ① preprocess (c0 = pp_first + vp_first)：
#     len(input[0])=1, len(output[0])=0  →  1 ≠ 0  →  不 append（初始 None 够用）
#     offset = 0（warmup 恒为 0）
#     mb_id = mb_id_table[0] = 0
#     input_tensor = input[0][0-0] = None  ← READ
# ② forward_step，以 None 驱动，从 data_iterator 读数据
# ③ output[0].append(out_c0m0)
# ④ recv_tensor_from_previous_stage(0, fwd) → recv=False（0 < pp-1=1）
# 状态 ▶ input[0]=[None]  input[1]=[]  output[0]=[out0]  output[1]=[]
#   ◀── len(input[0])=1 == len(output[0])=1，下次 c0 forward 前会触发 append ──▶

# ─── Warmup k=1 : fid=1 → (c0, m1) ─────────────────────────────────────────
# ① preprocess：
#     len(input[0])=1 == len(output[0])=1  →  条件成立  →  APPEND None！
#     input[0] = [None, None]              ← 现在有 2 个 None 槽位
#     offset=0, mb_id=mb_id_table[1]=1
#     input_tensor = input[0][1-0] = input[0][1] = None  ← READ
# ② forward_step(c0,m1)，以 None 驱动，从 data_iterator 读第 2 条数据
# ③ output[0].append(out_c0m1)
# ④ recv_tensor_from_previous_stage(1, fwd) → recv=True, next_chunk=c1
#     GPU1 刚完成 vid=0(c0,m0)，将激活值发回 GPU0（VPP 循环绕回）
#     input[1].append(recv_c1m0)
# 状态 ▶ input[0]=[None,None]  input[1]=[recv_c1m0]
#         output[0]=[out0,out1]  output[1]=[]

# ─── Warmup k=2 : fid=2 → (c1, m0) ─────────────────────────────────────────
# ① preprocess：c1 不是 vp_first_stage → 跳过 append 分支
#     offset=0, mb_id=mb_id_table[2]=0
#     input_tensor = input[1][0-0] = recv_c1m0  ← READ
# ② forward_step(c1,m0)
# ③ output[1].append(out_c1m0)
# ④ recv_tensor_from_previous_stage(2, fwd) → recv=True, next_chunk=c1
#     GPU1 完成 vid=1(c0,m1) → input[1].append(recv_c1m1)
# 状态 ▶ input[0]=[None,None]  input[1]=[recv_c1m0,recv_c1m1]
#         output[0]=[out0,out1]  output[1]=[out_c1m0]

# ─── Warmup k=3 : fid=3 → (c1, m1) ─────────────────────────────────────────
# ① offset=0, mb_id=1 → READ input[1][1] = recv_c1m1
# ③ output[1].append(out_c1m1)
# ④ recv_tensor_from_previous_stage(3, fwd) → recv=False
#     GPU1 完成 vid=2 → chunk_table[2]=c1=last_chunk → 不发激活值（last stage 产生 loss）
# 状态 ▶ input[0]=[None,None]  input[1]=[recv_c1m0,recv_c1m1]
#         output[0]=[out0,out1]  output[1]=[out_c1m0,out_c1m1]

# ─── Steady k=0 : fid=4(c0,m2) + bid=0(B_c1,m0) ────────────────────────────
# ==== forward fid=4 ====
# ① preprocess (c0)：
#     len(input[0])=2 == len(output[0])=2  →  APPEND None
#     input[0] = [None, None, None]
#     offset = num_released(fid=4, c0):
#       backward_mid = 4-4 = 0，chunk_table[:0].count(V-1-c0=1) = 0  →  offset=0
#     mb_id=mb_id_table[4]=2
#     input_tensor = input[0][2-0] = None  ← READ
# ③ output[0].append(out_c0m2)
# ④ recv=True, next=c1 → input[1].append(recv_c1m2)
#    [input[1] = [recv_c1m0, recv_c1m1, recv_c1m2]]
#
# ==== backward bid=0 → B(c1,m0) ====
# output_tensor_grads[1].append(None)   ← last stage，loss 提供梯度
# input[1].pop(0) = recv_c1m0           ← 释放！input[1]=[recv_c1m1,recv_c1m2]
# output[1].pop(0) = out_c1m0           ← 释放！output[1]=[out_c1m1]
# output_tensor_grads[1].pop(0) = None
# 状态 ▶ input[0]=[None,None,None]  input[1]=[recv_c1m1,recv_c1m2]
#         output[0]=[out0,out1,out2]  output[1]=[out_c1m1]

# ─── 快进到 Steady k=2 : fid=6(c1,m2) ── offset 生效场景 ────────────────────
# 此时 bid=0,1(B_c1,m0 和 B_c1,m1)已完成，input[1] 已 pop 两次
# ① preprocess (c1)：跳过 append 分支
#     offset = num_released(fid=6, c1):
#       backward_mid = 6-4 = 2
#       flipped_chunk = V-1-c1 = 0
#       chunk_table[:2].count(0) = [0,0].count(0) = 2  →  offset=2
#     mb_id = mb_id_table[6] = 2
#     index = 2 - 2 = 0
#     此时 input[1] = [recv_c1m2, recv_c1m3]（前两项已被 pop 掉）
#     input[1][0] = recv_c1m2  ✓
# ── offset 的本质：pop(0) 每次将整个列表前移，mb_id 减 offset 得到当前实际下标 ──

# p2p_communication.py:146-163
class P2PCommunicator:
    def __init__(self, pp_group, config):
        world_size = self.pp_group.size()          # pp=2 → world_size=2
        curr_rank_in_pg = self.pp_group.rank()     # GPU0→0, GPU1→1

        # 关键：模运算形成环形拓扑
        next_rank_pg = (curr_rank_in_pg + 1) % world_size  # GPU0→1, GPU1→0(绕回!)
        prev_rank_pg = (curr_rank_in_pg - 1) % world_size  # GPU0→1(绕回!), GPU1→0

        self.next_rank = dist.get_global_rank(self.pp_group, next_rank_pg)
        self.prev_rank = dist.get_global_rank(self.pp_group, prev_rank_pg)

# p2p_communication.py:465-483 — send_forward
def send_forward(self, output_tensors, is_last_stage):
    """Send tensor to next rank in pipeline (forward send)."""
    for output_tensor in output_tensors:
        if not is_last_stage:
            self._communicate(
                tensor_send_next=output_tensor,  # ← 发给 self.next_rank
                tensor_send_prev=None,
                recv_prev=False,
                recv_next=False,
            )

# p2p_communication.py:402-431 — recv_forward
def recv_forward(self, tensor_shapes, is_first_stage):
    """Receive tensor from previous rank in pipeline (forward receive)."""
    for tensor_shape in tensor_shapes:
        if is_first_stage:
            input_tensor = None          # ← first stage 不 recv，用 None
        else:
            input_tensor, _, _ = self._communicate(
                tensor_send_next=None,
                tensor_send_prev=None,
                recv_prev=True,          # ← 从 self.prev_rank 收
                recv_next=False,
            )

# p2p_communication.py:376-384 — _communicate 内部实际使用的 rank
p2p_reqs = p2p_func(
    tensor_send_prev=tensor_send_prev,
    tensor_recv_prev=tensor_recv_prev,
    tensor_send_next=tensor_send_next,       # → self.next_rank
    tensor_recv_next=tensor_recv_next,       # ← self.next_rank
    prev_pipeline_rank=self.prev_rank,       # 环形拓扑的 prev
    next_pipeline_rank=self.next_rank,       # 环形拓扑的 next
)