동기 RLHF 학습에서 응답 길이 편차로 긴 응답이 완료되기를 기다리며 GPU가 낭비되는 문제를 해결하는 PAT(Adaptive Tensor Parallelism)를 제안한다. PAT는 오프라인 프로파일링 기반 예측기로 재구성 시점과 목표 TP 설정을 결정하고, KV 캐시 마이그레이션·인플레이스 가중치 리샤딩·통신 그룹 재사용으로 재구성 오버헤드를 최소화한다. SGLang + VeRL 프레임워크에 구현해 LLaMA3.1-8B·Qwen3-14B로 평가한 결과, 생성 지연을 최대 34.6%, 엔드투엔드 RLHF 반복 지연을 최대 27.2% 감소시켰다.
- •RLHF 생성 단계에서 응답 길이 편차로 인한 GPU 낙비 문제를 동적 텐서 병렬성 재구성으로 해결하는 PAT 제안.
- •예측기 기반 온라인 재구성 + KV 캐시 마이그레이션·인플레이스 리샤딩으로 재구성 오버헤드를 최소화.
- •LLaMA3.1-8B·Qwen3-14B 평가에서 생성 지연 최대 34.6%, 엔드투엔드 RLHF 지연 최대 27.2% 감소.
Accelerating Long-Tail Generation in Synchronous RLHF Training via Adaptive Tensor Parallelism
- 1.PAT는 RLHF 생성 단계에서 텐서 병렬성(TP)을 동적으로 재구성해 긴 응답 생성 병목을 제거
- 2.LLaMA3.1-8B·Qwen3-14B 평가에서 생성 지연 최대 34.6%, RLHF 전체 반복 지연 최대 27.2% 단축
- 3.오프라인 프로파일링 기반 예측기로 재구성 비용 대비 이득을 분석해 재구성 시점·목표 TP를 자동 결정
- 4.SGLang 기반 구현, VeRL 프레임워크와 통합; KV-캐시 마이그레이션·재계산 비용 모델로 경량 전환 구현
왜 중요한가?
RLHF 학습의 주요 병목인 긴 응답 생성 시 GPU 유휴 문제를 동적 TP 재구성으로 해결, 실제 LLM 후훈련 속도를 27% 이상 높여 학습 비용 절감에 직접 기여한다.
본문 미리보기
arXiv:2605.23945v1 Announce Type: new Abstract: Reinforcement Learning from Human Feedback (RLHF) has become a key post-training paradigm for improving model quality. However, the synchronous three-stage RLHF pipeline is often bottlenecked by the generation stage, where response-length skew causes the effective batch size to shrink rapidly during decoding, leaving GPUs underutilized while a few long responses remain unfinished. Mainstream frameworks employ a static tensor parallelism (TP) confi
전체 내용이 궁금하다면?
원문을 직접 읽어보세요
