런타임 Draft Model 학습: Speculative Decoding을 실제 워크로드에 적용하기

개요

대규모 언어 모델(LLM)의 규모와 복잡성이 지속적으로 증가함에 따라, AI 데이터 센터와 서비스 제공업체에게 추론 비용 절감은 핵심 과제가 되었습니다. 이러한 모델을 대규모로 서빙하는 데 따르는 연산 비용은 모델 품질을 저하시키지 않으면서도 상당한 비용 절감을 달성할 수 있는 효율적인 최적화 기법의 필요성을 높이고 있습니다.

추론 최적화를 위해 disaggregation, KV cache 인지 라우팅, test-time 라우팅, quantization, speculative decoding 등 다양한 접근 방식이 등장했습니다. 이러한 기법 중 speculative decoding은 주요 클라우드 서비스 제공업체(CSP)에서 상당한 주목을 받고 있습니다. 이는 모델 품질 보존을 보장하면서 다른 최적화 방법과 결합할 수 있고, 대부분의 실용적인 시나리오에서 성능 향상을 제공하기 때문입니다.

Speculative decoding은 작고 빠른 모델(draft model)을 사용하여 draft token을 생성한 후, 원래 모델(target model)이 이를 병렬로 검증하는 방식으로 추론을 가속화합니다. 이 접근 방식은 target model만으로 순차적으로 토큰을 생성하는 것보다 근본적으로 더 효율적입니다.

위 그림은 4개의 draft token을 사용한 speculative decoding과 표준 decoding으로 10개의 토큰을 생성하는 타임라인을 보여줍니다. 평균적으로 2개의 draft token만 수락되더라도, speculative decoding은 약 2배의 속도 향상을 달성합니다.

이러한 효율성 향상이 가능한 이유는 target model이 여러 토큰을 병렬로 검증하는 데 필요한 시간이 단일 토큰을 생성하는 데 필요한 시간과 거의 동일하기 때문입니다. 이는 LLM 추론이 연산 바운드가 아닌 메모리 바운드이기 때문에 발생합니다. 병목은 실제 연산이 아닌 메모리에서 모델 가중치를 로딩하는 데 있습니다. 하나의 토큰을 검증하든 여러 토큰을 동시에 검증하든, target model은 동일한 가중치를 로딩해야 하므로 비슷한 지연 시간이 발생합니다. 이 메모리 접근 비용을 여러 토큰 검증에 걸쳐 분산시킴으로써, speculative decoding은 시퀀스 생성에 필요한 비용이 큰 target model forward pass의 총 횟수를 크게 줄입니다.

그러나 기존의 speculative decoding 접근 방식에는 핵심적인 한계가 있습니다: draft model은 일반적으로 실제 프로덕션 트래픽의 분포와 일치하지 않을 수 있는 범용 워크로드에 대해 사전 학습됩니다. 워크로드 분포는 서비스마다 크게 다르고 동일한 서비스 내에서도 시간이 지남에 따라 변화하기 때문에, 범용 데이터로 학습된 draft model은 특수하거나 변화하는 사용 사례에 대해 최적이 아닌 성능을 보이는 경우가 많습니다.

자가 개선 LLM 추론을 위한 Temporal Incremental Draft Engine (TIDE)

TIDE는 런타임 draft model 학습을 통해 이 문제를 해결합니다. TIDE는 추론 엔진으로 SGLang을 사용하고, draft model 학습을 위해 SpecForge를 활용하며, EAGLE3 speculative decoding 기법을 기반으로 합니다. 실시간 서비스 워크로드를 기반으로 draft model을 지속적으로 적응시킴으로써, TIDE는 수동 개입 없이 자동으로 추론 성능을 향상시킵니다.

EAGLE3: Hidden State로부터의 Draft Model 학습

EAGLE3는 draft model 학습에 독특한 접근 방식을 취합니다. 독립적인 언어 모델로 기능하는 기존의 draft model과 달리, EAGLE3의 draft model은 target model의 여러 중간 레이어에서 나온 hidden state를 입력으로 받아 target model의 출력 분포를 예측하도록 학습합니다.

TIDE 시스템 설계

TIDE의 아키텍처는 오버헤드와 복잡성을 최소화하면서 런타임 학습을 프로덕션 추론 시스템에 원활하게 통합하도록 설계되었습니다.

1. Hidden State 로깅이 포함된 추론 엔진

추론 엔진(SGLang)은 사용자 요청을 처리하기 위해 표준 prefill 및 decoding 연산을 수행합니다. 중요한 것은, 이러한 연산 중에 target model의 중간 레이어에서 hidden state를 캡처하여 스토리지에 저장한다는 점입니다. 이 hidden state는 실제 프로덕션 추론 중에 계산된 내부 표현을 나타내며, draft model의 학습 데이터로 사용됩니다.

2. 학습 및 모델 업데이트

추론 엔진이 계속 요청을 처리하는 동안, hidden state가 스토리지에 축적됩니다. 현재 워크로드 분포의 의미 있는 샘플을 나타내는 충분한 양의 데이터가 수집되면, 학습 프로세스가 자동으로 시작됩니다. 그러면 학습 엔진이 축적된 hidden state를 로드하고, 최근 워크로드를 기반으로 target model의 출력 분포를 더 잘 예측하도록 EAGLE3 draft model을 학습시킵니다. 학습이 수렴되면, 업데이트된 draft model이 이전 버전을 대체하여 추론 엔진에 다시 배포됩니다. 이로써 하나의 적응 사이클이 완료되며, 새로운 hidden state가 축적됨에 따라 프로세스가 계속되어 변화하는 워크로드 패턴에 대한 지속적인 적응을 보장합니다.

