[논문] RollMux: Phase-Level Multiplexing for Disaggregated RL Post-Training

발행: (2025년 12월 12일 오후 03:03 GMT+9)
9 min read
원문: arXiv

Source: arXiv - 2512.11306v1

Overview

RollMux는 강화학습(RL) 워크로드가 분산형 아키텍처—롤아웃(데이터 생성) 단계와 학습(모델 업데이트) 단계가 별도 클러스터에 존재하는 환경—로 이동하면서 나타난 병목 현상을 해결합니다. 온‑폴리시 알고리즘은 두 단계 간의 엄격한 동기화가 필요하기 때문에, 한 클러스터가 바쁠 때 다른 클러스터는 대기하게 되어 고가의 GPU 자원이 낭비됩니다. RollMux는 이러한 유휴 시간을 “채우는” 교차‑클러스터 스케줄링 프레임워크를 도입해, 실제 규모 GPU 테스트베드에서 비용 효율성을 최대 1.8배 향상시킵니다.

Key Contributions

  • Co‑execution Group Abstraction: 전체 하드웨어 풀을 격리된 로컬리티 도메인으로 분할하여, 작업들이 서로의 메모리 사용량에 방해받지 않으면서 자원을 공유할 수 있게 합니다.
  • Two‑Tier Scheduler:
    • Inter‑group scheduler는 보수적인 확률적 계획을 사용해 각 RL 작업을 어느 그룹에 배치할지(롤아웃 vs. 학습) 결정합니다.
    • Intra‑group scheduler는 GPU 활용도를 최대화하는 증명 가능한 최적 라운드‑로빈 방식을 구현합니다.
  • Warm‑Star Context Switching: 대규모 모델 상태가 호스트 메모리에 캐시된 채로 유지되도록 거주 제약을 강제하여, 롤아웃과 학습 단계 간 전환을 거의 즉시 수행할 수 있게 합니다.
  • Production‑Scale Evaluation: 656‑GPU(328 H20 + 328 H800) 클러스터에서 기존 분산 방식 대비 1.84×, 최선의 공동 배치 기준 대비 1.38×의 비용 효율성 향상을 보였으며, 100 % 서비스 수준 목표(SLO) 준수를 달성했습니다.

Methodology

  1. Problem Modeling: 저자들은 RL 파이프라인을 두 개의 교대로 진행되는 자원 집약적 단계(롤아웃 = 메모리‑바운드, 학습 = 컴퓨트‑바운드)로 모델링하고, 이 단계들이 동기화되어야 함을 강조합니다.
  2. Group Formation: 하드웨어 풀을 co‑execution groups으로 나눕니다—각 그룹은 하나의 작업 전체 수명 주기에 예약될 수 있는 GPU와 해당 호스트 메모리 집합을 포함합니다. 이는 거대한 모델 상태를 격리하고 비용이 큰 데이터 이동을 방지합니다.
  3. Inter‑Group Scheduling: 확률적 플래너는 후보 그룹에서 각 단계가 생성할 예상 유휴 시간(“버블”)을 평가하고, 전체 버블 비용을 최소화하는 그룹에 작업을 할당합니다. 플래너는 보수적이며, 워크로드 변동이 있더라도 SLO를 보장하는 배치를 선호합니다.
  4. Intra‑Group Scheduling: 그룹 내부에서는 RollMux가 라운드‑로빈 스케줄을 실행해 서로 다른 작업의 롤아웃 및 학습 작업을 교대로 배치함으로써 GPU를 “다중화”합니다. 저자들은 고정된 그룹 크기와 거주 제약 하에서 이 스케줄이 활용도를 최대로 만든다는 것을 증명했습니다.
  5. Implementation & Integration: 이 프레임워크는 기존 RL 오케스트레이션 스택(예: Ray RLlib)에 연결되며 표준 컨테이너 런타임을 활용합니다. 그룹 경계를 강제하고 스케줄링 결정을 수행하는 가벼운 데몬만 필요합니다.

Results & Findings

MetricBaseline (plain disaggregation)State‑of‑the‑art co‑locatedRollMux
비용 효율성 (수익 대비 $)1.0×1.38×1.84×
GPU 활용도 (평균)~45 %~60 %~82 %
SLO 달성률 (마감 시간 준수)96 %98 %100 %
Warm‑star 지연시간 (단계 전환)120 ms95 ms≈30 ms

Key takeaways

  • 한 단계의 유휴 “버블”을 다른 단계의 활성 단계와 겹치게 함으로써, RollMux는 온‑폴리시 RL 파이프라인을 괴롭히는 대부분의 대기 시간을 제거합니다.
  • 거주 제약이 모델을 호스트 메모리에 유지시켜 컨텍스트 전환 오버헤드를 3배 이상 감소시킵니다.
  • 656‑GPU 전체 클러스터에 높은 부하가 걸려도 스케줄러는 결정론적 SLO 보장을 유지하는데, 이는 프로덕션 RL 서비스에 필수적인 요구사항입니다.

Practical Implications

  • 클라우드 비용 절감: 대규모 RL(예: 로보틱스, 추천 시스템, 자율주행 시뮬레이터)을 운영하는 기업은 동일한 GPU 비용으로 거의 두 배에 달하는 처리량을 얻을 수 있습니다.
  • 클러스터 운영 간소화: 그룹 추상화 덕분에 운영팀은 RL 작업당 고정 “슬롯”을 할당하면 되며, 임시 메모리 고정 트릭을 사용할 필요가 없어 메모리 초과 오류 위험이 감소합니다.
  • 실험 회전 시간 단축: Warm‑star 컨텍스트 스위칭 덕분에 개발자는 정책 업데이트를 긴 데이터 생성 단계 없이 바로 반복할 수 있어 연구‑프로덕션 사이클이 가속됩니다.
  • 호환성: RollMux는 인기 RL 프레임워크 위에 플러그인 형태로 동작하므로 기존 코드베이스는 최소한의 변경(주로 그룹 크기와 거주 정책 설정)만 필요합니다.

Limitations & Future Work

  • 온‑폴리시 중심: 현재 설계는 엄격한 롤아웃‑학습 동기화를 전제로 하며, 오프‑폴리시 혹은 비동기 RL 알고리즘은 큰 이점을 얻지 못할 수 있습니다.
  • 정적 그룹 크기: 그룹은 작업 시작 시 정의되며, 작업 부하 급증 시 동적 확장(예: 스케일‑아웃)은 아직 지원되지 않습니다.
  • 하드웨어 다양성: 평가가 동질적인 NVIDIA H20/H800 GPU에서 수행되었으며, 이종 가속기(TPU, AMD GPU)에서는 새로운 스케줄링 과제가 발생할 수 있습니다.
  • 향후 방향: 이종 자원을 다루는 확률적 플래너 확장, 동적 그룹 재분할 지원, 그리고 비디오 트랜스코딩·대규모 데이터 전처리와 같은 파이프라인형 워크로드에 대한 적용 가능성 탐색 등을 계획하고 있습니다.

Authors

  • Tianyuan Wu
  • Lunxi Cao
  • Yining Wei
  • Wei Gao
  • Yuheng Zhao
  • Dakai An
  • Shaopan Xiong
  • Zhiqiang Lv
  • Ju Huang
  • Siran Yang
  • Yinghao Yu
  • Jiamang Wang
  • Lin Qu
  • Wei Wang

Paper Information

  • arXiv ID: 2512.11306v1
  • Categories: cs.DC
  • Published: December 12, 2025
  • PDF: Download PDF
Back to Blog

관련 글

더 보기 »