기술 요약: DOPPLER – 비동기 데이터플로우 그래프에서의 장치 할당을 위한 이중 정책 학습

1. 문제 정의

본 논문은 복잡한 머신러닝 워크로드(예: Transformer, Mixture-of-Experts)를 정적 데이터플로우 그래프로 실행하는 멀티 GPU 시스템에서의 **장치 할당 문제(device assignment problem)**를 다룹니다. 목표는 연산(vertex)을 디바이스에 매핑하여 전체 실행 시간을 최소화하는 것이며, 이는 Work-Conserving (WC) 시스템 내에서 수행됩니다.

모든 디바이스가 다음 단계로 진행하기 전 가장 느린 연산이 끝날 때까지 대기하는 기존의 벌크 동기식(bulk-synchronous) 프레임워크(예: 표준 PyTorch 또는 TensorFlow 실행)와 달리, WC 시스템은 데이터 의존성이 충족되는 즉시 연산을 비동기적으로 스케줄링합니다. WC 시스템은 통신과 연산을 중첩함으로써 상당한 속도 향상 잠재력을 제공하지만, 다음과 같은 고유한 과제를 야기합니다:

• 확률적 실행 순서 (Stochastic Execution Order): 전역 동기화 없이 실행되므로, 실행 순서가 동적이며 하드웨어의 이질성 및 자원 경합에 민감합니다.
• 복잡한 최적화: 이러한 확률적 특성 때문에 실행 시간을 할당의 폐쇄형 함수(closed-form function)로 표현할 수 없습니다.
• 상충하는 목표: 효과적인 할당은 GPU 부하 분산과 디바이스 간 통신 최소화를 동시에 달야 하며, 비동기 환경에서는 부하 분산 자체가 본질적으로 시간적(temporal) 성격을 띠기 때문에 이 작업은 매우 어렵습니다.

기존의 학습 기반 접근 방식들은 벌크 동기식 가정에 의존하고, 시스템 역학을 모델링하지 않는 단일 배치 정책을 사용하며, 실제 배포 중 최적화가 부족하다는 한계가 있습니다.

2. 방법론: DOPPLER 프레임워크

저자들은 **이중 정책 순차 결정 방식(dual-policy sequential decision scheme)**과 3단계 훈련 파이프라인을 사용하는 강화 학습(RL) 프레임워크인 DOPPLER를 제안합니다.

2.1 이중 정책 정식화 (Dual-Policy Formulation)

단일한 모놀리식 정책을 학습하는 대신, DOPPLER는 할당 과정을 에피소드형 마르코프 결정 과정(MDP)으로 모델링된 두 가지 별개의 정책으로 분해합니다:

1. SEL (Selection Policy): 후보 노드 집합에서 다음에 할당할 연산(vertex)을 선택합니다. 이 정책은 부분적으로 할당된 그래프를 탐색하여 비결정론적인 "시간"의 흐름을 근사하며, 효과적으로 실행 순서를 학습합니다.
2. PLC (Placement Policy): 선택된 연산에 대한 구체적인 GPU 디바이스를 결정하며, 부하를 균형 있게 맞추고 통신을 최소화하는 것을 목표로 합니다.

아키텍처:

• 두 정책 모두 메시지 패싱(message passing)을 사용하는 **그래프 신경망(GNN)**을 활용하여 데이터플로우 그래프의 구조, 의존성 및 임계 경로(critical paths)를 인코딩합니다.
• SEL은 임계 경로(하위 레벨 및 상위 레벨 경로)와 노드 특징으로부터 정보를 집계하여 다음 노드를 선택합니다.
• PLC는 노드 표현, 디바이스 특징(예: 현재 부하, 종료 시간), 그리고 이미 배치된 노드 집합을 결합하여 최적의 디바이스를 선택합니다 매 단계마다 메시지 패싱을 다시 실행하는 데 따르는 과도한 비용을 피하기 위해, DOPPLER는 에피소드당 한 번의 메시지 패싱을 수행하고 디바이스 특징을 동적으로 업데이트하며 그래프 임베딩을 매번 재계산하지 않습니다.

2.2 3단계 훈련 파이프라인

DOPPLER는 수렴과 적응성을 보장하기 위해 "실행을 통한 학습(learning-by-doing)" 패러다임을 따르는 세 단계를 채택합니다:

1. 1단계 (모방 학습 / 오프라인): 정책이 휴리스틱 교사(특히 Critical Path 휴리스틱)를 모방하도록 지도 학습을 통해 사전 훈련됩니다. 이는 탐색이 시작되기 전 강력한 초기 정책을 제공함으로써 수렴을 가속화합니다.
2. 2단계 (시뮬레이션 기반 RL / 오프라인): (WC 실행 로직을 구현한) 소프트웨어 시뮬레이터 내에서 강화 학습을 사용하여 정책을 정교화합니다. 보상은 시뮬레이션된 실행 시간으로부터 도출됩니다. 이 단계는 실제 하드웨어의 지연 시간 없이 빠른 반복 학습을 가능하게 합니다.
3. 3단계 (실제 시스템 RL / 온라인): 훈련된 정책을 실제 WC 시스템에 배포합니다. 실제 실행 시간을 보상으로 사용하여 지속적으로 정교화됩니다. 이 단계는 시뮬레이터가 완벽하게 포착할 수 없는 특정 하드웨어 특성(예: NVLink 지터, 경합)에 정책을 적응시키며, 사용자에게 추가적인 오버헤드를 발생시키지 않습니다.

