🤖 machine learning

Runtime-Orchestrated Second-Order Optimization for Scalable LLM Training

O artigo apresenta a Asteria, um sistema de tempo de execução que viabiliza a otimização de segunda ordem prática e escalável para grandes modelos de linguagem, distribuindo dinamicamente os estados do otimizador através de hierarquias de memória heterogêneas e desacoplando os cálculos caros de pré-condicionamento do caminho crítico de treinamento em GPU, a fim de reduzir a sobrecarga e acelerar a convergência.

Autores originais: Yishun Lu, Junhao Zhang, Zeyu Yang, Wes Armour

Publicado 2026-05-18

📖 5 min de leitura🧠 Leitura aprofundada

CC BY 4.0

Autores originais: Yishun Lu, Junhao Zhang, Zeyu Yang, Wes Armour

Artigo original sob licença CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ Esta é uma explicação gerada por IA do artigo abaixo. Não foi escrita nem endossada pelos autores. Para precisão técnica, consulte o artigo original. Ler aviso legal completo

Imagine que você está tentando ensinar um robô gigante e superinteligente (um Modelo de Linguagem de Grande Escala) a escrever histórias. Para fazer isso, você precisa de um "professor" (um otimizador) que observe os erros do robô e lhe diga como melhorar.

Atualmente, a maioria dos professores usa um método simples chamado AdamW. É como um aluno que verifica seu dever de casa, vê algumas respostas erradas e faz correções pequenas e rápidas. É rápido e funciona bem, mas exige muita prática (milhões de exemplos) para ficar realmente bom.

Existe um método de professor mais inteligente e avançado chamado Otimização de Segunda Ordem (como SOAP ou Shampoo). Este professor não olha apenas o que estava errado; ele analisa a forma dos erros para entender o problema profundamente. Aprende muito mais rápido e precisa de menos exemplos. No entanto, há um grande problema: este professor inteligente é incrivelmente pesado. Requer quantidades massivas de memória e poder de computação para fazer sua matemática complexa, frequentemente fazendo o computador travar ou funcionar tão lentamente que o professor simples termina o trabalho primeiro.

Apresentamos a Asteria: O "Gerente Logístico Inteligente"

Os pesquisadores de Oxford construíram um novo sistema chamado Asteria. Pense na Asteria não como um novo professor, mas como um brilhante gerente logístico que reorganiza toda a sala de aula para que o professor inteligente finalmente possa fazer seu trabalho sem quebrar a escola.

Veja como a Asteria resolve os três maiores problemas, usando analogias simples:

1. O Problema dos "Móveis em um Quarto Minúsculo" (Memória)

O Problema: O professor inteligente precisa manter enormes e complexos gráficos (matrizes) em sua cabeça (a memória da GPU). Em um computador padrão, não há espaço suficiente. É como tentar caber uma cama de casal, uma mesa de jantar e um guarda-roupa em um apartamento estúdio. O quarto fica cheio e o professor é obrigado a desistir.
A Solução da Asteria: A Asteria age como um especialista em móveis dobráveis. Ela percebe que o professor não precisa segurar tudo nas mãos ao mesmo tempo.

Mantém as partes mais ativas dos gráficos na GPU (a mesa).
Move as partes pesadas e menos frequentemente usadas para a CPU (o corredor) ou até mesmo para um disco rígido (a garagem).
Crucialmente, ela traz as partes pesadas de volta à mesa exatamente quando são necessárias. Isso permite que o professor inteligente trabalhe em um laptop pequeno tão bem quanto em um supercomputador.

2. O Problema do "engarrafamento" (Velocidade)

O Problema: A cada poucos passos, o professor inteligente precisa realizar um cálculo massivo e complexo (como resolver um quebra-cabeça gigante) para atualizar seus gráficos. Isso leva muito tempo e interrompe toda a classe de trabalhar. É como um caminhão de entregas bloqueando toda a rodovia enquanto descarrega um único pacote. A GPU (o cérebro do computador) fica ociosa, esperando.
A Solução da Asteria: A Asteria introduz uma linha de montagem paralela.

