Resumo Técnico: DOPPLER – Aprendizado de Política Dupla para Atribuição de Dispositivos em Grafos de Fluxo de Dados Assíncronos

1. Definição do Problema

O artigo aborda o problema de atribuição de dispositivos em sistemas multi-GPU que executam cargas de trabalho complexas de aprendizado de máquina (ex: Transformers, Mixture-of-Experts) representadas como grafos de fluxo de dados estáticos. O objetivo é mapear operações (vértices) para dispositivos a fim de minimizar o tempo total de execução dentro de um sistema de Consumo de Trabalho (Work-Conserving - WC).

Ao contrário das estruturas tradicionais de processamento em lote síncrono (ex: PyTorch ou TensorFlow padrão), onde todos os dispositivos esperam pela operação mais lenta antes de prosseguir, os sistemas WC agendam operações assincronamente assim que suas dependências de dados são atendidas. Embora os sistemas WC ofereçam um potencial significativo de aceleração ao sobrepor comunicação e computação, eles introduzem desafios únicos:

• Ordem de Execução Estocástica: Sem sincronização global, a ordem de execução é dinâmica e sensível à heterogeneidade do hardware e à contenção de recursos.
• Otimização Complexa: O tempo de execução não pode ser expresso como uma função de forma fechada da atribuição devido a essa estocasticidade.
• Objetivos Conflitantes: Uma atribuição eficaz deve equilibrar simultaneamente a carga da GPU e minimizar a comunicação entre GPUs, uma tarefa dificultada pelo fato de que o balanceamento de carga é inerentemente temporal em configurações assíncronas.

Abordagens anteriores baseadas em aprendizado são limitadas por sua dependência de suposições de processamento em lote síncrono, pelo uso de uma política de posicionamento única sem modelar a dinâmica do sistema e pela falta de otimização durante a implantação real.

2. Metodologia: O Framework DOPPLER

Os autores propõem o DOPPLER, um framework de aprendizado por reforço (RL) que utiliza um esquema de decisão sequencial de política dupla e um pipeline de treinamento de três estágios.

2.1 Formulação de Política Dupla

Em vez de aprender uma única política monolítica, o DOPPLER decompõe o processo de atribuição em duas políticas distintas, modeladas como um Processo de Decisão de Markov (MDP) episódico:

1. SEL (Política de Seleção): Seleciona a próxima operação (vértice) para atribuir a partir do conjunto de nós candidatos. Esta política percorre o grafo parcialmente atribuído para aproximar o fluxo não determinístico do "tempo", efetivamente aprendendo uma ordem de execução.
2. PLC (Política de Posicionamento): Determina o dispositivo GPU específico para a operação selecionada, visando equilibrar a carga e minimizar a comunicação.

Arquitetura:

• Ambas as políticas utilizam Redes Neurais de Grafos (GNNs) com passagem de mensagens para codificar a estrutura do grafo de fluxo de dados, as dependências e os caminhos críticos.
• A SEL agrega informações de caminhos críticos (caminhos de nível inferior e superior) e características dos nós para selecionar o próximo nó.
• A PLC combina representações de nós, características do dispositivo (ex: carga atual, tempo de término) e o conjunto de nós já posicionados para selecionar o dispositivo ideal.
• Otimização de Eficiência: Para evitar custos proibitivos de treinamento associados à reexecução da passagem de mensagens em cada etapa, o DOPPLER realiza a passagem de mensagens uma única vez por episódio e atualiza as características do dispositivo dinamicamente sem recomputar as incorporações (embeddings) do grafo a cada passo.

2.2 Pipeline de Treinamento de Três Estágios

O DOPPLER emprega um paradigma de "aprender fazendo" através de três estágios para garantir convergência e adaptabilidade:

1. Estágio I (Aprendizado por Imitação / Offline): As políticas são pré-treinadas usando aprendizado supervisionado para imitar um professor heurístico (especificamente, a heurística de Caminho Crítico). Isso acelera a convergência ao fornecer uma política inicial forte antes do início da exploração.
2. Estágio II (RL Baseado em Simulação / Offline): As políticas são refinadas usando aprendizado por reforço dentro de um simulador de software (implementando a lógica de execução WC). A recompensa é derivada do tempo de execução simulado. Este estágio permite iterações rápidas sem a latência do hardware real.
3. Estágio III (RL de Sistema Real / Online): As políticas treinadas são implantadas em um sistema WC real. Elas são continuamente refinadas via RL online usando tempos de execução reais como recompensas. Este estágio adapta a política a características específicas do hardware (ex: jitter de NVLink, contenção) que simuladores não conseguem capturar perfeitamente, sem adicionar sobrecarga adicional ao usuário.

