Deep Optimizer States: Towards Scalable Training of Transformer Models Using Interleaved Offloading

O artigo apresenta o "Deep Optimizer States", uma técnica inovadora que divide modelos de linguagem grandes em subgrupos para alternar dinamicamente a atualização dos estados do otimizador entre CPU e GPU, explorando as flutuações de uso de memória e superando o gargalo de memória para acelerar o treinamento em 2,5 vezes em comparação com abordagens existentes.

O essencial

Imagine que você está tentando ensinar um robô superinteligente (um modelo de IA, como o ChatGPT) a escrever poemas, diagnosticar doenças ou prever o clima. Para isso, o robô precisa de um "cérebro" gigantesco, com bilhões de conexões (parâmetros).

O problema é que esse cérebro é tão grande que não cabe na memória de trabalho do computador (a GPU), que é como a mesa de trabalho do robô. A mesa é pequena, mas o cérebro é enorme.

O Problema: A "Parede da Memória"

Antes desta pesquisa, os cientistas tinham duas opções ruins:

1. Parar tudo: Tentar encaixar o cérebro inteiro na mesa. Como não cabe, o treinamento trava ou exige supercomputadores caríssimos.
2. Usar o armário (CPU): Guardar a parte do cérebro que não cabe na mesa no armário (memória do computador, chamada de CPU). A cada passo de aprendizado, o robô tem que correr até o armário, pegar um pedaço do cérebro, trazer para a mesa, trabalhar, e levar de volta.

O gargalo: O corredor entre a mesa e o armário (o cabo PCIe) é estreito e lento. Além disso, o robô na mesa (GPU) é um atleta olímpico, super rápido, enquanto o ajudante no armário (CPU) é um pouco mais lento. Quando o robô precisa esperar o ajudante trazer os dados, ele fica parado, perdendo tempo valioso.

A Solução: "Deep Optimizer States" (Estados Profundos do Otimizador)

Os autores deste artigo criaram uma nova técnica chamada Deep Optimizer States. Eles não tentaram apenas "levar e trazer" de forma tradicional. Eles mudaram a estratégia de como o trabalho é dividido.

Aqui está a analogia simples:

1. O Antigo Método (DeepSpeed TwinFlow)

Imagine que você tem 100 caixas de ferramentas (o cérebro do robô) para montar um móvel.

• A mesa (GPU) tem espaço para apenas 2 caixas.
• O armário (CPU) guarda as outras 98.
• O método antigo: Você decide que as caixas 1 e 2 ficam na mesa para sempre. As caixas 3 a 100 ficam no armário. O ajudante (CPU) pega a caixa 3, trabalha nela, leva de volta, pega a 4, trabalha, leva de volta...
• O problema: Enquanto o ajudante trabalha na caixa 3, a mesa (o robô rápido) fica parada esperando. E quando o ajudante termina, ele tem que correr até a mesa para entregar o resultado, e a mesa fica parada esperando ele voltar. É um ritmo de "um de cada vez", cheio de esperas.

2. O Novo Método (Deep Optimizer States)

A equipe percebeu algo inteligente: A mesa fica vazia em certos momentos!
Quando o robô está "pensando" (fase de frente) ou "analisando erros" (fase de trás), ele usa muita memória. Mas, quando chega a hora de "atualizar o aprendizado" (atualizar os pesos), a mesa libera muito espaço.

A nova técnica faz o seguinte:

• Divisão Dinâmica: Em vez de deixar apenas 2 caixas na mesa, eles dividem as 100 caixas em grupos menores.
• Trabalho em Paralelo (Interleaved):
  • Enquanto o ajudante (CPU) está trabalhando na caixa 3, o robô (GPU) já está trabalhando na caixa 4 (que foi trazida para a mesa no momento em que ela ficou vazia).
  • Enquanto o robô trabalha na caixa 4, o ajudante já está trazendo a caixa 5.
  • Eles se cruzam! O ajudante entrega o trabalho da caixa 3 enquanto o robô termina a caixa 4.