3. 주요 기여

1. 문제 정식화: 본 논문은 비동기적, 워크 컨서빙(work-conserving) 멀티 GPU 시스템의 맥락에서 실행 순서의 확률적 특성을 다루는 장치 할당 문제를 조사합니다.
2. 이중 정책 접근 방식: 하나의 정책은 노드의 근사적 순차 탐색을 학습하고 다른 하나는 디바이스 배치를 처리하도록 책임을 분리하여, 실행 역학(dynamics)과 하드웨어 제약을 모두 포착합니다.
3. 3단계 프레임워크: 모방 학습, 시뮬레이션 기반 RL, 그리고 온라인 실시스템 RL을 결ari한 하이브리드 훈련 전략을 제안합니다. 이 방식은 빠른 수렴을 촉진하고 배포 중 지속적인 최적화를 가능하게 합니다.
4. 확장 가능한 GNN 구현: 저자들은 수렴을 저해하지 않으면서 훈련 시간을 크게 줄이는 효율적인 메시지 패싱 근사법을 제안합니다.

4. 실험 결과

NVIDIA Tesla P100 GPU 4개를 사용하여 다양한 워크로드(Chain Matrix Multiplications (CHAINMM), Feedforward Neural Networks (FFNN), Llama 트랜스포머 블록/레이어)에 대해 실험을 진행했습니다.

• 베이스라인 대비 성능: DOPPLER(특히 완전히 훈련된 DOPPLER-SYS)는 Critical Path 휴리스틱, PLACETO(단일 정책 RL), GDP(그래프 임베딩 RL), 그리고 커스텀 열거형 최적화 도구를 포함한 베이스라인들을 일관되게 능가했습니다.
  • DOPPLER는 최상의 베이스라인인 GDP와 비교했을 때 최대 52.7% 낮은 실행 시간을 달려냈습니다.
  • Critical Path 휴리스틱과 비교하여 실행 시간을 최대 78.2% 단축했습니다.
  • 커스텀 열거형 최적화 도구보다 최대 13.8% 더 우수한 성능을 보였습니다.
• 절제 연구 (Ablation Studies):
  • 이중 정책: SEL과 PLC 정책을 결합했을 때 특히 복잡한 모델에서 가장 좋은 결과를 얻었으며, 이는 순서와 배치를 모두 학습하는 것의 필요성을 입증합니다.
  • 훈련 단계: 실제 시스템에서만 훈련할 경우 수렴이 느리고 불안정했습니다. 모방 학습(1단계)과 시뮬레이션(2단계)의 조합이 온라인 정교화(3단계) 전의 안정적이고 빠른 수렴에 결정적이었습니다.
• 전이 학습 (Transfer Learning): DOPPLER는 강력한 일반화 능력을 보여주었습니다. 단순한 그래프(FFNN)에서 훈련된 정책은 몇 번의 에피소드(2k–4k)만으로도 보지 못한 아키텍처(Llama)나 다른 하드웨어 구성(예: 4개의 P100에서 8개의 V100으로의 전이)에서 잘 작동하도록 미세 조정될 수 있었습니다.
• 확장성: DOPPLER의 훈련 및 추론 시간은 데이터플로우 그래프의 크기에 따라 선형적으로 확장되며, 다른 RL 기반 베이스라인들에 비해 효율성이 뛰어납니다.

5. 의의 및 주장

본 논문은 DOPPLER가 정적, 벌크 동기식 가정을 넘어 현대적인 ML 워크로드를 최적화하는 중요한 진전임을 주장합니다. 그 의의는 다음과 같습니다:

• 실용적 효율성: 워크 컨서빙 실행을 가능하게 함으로써, DOPPLER는 대규모 배포(예: ChatGPT 규모의 워크로드에서 연간 수백만 GPU 시간 절감)에서 상당한 GPU 시간을 절약할 수 있습니다.
• 적응성: 정적 휴리스틱이나 순수 시뮬레이션 기반 모델과 달리, DOPPLLE의 온라인 학습 능력은 실제 운영 하드웨어 환경의 노이즈가 있는 현실에 적응할 수 있게 합니다.
• 방법론적 혁신: 이중 정책 분해는 장치 할당의 조합 폭발 문제를 해결하기 위해 "무엇을 다음에 실행할 것인가"와 "어디에서 실행할 것인가"를 분리하는 더 미묘한 접근 방식을 제공하며, 이는 비동기 실행의 시간적 특성과 더 잘 부합합니다.

저자들은 이 접근 방식이 기존 ML 프레임워크의 변경을 요구하지 않고 스케줄링 레이어를 최적화하는 것이므로, 기존의 멀티 GPU 인프라에 대한 실용적인 솔루션임을 강조합니다.