이기종 GPU 환경에서 LLM 강화학습을 가속하는 HetRL

초록

HetRL은 이기종 GPU와 네트워크로 구성된 분산 환경에서 대규모 언어 모델(LLM)의 강화학습(RL) 훈련 효율을 극대화하는 시스템입니다. 복잡한 RL 워크플로우의 스케줄링을 다단계 검색 프레임워크로 최적화하여, 기존 최신 시스템 대비 평균 3.17배의 처리량 향상을 달성했습니다.

상세 분석

HetRL의 핵심 기여는 이기종 환경에서의 RL 훈련 스케줄링을 제약 조건이 있는 결합 최적화 문제로 공식화하고, 이를 해결하기 위한 혁신적인 알고리즘을 제안한 점입니다. 기존 RL 훈련 시스템(예: verl, OpenRLHF)은 높은 대역폭의 네트워크로 연결된 동종 GPU 클러스터에 최적화되어 있어, 다양한 성능의 GPU와 네트워크로 구성된 현실적 인프라에서는 비효율적이었습니다. 반면, HetRL은 (1) 다단계 검색 프레임워크를 통해 복잡한 검색 공간을 분해하고, (2) 연속적 반감(Successive Halving)을 통해 검색 예산을 할당함으로써 실용적인 시간 내에 근사 최적 해를 찾습니다.

구체적으로, 다단계 검색은 세 단계로 진행됩니다. 첫째, RL 워크플로우 내 여러 태스크(액터 생성, 크리틱 추론, 보상 추론, 액터/크리틱 학습 등)를 그룹화합니다. 둘째, 이러한 태스크 그룹을 서로 다른 성능의 GPU로 구성된 ‘Coarse-grained GPU 그룹’에 할당합니다. 이는 서로 다른 계산 특성(메모리 집약적 생성 작업 vs 계산 집약적 학습 작업)을 가진 태스크를 적절한 GPU 유형에 배치하는 전략입니다. 셋째, 각 태스크 그룹 내에서 모델의 텐서, 파이프라인, 데이터 병렬화(TP, PP, DP) 전략을 결정하고, 최종적으로 개별 작업 단위(tasklet)를 세부 GPU에 매핑하는 ‘Fine-grained 할당’을 수행합니다.

이러한 계층적 접근은 검색 공간을 효과적으로 축소하며, 각 단계에서 비용 모델을 활용해 후보 배치 계획의 실행 시간을 빠르게 추정합니다. 20,000 GPU-hour에 달하는 대규모 평가는 HetRL의 효용성을 입증합니다. 특히 GPU 유형(A100, L4, L40S, H100)과 네트워크 대역폭이 혼합된 지리적으로 분산된 환경에서 최대 9.17배의 성능 향상을 보였으며, 이는 미드레인지나 이전 세대의 유휴 GPU 자원을 LLM RL 훈련에 효과적으로 통합할 수 있는 길을 제시합니다.