유휴 GPU로 LLM 강화학습 롤아웃을 가속하는 MIT TLT

롤아웃이 강화학습 클러스터를 멈춰 세우는 방식

LLM의 추론 능력을 강화하는 RL 훈련에서는 롤아웃이 전체 실행 시간의 최대 85%를 차지한다. 모델이 여러 답변을 순차적으로 생성하는 동안, 복잡한 쿼리를 처리하는 일부 프로세서와 달리 나머지 GPU는 대기하게 된다.

문제는 입력 난이도가 고르게 분포하지 않는다는 데 있다. 소수의 긴 작업이 배치 전체의 종료 시점을 결정하는 롱테일 특성 때문에, 단순 병렬화나 배치 크기 조정만으로는 병목을 없애기 어렵다. 이 비효율은 LLM 학습 비용을 수백만 달러 단위로 높이고 에너지 소비와 탄소 배출에도 영향을 준다.

MIT HAN Lab의 Song Han 교수 연구팀은 이 대기 시간을 연산 자원으로 전환하는 TLT(Taming the Long-Tail)를 제안했다. 이 연구는 정확도를 유지하면서 학습 속도를 70~210% 높였으며, ACM ASPLOS 학회에서 발표될 예정이다.

유휴 자원에서 드래프터를 따라잡게 만든다

TLT의 중심에는 적응형 드래프터 트레이너와 적응형 롤아웃 엔진이 있다.

적응형 드래프터 트레이너는 롤아웃 대기 중인 유휴 프로세서에서 소형 드래프터 모델을 실시간으로 학습한다. 타겟 모델이 RL 훈련 중 수천 번 업데이트되는 상황에서도 드래프터가 함께 적응하도록 구성한 것이다. 기존에 쓰이지 않던 시간만 활용하므로 추가 GPU 자원이 필요하지 않으며, 이렇게 학습된 드래프터는 추론 단계의 효율적 배포용 모델로도 바로 활용할 수 있다.

롤아웃 엔진은 입력 배치의 특성에 맞춰 투기적 디코딩 전략을 조정한다. 드래프터가 먼저 만든 출력을 타겟 모델이 한 번에 병렬 검증하고, 검증을 통과한 토큰은 즉시 채택한다. 실패한 토큰은 다시 생성한다. 쿼리 복잡도에 따라 드래프트 토큰 수도 바꿔 속도와 정확도 사이의 균형을 맞춘다.

⤢✕

정적 드래프터가 학습 단계에서 맞지 않는 이유

투기적 디코딩은 이미 LLM 추론 가속에 쓰이는 방식이다. 다만 RL 훈련에서는 타겟 모델 자체가 계속 갱신되므로, 사전에 학습한 드래프터를 그대로 쓰기 어렵다. 몇 스텝이 지나면 정적 드래프터와 타겟 모델의 정합성이 급격히 낮아져 실질적인 역할을 잃는다.

구분	기존 투기적 디코딩	TLT 적응형 투기적 디코딩
드래프터 모델	사전 학습된 정적 모델	훈련 중 실시간 적응하는 동적 모델
타겟 모델 변화	고정된 추론 모델 대상	수천 번 업데이트되는 RL 학습 모델 대상
드래프터 정합성	시간 경과에 따라 급격히 저하	유휴 GPU 활용으로 지속적 정합성 유지
적용 단계	추론(inference) 전용	학습(training) 단계까지 확장
연산 비용	드래프터 학습에 별도 자원 필요	유휴 자원만 사용하여 추가 비용 제로

TLT는 유휴 GPU에서 드래프터를 계속 갱신해 이 간극을 줄인다. 학습용 가속과 추론용 모델이라는 두 결과물을 같은 과정에서 얻는 구조다.

검증 통과율을 피드백으로 쓰는 롤아웃 제어

TLT는 제어 이론의 폐루프 구조를 설계에 적용한다. 드래프터의 검증 통과율을 관측하고, 그 결과에 따라 드래프트 길이와 학습 강도를 바꾼다.

타겟 모델이 빠르게 업데이트될 때 드래프터가 발산하지 않도록 적응 속도를 제어하는 안정화 메커니즘도 포함한다. 모든 쿼리에 같은 계산량을 할당하는 대신, 쿼리 복잡도에 비례해 연산을 배정한다. 유휴 자원을 남기지 않는 스케줄링 최적화와도 연결된다.

⤢✕

실험에서 확인한 가속과 재활용 가능성

연구팀은 여러 규모의 추론 LLM을 대상으로 TLT를 검증했다. 모델과 태스크에 따라 학습 속도는 70~210% 향상됐고, 가속 과정에서도 원래 모델 대비 정확도 손실은 없었다.

금융 트렌드 예측과 전력망 리스크 감지를 포함한 응용 분야에서도 효과를 확인했다. 학습 중 생성한 드래프터 모델은 추론 효율화에 재활용할 수 있으며, 같은 성능을 달성하는 데 필요한 총 GPU 시간은 절반 이하로 감소했다.

연구에는 MIT 포스트닥 연구원 Qinghao Hu와 MIT EECS 대학원생 Shang Yang이 공동 제1저자로 참여했다. Song Han은 MIT EECS 부교수이자 선임 저자이며 NVIDIA Distinguished Scientist다.

학습 인프라에 적용될 때의 변화

대규모 LLM 학습 파이프라인에 TLT를 도입하면 GPT-4급 모델의 학습 비용은 수천만 달러에서 절반 이하로 감소할 가능성이 있다. 같은 성능을 위한 에너지 소비와 탄소 배출도 줄어들고, 학습 비용 장벽이 낮아지면서 중소 규모 연구팀과 기업의 참여가 확대될 수 있다.

드래프터를 학습 부산물로 남겨 추론 단계 효율화까지 연결하는 점도 특징이다. 기존 GPU 클러스터의 유휴 시간을 생산적인 작업으로 바꾸면 하드웨어 활용률과 인프라 ROI를 함께 개선할 수 있다. RL 기반 추론 모델 학습이 주류로 자리잡은 현시점에서, 유휴 자원을 활용하는 TLT의 접근은 대규모 AI 학습 인프라 설계에 확산될 전망이다.