Technische Samenvatting: DOPPLER – Dual-Policy Learning voor Device Assignment in Asynchrone Dataflow-grafen

1. Probleemstelling

Het artikel behandelt het device assignment probleem in multi-GPU-systemen die complexe machine learning-workloads (bijv. Transformers, Mixture-of-Experts) uitvoeren die worden gerepresenteerd als statische dataflow-grafen. Het doel is om operaties (vertices) toe te wijzen aan devices om de totale executietijd te minimaliseren binnen een Work-Conserving (WC) systeem.

In tegenstelling tot traditionele bulk-synchrone frameworks (bijv. standaard PyTorch of TensorFlow executie) waarbij alle devices wachten op de traagste operatie voordat ze verder gaan, plannen WC-systemen operaties asynchroon zodra hun data-afhankelijkheden zijn voldaan. Hoewel WC-systemen aanzienlijk potentieel bieden voor snelheidswinsten door communicatie en berekening te overlappen, introduceren ze unieke uitdagingen:

• Stochastische Executievolgorde: Zonder globale synchronisatie is de volgorde van executie dynamisch en gevoelig voor hardware-heterogeniteit en resource-contention.
• Complexe Optimalisatie: De executietijd kan vanwege deze stochastiek niet worden uitgedrukt als een gesloten functie van de toewijzing.
• Tegenstrijdige Doelen: Een effectieve toewijzing moet tegelijkertijd de GPU-belasting balanceren en de communicatie tussen GPU's minimaliseren, een taak die wordt bemoeilijkt omdat load balancing inherent temporeel is in asynchrone settings.

Eerdere learning-gebaseerde benaderingen worden beperkt door hun afhankelijkheid van bulk-synchrone assumpties, het gebruik van een enkele placement policy zonder systeemdynamiek te modelleren, en een gebrek aan optimalisatie tijdens de daadwerkelijke deployment.

2. Methodologie: Het DOPPLER-framework

De auteurs stellen DOPPLER voor, een reinforcement learning (RL) framework dat een dual-policy sequential decision scheme en een driedelige trainingspipeline gebruikt.

2.1 Dual-Policy Formulering

In plaats van één monolithische policy te leren, deelt DOPPLER het toewijzingsproces op in twee afzonderlijke policies, gemodelleerd als een episodische Markov Decision Process (MDP):

1. SEL (Selection Policy): Selecteert de volgende operatie (vertex) om toe te wijzen uit de set van kandidaat-nodes. Deze policy doorloopt de gedeeltelijk toegewezen graaf om de niet-deterministische flow van "tijd" te benaderen, waardoor effectief een executievolgorde wordt geleerd.
2. PLC (Placement Policy): Bepaalt de specifieke GPU-device voor de geselecteerde operatie, met als doel de belasting te balanceren en communicatie te minimaliseren.

Architectuur:

• Beide policies maken gebruik van Graph Neural Networks (GNNs) met message passing om de structuur, afhankelijkheden en kritieke paden van de dataflow-graaf te encoderen.
• SEL aggregeert informatie van kritieke paden (bottom-level en top-level paden) en node-features om de volgende node te selecteren.
• PLC combineert node-representaties, device-features (bijv. huidige belasting, eindtijd) en de set van reeds geplaatste nodes om de optimale device te selecteren.
• Efficiëntie-optimalisatie: Om de prohibitieve kosten te vermijden die gepaard gaan met het telkens opnieuw uitvoeren van message passing bij elke stap, voert DOPPLER de message passing één keer per episode uit en werkt de device-features dynamisch bij zonder de graph embeddings bij elke stap opnieuw te berekenen.

2.2 Driedelige Trainingspipeline

DOPPLER hanteert een "learning-by-doing" paradigma over drie stadia om convergentie en adaptiviteit te waarborgen:

1. Fase I (Imitation Learning / Offline): De policies worden vooraf getraind met behulp van supervised learning om een heuristische docent (specifiek de Critical Path heuristiek) na te bootsen. Dit versnelt de convergentie door een sterke initiële policy te bieden voordat exploratie begint.
2. Fase II (Simulatie-gebaseerde RL / Offline): De policies worden verfijnd met behulp van reinforcement learning binnen een software-simulator (die de WC-executielogica implementeert). De reward wordt afgeleid van de gesimuleerde executietijd. Deze fase maakt snelle iteratie mogelijk zonder de latentie van echte hardware.
3. Fase III (Real-System RL / Online): De getrainde policies worden ingezet in een echt WC-systeem. Ze worden continu verfijnd via online RL met werkelijke executietijden als rewards. Deze fase past de policy aan aan specifieke hardware-kenmerken (bijv. NVLink jitter, contention) die simulators niet perfect kunnen vatten, zonder extra overhead voor de gebruiker.