Em vez de parar a aula principal para fazer a matemática pesada, a Asteria envia o "trabalho pesado" para uma equipe de trabalhadores no escritório de trás (a CPU).
Enquanto a aula principal (a GPU) continua ensinando o robô, os trabalhadores do escritório de trás resolvem silenciosamente os quebra-cabeças complexos em segundo plano.
Quando a aula precisa dos novos gráficos, o escritório de trás já os terminou e os deslizou para lá. A aula principal nunca precisa parar. O "engarrafamento" desaparece.

3. O Problema do "Chat de Grupo" (Treinamento Distribuído)

O Problema: Ao treinar com muitos computadores ao mesmo tempo, todos geralmente precisam parar e concordar com exatamente as mesmas informações antes de avançar para o próximo passo. É como um chat de grupo onde todos precisam esperar a pessoa mais lenta responder antes que alguém possa digitar novamente.
A Solução da Asteria: A Asteria usa uma política de "Bom o Suficiente".

Em vez de esperar que todos estejam perfeitamente sincronizados, ela permite que os computadores trabalhem com informações ligeiramente "desatualizadas" por um curto período.
Ela envia atualizações apenas quando são verdadeiramente necessárias.
Isso mantém o grupo avançando rapidamente, como uma equipe que confia uns nos outros para continuar trabalhando enquanto atualizações são recuperadas depois, em vez de parar todo o projeto a cada cinco minutos.

Os Resultados: O Que Eles Descobriram?

Os pesquisadores testaram a Asteria em hardware real, incluindo um único computador poderoso (Nvidia DGX Spark) e um grande cluster de computadores (Nvidia GH200).

Funciona em máquinas pequenas: Eles conseguiram treinar um modelo de linguagem grande (1 bilhão de parâmetros) em uma única máquina que anteriormente não conseguia lidar com os requisitos de memória deste professor inteligente.
É mais rápido em tempo real: Como esconde os "engarrafamentos", o treinamento termina em menos tempo de relógio real, mesmo que a matemática seja mais difícil.
Economiza energia: Ao manter o cérebro do computador (GPU) ocupado e não deixá-lo ocioso esperando por cálculos, a Asteria na verdade usa menos energia total para alcançar o mesmo resultado em comparação com os métodos antigos e pesados.
Não perde qualidade: O modelo aprende tão bem quanto aprenderia com o professor inteligente "nativo" (não otimizado), mas chega lá muito mais rápido e barato.

Em Resumo:
A Asteria não inventa uma nova fórmula matemática. Em vez disso, ela repensa como o computador executa essa matemática. Ela separa o trabalho pesado do trabalho principal, usa todo o espaço de armazenamento disponível com sabedoria e permite que os computadores trabalhem de forma assíncrona. Isso transforma um método "teoricamente ótimo, mas praticamente impossível" em uma ferramenta prática para treinar a próxima geração de IA.

Resumo Técnico: Asteria – Otimização de Segunda Ordem Orquestrada em Tempo de Execução para Treinamento Escalável de LLMs

1. Declaração do Problema

Métodos de otimização de segunda ordem (por exemplo, Shampoo, SOAP, KL-Shampoo) oferecem eficiência estatística superior e convergência mais rápida para o treinamento de Modelos de Linguagem de Grande Escala (LLMs) em comparação com otimizadores de primeira ordem como AdamW. No entanto, sua implantação prática em escala é impedida por barreiras severas ao nível do sistema:

