s10批处理与调度24 分钟阅读

Overlap the Pipeline

让执行与状态处理重叠

读懂 ZOS 的前向、后处理、采样顺序,以及滞后一轮的状态修正。

异步首先是一组依赖关系:先发出可执行的工作,再在需要结果时等待。

01串行等待会在设备时间线上留下空洞

若每轮都等待采样回传、更新请求、重新组批,再发出下一次前向,执行链路就会在控制工作期间等待。异步调度尝试重叠这些阶段。关键是明确哪些数据已可用、哪些状态可先修正,以及何时必须检查完成事件,而不是仅给函数加 async。

s10_async_step.py
# Framework pseudocode for ZOS dependencies and deferred state.
previous = None

for plan in scheduler:
    metadata = prepare_inputs(plan)
    if previous is not None:
        metadata = correct_on_device(metadata, previous)
    forward = launch_forward(metadata)  # preserve execution dependencies
    if previous is not None:
        wait_for(previous.sample_end)
        process_outputs(previous)
    sample = launch_sample(forward, metadata)
    current = StepState(plan, sample, record_event())
    if plan.is_prefill or not zos_enabled:
        wait_for(current.sample_end)
        process_outputs(current)
        previous = None
    else:
        previous = current

flush_if_present(previous)

02先发本轮前向,再处理上一轮结果

DeepInfer Worker 的统一顺序是当前前向、上一轮后处理、当前采样。ZOS decode 把当前后处理保留到下一轮,从而让它与下一次前向重叠。prefill 或关闭 ZOS 时则当轮处理。最后没有下一轮可借用,必须显式 flush,否则最后一批已经生成的结果可能还没有交回请求层。

03滞后的长度不能直接用于计算

上一轮尚未后处理时,控制侧序列长度可能落后一拍,尤其草稿接受数还未落账。上游 _correct_seqlens_qk 根据采样结果修正长度与输入,并约束前向路径的同步边界。伪代码用于阅读依赖关系,具体事件、元数据和状态所有权应沿 Worker 的实现核对。

04跨轮重叠要保留完成事件与状态所有权

官方生产设计:官方部署文章把调度、执行与结果处理放到跨轮流水线中重叠。理解这条时间线时,每个阶段都要标出它消费哪一轮数据、由谁持有结果,以及哪些状态尚未提交。下一轮前向已经发出,不代表上一轮结果已经对请求层可见。 开源核对:WorkerStepState 的构造把采样结果与完成事件一起保留;_post_process_step等待该事件后才消费结果并回调引擎。ZOS decode 将这份状态交给下一轮,结束时再由 flush 排空。因此重叠既需要前向与后处理的先后安排,也需要明确的就绪信号和在途状态生命周期。

TRY IT YOURSELF

动手想一想

画出三轮控制与执行时间线,标注 forward、sample、postprocess 和 flush。假设上一轮只接受一个草稿,指出下一轮哪些元数据需要修正。

把这个机制,加入你的工具箱。