[논문 리뷰] Orca: A Distributed Serving System forTransformer-Based Generative Models

ML/논문

[논문 리뷰] Orca: A Distributed Serving System forTransformer-Based Generative Models

minturtle 2025. 3. 25. 23:25

안녕하세요, 오늘은 이전에 발표한 vLLM이 사용하는 배치 최적화 기법인 Continous Batching(Iteration Level Batching)을 제시한 논문인 Orca: A Distributed Serving System forTransformer-Based Generative Models에 대한 리뷰를 해보도록 하겠습니다.

논문 원본 PDF는 아래와 같습니다.

https://www.usenix.org/system/files/osdi22-yu.pdf

본 논문은 22년 07월에 발표되었습니다.

Introduction

생성 모델을 사용하기 위해 별도의 서비스에 추론 절차를 위임하게 됨.
위임받은 서비스는 낮은 대기 시간과 높은 처리량으로 추론 결과를 제공해야 하며, Triton Inference Server와 Tensorflow Serving과 같은 추론 시스템의 개발을 촉진함
이러한 시스템은 실제 텐서 연산을 추론 시스템 내의 실행 엔진으로 수행할 수 있음
- 예를 들어, Triton과 FasterTransformer 를 결합하여 배포할 수 있음. 이 경우, Triton은 여러 클라이언트 요청을 배치로 그룹화하는 주된 책임이 있으며, FasterTransformer는 Triton으로부터 배치를 받아 배치 방식으로 추론 절차를 수행
다만 자동회귀 모델은 다음 토큰을 하나씩 생성하기 때문에 여러번 반복 수행
- 처리량을 늘리기 위해 배치로 처리하는데, 기존 시스템은 배치 내에 먼저 끝난 요청이 있어도 배치가 전부 끝날때 까지 기달려야함 → 대기 시간 증가
이 논문에서는 요청 대신 반복 단위로 엔진 실행을 예약 → 서비스 시스템은 엔진에서 반환될 때마다 요청의 처리가 완료되었는지 확인하고, 완료되면 즉시 반환
다만 반복 단위 배치를 적용하기 위해선 input token의 길이를 맞춰줘야함 → 선택적 배칭(특정 작업 집합에 대해서만 배칭)을 적용, 선택적 배칭은 작업의 다양한 특성을 고려하여 배치를 나누고 Attention1 작업에 대해 각 요청을 개별적으로 처리하며, Transformer 모델의 다른 작업에는 배칭

Background

Transformer & KV Cache

AutoRegressive Transformer Model은 prefill 단계에서 입력 분석 및 첫번째 출력 토큰 생성

→ 다음 토큰 생성, 끝날 때까지 반복(decoding 단계)

어텐션은 모든 이전 토큰의 KV값이 필요 → 다시계산하는 것보다 KV Cache를 적용해 재사용 가능

ML Inference Serving System

사용자가 추론 서비스에 추론 요청 → 서비스는 미리 정의된 ML 모델을 기반으로 제공된 리소스를 사용하여 요청에 대한 응답을 제공
일반적으로 GPU나 TPU 같은 HW 가속기를 사용하며, 일반적으로 추론 프레임 워크의 아키텍처는 아래와 같음

Endpoint: 외부로부터 request를 받아들이고, 처리된 response을 반환하는 인터페이스
Request Queue: Endpoint로부터 전달받은 요청들이 대기하는 큐. 스케줄러는 이 큐에서 요청을 가져와 batch를 생성
Scheduler: 요청 큐에서 요청을 가져와 엔진에 전달할 배치를 구성하고, 실행 엔진에게 전달
Execution Engine: 스케줄러로부터 전달받은 배치를 실제로 처리하여 토큰을 생성

여기서 추론 서비스는 Execution Engine과 구분되어, 추론 과정은 Engine에게 위임하고, 그 외의 실행의 배치, 여러 모델 중 모델 선택, 모델의 배포하는 등의 측면에 집중한다.
하지만 Triton과 같은 기존의 스케줄러는 배치 내의 토큰이 먼저 끝나더라도 Early-Return이 불가능하다.

위 그림을 예시로 들어보자. x2가 iter2에서 <EOS>를 생성했지만, x1이 아직 끝나지 않았기 때문에 x1이 끝날때 까지 x2도 응답하지 못하고 대기해야 한다.

Challenges and Proposed Solutions

Transformer 기반 생성 모델의 추론 프레임워크를 개발하는데 필요한 과제와, 이를 해결하기 위한 teration-level scheduling과 selective batching을 제시

Challange 1 : Early Return 및 늦게 들어온 요청

기존 시스템은 요청들을 배치(그룹) 단위로 처리
다만 배치 내의 모든 요청이 다 끝날 때 까지 먼저 끝나도 기다림
위에서 언급한 x2의 대기 문제

Solution 1 : Iteration-level Scheduling

Early Return을 구현하기 위해 반복 단위 스케줄링을 수행
스케줄러는 매번 반복 단위로 작업을 처리하며, 아래의 과정을 거침
1. 실행할 요청 선택(n개)
2. 요청에 대해 하나의 반복을 실행(step, 하나의 토큰 생성)
3. 생성된 토큰이 <EOS>이면 해당 요청 early return

