Spotlight: Synergizing Seed Exploration and Spot GPUs for DiT RL Post-Training

기술 요약: SPOTLIGHT

문제 정의

Diffusion Transformer(DiT)를 위한 강화 학습(RL) 사후 학습(post-training)은 종종 수천 개의 고성능 GPU를 필요로 하며, 비용이 매우 많이 듭니다. 표준적인 GRPO 스타일의 워크플로우는 두 단계로 구성됩니다: 롤아웃 단계(무작위 시드를 사용하여 샘플 생성)와 학습 단계(그래디언트 계산 및 가중치 업데이트).

기존의 비용 절감 전략들은 DiT에 적용될 때 다음과 같은 두 가지 주요 병목 현상에 직면합니다:

1. 시드 탐색 오버헤드(Seed Exploration Overhead): DanceGRPO나 TreeGRPO와 같은 방법들은 수렴을 가속화하기 위해 여러 무작위 시드를 탐색하여 고대비 샘플(높은 보상과 낮은 보상)을 찾습니다. 그러나 이 탐색 단계는 임계 경로(critical path) 상에 위치하여, 매 반복(iteration)마다 계산 시간을 추가함으로써 잠재적인 비용 절감 효과를 상쇄합니다.
2. 스팟 GPU 유휴 상태 및 파편화(Spot GPU Idleness and Fragmentation): 스팟 GPU는 69~77% 더 저렴하지만, 선점(preemptible)될 수 있으며 노드 간 대역폭이 제한적입니다.
   • 유휴 상태: LLM에서는 롤아웃 지연 시간이 다양하여 파이프라이닝이 가능하지만, DiT 롤아웃은 거의 균일한 시간에 완료됩니다. 이는 스팟 GPU가 학습 단계 동안 전체 시간 동안 유휴 상태로 머물게 되는 동기화 장벽을 생성합니다.
   • 파편화: 빈번한 스팟 선점은 대규모 DiT 모델 실행에 필수적인 시퀀스 병렬성(Sequence Parallelism, SP) 그룹을 깨뜨립니다. 이러한 그룹을 재구성하는 작업은 일반적으로 모델을 다시 로드하고 엔진을 초기화하는 데 수 분이 소요되어 상당한 자원을 낭비합니다.

방법론: SPOTLIGHT

저자들은 시드 탐색과 스팟 자원을 결능합하여 DiT RL 사후 학습을 위해 스팟 GPU를 활용하도록 설계된 최초의 시스템인 SPOTLIGHT를 제시합니다. 이 시스템은 두 가지 핵심 통찰력과 세 가지 핵심 기술을 기반으로 구축되었습니다.

핵심 통찰력

1. 탐색은 오래된 가중치를 허용함(Exploration Tolerates Stale Weights): 저자들은 무작위 시드의 상대적 순위(어떤 시드가 높은 보상 또는 낮은 보상을 생성하는지)가 현재 업데이트된 가중치가 아닌 이전 반복의 모델 가중치를 사용하더라도 매우 일관되게 유지된다는 것을 입증했습니다. 이를 통해 예약된 클러스터가 학습 단계를 수행하는 동안 스팟 GPU에서 오래된 가중치를 사용하여 시드 탐색을 실행할 수 있으며, 결과적으로 탐색 작업을 임계 경로에서 제거할 수 있습니다.
2. SP 재구성은 상태를 재사용할 수 있음(SP Reconfiguration Can Reuse State): SP 그룹 재구성 비용이 높은 이유는 CPU 스케줄러 초기화 및 원격 모델 로딩 때문입니다. 스케줄러를 분리하고 동일 노드 내의 피어(peer)로부터 인트라 노드 가중치 복사(NVLink 활용)를 사용함으로써, SP 재구성 시간을 수 분에서 수 초 미만으로 단축할 수 있습니다.

시스템 아키텍처 및 기술

SPOTLIGHT는 오픈 소스 RL 플랫폼인 ROLL 상에 구현되었으며, 세 가지 주요 기술을 채택합니다:

1. 밴딧 기반 탐색 플래너(Bandit-Based Exploration Planner):

   • 시스템은 탐색 계획을 멀티 암드 밴딧(multi-armed bandit) 문제로 정식화합니다.
   • 학습 단계에서 제공되는 시간 예산 내에서 보상 분산을 최대화하기 위해 시퀀스 수와 유효 디노이징 단계(Teacache와 같은 캐싱 전략 사용)를 동적으로 선택합니다.
   • 새로운 구성을 탐색하는 것과 알려진 고분산 설정을 활용하는 것 사이의 균형을 맞추기 위해 UCB(Upper Confidence Bound) 전략을 사용합니다.

