Post

ORCA: continuous batching은 어디서 왔는가 (2) — 분산 아키텍처와 예약 스케줄링

ORCA: continuous batching은 어디서 왔는가 (2) — 분산 아키텍처와 예약 스케줄링

지난 글에서는 ORCA의 두 핵심 아이디어인 iteration-level scheduling과 selective batching을 정리했습니다. 요청을 iteration 단위로 스케줄링하고, Attention만 개별 처리하고 나머지는 평탄화해 배칭하는 기법이었죠. 그런데 이 두 가지만으로는 실제 서비스가 되지 않습니다. GPT-3 175B 같은 모델은 GPU 하나에 올라가지도 않고, 매 iteration마다 K/V가 쌓이는 메모리를 관리하지 않으면 시스템이 멈춰버립니다.

이번 글에서는 ORCA가 이 아이디어를 실제 대규모 분산 시스템으로 만들기 위해 도입한 설계 — 분산 아키텍처, 메모리 예약 스케줄링, 파이프라이닝 — 을 정리합니다.

모델을 여러 GPU에 어떻게 쪼개는가

수천억 파라미터 모델은 GPU 하나에 들어가지 않습니다. 파라미터와 그에 딸린 연산을 여러 GPU와 여러 머신에 나눠야 합니다. ORCA는 학습 시스템에서 먼저 개발된 두 가지 모델 병렬화를 조합합니다.

  • intra-layer parallelism: 레이어 하나 안의 행렬곱(Linear, Attention)과 그 파라미터를 여러 GPU에 쪼갭니다.
  • inter-layer parallelism: Transformer 레이어들을 여러 GPU에 나눠 배치합니다. ORCA는 각 GPU에 같은 수의 레이어를 할당합니다.

예를 들어 4-layer 모델을 2개의 inter-layer 파티션으로 나누고, 각 파티션을 다시 3개의 intra-layer 파티션으로 쪼개면 총 6개의 GPU를 사용하게 됩니다. 이 자체는 FasterTransformer도 쓰는, 새로울 것 없는 기법입니다.

진짜 최적화는 control/data plane 분리에 있다

ORCA 분산 아키텍처의 핵심은 병렬화 전략이 아니라, 통신 채널을 둘로 나눈 것입니다. 이 부분이 ORCA를 기존 시스템과 갈라놓습니다.

엔진은 여러 worker로 구성되고, 각 worker가 inter-layer 파티션 하나를 맡습니다. worker는 서로 다른 머신에 있을 수 있고, 내부에는 GPU를 제어하는 CPU 스레드가 있습니다. 한 iteration이 스케줄되면 engine master가 배치 정보(토큰, 그리고 요청 id·토큰 index 등이 담긴 control message)를 첫 worker에게 넘기고, 이 정보가 worker들을 따라 흘러갑니다.

여기서 기존 시스템의 문제가 드러납니다. FasterTransformer나 Megatron-LM은 이 control message를 GPU 간 통신 채널인 NCCL로 주고받습니다. 문제는 CPU가 발행해야 하는 control message를 GPU 채널로 보내다 보니, 메시지를 주고받을 때마다 CPU-GPU 동기화가 강제된다는 점입니다. 이 동기화가 매 iteration마다 발생하니, 쌓이면 무시할 수 없는 오버헤드가 됩니다.

ORCA는 채널을 성격에 따라 분리합니다.

  • control message + 토큰 → gRPC 같은 GPU를 거치지 않는 별도 채널. CPU 스레드가 GPU 커널을 발행하는 데 쓰는 데이터입니다.
  • 중간 텐서 데이터 → NCCL. GPU가 생산하고 소비하는 데이터입니다.

이렇게 나누면 CPU가 GPU의 커널 완료를 기다리지 않고, 다음 worker에게 control message를 먼저 넘길 수 있습니다. GPU를 놀리는 시간이 줄어듭니다. 논문 실험에서 inter-layer 파티션을 쓰는 175B 모델 기준으로, 이 control/data plane 분리만으로 FasterTransformer 대비 최대 47% 성능 향상이 나옵니다.

