본문으로 건너뛰기

[vLLM] Pipeline Parallelism: 파이프라인 병렬화

들어가며

단일 GPU에 올라가지 않는 대규모 모델을 서빙하려면 모델을 여러 GPU에 분산해야 한다. Pipeline Parallelism(PP)은 모델의 레이어를 여러 GPU에 나누어 배치하는 방식이다. Tensor Parallelism(TP)이 하나의 레이어를 쪼개는 것과 달리, PP는 레이어 그룹 단위로 분할한다.

논문: GPipe: Efficient Training of Giant Neural Networks using Pipeline Parallelism

공식 문서: https://docs.vllm.ai/en/latest/serving/distributed_serving.html

공식 문서

vLLM 공식 문서: Parallelism & Scaling

핵심 구조/코드 분석

PP 지원 여부와 Executor 선택

PP를 사용하려면 supports_pp = True인 Executor가 필요하다.

class Executor(ABC):
    uses_ray: bool = False
    supports_pp: bool = False

class MultiprocExecutor(Executor):
    supports_pp: bool = True  # PP 지원

UniProcExecutor는 PP를 지원하지 않는다. PP는 최소 2개의 프로세스가 필요하므로 MultiprocExecutor나 RayDistributedExecutor를 사용해야 한다.

World Size 계산

def __init__(self, vllm_config: VllmConfig) -> None:
    tp_size, pp_size, pcp_size = self._get_parallel_sizes()
    assert self.world_size == tp_size * pp_size * pcp_size

world_size = TP * PP * PCP로 계산된다. PCP(Prefill Context Parallel)는 프리필 단계에서의 추가 병렬화이다. 예를 들어 TP=4, PP=2이면 8개의 GPU 워커가 필요하다.

AsyncIntermediateTensors: 스테이지 간 비동기 통신

PP에서 가장 중요한 것은 스테이지 간 텐서 전달이다.

class AsyncIntermediateTensors(IntermediateTensors):
    def __init__(
        self,
        tensors: dict[str, torch.Tensor],
        comm_handles: list[Handle] | None = None,
        comm_postprocess: list[Callable[[], None]] | None = None,
    ) -> None:
        super().__init__(tensors)
        self._comm_handles = comm_handles
        self._comm_waited = False

    def wait_for_comm(self) -> None:
        if self._comm_waited:
            return
        if self._comm_handles:
            for handle in self._comm_handles:
                handle.wait()
        if self._comm_postprocess:
            for fn in self._comm_postprocess:
                fn()
        self._comm_waited = True

    def __getattribute__(self, name: str):
        if name == "tensors" and not object.__getattribute__(self, "_comm_waited"):
            object.__getattribute__(self, "wait_for_comm")()
        return object.__getattribute__(self, name)

이 클래스의 설계가 매우 인상적이다:

  1. 지연 대기(Lazy Wait): .tensors에 접근할 때만 통신 완료를 기다린다. 통신이 진행되는 동안 다른 연산을 수행할 수 있다.
  2. __getattribute__ 후킹: Python의 속성 접근 메커니즘을 활용하여, 사용자 코드 변경 없이 자동으로 통신 동기화를 수행한다.
  3. 한 번만 대기: _comm_waited 플래그로 중복 대기를 방지한다.

Worker에서의 PP Send 관리

class Worker(WorkerBase):
    def __init__(self, ...):
        # pending non-blocking PP send work from the previous iteration
        self._pp_send_work: list[Handle] = []

이전 이터레이션의 비동기 전송이 아직 완료되지 않았을 수 있으므로, 핸들을 보관하고 있다가 다음 이터레이션 시작 시 완료를 확인한다.

execute_model과 PP

# Executor의 execute_model
def execute_model(
    self, scheduler_output: SchedulerOutput, non_block: bool = False
) -> ModelRunnerOutput | None | Future[ModelRunnerOutput | None]:
    output = self.collective_rpc(
        "execute_model", args=(scheduler_output,), non_block=non_block
    )
    return output[0]  # 첫 번째 워커(driver)의 결과만 반환

PP에서 collective_rpc는 모든 워커에 execute_model을 호출하지만, 최종 출력은 마지막 스테이지의 driver 워커에서만 생성된다. output[0]은 driver 워커의 결과이다.

MultiprocExecutor의 워커 생성

for local_rank in range(self.local_world_size):
    global_rank = global_start_rank + local_rank
    is_driver_worker = self._is_driver_worker(global_rank)

    unready_worker_handle = WorkerProc.make_worker_process(
        vllm_config=self.vllm_config,
        local_rank=local_rank,
        rank=global_rank,
        distributed_init_method=distributed_init_method,
        input_shm_handle=scheduler_output_handle,
        shared_worker_lock=shared_worker_lock,
        is_driver_worker=is_driver_worker,
    )

각 PP 스테이지에 해당하는 워커 프로세스가 생성된다. shared_worker_lock으로 동시에 하나의 워커만 GPU 연산을 수행하도록 제어할 수 있다.

MessageQueue 브로드캐스트

self.rpc_broadcast_mq = MessageQueue(
    self.world_size,
    self.local_world_size,
    max_chunk_bytes=max_chunk_bytes,
    connect_ip=mq_connect_ip,
)
scheduler_output_handle = self.rpc_broadcast_mq.export_handle()

스케줄러 출력(SchedulerOutput)을 모든 PP 스테이지에 브로드캐스트하기 위해 공유 메모리 기반 MessageQueue를 사용한다. 각 스테이지가 동일한 배치 정보를 받아야 정확한 실행이 보장된다.

PP vs TP 비교

특성 Pipeline Parallelism Tensor Parallelism
분할 단위 레이어 그룹 개별 레이어
통신 패턴 스테이지 간 순차 전달 All-Reduce
통신 빈도 1회/레이어그룹 1회/레이어
노드 간 사용 적합 (낮은 대역폭 OK) 부적합 (높은 대역폭 필요)
버블 오버헤드 있음 (마이크로배치로 완화) 없음

PP는 노드 간 연결(InfiniBand 등)의 대역폭이 제한적일 때 유리하다. 스테이지 간 통신이 All-Reduce보다 적기 때문이다.

왜 이 설계인가

  1. 비동기 통신 오버랩: AsyncIntermediateTensors로 연산과 통신을 겹친다. PP의 주요 병목인 "파이프라인 버블"을 줄이는 핵심 최적화이다.

  2. 공유 메모리 브로드캐스트: 프로세스 간 SchedulerOutput 전달에 직렬화 오버헤드가 없다. PP 스테이지가 늘어나도 브로드캐스트 비용이 증가하지 않는다.

  3. TP + PP 결합: world_size = TP * PP 공식으로 두 병렬화를 자유롭게 결합할 수 있다. 8 GPU에서 TP=4, PP=2나 TP=2, PP=4를 선택할 수 있다.

  4. Executor 수준 추상화: PP의 복잡성이 MultiprocExecutor 내부에 캡슐화되어, 스케줄러와 엔진 코어는 PP 여부를 알 필요가 없다.

정리

vLLM의 Pipeline Parallelism은 대규모 모델 서빙의 확장성을 담보하는 핵심 기능이다. AsyncIntermediateTensors의 지연 대기 패턴과 공유 메모리 기반 브로드캐스트로 PP의 고유한 오버헤드를 최소화한다. Tensor Parallelism과 결합하여 수백 GPU 규모의 클러스터에서도 효율적인 서빙이 가능하다.

댓글

관련 포스트

vLLM 의 다른글