RollMux 강화학습 효율을 극대화하는 차세대 클러스터 스무딩 기술

초록

강화학습(RL) 사후 학습 과정에서 발생하는 클러스터의 유휴 시간(Dependency Bubble)을 제거하기 위해, 분리된 자원을 교차 실행하여 효율성을 극대화하는 RollMux 스케줄링 프레임워크를 제안합니다.

상세 분석

RollMux의 핵심 기술적 가치는 ‘의존성 버블(Dependency Bubble)‘이라 불리는 자원 낭비 문제를 해결하기 위해 도입된 ‘상호 배타적이지 않은 자원 공유’ 메커니즘에 있습니다. 최근 대규모 언어 모델(LLM)의 RLHF와 같은 강화학습 사후 학습은 메모리 집약적인 롤아웃(Rollout) 단계와 연산 집약적인 트레이닝(Training) 단계를 물리적으로 분리된 클러스터에서 수행하는 ‘분리형(Disaggregated) 아키텍처’로 진화하고 있습니다. 그러나 온-폴리시(On-policy) 알고리즘의 특성상 두 단계 사이에는 엄격한 동기화가 필요하며, 이로 인해 한쪽 클러스터가 작업을 수행하는 동안 다른 쪽은 데이터가 준비될 때까지 아무런 작업도 하지 못하는 구조적 한계가 존재합니다.

RollMux는 이를 해결하기 위해 ‘공동 실행 그룹(Co-execution group)‘이라는 추상화 계층을 도입하여 클러스터를 독립적인 로컬리티 도메인으로 분할합니다. 이 구조 내에서 작동하는 2단계 스케줄링 시스템은 매우 정교합니다. 먼저 그룹 간 스케줄러(Inter-group scheduler)는 보수적인 확률적 계획(Conservative stochastic planning)을 사용하여 작업 배치 최적화를 수행하며, 그룹 내 스케줄러(Intra-group scheduler)는 증명 가능한 최적의 라운드 로빈(Round-robin) 방식을 사용하여 처리량을 극대화합니다.

특히 주목할 기술적 혁신은 ‘상주 제약(Residency constraint)‘입니다. 이는 거대한 모델 상태를 호스트 메모리에 캐싱하여, 작업 전환 시 모델을 다시 로드하는 오버헤드를 최소화하는 ‘웜 스타(Warm-star)’ 컨텍스 스위칭을 가능하게 합니다. 이는 단순한 자원 배분을 넘어, 하드웨어의 물리적 한계를 소프트웨어적 지능으로 극복하여 GPU 자원의 유효 활용도를 극적으로 끌어올린 사례라고 평가할 수 있습니다.

본 논문은 강화학습(RL) 사후 학습의 효율성을 혁신적으로 개선하기 위한 새로운 클러스터 스케줄링 프레임워크인 RollMux를 제안합니다. 현대의 대규모 AI 학습 환경은 하드웨어 효율을 극대화하기 위해 메모리 대역폭에 민감한 롤아웃(Rollout) 작업과 연산 성능에 민감한 트레이닝(Training) 작업을 서로 다른 전용 클러스터에 배치하는 ‘분리형 아키텍처’를 채택하고 있습니다. 하지만 이러한 분리형 구조는 온-폴리시 알고리즘의 필수적인 동기화 요구사항 때문에 심각한 병목 현상을 야기합니다. 즉, 롤아웃 단계가 진행되는 동안 트레이닝 클러스터는 유휴 상태가 되고, 반대로 트레이닝 단계가 진행될 때는 롤아웃 클러스터가 대기 상태에 머무는 ‘의존성 버블’이 발생하게 됩니다.

RollMux는 이러한 유휴 시간을 재활용하기 위해 ‘단계 수준의 멀티플렉싱(Phase-Level Multiplexing)’ 전략을 사용합니다. RollMux의 핵심 아이디어는 한 작업의 유휴 상태를 다른 작업의 활성 단계로 채우는 것입니다. 이를 실현하기 위해 저자들은 ‘공동 실행 그룹(Co-execution group)‘이라는 개념을 도입했습니다. 이 추상화는 클러스터를 격리된 도메인으로 나누어, 서로 다른 작업들이 서로의 자원을 침범하지 않으면서도 효율적으로 교차 실행될 수 있는 환경을 제공합니다.

RollMux의 스케줄링 아키텍처는 두 개의 계층으로 구성됩니다. 첫 번째 계층인 ‘그룹 간 스케줄러(Inter-group scheduler)‘는 작업의 배치와 위치를 결정하기 위해 보수적인 확률적 계획법을 사용합니다. 이는 예측 불가능한 작업 완료 시간을 고려하여 자원 충돌을 방지하고 최적의 배치를 찾아냅니다. 두 번째 계층인 ‘그룹 내 스케줄러(Intra-group scheduler)‘는 각 그룹 내부에서 작업들을 라운드 로빈 방식으로 배치하며, 이 스케줄링 방식은 수학적으로 최적임이 증명되었습니다.

또한, RollMux는 대규모 모델 전환 시 발생하는 막대한 오버헤드를 관리하기 위해 ‘상주 제약(Residency constraint)‘을 적용합니다. 이는 모델의 가중치와 상태를 호스트 메모리에 지속적으로 캐싱하여, 작업 전환 시 모델을 디스크나 원격 저장소에서 다시 불러올 필요 없이 즉각적으로 전환할 수 있는 ‘웜 스타(Warm-star) 컨텍스 스위칭’을 가능하게 합니다. 이를 통해 컨텍스 스위칭에 소요되는 지연 시간을 획기적으로 줄였습니다.

연구진은 328개의 H20 GPU와 328개의 H800 GPU로 구성된 실제 운영 규모의 테스트베드에서 RollMux를 평가했습니다. 실험 결과, RollMux는 기존의 표준 분리형 방식 대비 1.84배, 그리고 기존의 최첨단(SOTA) 공동 배치(Co-located) 베이스라인 대비 1.38배의 비용 효율성 향상을 달성했습니다. 더욱 놀라운 점은 이러한 높은 효율성을 달성하면서도 서비스 수준 목표(SLO)를 100% 준수했다는 것입니다. 이는 RollMux가 대규모 AI 인프라 운영에 있어 성능과 안정성을 동시에 확보할 수 있는 실질적인 솔루션임을 입증합니다.