• O Corredor Inteligente: Eles também mudaram como os dados viajam pelo corredor. Em vez de levar o material "desmontado" (formato de baixa precisão) e montar na mesa, eles levam o material "montado" (alta precisão) direto, o que é mais rápido e evita erros de conversão.

O Resultado: Uma Dança Perfeita

Com essa nova coreografia:

• Ninguém fica parado: O robô rápido (GPU) e o ajudante (CPU) trabalham ao mesmo tempo, sem esperar um pelo outro.
• O corredor é usado melhor: O tráfego de dados é constante e eficiente.
• Velocidade: O treinamento ficou 2,5 vezes mais rápido do que os métodos anteriores.

Por que isso importa?

Antes, para treinar modelos gigantes, você precisava de um data center inteiro e caro. Com essa técnica, você pode treinar modelos muito grandes em um único computador potente (ou em computadores menores), economizando dinheiro e tempo. É como se você pudesse construir um arranha-céu usando apenas uma escada de mão, porque aprendeu a subir e descer de forma tão eficiente que o tempo de construção cai pela metade.

Resumo da Ópera:
O papel propõe uma maneira inteligente de dividir o trabalho entre a memória rápida (GPU) e a memória lenta (CPU), fazendo com que elas trabalhem juntas em ritmo de dança, em vez de ficarem esperando uma pela outra. Isso quebra a "parede da memória" e acelera a criação de Inteligência Artificial.

Resumo técnico

Resumo Técnico: Deep Optimizer States

1. O Problema: A "Parede de Memória" no Treinamento de LLMs

O treinamento de Modelos de Linguagem Grandes (LLMs) e Transformers enfrenta um gargalo crítico conhecido como "parede de memória". À medida que os modelos crescem para centenas de bilhões de parâmetros, o estado do otimizador (que inclui parâmetros FP32, momento e variância para otimizadores adaptativos como Adam) torna-se massivo, superando a capacidade de memória das GPUs, mesmo com técnicas de paralelismo 3D (dados, pipeline e tensor).

Para contornar isso, abordagens de ponta (como DeepSpeed Offload e ZeRO-Offload) descarregam (offload) o estado do otimizador para a memória do host (CPU). No entanto, isso introduz dois gargalos severos:

1. Largura de Banda Limitada (PCIe): A transferência de dados entre CPU e GPU é lenta (tipicamente 25-50 GB/s), criando um gargalo de I/O.
2. Capacidade Computacional Desigual: As CPUs são ordens de magnitude mais lentas que as GPUs para atualizar os parâmetros. Em testes, as GPUs atualizavam ~100 bilhões de parâmetros/segundo, enquanto as CPUs geriam apenas ~8 bilhões/segundo.
3. Subutilização de Recursos: Nas abordagens atuais, a memória da GPU e a largura de banda da PCIe ficam subutilizadas durante as fases de atualização, pois os dados são transferidos de forma bloqueante e sequencial, deixando a GPU ociosa enquanto a CPU trabalha.

2. Metodologia e Design do Sistema

Os autores propõem o Deep Optimizer States, uma técnica de middleware que utiliza o descarregamento intercalado (interleaved offloading) para otimizar o treinamento híbrido CPU-GPU. A solução baseia-se na observação de que a utilização da memória da GPU flutua drasticamente entre as fases de forward, backward e update.

Princípios Chave de Design:

• Atualizações Intercaladas (GPU e CPU): Em vez de manter um subconjunto estático do otimizador na GPU e o resto na CPU, o sistema divide o estado do otimizador em subgrupos menores. Durante a fase de atualização, alguns subgrupos são atualizados na GPU enquanto outros são processados na CPU, alternando dinamicamente.
• Sobreposição (Overlapping) de Computação e Transferência: O sistema utiliza transferências assíncronas (H2D e D2H) para sobrepor o movimento de dados com a computação. Enquanto a CPU atualiza um subgrupo, a GPU pode estar atualizando outro, e a transferência de dados para o próximo subgrupo ocorre simultaneamente em streams dedicados.
• Gerenciamento Eficiente de Gradientes: Os gradientes gerados na fase backward (FP16) são convertidos para FP32 diretamente na GPU (a alta velocidade) antes de serem transferidos para a CPU, evitando conversões lentas na memória do host e alocação de memória não fixada (unpinned).
• Modelo de Desempenho Dinâmico: Os autores desenvolveram um modelo matemático para calcular a proporção ideal de subgrupos a serem atualizados na GPU (o "stride" de atualização). O modelo equilibra a velocidade de atualização da CPU vs. GPU e a largura de banda da PCIe para maximizar a sobreposição e evitar gargalos de I/O.