매 iteration마다 어떤 요청을 고를 것인가

이제 스케줄러 이야기입니다. selective batching 덕분에 스케줄러는 아무 요청 집합이나 배치로 묶을 수 있으니, 남은 질문은 “매 iteration마다 request pool에서 무엇을 고르는가”입니다.

ORCA의 답은 의외로 단순합니다. 먼저 도착한 요청을 먼저 처리하는 FCFS입니다. 정확히는 iteration 수준의 FCFS를 보장합니다 — pool 안의 두 요청 xi, xj에 대해 xi가 먼저 도착했다면, xi는 xj와 같거나 더 많은 iteration을 돈 상태여야 합니다. (단, 늦게 온 요청이라도 필요한 iteration 수가 적으면 클라이언트에는 먼저 응답이 나갈 수 있습니다.) 스케줄러는 그냥 도착 순서대로 정렬해 max batch size만큼 뽑습니다.

여기에 두 가지 제약이 붙습니다.

첫째, max batch size. 배치를 키우면 처리량은 늘지만 지연도 늘고, 어느 지점부터는 처리량 증가폭이 줄어듭니다(diminishing returns). 그래서 운영자가 지연 예산에 맞춰 조정하는 knob으로 둡니다.

둘째, GPU 메모리 제약. 이게 진짜 까다로운 부분입니다.

K/V 메모리가 데드락을 부른다

지난 글에서 K/V manager가 요청별 key/value를 요청이 끝날 때까지 들고 있다고 했습니다. 여기서 함정이 있습니다. 중간 결과 버퍼는 연산이 끝나면 바로 재활용되지만, K/V 버퍼는 해당 요청이 완전히 끝나야만 회수됩니다. 그리고 increment phase가 진행될수록 토큰이 쌓이니 K/V도 계속 커집니다.

순진하게 구현하면 어떻게 될까요? 다음 토큰의 K/V를 저장할 공간이 없는데, 공간을 비워줄 요청도 아직 안 끝난 상황 — 즉 pool 안의 어떤 요청도 실행할 수 없는 데드락에 빠질 수 있습니다. 실행을 해야 요청이 끝나고, 요청이 끝나야 메모리가 나는데, 메모리가 없어서 실행을 못 하는 교착 상태입니다.

ORCA의 해법은 사전 예약(reservation)입니다. 요청이 처음 스케줄될 때, 그 요청이 앞으로 생성할 최대 토큰 수(max_tokens)만큼의 K/V 슬롯을 미리 예약합니다. 예약이 가능한지는 현재 예약된 슬롯 수(n_rsrv)와 K/V manager 전체 용량(n_slots)을 비교해 판단하고, 공간이 부족하면 그 요청은 이번 iteration에 아예 뽑지 않습니다.

한번 예약에 성공하면, 그 요청이 끝날 때까지 K/V 버퍼 할당이 보장됩니다. 요청의 토큰 수가 max_tokens를 넘을 수 없으니까요. 실행 중에 메모리가 부족해질 일이 없어지고, 따라서 데드락도 없습니다.

이 설계에는 실용적인 장점도 있습니다. max_batch_size는 지연과 처리량의 트레이드오프를 실험으로 측정해가며 튜닝해야 하지만, n_slots는 그냥 GPU 메모리 한도 안에서 가능한 최댓값으로 잡으면 됩니다. 실험 없이 설정할 수 있는 knob인 셈입니다.

배치를 쪼개지 않고 파이프라인을 채운다

마지막은 파이프라이닝입니다. ORCA 스케줄러는 배치 하나의 결과를 기다리지 않습니다. 현재 실행 중인 배치 수가 worker 수(n_workers)에 도달할 때까지 계속 새 배치를 밀어 넣습니다. 그러면 모든 worker가 놀지 않고 각자 다른 배치를 동시에 처리합니다.

