pipeline parallelism

在这里插入图片描述
 However, because each device depends on the output of the previous one, some devices may be idle at times, which means resource underutilization. To reduce these idle periods, the input batch can be split into smaller microbatches. Each microbatch flows through the pipeline one by one, and gradients are accumulated at the end. This microbatching improves GPU utilization, though it does not completely eliminate idle time.
在这里插入图片描述
pic source

vllm Executor 和 Worker

在这里插入图片描述
 图例中集群中有两个节点,node0(GPU0,GPU1)和 node1(GPU2,GPU3)。假设的并行配置:tp = 2,pp = 2。
 EngineCore does not run forward passes itself. It hands a SchedulerOutput to an Executor, which fans out to one or more Worker processes. Each worker holds part of the model and is driven by a ModelRunner. This separation is what lets vLLM target tensor/pipeline/data parallelism. from open-source-wikis
 A worker class that executes (a partition of) the model on a GPU. Each worker is associated with a single GPU. The worker is responsible for maintaining the KV cache and executing the model on the GPU. In case of distributed inference, each worker is assigned a partition of the model. from vllm doc
 我之前的博客分析,vllm分析(三)——EngineCore 到 Worker 初始化流程分析

vllm中的 pp 处理流程

 流水线并行(Pipeline Parallelism, PP) 通过将模型按层切分到多个 GPU,有效降低了单卡显存占用,是大规模推理的必备技术。然而,PP 带来的核心挑战——流水线气泡(Pipeline Bubbles)——会显著降低 GPU 利用率。
vLLM 针对pipeline parallelism的优化策略:

  • 批次队列(Batch Queue):允许 CPU 调度与 GPU 执行重叠,用多批次填充流水线。
  • 异步通信(Async Send/Recv):利用非阻塞点对点通信,将数据传输与计算重叠。
  • 动态分块(Chunked Pipeline Parallel, CPP):根据各 stage 计算负载动态切分输入序列,平衡各 stage 的计算时间,进一步减少气泡。

step_with_batch_queue

 批次队列的结构: self.batch_queue: deque = deque(maxlen=self.batch_queue_size)

    @property
    def max_concurrent_batches(self) -> int:
        # PP requires PP-size concurrent batches to fill the pipeline.
        # Async scheduling requires 2 concurrent batches to overlap.
        pp_size = self.parallel_config.pipeline_parallel_size
        if self.scheduler_config.async_scheduling:
            if self.use_v2_model_runner:
                return pp_size + 1
            # V1 Model Runner does not fully support async scheduling with PP.
            if pp_size <= 1:
                return 2
        return pp_size

 当 pipeline_parallel_size > 1 且 max_concurrent_batches > 1 时,Engine Core 会启用 step_with_batch_queue 作为主步进函数(替代简单的同步 step)。该函数的核心思想是:优先填满批次队列,再阻塞等待结果,以此提高流水线并发度。