Algoritmo de Agendamento:
O sistema utiliza um algoritmo que decide, para cada subgrupo, se ele deve ser atualizado na GPU ou na CPU com base na razão ótima ($k$) derivada do modelo de desempenho. Isso permite que a GPU processe atualizações enquanto a CPU prepara os próximos dados, eliminando tempos de espera bloqueantes.

3. Principais Contribuições

1. Análise de Comportamento: Demonstração de que a fragmentação fina do estado do otimizador em subgrupos não impacta o tempo de treinamento, mas revela flutuações significativas na utilização da memória da GPU e subutilização da PCIe, criando oportunidades para otimização.
2. Técnica de Interleaving: Proposta de uma nova abordagem para mover e atualizar subgrupos do otimizador dinamicamente entre CPU e GPU, superando as limitações das soluções estáticas (como TwinFlow/ZeRO-Offload++).
3. Modelo de Desempenho: Desenvolvimento de um modelo analítico para determinar a frequência ideal de offloading na GPU, maximizando a sobreposição entre computação e transferência de dados.
4. Implementação em Middleware: Integração da técnica no DeepSpeed e Megatron-LM, permitindo treinamento híbrido eficiente sem exigir alterações profundas nos frameworks existentes.
5. Validação Experimental: Avaliação extensiva em modelos de até 20 bilhões de parâmetros em configurações com recursos limitados.

4. Resultados Experimentais

Os experimentos foram realizados em um nó com 4x GPUs H100 (80GB) e 192 núcleos de CPU.

• Aceleração de Iteração: O Deep Optimizer States alcançou 2,5x mais rápido nas iterações de treinamento em comparação com a abordagem de ponta (DeepSpeed ZeRO-3 com offload total para CPU).
• Throughput de Atualização: Houve um aumento de até 3x na velocidade de atualização dos parâmetros do modelo em várias configurações.
• Eficiência de Memória: O método permite treinar modelos de 20B parâmetros em um único nó com a mesma velocidade que modelos de 7B parâmetros em sistemas de ponta atuais, ou seja, oferece o desempenho de modelos maiores com menos memória GPU.
• Otimização de Recursos: O sistema alcançou utilização próxima de 100% da GPU e ~40% da largura de banda da PCIe (próximo ao máximo prático para transferências D2H/H2D), enquanto as soluções estáticas deixavam a GPU ociosa durante as atualizações da CPU.
• Escalabilidade: O método manteve ganhos de desempenho (até 2,5x) mesmo com o aumento do grau de paralelismo de dados, demonstrando eficiência em escala.

5. Significado e Impacto

O Deep Optimizer States representa um avanço significativo no treinamento escalável de LLMs em hardware com restrições de memória. Ao transformar o problema de offloading de um gargalo de I/O estático em um processo dinâmico e sobreposto, a técnica:

• Democratiza o Treinamento: Permite que pesquisadores e empresas treinem ou fine-tunem modelos grandes (20B+) em nós de computação mais acessíveis (single-node), sem depender de clusters massivos.
• Maximiza Investimentos em Hardware: Extrai o máximo desempenho de sistemas híbridos existentes, melhorando a relação custo-benefício de clusters com CPUs e GPUs.
• Prepara para o Futuro: A abordagem é particularmente relevante para futuras arquiteturas (como Grace Hopper) que possuem interconexões de alta largura de banda entre CPU e GPU, onde o gerenciamento dinâmico de dados será ainda mais crítico.

Em suma, o trabalho resolve o desequilíbrio entre a capacidade computacional da GPU e a lentidão das atualizações baseadas em CPU, preenchendo a lacuna deixada pelas soluções de offload estático atuais.