Riepilogo Tecnico: DOPPLER – Apprendimento a Doppia Politica per l'Assegnazione dei Dispositivi in Grafi di Dataflow Asincroni

1. Definizione del Problema

Il documento affronta il problema dell'assegnazione dei dispositivi in sistemi multi-GPU che eseguono carichi di lavoro complessi di machine learning (ad es. Transformer, Mixture-of-Experts) rappresentati come grafi di dataflow statici. L'obiettivo è mappare le operazioni (vertici) ai dispositivi per minimizzare il tempo totale di esecuzione all'interno di un sistema Work-Conserving (WC).

A differenza dei tradizionali framework bulk-synchronous (ad es. lo standard PyTorch o TensorFlow), dove tutti i dispositivi attendono l'operazione più lenta prima di procedere, i sistemi WC pianificano le operazioni in modo asincrono non appena le loro dipendenze di dati sono soddisfatte. Sebbene i sistemi WC offrano un potenziale significativo di accelerazione grazie alla sovrapposizione di comunicazione e computazione, essi introducono sfide uniche:

• Ordine di Esecuzione Stocastico: Senza sincronizzazione globale, l'ordine di esecuzione è dinamico e sensibile all'eterogeneità dell'hardware e alla contesa delle risorse.
• Ottimizzazione Complessa: Il tempo di esecuzione non può essere espresso come una funzione a forma chiusa dell'assegnazione a causa di questa stocasticità.
• Obiettivi Conflittuali: Un'assegnazione efficace deve bilanciare simultaneamente il carico della GPU e minimizzare la comunicazione inter-GPU, un compito reso difficile dal fatto che il bilanciamento del carico è intrinsecamente temporale in contesti asincroni.

Gli approcci precedenti basati sull'apprendimento sono limitati dalla loro dipendenza da assunzioni bulk-synchronous, dall'uso di una singola politica di posizionamento senza modellare la dinamica del sistema e dalla mancanza di ottimizzazione durante l'effettiva implementazione.

2. Metodologia: Il Framework DOPPLER

Gli autori propongono DOPPLER, un framework di reinforcement learning (RL) che impiega uno schema decisionale sequenziale a doppia politica e una pipeline di addestramento in tre fasi.

2.1 Formulazione a Doppia Politica

Invece di apprendere una singola politica monolitica, DOPPLER decompone il processo di assegnazione in due politiche distinte, modellate come un processo decisionale di Markov (MDP) episodico:

1. SEL (Politica di Selezione): Seleziona la successiva operazione (vertice) da assegnare dall'insieme dei nodi candidati. Questa politica attraversa il grafo parzialmente assegnato per approssimare il flusso non deterministico del "tempo", apprendendo effettivamente un ordine di esecuzione.
2. PLC (Politica di Posizionamento): Determina il dispositivo GPU specifico per l'operazione selezionata, con l'obiettivo di bilanciare il carico e minimizzare la comunicazione.

Architettura:

• Entrambe le politiche utilizzano Reti Neurali su Grafi (GNN) con passaggio di messaggi per codificare la struttura del grafo di dataflow, le dipendenze e i percorsi critici.
• SEL aggrega informazioni dai percorsi critici (percorsi di livello inferiore e superiore) e dalle caratteristiche dei nodi per selezionare il nodo successivo.
• PLC combina le rappresentazioni dei nodi, le caratteristiche del dispositivo (ad es. carico attuale, tempo di fine) e l'insieme dei nodi già posizionati per selezionare il dispositivo ottimale.
• Ottimizzazione dell'Efficienza: Per evitare costi di addestramento proibitivi associati alla riesecuzione del passaggio di messaggi ad ogni step, DOPPLER esegue il passaggio di messaggi una sola volta per episodio e aggiorna dinamicamente le caratteristiche dei dispositivi senza ricalcolare le embedding del grafo ad ogni step.

2.2 Pipeline di Addestramento in Tre Fasi

DOPPLER impiega un paradigma di "apprendimento facendo" attraverso tre fasi per garantire convergenza e adattabilità:

1. Fase I (Imitation Learning / Offline): Le politiche sono pre-addestrate utilizzando l'apprendimento supervisionato per imitare un insegnante euristico (specificamente, l'euristica del Percorso Critico). Ciò accelera la convergenza fornendo una politica iniziale forte prima dell'inizio dell'esplorazione.
2. Fase II (RL basato su Simulazione / Offline): Le politiche vengono perfezionate utilizzando il reinforcement learning all'interno di un simulatore software (che implementa la logica di esecuzione WC). La ricompensa è derivata dal tempo di esecuzione simulato. Questa fase consente iterazioni rapide senza la latenza dell'hardware reale.
3. Fase III (RL su Sistema Reale / Online): Le politiche addestrate vengono implementate in un sistema WC reale. Esse vengono continuamente perfezionate tramite RL online utilizzando i tempi di esecuzione reali come ricompense. Questa fase adatta la politica alle specifiche caratteristiche dell'hardware (ad es. jitter di NVLink, contesa) che i simulatori non possono catturare perfettamente, senza alcun overhead aggiuntivo per l'utente.