Restrições de Memória: Métodos de segunda ordem exigem o armazenamento de grandes matrizes de covariância ( $L$ e $R$ ) e a realização de decomposições em autovalores densas, resultando em uma pegada de memória de $O(N^2)$ e um custo computacional de $O(N^3)$ . Em ambientes com restrições de memória (por exemplo, estações de trabalho de nó único com Arquiteturas de Memória Unificada), esses estados esgotam a memória física, causando falhas de Memória Insuficiente (OOM) ou forçando reduções no tamanho do modelo e no tamanho do lote.
Picos de Latência: O cálculo periódico de raízes inversas (por exemplo, $L^{-1/4}$ ) envolve operações matriciais de custo cúbico. Em implementações existentes, essas atualizações síncronas bloqueiam o pipeline de treinamento da GPU, causando picos de latência periódicos severos que interrompem a sobreposição de granularidade fina entre computação e comunicação essencial para treinamento distribuído eficiente (por exemplo, FSDP).
Sobrecarga de Sincronização: Métodos distribuídos atuais de segunda ordem dependem de atualizações globais de estado rígidas e síncronas (AllGather) em todas as ranks. Isso acopla a comunicação rápida intra-nó com a sincronização lenta inter-nó, impedindo que o sistema explore limites de desatualização assíncronos e conscientes da topologia.

Soluções existentes, como particionamento baseado em DTensor, falham em abordar esses problemas em configurações de nó único com recursos limitados ou quando a sobrecarga de comunicação da sincronização global supera os benefícios.

2. Metodologia: O Tempo de Execução Asteria

O Asteria é um sistema de tempo de execução projetado em co-criação entre hardware e software, construído sobre o PyTorch FSDP. Ele desacopla a lógica de otimização de segunda ordem do caminho crítico de treinamento da GPU através de três mecanismos principais:

A. Tiering de Memória Assimétrico Adaptativo à Arquitetura

Diferentemente dos otimizadores de primeira ordem que tratam o estado como vetores homogêneos, o Asteria reconhece o ciclo de vida assimétrico dos estados de segunda ordem:

GPU: Armazena factor_matrices (fatores de Kronecker) que são atualizados frequentemente durante o treinamento.
Host (CPU/UVM): Armazena inv_factor_matrices (fatores inversos). Estes são mantidos em memória acessível à CPU para evitar pressão na memória da GPU.
NVMe (Opcional): Fornece um nível de staging para execuções com capacidade limitada, permitindo que os estados sejam trocados via computação serializada ou paginação JIT.
Mecanismo: O tempo de execução captura snapshots do lado do host das matrizes de fatores, despacha trabalhadores assíncronos da CPU para calcular raízes inversas e escreve os resultados de volta na memória do host. Quando necessário, a GPU acessa esses dados através de visualizações de Memória Virtual Unificada (UVM) ou pré-busca explícita, evitando cópias redundantes host-dispositivo.

B. Pipeline de Estado Sombra Orquestrado por Hooks

Para eliminar picos de latência, o Asteria remove o cálculo de raiz inversa de $O(N^3)$ do caminho crítico da GPU:

Stream Sombra: Um stream CUDA de baixa prioridade lida com o staging de estado no lado da GPU (pré-busca e paginação) sem bloquear o stream principal de treinamento.
Execução Assíncrona: Atualizações caras de raiz inversa são executadas assincronamente em threads de trabalhadores da CPU.
Integração de Hooks: Hooks leves do PyTorch (register_forward_hook, register_full_backward_pre_hook) atuam como sinais de agendamento. Eles acionam o esvaziamento de transferências de estado prontas e a pré-busca de fatores inversos que serão consumidos em breve, garantindo que a GPU tenha pré-condicionadores atualizados sem interromper o loop principal de treinamento.

C. Coerência de Desatualização Limitada Consciente da Topologia

Para treinamento distribuído, o Asteria substitui a sincronização global rígida por um protocolo de coerência seletivo:

Rastreamento de Desatualização: Um CoherenceRegistry rastreia a frescura de cada bloco de pré-condicionador. Apenas blocos que excedem um orçamento de desatualização configurado são sincronizados.
Sincronização Hierárquica: A sincronização ocorre em dois estágios: primeiro dentro de grupos de processos locais ao nó, depois através de representantes de nó. Isso minimiza o tráfego entre nós e evita viagens de ida e volta desnecessárias da CPU para a GPU mantendo a comunicação na memória acessível ao host.
Barreira Limitada: O tempo de execução impõe uma barreira local garantindo que, se uma versão do pré-condicionador se tornar muito desatualizada enquanto um trabalho da CPU estiver em andamento, o otimizador aguarda a conclusão, equilibrando frescura com vazão.