3. Principais Contribuições

1. Formulação do Problema: O artigo investiga a atribuição de dispositivos especificamente no contexto de sistemas multi-GPU assíncronos e de consumo de trabalho (work-conserving), abordando a natureza estocástica da ordem de execução.
2. Abordagem de Política Dupla: Introduz uma separação de responsabilidades onde uma política aprende a ordem de travessia dos nós e a outra lida com o posicionamento do dispositivo, capturando tanto a dinâmica de execução quanto as restrições de hardware.
3. Framework de Três Estágios: O DOPPLER propõe uma estratégia de treinamento híbrida combinando aprendizado por imitação, RL baseado em simulação e RL online em sistema real. Essa abordagem acelera a convergência e permite a otimização contínua durante a implantação.
4. Implementação Escalável de GNN: Os autores propõem uma aproximação eficiente de passagem de mensagens que reduz significativamente o tempo de treinamento sem comprometer a convergência.

4. Resultados Experimentais

Experimentos foram conduzidos em quatro GPUs NVIDIA Tesla P100 usando várias cargas de trabalho: Multiplicações de Matrizes em Cadeia (CHAINMM), Redes Neurais Feedforward (FFNN) e blocos/camadas do transformador Llama.

• Desempenho vs. Baselines: O DOPPLER (especificamente o DOPPLER-SYS totalmente treinado) superou consistentemente os baselines, incluindo a heurística de Caminho Crítico, o PLACETO (RL de política única), o GDP (RL de incorporação de grafo) e um otimizador enumerativo customizado.
  • O DOPPLER alcançou até 52,7% menos tempos de execução em comparação com o melhor baseline (GDP).
  • Reduziu o tempo de execução em até 78,2% em comparação com a heurística de Caminho Crítico.
  • Superou o otimizador enumerativo customizado em até 13,8%.
• Estudos de Ablação:
  • Política Dupla: A combinação de ambas as políticas SEL e PLC rendeu os melhores resultados, particularmente para modelos complexos, demonstrando a necessidade de aprender tanto a ordem quanto o posicionamento.
  • Estágios de Treinamento: Treinar apenas no sistema real levou a uma convergência lenta e instabilidade. A combinação de aprendizado por imitação (Estágio I) e simulação (Estágio II) foi crucial para uma convergência estável e rápida antes do refinamento online (Estágio III).
• Transfer Learning: O DOPPLER demonstrou forte generalização. Políticas treinadas em grafos simples (FFNN) puderam ser ajustadas com poucos episódios (2k–4k) para performar bem em arquiteturas não vistas (Llama) e diferentes configurações de hardware (ex: transferindo de 4 P100s para 8 V100s).
• Escalabilidade: O treinamento e a inferência do DOPPLER escalam linearmente com o tamanho do grafo de fluxo de dados, superando outros baselines de RL em eficiência.

5. Significância e Alegações

O artigo afirma que o DOPPLER representa um passo significativo na otimização de cargas de trabalho modernas de ML ao ir além das suposições estáticas de processamento em lote síncrono. Sua significância reside em:

• Eficiência Prática: Ao permitir a execução de consumo de trabalho (work-conserving), o DOPPLO pode economizar substanciais horas de GPU em implantações de larga escala (ex: economizando milhões de horas de GPU anualmente para cargas de trabalho do nível do ChatGPT).
• Adaptabilidade: Diferente de heurísticas estáticas ou modelos puramente baseados em simulação, a capacidade de aprendizado online do DOPPLER permite que ele se adapte às realidades ruidosas dos ambientes de hardware de produção.
• Inovação Metodológica: A decomposição de política dupla oferece uma abordagem mais matizada para a explosão combinatória dos problemas de atribuição de dispositivos, separando o "o que executar a seguir" de "onde executá-lo", o que se alinha melhor com a natureza temporal da execução assíncrona.

Os autores enfatizam que sua abordagem não requer mudanças nas estruturas de ML subjacentes, mas sim otimiza a camada de agendamento, tornando-a uma solução prática para as infraestruturas multi-GPU existentes.