step_with_batch_queue

    def step_with_batch_queue(
        self,
    ) -> tuple[dict[int, EngineCoreOutputs] | None, bool]:
        """Schedule and execute batches with the batch queue.
        Note that if nothing to output in this step, None is returned.

        The execution flow is as follows:
        1. Try to schedule a new batch if the batch queue is not full.
        If a new batch is scheduled, directly return an empty engine core
        output. In other words, fulfilling the batch queue has a higher priority
        than getting model outputs.
        2. If there is no new scheduled batch, meaning that the batch queue
        is full or no other requests can be scheduled, we block until the first
        batch in the job queue is finished.
        3. Update the scheduler from the output.
        """

        batch_queue = self.batch_queue
        assert batch_queue is not None

        # Try to schedule a new batch if the batch queue is not full, but
        # the scheduler may return an empty batch if all requests are scheduled.
        # Note that this is not blocking.
        assert len(batch_queue) < self.batch_queue_size

        model_executed = False
        deferred_scheduler_output = None
        if self.scheduler.has_requests():
            scheduler_output = self.scheduler.schedule(self._should_throttle_prefills())
            with self.log_error_detail(scheduler_output):
                exec_future = self.model_executor.execute_model(
                    scheduler_output, non_block=True
                )
            if self.is_ec_consumer:
                model_executed = scheduler_output.total_num_scheduled_tokens > 0

            if self.is_pooling_model or not model_executed:
                # No sampling required (no requests scheduled).
                future = cast(Future[ModelRunnerOutput], exec_future)
            else:
                if not scheduler_output.pending_structured_output_tokens:
                    # We aren't waiting for any tokens, get any grammar output
                    # and sample immediately.
                    grammar_output = self.scheduler.get_grammar_bitmask(
                        scheduler_output
                    )
                    future = self.model_executor.sample_tokens(
                        grammar_output, non_block=True
                    )
                else:
                    # We need to defer sampling until we have processed the model output
                    # from the prior step.
                    deferred_scheduler_output = scheduler_output

            if not deferred_scheduler_output:
                # Add this step's future to the queue.
                batch_queue.appendleft((future, scheduler_output, exec_future))
                if len(batch_queue) < self.batch_queue_size and (
                    model_executed or self.scheduler.has_requests()
                ):
                    # Don't block on next worker response unless the queue is full
                    # or there are no more requests to schedule.
                    return None, model_executed

        elif not batch_queue:
            # Queue is empty. We should not reach here since this method should
            # only be called when the scheduler contains requests or the queue
            # is non-empty.
            return None, False

        # Block until the next result is available.
        future, scheduler_output, exec_model_fut = batch_queue.pop()
        with (
            self.log_error_detail(scheduler_output),
            self.log_iteration_details(scheduler_output),
        ):
            model_output = future.result()
            if model_output is None:
                # None from sample_tokens() implies that the original execute_model()
                # call failed - raise that exception.
                exec_model_fut.result()
                raise RuntimeError("unexpected error")

        # Before processing the model output, process any aborts that happened
        # during the model execution.
        self._process_aborts_queue()
        engine_core_outputs = self.scheduler.update_from_output(
            scheduler_output, model_output
        )

        # NOTE(nick): We can either handle the deferred tasks here or save
        # in a field and do it immediately once step_with_batch_queue is
        # re-called. The latter slightly favors TTFT over TPOT/throughput.
        if deferred_scheduler_output:
            # When draft tokens are used with structured output, validate them
            # before computing the grammar bitmask for the deferred request.
            if self.check_for_draft_tokens:
                draft_token_ids = self.model_executor.take_draft_token_ids()
                if draft_token_ids is not None:
                    # Update the draft token ids in the scheduler output to
                    # filter out the invalid spec tokens, which will be padded
                    # with -1 and skipped by the grammar bitmask computation.
                    self.scheduler.update_draft_token_ids_in_output(
                        draft_token_ids, deferred_scheduler_output
                    )
            # We now have the tokens needed to compute the bitmask for the
            # deferred request. Get the bitmask and call sample tokens.
            grammar_output = self.scheduler.get_grammar_bitmask(
                deferred_scheduler_output
            )
            future = self.model_executor.sample_tokens(grammar_output, non_block=True)
            batch_queue.appendleft((future, deferred_scheduler_output, exec_future))

        return engine_core_outputs, model_executed

 非阻塞地调度一个批次,并将其放入队列头部(appendleft):batch_queue.appendleft((future, scheduler_output, exec_future))。
 当不满足提前返回条件时(即队列已满,或者没有更多请求可以调度但队列非空),代码进入阻塞阶段:model_output = future.result()。

Worker 的 execute_model的处理过程

 每个 PP stage 的 Worker 进程在收到 execute_model 命令后,按照“等待前序发送完成 → 异步接收前序数据 → 本地模型计算 → 异步发送后序数据”的模式执行。关键代码在Worker.execute_model

确保上一轮非阻塞发送已完成
 Worker 维护一个 _pp_send_work 列表,存储上一轮发起的 isend 操作句柄。进入新一步时,首先对这些句柄调用 wait(),以确保前序 stage 的数据已被完全接收,避免发送缓冲区冲突。这是异步发送的必要清理工作。

# 确保上一轮的非阻塞发送已完成
if self._pp_send_work:
    for handle in self._pp_send_work:
        handle.wait()
    self._pp_send_work = []