Challange 2. 배치 처리의 어려움

배치 처리가 어려운 경우가 존재할 수 있다. 그 경우는 아래와 같다.
1. 두 요청이 서로 다른 입력 크기를 가지는 경우
2. 두 요청이 서로 다른 위치(인덱스)에서 토큰을 처리하는 경우
3. 각 요청이 서로 다른 decoding, prefill 단계에 속하는 경우

Solution 2. Selective Batching

보통 배치 처리에서는 여러 요청을 함께 묶어서 처리하는데, 선택적 배치는 모든 연산에서 배치를 적용하는 대신, 일부 연산에만 배치를 적용
어텐션 외의 Norm 연산, GeLU 등은 입력 텐서의 크기를 맞춰줄 필요가 없으므로 배치 처리가 가능
어텐션 연산은 크기를 맞춰줘야 하므로 배치를 하지않고 요청마다 각각 수행

Orca Design

Orca에 대한 아키텍쳐 설명

Distributed Arch.

모델을 여러 GPU에 분산 실행해야 할 필요가 생김
ORCA는 Transformer 모델 병렬화 기술인 층 내부 병렬화(intra-layer parallelism)와 층 간 병렬화(inter-layer parallelism)를 조합하여 사용
- intra-layer parallelism는 행렬 곱셈(Linear 및 Attention 연산)과 연관된 매개변수를 여러 개의 GPU에 분산
- inter-layer parallelism는 전체 모델을 여러 계층으로 나누어 각 GPU에 분산

즉 두 병렬화 기법을 적용한다는 것은, 모델을 층별로, 층을 GPU별로 분산한다는 것을 의미

Execution Engine

Orca는 추론 서비스 외에 실행 엔진 또한 구현했으며, 그 구조는 아래와 같음.

전체 구조 : 여러개의 Worker로 구성, Worker는 하나이상의 GPU를 제어
엔진 마스터: 스케줄러로부터 스케줄 정보를 받아 Worker들에게 전달하고, Worker들은 GPU를 사용하여 실제 텐서 연산을 수행
Control Plane: 엔진 마스터와 Worker들의 컨트롤러 간의 통신을 담당, gRPC와 같은 별도의 통신 채널을 사용(메시지 내용, 토큰 정보)
Data Plane : GPU 간의 데이터 전송을 담당, NCCL 사용
Worker : 모델의 일부 레이어 담당, GPU 스레드 관리 및 GPU 커널 실행
파이프라인 : 엔진 마스터는 스케줄된 배치에 대한 정보를 첫 번째 Worker에게 전달하고, 각 Worker는 GPU 연산 완료를 기다리지 않고 다음 Worker에게 제어 메시지를 전달하다가, 마지막 Worker는 GPU 연산이 완료될때 까지 기다리고 결과를 반환

기존의 추론 시스템인 FasterTransformer이나 Megatron-LM과는 다르게, Orca는 Data Plane과 Control Plane이 구분되어 있어 GPU-CPU 간의 통신 오버헤드가 발생하지 않아 더욱 효율적임

Scheduling Algorithm

Orca는 매 Iteration마다 어떤 요청을 선택할 것인지 결정
이때 Iteration 수준의 FIFO를 보장
배치 크기, GPU 메모리 제약(KV 캐시)을 고려해야 함.
- ORCA 시스템은 최적의 처리량을 유지하면서 대기 시간을 적절히 관리하기 위해, 최대 배치 크기라는 개념을 도입
- 스케줄러가 처리가 완료되기 전에 다른 요청의 이터레이션을 시작하려고 하면 DeadLock 발생가능, 이를 방지하려면 Orca 스케줄러가 미리 할당된 메모리 영역을 잘 관리해야하며, 남아있는 메모리 크기를 인식하고 새로운 요청을 처리할 수 있는지 여부 판단이 중요
ORCA 스케줄러는 배치 크기와 GPU 메모리 제약을 모두 고려하여 아래의 방식으로 요청 처리
- 스케줄러는 요청의 도착 시간을 기준으로 요청을 선택. 가장 먼저 도착한 요청을 먼저 처리하도록 하며, 동시에 배치 크기를 최대 배치 크기로 제한
- 요청이 처음 스케줄링될 때, 키와 값을 저장할 충분한 공간을 예약. 이 예약은 Attention K/V 관리자가 키와 값을 저장하는 데 사용

1. Select 함수

목표: max batch size(max_bs)가 넘지 않는 선에서 FIFO 순으로, GPU 메모리 제약 조건을 준수하면서 Attention K/V manager에 필요한 메모리 슬롯을 예약
입력:
- pool : 현재 요청들의 집합
- n_rsrv : 현재 예약된 Attention KV 슬롯의 수
- max_bs : max batch size
- n_slots : 최대 등록 가능한 Attention KV 슬롯의 수
알고리즘:
1. 현재 실행중이지 않은(!RUNNING) req를 도착순으로 오름차순 정렬
2. 정렬한 요청을 하나씩 탐색하면서, 아직 초기화 단계의 req라면 해당 요청의 최대 토큰 수만 큼 KV 슬롯을 예약 시도, 이 때 n_slots (최대 가용 슬롯수)를 초과하면, 메모리 부족으로 인한 순회 종료
3. 순회 종료는 batch_size가 max_bs를 채웠을때도 종료