이 지점에서 기존 방식과의 차이가 극명합니다. FasterTransformer는 request-level 스케줄링이라, 현재 배치가 끝나기 전에는 다음 배치를 넣을 수 없습니다. 그래서 파이프라인을 채우려면 배치를 여러 개의 microbatch로 쪼개 파티션들에 흘려보내야 합니다. 그런데 이러면 각 파티션이 원래 배치보다 작은 microbatch를 처리하게 되어 배칭 이득이 줄고, microbatch가 너무 크면 파이프라인에 버블(빈 구간)이 생깁니다. 결국 배칭 효율(큰 microbatch) 대 파이프라인 효율(적은 버블) 사이의 트레이드오프에 걸립니다.

ORCA는 이 트레이드오프 자체가 없습니다. iteration-level 스케줄링 덕분에 배치를 쪼갤 필요 없이, iteration 단위로 자연스럽게 파이프라이닝됩니다.

그래서 얼마나 빨라졌는가

논문은 GPT-3 계열 모델(최대 341B)로 FasterTransformer와 비교합니다. 핵심 결과는 같은 수준의 지연에서 처리량 36.9배 향상입니다.

이 숫자가 나오는 이유는 지금까지 정리한 설계가 겹겹이 맞물린 결과입니다. iteration-level scheduling이 먼저 끝난 요청을 즉시 내보내고 새 요청을 바로 합류시켜 불필요한 연산과 대기를 없애고, selective batching이 그 유연한 스케줄링을 실제 배칭으로 연결하고, control/data plane 분리가 iteration마다의 동기화 오버헤드를 걷어내고, 예약 스케줄링과 파이프라이닝이 메모리 안정성과 worker 활용률을 챙깁니다.

흥미로운 점은, 요청들이 모두 같은 길이인 동질적(homogeneous) 워크로드에서는 FasterTransformer도 배치 크기를 키우면 성능이 오른다는 것입니다. 먼저 끝나는 요청 문제가 애초에 안 생기니까요. ORCA의 진짜 강점은 요청마다 도착 시점도, 입력 길이도, 생성 토큰 수도 제각각인 현실적인 워크로드에서 드러납니다. 바로 실제 LLM 서빙이 마주하는 상황이죠.

정리하며 — vLLM으로 가는 길

두 편에 걸쳐 ORCA를 정리했습니다. 다시 한 줄로 요약하면 이렇습니다.

  • iteration-level scheduling: request가 아니라 iteration 단위로 스케줄링해 끝난 요청은 즉시 반환, 새 요청은 즉시 합류.
  • selective batching: Attention만 개별 처리하고 나머지는 평탄화해 배칭, 임의 요청 집합의 배칭을 가능케 함.
  • 분산 설계: control/data plane 분리로 동기화 오버헤드 제거, K/V 예약 스케줄링으로 데드락 방지, 배치를 쪼개지 않는 파이프라이닝.

오늘날 vLLM의 continuous batching은 ORCA의 iteration-level scheduling을 정확히 계승한 것입니다. 다만 ORCA에는 한 가지 아쉬운 지점이 남아 있었습니다. 바로 max_tokens만큼 K/V를 미리 통째로 예약하는 방식입니다. 요청이 실제로는 짧게 끝나도 최대치를 예약해두니, 메모리를 낭비하게 됩니다. 이 문제를 페이지 단위 동적 할당으로 풀어낸 것이 바로 vLLM의 PagedAttention입니다. ORCA가 “무엇을 언제 실행할지”의 스케줄링 혁신이었다면, vLLM은 그 위에 “메모리를 어떻게 잘게 쪼개 쓸지”를 얹은 셈입니다.

두 논문을 나란히 놓고 보면, 지금 우리가 쓰는 LLM 서빙 스택이 어떤 문제들을 하나씩 풀어오며 만들어졌는지가 선명하게 보입니다.


📚 시리즈: LLM 서빙 인프라

  1. OS의 페이징을 KV Cache에 가져온 PagedAttention
  2. ORCA: continuous batching은 어디서 왔는가 (1) — iteration-level scheduling과 selective batching
  3. ORCA: continuous batching은 어디서 왔는가 (2) — 분산 아키텍처와 예약 스케줄링 ← 지금 글
  4. GPU, TPU, NPU — AI 칩은 무엇이 다르고 언제 무엇을 쓰는가
This post is licensed under CC BY 4.0 by the author.