异步接收前序数据(非首 Rank)
 若当前 stage 不是第一个 PP rank,则调用 get_pp_group().irecv_tensor_dict() 发起非阻塞接收,返回一个包含张量字典、通信句柄和后处理函数的 AsyncIntermediateTensors 对象。该对象重载了 getattribute,在首次访问其 tensors 属性时,会自动调用 wait_for_comm() 等待接收完成。这一设计将通信等待延迟到计算真正需要输入数据的那一刻,从而最大化通信与计算的重叠。

if forward_pass and not get_pp_group().is_first_rank:
    tensor_dict, comm_handles, comm_postprocess = \
        get_pp_group().irecv_tensor_dict(...)
    intermediate_tensors = AsyncIntermediateTensors(
        tensor_dict, comm_handles=comm_handles, ...
    )

本地模型计算
 调用 model_runner.execute_model(scheduler_output, intermediate_tensors) 执行前向传播。对于首 rank,intermediate_tensors 为空,输入来自嵌入层;对于其他 rank,输入则来自前序 stage 的激活值(通过异步接收获得)。

非阻塞发送后序数据(非尾 Rank)
 若当前 stage 不是最后一个 PP rank,则从 output.tensors 中提取中间激活张量,通过 get_pp_group().isend_tensor_dict() 发起非阻塞发送,并将返回的通信句柄保存到 _pp_send_work 中供下一轮等待。随后返回 None(不返回最终结果),因为最终输出只由尾 rank 产生。
 若不是最后一个 PP Stage,将计算得到的中间张量通过 isend_tensor_dict 非阻塞发送:

if not get_pp_group().is_last_rank:
    self._pp_send_work = get_pp_group().isend_tensor_dict(
        output.tensors, ...
    )
    return None  # 不返回最终结果

这种显式的异步通信管理,使得每个 stage 在计算当前批次的同时,能够重叠地传输上一批次的输出和下一批次的输入,极大降低了通信气泡。

PP场景,Scheduler的行为

 大致分析PP场景Scheduler的行为。Scheduler.schedule
 PP使用异步AsyncScheduler,在首rank,假设request a在current_step放入调度队列,进入推理过程。request a必须等待尾rank 返回结果,才可以再次进入推理。
 在异步 AsyncScheduler模式下,设置步调:在AsyncScheduler._update_after_schedule中,为每个被调度的请求设置下一次可调度步数:

request.next_decode_eligible_step = self.current_step + self.pp_size

 检查步调:在基类 Scheduler.schedule() 遍历 running 队列时,检查此约束:

if self.current_step < request.next_decode_eligible_step:
    req_index += 1
    continue  # 未到期,本轮跳过该请求

 PP使用同步Scheduler,之前的批次调度的 request在没有更新结果之前,num_new_tokens = 0, 本轮调度跳过该请求。对应的代码逻辑

num_new_tokens = (
    request.num_tokens_with_spec
    + request.num_output_placeholders
    - request.num_computed_tokens
)

if num_new_tokens == 0:
    # The request cannot be scheduled because one of the following
    # reasons:
    # 1. No new tokens to schedule. This may happen when
    #    (1) PP>1 and we have already scheduled all prompt tokens
    #    but they are not finished yet.
    #    (2) Async scheduling and the request has reached to either
    #    its max_total_tokens or max_model_len.
    # 2. The encoder budget is exhausted.
    # 3. The encoder cache is exhausted.
    # 4. Insufficient budget for a block-aligned chunk in hybrid
    #    models with mamba cache mode \"align\".
    # NOTE(woosuk): Here, by doing `continue` instead of `break`,
    # we do not strictly follow the FCFS scheduling policy and
    # allow the lower-priority requests to be scheduled.
    req_index += 1
    continue

#Reference
[1] vllm v0.6.0代码走读(三)–pipeline parallelism
[2] pp下的step:self.step_with_batch_queue
[3] Pipeline Parallelism in SGLang: Scaling to Million-Token Contexts and Beyond
[4] Chunked Pipeline Parallel (CPP)
[5] [RFC]: Pipeline-Parallelism for vLLM V1 #11945
[6] [Feature]: Pipeline Parallel Features & Performance Optimizations
[7] [Core] Pipeline Parallel support for Model Runner V2
[8] [Distributed] Add async P2P overlap for pipeline parallelism

Logo

免费领 150 小时云算力,进群参与显卡、AI PC 幸运抽奖

更多推荐