2. 스케줄링 알고리즘

Select 함수에서 req 선택
추론 엔진에 추론 명령 전달
for req in batch do
- req의 상태를 RUNNING으로 변경
현재 스케줄된 작업의 수를 나타내는 n_scheduled 변수를 1 증가
n_scheduled가 n_worker과 같다는 것은 모든 워커가 작업을 수행 중임을 의미, 더 작으면 req를 더 실행할 수 있음(6번을 실행하지 않고 다시 1번으로)
n_scheduled == n_worker인 경우, 다음 작업을 위해선 스케줄된 배치 중 하나가 완료될때까지 대기
- foreach req in the returned batch do
  - 요청의 상태를 INCREMENT, 즉 다음 iter를 위한 준비가 됨으로 변경
  - 요청이 완료되었다면, 해당 요청이 사용한 KV 캐시 슬롯을 해제하고, n_scheduled를 1감소

Evaluation

FasterTransformer과 비교
Megatron-LM이나 DeepSpeed도 있지만 추론 최적보다는 학습 최적
시나리오 2가지
- ORCA 엔진의 성능만을 테스트 하는 방법
  - ORCA 스케줄러는 사용하지 않고, 테스트 스크립트가 요청 배치를 받아서 ORCA 엔진에 반복적으로 주입합니다. 이를 통해 ORCA가 각 요청을 어떻게 처리하는지 살펴봄
  - 가정 : 실험에서는 모든 요청이 동일한 입력 토큰과 출력 토큰을 가진다고 가정합니다. 즉, 각 요청의 복잡도가 같다고 보고 성능을 평가
  - 비교 대상 : 이렇게 얻은 배치 처리 시간을 다른 시스템인 FasterTransformer와 비교하여 ORCA의 성능을 분석
- ORCA의 엔드 투 엔드 성능을 테스트
  - 실제 클라이언트 요청 트레이스가 없으므로, 합성된 요청 트레이스를 사용
  - 합성된 요청 트레이스는 무작위로 생성된 가상 요청
  - 모든 요청에 대해 <eos>를 생성하지 않고 max_gen_tokens 수만큼 생성한다고 가정
  - 각 요청이 언제 도착할지는 포아송 분포를 사용하여 결정
  - 성능 측정은 지연시간과 처리량으로 측정하며, FasterTransformer와 비교함

시나리오 1. Orca 엔진 테스트

배치의 모든 요청이 동일한 수의 입력 토큰(32 또는 128)을 가지며, 32개의 토큰을 생성하도록 설정
실험 세트에서 배치 내 모든 요청은 동시에 처리 시작과 종료가 이루어짐
13B, 101B, 175B의 세 가지 서로 다른 모델을 사용하여 실험을 진행

13B 1GPU에서는 fastertransformer가 batch크기가 8이상에서는 실행이 불가능
101B 8GPU에서는 16 batch size이상에는 처리가 불가
175B 16GPU 역시 16 batch size이상은 처리가 불가했으며, 전체적으로 Orca 엔진이 최대 47 %더 빠름

전반적으로 Orca가 FasterTransformer에 비해 더 높은 처리량과 낮은 지연시간을 제공하며, 모델이 커질수록 성능 향상이 두드러짐
ORCA의 max batch size (max_bs)가 증가함에 따라 지연 시간에 큰 영향을 주지 않으면서 처리량이 증가하는데, 이는 Iter-level Scheduling이 early-finished 및 late-joining 요청 문제를 해결하기 때문

엔드 투 엔드

ORCA는 배치 크기 증가로 처리량을 효율적으로 개선할 수 있지만, FasterTransformer는 특정 배치 크기에서만 성능 향상이 나타남

부하가 증가함에 따라 ORCA는 지연 시간의 증가가 상대적으로 작으면서도 더 높은 처리량을 제공
- ORCA의 스케줄러가 늦게 도착한 요청을 현재 진행 중인 배치에 효과적으로 통합할 수 있기 때문
다만 FasterTransformer는 아래의 시나리오를 효율적으로 처리하지 못함
- 서로 다른 시간에 도착하는 요청
- 완료하는 데 필요한 반복 횟수가 다른 요청
- 서로 다른 수의 입력 토큰으로 시작하는 요청
특히, inter-layer partitions을 두 개 이상 사용하는 175B 또는 341B 모델에서는 비슷한 수준의 지연 시간에서 ORCA가 훨씬 더 높은 처리량을 달성(약 37배)
ORCA의 반복 수준 스케줄링의 early-return 문제와 late-joining문제를 해결하는 덕분에 배치 크기를 올릴수록 지연시간에 큰 영향을 주지 않으면서 처리량을 높일 수 있었음.