3. Belangrijkste Bijdragen

1. Probleemformulering: Het artikel onderzoekt device assignment specifiek binnen de context van asynchrone, work-conserving multi-GPU-systemen, waarbij de stochastische aard van de executievolgorde wordt geadresseerd.
2. Dual-Policy Benadering: Het introduceert een scheiding van verantwoordelijkheden waarbij één policy de benaderde traversatievolgorde van nodes leert en de andere de device-plaatsing afhandelt, wat zowel de executiedynamiek als de hardwarebeperkingen vangt.
3. Driedelig Framework: DOPPLER stelt een hybride trainingsstrategie voor die imitation learning, simulatie-gebaseerde RL en online real-system RL combineert. Deze aanpak versnelt convergentie en maakt continue optimalisatie tijdens deployment mogelijk.
4. Schaalbare GNN Implementatie: De auteurs stellen een efficiënte message-passing benadering voor die de trainingstijd aanzienlijk vermindert zonder de convergentie in gevaar te brengen.

4. Experimentele Resultaten

Experimenten werden uitgevoerd op vier NVIDIA Tesla P100 GPU's met diverse workloads: Chain Matrix Multiplications (CHAINMM), Feedforward Neural Networks (FFNN) en Llama transformer blocks/layers.

• Prestaties vs. Baselines: DOPPLER (specifiek de volledig getrainde DOPPLER-SYS) presteerde consistent beter dan de baselines, inclusief de Critical Path heuristiek, PLACETO (single-policy RL), GDP (graph embedding RL) en een custom enumeratieve optimizer.
  • DOPPLER behaalde tot wel 52,7% lagere executietijden vergeleken met de beste baseline (GDP).
  • Het reduceerde de runtime met tot wel 78,2% vergeleken met de Critical Path heuristiek.
  • Het presteerde tot wel 13,8% beter dan de custom enumeratieve optimizer.
• Ablatie-studies:
  • Dual Policy: Het combineren van zowel de SEL als de PLC policy leverde de beste resultaten op, vooral voor complexe modellen, wat de noodzaak aantoont om zowel de volgorde als de plaatsing te leren.
  • Trainingsstadia: Alleen trainen op het echte systeem leidde tot trage convergentie en instabiliteit. De combinatie van imitation learning (Fase I) en simulatie (Fase II) was cruciaal voor een stabiele en snelle convergentie vóór online verfijning (Fase III).
• Transfer Learning: DOPPLER demonstreerde sterke generalisatie. Policies getraind op eenvoudige grafen (FFNN) konden worden verfijnd met few-shot episodes (2k–4k) om goed te presteren op ongeziene architecturen (Llama) en verschillende hardwareconfiguraties (bijv. transfereren van 4 P100's naar 8 V100's).
• Schaalbaarheid: De training- en inferentietijden van DOPPLER schalen lineair met de grootte van de dataflow-graaf, wat andere RL-gebaseerde baselines efficiënter maakt.

5. Betekenis en Claims

Het artikel claimt dat DOPPLER een significante stap voorwaarts is in het optimaliseren van moderne ML-workloads door verder te gaan dan statische, bulk-synchrone assumpties. De betekenis ligt in:

• Praktische Efficiëntie: Door work-conserving executie mogelijk te maken, kan DOPPLER aanzienlijke hoeveelheden GPU-uren besparen in grootschalige deployments (bijv. het besparen van miljoenen GPU-uren jaarlijks voor ChatGPT-schaal workloads).
• Adaptiviteit: In tegenstelling tot statische heuristieken of puur simulatie-gebaseerde modellen, stelt de online leerfunctionaliteit van DOPPLER het systeem in staat om zich aan te passen aan de specifieke, ruisgevoelige realiteit van productie-hardwareomgevingen.
• Methodologische Innovatie: De dual-policy decompositie biedt een meer genuanceerde aanpak van de combinatorische explosie van device assignment problemen, door de "wat draai ik nu eerst" te scheiden van "waar draai ik het", wat beter aansluit bij het temporele karakter van asynchrone executie.

De auteurs benadrukken dat hun aanpak geen wijzigingen vereist in de onderliggende ML-frameworks, maar eerder de scheduling-laag optimaliseert, wat het een praktische oplossing maakt voor bestaande multi-GPU infrastructuren.