3. Contributi Chiave

1. Formulazione del Problema: Il documento investiga l'assegnazione dei dispositivi specificamente nel contesto di sistemi multi-GPU asincroni e work-conserving, affrontando la natura stocastica dell'ordine di esecuzione.
2. Approccio a Doppia Politica: Introduce una separazione delle responsabilità in cui una politica apprende l'ordine di attraversamento dei nodi e l'altra gestisce il posizionamento dei dispositivi, catturando sia la dinamica di esecuzione che i vincoli hardware.
3. Framework a Tre Fasi: DOPPLER propone una strategia di addestramento ibrida che combina l'apprendimento per imitazione, il RL basato su simulazione e il RL online su sistema reale. Questo approccio accelera la convergenza e consente l'ottimizzazione continua durante l'implementazione.
4. Implementazione GNN Scalabile: Gli autori propongono un'approssimazione efficiente del passaggio di messaggi che riduce significativamente i tempi di addestramento senza compromettere la convergenza.

4. Risultati Sperimentali

Gli esperimenti sono stati condotti su quattro GPU NVIDIA Tesla P100 utilizzando vari carichi di lavoro: Moltiplicazioni di Matrici in Catena (CHAINMM), Reti Neurali Feedforward (FFNN) e blocchi/layer Transformer Llama.

• Prestazioni rispetto ai Baseline: DOPPLER (specificamente il modello completamente addestrato DOPPLER-SYS) ha superato costantemente i baseline, inclusa l'euristica del Percorso Critico, PLACETO (RL a singola politica), GDP (RL basato su embedding di grafi) e un ottimizzatore enumerativo personalizzato.
  • DOPPLER ha ottenuto tempi di esecuzione fino al 52,7% inferiori rispetto al miglior baseline (GDP).
  • Ha ridotto il tempo di esecuzione fino al 78,2% rispetto all'euristica del Percorso Critico.
  • Ha superato l'ottimizzatore enumerativo personalizzato fino al 13,8%.
• Studi di Ablazione:
  • Doppia Politica: Combinare entrambe le politiche SEL e PLC ha prodotto i risultati migliori, in particolare per i modelli complessi, dimostrando la necessità di apprendere sia l'ordine che il posizionamento.
  • Fasi di Addestramento: L'addestramento sul sistema reale da solo ha portato a una convergenza lenta e instabile. La combinazione di apprendimento per imitazione (Fase I) e simulazione (Fase II) è stata cruciale per una convergenza stabile e veloce prima del raffinamento online (Fase III).
• Transfer Learning: DOPPLER ha dimostrato una forte capacità di generalizzazione. Le politiche addestrate su grafi semplici (FFNN) potevano essere perfezionate con pochi episodi (2k–4k) per funzionare bene su architetture inedite (Llama) e diverse configurazioni hardware (ad es. trasferimento da 4 P100 a 8 V100).
• Scalabilità: L'addestramento e l'inferenza di DOPPLER scalano linearmente con la dimensione del grafo di dataflow, superando in efficienza altri baseline basati su RL.

5. Significato e Rivendicazioni

Il documento sostiene che DOPPLER rappresenti un passo avanti significativo nell'ottimizzazione dei moderni carichi di lavoro ML, andando oltre le assunzioni statiche e bulk-synchronous. Il suo significato risiede in:

• Efficienza Pratica: Consentendo l'esecuzione work-conserving, DOPPLER può risparmiare sostanziali ore di GPU in implementazioni su larga scala (ad es. risparmiando milioni di ore di GPU annualmente per carichi di lavoro su scala ChatGPT).
• Adattabilità: A differenza delle euristiche statiche o dei modelli puramente basati sulla simulazione, la capacità di apprendimento online di DOPPLER gli consente di adattarsi alle realtà rumorose degli ambienti hardware di produzione.
• Innovazione Metodologica: La decomposizione a doppia politica offre un approccio più sfumato al problema dell'esplosione combinatoria dell'assegnazione dei dispositivi, separando il "cosa eseguire dopo" dal "dove eseguirlo", il che si allinea meglio con la natura temporale dell'esecuzione asincrona.

Gli autori sottolineano che il loro approccio non richiede modifiche ai framework di ML sottostanti, ma ottimizza piuttosto lo strato di scheduling, rendendolo una soluzione pratica per le esistenti infrastrutture multi-GPU.