이 아키텍처는 런타임 학습의 핵심 과제들을 우아하게 해결합니다:

• 거의 제로에 가까운 추론 오버헤드: Hidden state 로깅은 추론 연산과 비동기적으로 실행되어, 오버헤드가 거의 완전히 숨겨지며 요청 처리에 사실상 지연 시간을 추가하지 않습니다.
• 비동기 학습: 학습은 추론과 독립적으로 진행되므로, 요청 처리를 차단하거나 느리게 하지 않습니다.
• 자동 적응: 전체 사이클이 수동 개입 없이 자율적으로 실행되며, 워크로드 변화에 지속적으로 적응합니다.
• 리소스 효율성: 이기종 GPU 시스템에서 학습을 다른 하드웨어로 오프로드하고, 고성능 GPU는 추론에 집중할 수 있습니다.

이 아키텍처의 단순성과 자동화는 수동 튜닝과 개입이 비용이 많이 들고 비실용적인 프로덕션 배포에 적합합니다.

평가

런타임 draft model 학습의 효과를 평가하기 위해, TIDE의 성능을 범용 데이터로 사전 학습된 정적 draft model과 비교했습니다. 평가에는 한국어 대화 쿼리를 나타내는 dbdu/ShareGPT-74k-ko 데이터셋을 사용했으며, 워크로드별 적응을 위한 현실적인 테스트베드를 제공합니다. 실험에서는 lmsys/gpt-oss-120b-bf16을 target model로, lmsys/EAGLE3-gpt-oss-120b-bf16을 베이스라인 사전 학습 draft model로 사용했으며, 추론은 AMD MI300X 또는 NVIDIA H100 GPU에서 실행하고 draft model 학습은 AMD MI250 GPU에서 수행하여 TIDE가 이기종 GPU 리소스를 효과적으로 활용할 수 있음을 보여줍니다.

Figure 3에서 보듯이, TIDE는 동시성 수준에 따라 사전 학습된 draft model 대비 1.14×에서 1.35×의 출력 토큰 처리량 속도 향상을 달성합니다.

Figure 4는 런타임 학습의 핵심적인 장점을 보여줍니다: TIDE가 워크로드로부터 지속적으로 학습함에 따라 처리량이 향상됩니다. 사전 학습된 draft model은 사전 학습에 의해 성능이 고정되어 있으므로 실험 내내 비교적 일정한 처리량을 유지합니다. 반면, TIDE는 한국어 대화 데이터셋의 특정 패턴에 적응하면서 처리량이 점진적으로 향상되는 것을 보여줍니다.

TIDE의 성능 향상 메커니즘을 더 잘 이해하기 위해, acceptance length — 각 검증 단계에서 target model이 수락하는 draft token의 수 — 가 다양한 동시성 수준에서 시간에 따라 어떻게 변화하는지 분석했습니다.

Figure 5는 런타임 학습의 중요한 특성을 보여줍니다: 높은 동시성 수준은 acceptance length의 더 빠른 향상으로 이어집니다. 그 이유는 다음과 같습니다:

• 더 빠른 데이터 축적: 높은 동시성에서는 더 많은 요청이 동시에 처리되어, TIDE가 더 빠른 속도로 학습 데이터를 수집할 수 있습니다. 이는 학습 프로세스를 가속화하고 워크로드 분포에 대한 더 빠른 적응을 가능하게 합니다.
• 더 빈번한 모델 업데이트: 단위 시간당 더 많은 학습 샘플이 사용 가능하므로, draft model을 통계적으로 유의미한 배치로 더 자주 업데이트할 수 있어, 더 빠른 수렴과 더 넓은 패턴 커버리지를 달성합니다.

결론

TIDE는 런타임 draft model 학습이 프로덕션 환경에서 speculative decoding 성능을 크게 향상시킬 수 있음을 보여줍니다. 한국어 대화 데이터셋에 대한 실험 결과, 정적 사전 학습 draft model 대비 1.14×에서 1.35×의 출력 토큰 처리량 속도 향상을 보여주며, 시스템이 실시간 워크로드에 적응함에 따라 성능이 지속적으로 향상됩니다.

성능 향상 외에도, TIDE는 이기종 GPU 시스템에서 강력한 이점을 제공합니다. 이전 세대 GPU나 활용도가 낮은 하드웨어와 같은 유휴 또는 학습에 최적화된 리소스를 draft model 학습에 활용하고, 고성능 GPU는 추론에 전념하게 함으로써, TIDE는 리소스 활용도와 추론 효율성을 동시에 향상시킵니다. 이는 시스템 수준에서 더 나은 전체적인 비용 대비 성능을 제공합니다.

런타임 학습을 프로덕션 추론 시스템에 원활하게 통합하기 위해, 오픈소스 프로젝트에 기여하여 TIDE의 기능을 더 넓은 커뮤니티에 제공하고 있습니다:

• SpecForge: PR#1, PR#2
• SGLang: PR#3

AI 워크로드가 계속 다양화되고 진화함에 따라, 변화하는 패턴에 자동으로 적응하는 TIDE와 같은 시스템은 대규모에서 효율적이고 비용 효과적인 추론을 유지하는 데 점점 더 필수적이 될 것입니다.