2. 탄력적 시퀀스 병렬성(Elastic Sequence Parallelism, SP):

   • 지속적 스케줄러(Persistent Scheduler): CPU 스케줄러를 GPU 워커와 분리하여 SP 변경 시에도 노드에 상주하게 함으로써 반복적인 초기화 비용을 제거합니다.
   • 인트라 노드 가중치 로딩(Intra-Node Weight Loading): 새로운 GPU 워커가 동일 노드 내의 SP 그룹에 합류할 때, 원격 노드에서 가져오는 대신 NVLink를 통해 인접한 피어로부터 가중치를 복사하여 재구성 지연 시간을 획기적으로 줄입니다.

3. 선점 대응형 풀 기반 요청 스케줄러(Preemption-Aware Pull-Based Request Scheduler):

   • 롤아웃 워커는 작업을 정적으로 할당받는 대신 중앙 집중식 큐에서 요청을 가져오며(pull), 이를 통해 이질적인 워커 속도에 적응합니다.
   • 상태 커밋(State Commitment): 중간 롤아웃 상태(잠재 변수 및 메타데이터)는 예약된 노드의 지속적인 텐서 스토어(Tensor Store)에 커밋됩니다. 만약 스팟 GPU가 선점되면, 시스템은 진행 중인 상태를 복구하여 다른 워커에서 요청을 재개할 수 있어 진행 손실을 최소화합니다.

결과

시스템은 **Qwen-Image (20B)**를 대상으로 DeepSeek-OCR 및 Geneval 데이터셋의 512×512 및 1280×1280 해상도에서 평가되었습니다.

• 비용 절감: SPOTLIGHT는 최신 베이스라인(스팟 GPU를 최적화 없이 사용하는 시스템 및 예약 전용 시스템 포함) 대비 총 학습 비용을 1.4배에서 6.4배까지 절감합니다.
• 수렴 속도: 베이스라인보다 목표 검증 점수에 4배 빠르게 도달합니다. 이는 효과적인 시드 탐색 덕분에 필요한 반복 횟수를 줄이고, 탐색과 학습을 중첩(overlap)시킴으로써 달성되었습니다.
• 품질: SPOTLIGHT는 동적인 고대비 시드 선택 덕분에 모든 데이터셋에서 베이스라인보다 우수한 검증 점수를 달성했습니다.
• 효율성:
  • 탐색 오버헤드는 탐색을 학습 윈도우 내에 맞춤으로써 낮은 수준(반복 시간의 2~3%)으로 유지됩니다.
  • 탄력적 SP를 통해 시스템은 스팟 용량 변화를 거의 실시간으로 추적할 수 있는 반면, 일반적인 방식은 재구성 중에 수 분의 지연이 발생합니다.
  • 시스템은 스팟 GPU 수가 8개에서 64개로 증가함에 따라 처리량(throughput)에서 거의 선형적인 확장성을 보여줍니다.

의의 및 주장

본 논문은 SPOTLIGHT가 DiT RL 사후 학습을 위해 스팟 GPU를 성공적으로 활용하는 최초의 시스템이라고 주장합니다. 그 의의는 다음과 같습니다:

1. 임계 경로 의존성 타파: 시드 탐색이 오래된 가중치를 허용할 수 있음을 증명함으로써, 시스템은 빠른 수렴(탐색을 통한)과 반복 지연 시간 사이의 트레이드오프를 제거했습니다.
2. 스팟 유휴 상태 해결: LLM 중심의 기존 시스템들이 간과했던 시너지인, 학습 단계의 유휴 스팟 용량을 탐색에 활용함으로써 DiT 롤아웃의 특수한 유휴 문제를 해결했습니다.
3. 선점에 대한 견고성: 탄력적 SP와 상태 지속성을 통해, DiT 학습을 빈번한 선점이 발생하는 스팟 시장 환경에서도 회복 탄력성을 갖게 하여 파편화된 자원 풀을 실행 가능한 비용 효율적 학습 인프라로 전환했습니다.

저자들은 이러한 기여가 이미지 품질을 희생하지 않으면서도 비용 효율적인 DiT RL 사후 학습 시스템을 가능하게 한다고 결론짓습니다.