RLHFSpec 적응형 초안 선택으로 RLHF 생성 가속화

초록

RLHFSpec은 RLHF 훈련에서 가장 큰 병목인 생성 단계에 추측 디코딩(speculative decoding)을 도입하고, 동적 워크로드에 맞춰 초안 전략을 자동 선택하며, 샘플 재배치를 통해 GPU 활용도를 높인다. 이를 통해 생성 단계와 전체 파이프라인의 처리량을 크게 향상시킨다.

상세 분석

본 논문은 대규모 언어 모델(LLM)의 Reinforcement Learning from Human Feedback(RLHF) 훈련 과정에서 생성 단계가 전체 실행 시간의 68 % 이상을 차지한다는 실증적 분석으로 시작한다. 기존의 추측 디코딩은 온라인 서빙에서 작은 초안 모델(SSM)로 토큰을 미리 생성하고, 메인 LLM이 이를 한 번에 검증함으로써 레이턴시를 감소시켰지만, RLHF의 특수한 워크로드—고정된 샘플 수와 가변적인 응답 길이—에 그대로 적용하면 정적인 초안 토큰 수(n) 설정이 비효율을 초래한다. 초기에는 높은 워크로드로 인해 과도한 검증 비용이 발생하고, 후반부에는 짧은 샘플이 대부분이 되어 자원이 유휴 상태가 된다.

RLHFSpec은 두 가지 핵심 기법으로 이 문제를 해결한다. 첫째, 워크로드‑aware 초안 전략 선택기이다. 현재 인스턴스의 토큰 처리량과 남은 샘플 수를 실시간으로 측정하고, 검증 비용과 기대 수용 토큰 수를 모델링한 비용 함수에 따라 최적에 가까운 n 값을 동적으로 결정한다. 이 과정은 라인 수준의 경량 예측 모델을 사용해 오버헤드를 최소화한다. 둘째, 샘플 재배치 메커니즘이다. 긴 응답을 가진 샘플이 집중된 인스턴스와 짧은 샘플만 남은 인스턴스 간에 주기적으로 샘플을 이동시켜 GPU 활용률을 균등하게 만든다. 재배치는 두 단계로 이루어지는데, 먼저 재배치 정책을 결정하고, 이후 레이어‑레벨의 비동기 통신을 이용해 데이터를 옮김으로써 계산과 통신을 겹친다.

실험에서는 Llama‑2‑7B와 DeepSeek‑7B를 기반으로 RLHFSpec을 구현하고, 기존 최첨단 시스템인 SpecDec‑RLHF와 비교했다. 결과는 생성 단계에서 최대 2.3배, 전체 RLHF 파이프라인에서는 평균 1.8배의 처리량 향상을 보여준다. 특히, 워크로드‑aware 초안 선택이 없을 경우와 비교했을 때 평균 12 % 이상의 추가 속도 개선을 기록했으며, 샘플 재배치를 적용하지 않은 경우 대비 전체 GPU 사용률이 18 % 상승했다.

이 논문은 RLHF 훈련의 구조적 특성을 고려한 시스템 수준 최적화가 알고리즘 수준의 개선만큼 중요함을 입증한다. 추측 디코딩을 정적 파라미터가 아닌 동적 워크로드에 맞춰 조정하고, 샘플 레벨에서 부하를 균등화함으로써 GPU 자원을 최대한 활용한다는 설계 철학은 향후 다른 배치‑비의 작업(예: 대규모 파인튜닝, 멀티모달 생성)에도 확장 가능할 것으로 보인다. 다만, 초안 모델과 메인 모델 간의 메모리 격리, 재배치 시 발생할 수 있는 데이터 일관성 문제, 그리고 초안 전략 선택을 위한 비용 모델의 일반화 가능성 등에 대한 추가 연구가 필요하다.