3. Contribuições Principais

Desacoplamento em Tempo de Execução: O Asteria demonstra que a otimização de segunda ordem pode ser tornada prática não simplificando o algoritmo, mas repensando o gerenciamento de estado, a computação em segundo plano e a sincronização no nível do tempo de execução.
Escalabilidade de Memória: Permite treinamento de segunda ordem em hardware de nó único com restrições de memória (por exemplo, uma única GPU GB10 com 128GB de memória unificada) onde o particionamento tradicional é impossível.
Ocultação de Latência: Ao mover os cálculos de raiz inversa para o host e usar um pipeline sombra, o Asteria achata os picos de latência periódicos, alcançando um perfil de execução de hardware suave comparável aos métodos de primeira ordem.
Eficiência Energética: O sistema melhora o trade-off de perda energética, alcançando mais redução de perda por unidade de energia ao reduzir o tempo ocioso da GPU e otimizar a sobreposição da computação do lado do host.

4. Resultados Experimentais

Os autores avaliaram o Asteria em modelos da família OLMo (660M, 1B, 7B parâmetros) usando o framework de treinamento OLMo em duas plataformas:

Nó Único (DGX Spark, GPU GB10):
- Vazão: Asteria-SOAP e Asteria-KL-Shampoo eliminaram os picos periódicos de tempo de passo observados em implementações nativas (reduzindo o tempo de passo máximo de ~96s para ~66.5s).
- Energia: Enquanto métodos de segunda ordem nativos consumiram ~17-20% mais energia total de SoC do que AdamW, o Asteria reduziu significativamente essa lacuna (por exemplo, Asteria-KL-Shampoo consumiu apenas ~7.9% mais energia do que AdamW).
- Eficiência: Asteria-KL-Shampoo alcançou a maior eficiência normalizada de redução de perda (3.3010), superando tanto métodos nativos de segunda ordem quanto AdamW.
Multi-Nó (Cluster GH200):
- Convergência: Variantes do Asteria mantiveram trajetórias de convergência próximas às suas contrapartes nativas de segunda ordem, superando significativamente o AdamW no tempo de relógio para atingir níveis de perda alvo.
- Tolerância à Desatualização: Experimentos variando o orçamento de desatualização ( $S$ ) mostraram que aumentar $S$ de 1 para 5 reduziu substancialmente o tempo de treinamento com impacto negligenciável na qualidade final do modelo.
- Escalonamento Forte: Em modelos de 7B escalando de 2 para 16 nós, o Asteria alcançou maior aceleração e menor tempo por passo em comparação com bases nativas, provando-se eficaz mesmo à medida que o custo relativo da manutenção do estado aumenta.

5. Significado e Alegações

O artigo argumenta que o principal obstáculo ao treinamento prático de LLMs de segunda ordem não é o custo matemático das atualizações conscientes da curvatura, mas o descompasso de sistemas entre a manutenção de estados de segunda ordem e as pilhas de treinamento modernas.

Reenquadramento do Problema: Os autores afirmam que a otimização de segunda ordem deve ser vista como um problema de co-design envolvendo algoritmos, agendamento de tempo de execução e sistemas de memória heterogênea.
Viabilidade: Eles demonstram que grandes memórias de GPU agregadas não são um pré-requisito fundamental para treinamento de segunda ordem; com colocação consciente do ciclo de vida e acesso em etapas, esses métodos podem ser executados em orçamentos de memória apertados.
Praticidade: Ao tratar a frescura do pré-condicionador e a comunicação como recursos gerenciados pelo tempo de execução em vez de estados globais rígidos, o Asteria estreita a lacuna entre a promessa estatística da otimização de segunda ordem e sua implantabilidade no mundo real.

O artigo conclui que o Asteria torna a otimização de segunda ordem uma escolha de sistemas viável para treinamento de LLMs, oferecendo melhor convergência e eficiência energética sem exigir simplificação algorítmica.