Resumo Técnico: Kareus – Redução Conjunta de Energia Dinâmica e Estática no Treinamento de Grandes Modelos

Declaração do Problema

O crescimento rápido da demanda por computação de IA superou a oferta de energia, tornando a energia um recurso caro e disputado para o treinamento de grandes modelos. Os esforços de otimização atuais tipicamente abordam ou a energia dinâmica (consumida por computação ativa, redutível via escalonamento de frequência) ou a energia estática (consumida pelo hardware independentemente da atividade, redutível ao encurtar o tempo de execução).

As soluções existentes sofrem de uma limitação crítica: elas otimizam esses fatores isoladamente ou via combinações ingênuas.

• Abordagens de escalonamento de frequência (ex: Perseus) reduzem a energia dinâmica ao baixar as frequências das GPUs para caminhos não críticos, mas frequentemente seguem a execução sequencial de kernels, perdendo oportunidades de sobrepor comunicação e computação.
• Abordagens de escalonamento de kernel de grão fino (ex: Nanobatching) reduzem a energia estática ao sobrepor comunicação e computação, mas ignoram a energia dinâmica e o escalonamento de frequência.

O artigo identifica que a alocação de SM, o tempo de lançamento de kernel e a frequência da GPU são interdependentes. Alterar um fator desloca a configuração ideal dos outros. Por exemplo, reduzir a frequência torna os kernels relativamente mais limitados por computação (compute-bound), alterando a estratégia ideal de contenção de recursos e sobreposição. Consequentemente, otimizar subconjuntos desses fatores ou combiná-los de forma ingênua falha em levar a fronteira de compensação tempo–energia ao seu limite teórico.

Metodologia

Os autores propõem o Kareus, um sistema de treinamento projetado para otimizar conjuntamente a alocação de SM, o tempo de lançamento de kernel e a frequência da GPU para encontrar a fronteira de Pareto ótima entre tempo e energia.

1. Modelo de Execução de Sobreposição Particionada

Para tornar a otimização conjunta tratável, o Kareus decompõe o problema global em subproblemas locais usando um modelo de execução de sobreposição particionada:

• Particionamento: O grafo de execução das passagens forward e backward é dividido em partições fixas. Uma partição consiste em um kernel de comunicação (ex: AllReduce) e a sequência contígua mais longa de kernels de computação de um microbatch diferente que não possui dependências de dados com o kernel de comunicação.
• Generalização: Este modelo generaliza as técnicas de nanobatching, adicionando controle de grão fino sobre a alocação de SM e o tempo de lançamento para cada partição.
• Composição Hierárquica: O sistema otimiza cada partição de forma independente e, em seguida, compõe hierarquicamente essas fronteiras locais em uma fronteira global de microbatch e de iteração.

2. Otimização Bayesiana de Múltiplos Objetivos e Múltiplas Passagens (MBO)

Como a busca exaustiva no espaço de configuração (envolvendo contagens de SM, tempos de lançamento e frequências) é proibitivamente cara (exigindo milhares de horas de GPU e causando interferência térmica), o Kareus emprega um algoritmo MBO especializado:

• Modelos Substitutos (Surrogate Models): Utiliza dois modelos distintos de Árvore de Decisão Potencializada por Gradiente (XGBoost) para prever o tempo de execução e a energia dinâmica, respectivamente.
• Seleção de Candidatos de Múltiplas Passagens: Para explorar a fronteira de Pareto de forma eficaz, o algoritmo utiliza múltiplas funções de aquisição em paralelo:
  • Melhoria de Hipervolume (HVI): Três funções HVI distintas guiam a busca para menor energia total, menor energia dinâmica (favorecendo frequências mais baixas) e menor energia estática (favorecendo execução mais rápida).
  • Exploração de Incerteza: Um conjunto bootstrap de modelos substitutos estima a incerteza da previsão, priorizando candidatos em regiões subexploradas do espaço de busca.
• Perfilamento Termicamente Estável: Para garantir medições de energia precisas, o sistema executa partições sobre uma janela de 5 segundos e insere um período de resfriamento de 5 segundos entre as configurações para evitar viés térmico.

3. Implementação do Sistema

O Kareus é implementado integrando-se ao Megatron-LM e ao Perseus.

• Utiliza MSCCL++ para controle de grão fino da alocação de SM.
• Emprega uma torch.autograd.Function customizada para envolver partições, permitindo cronogramas distintos para as passagens forward e backward.
• Em tempo de execução, seleciona um cronograma de execução de uma fronteira pré-computada com base nas restrições do usuário (ex: prazo de tempo ou orçamento de energia) e ajusta dinamicamente as frequências da GPU via um controlador adaptado do Perseus.

Principais Resultados

O Kareus foi avaliado em 14 cargas de trabalho, incluindo treinamento real no Llama 3.2 3B e Qwen 3 1.7B, e emulação de larga escala do Llama 3.3 70B em 16 GPUs NVIDIA A100.

• Compensação Tempo–Energia: Comparado à linha de base de estado da arte (Megatron-LM + Perseus), o Kareus alcança:
  • Até 28,3% de redução na energia de treinamento para o mesmo tempo de treinamento.
  • Até 27,5% de redução no tempo de treinamento para o mesmo consumo de energia.
• Desempenho de Throughput Máximo: No modo de throughput máximo, o Kareus reduz o tempo de iteração em até 14,9% e a energia em até 22,1% em comparação com a linha de base Megatron-LM.
• Estudos de Ablação: Remover o escalonamento de frequência ou o escalonamento de kernel resulta em degradação significativa do desempenho (ex: remover o escalonamento de kernel leva a um aumento de 10,8% na energia), confirmando a necessidade da otimização conjunta.
• Escalabilidade: Resultados de emulação em larga escala para o Llama 3.3 70B mostram tendências consistentes com os resultados do banco de testes físico, demonstrando a aplicabilidade do sistema a modelos massivos.

Significância e Alegações

O artigo alega que o Kareus empurra a fronteira de compensação tempo–energia para o treinamento de grandes modelos, demonstrando que a energia dinâmica e a estática não podem ser otimizadas independentemente. A contribção central é o insight de que os cronogramas de execução (alocação de SM e tempo de lançamento) e a frequência da GPU são conjunta e interdependentemente determinados.

Os autores enfatizam que:

1. O Controle Conjunto é Superior: Otimizar conjuntamente a frequência e o escalonamento é comprovadamente mais eficiente energeticamente do que deixar o controle de frequência para controladores de potência de hardware ou otimizar o escalonamento sem consciência da frequência.
2. A Interdependência é a Chave: A alocação de SM e o tempo de lançamento ideais são específicos para cada frequência; um cronograma ideal em uma frequência pode ser subótimo em outra.
3. Eficiência Prática: Ao decompor o problema em partições e usar MBO, o Kareus alcança esses ganhos sem o custo proibitivo de uma busca exaustiva, tornando-o uma solução prática para as atuais e futuras cargas de trabalho de treinamento de IA em larga escala.