Decoupled DiLoCo for Resilient Distributed Pre-training

1. O Problema

O pré-treinamento de Grandes Modelos de Linguagem (LLMs) em escala moderna depende quase exclusivamente do paradigma SPMD (Single Program, Multiple Data). Neste modelo, todos os aceleradores (TPUs/GPUs) devem estar rigidamente acoplados e sincronizados globalmente a cada passo de treinamento.

• Gargalo de Confiabilidade: Qualquer falha de hardware transitória, "straggler" (acelerador lento) ou atraso de rede paralisa todo o cluster, desperdiçando tempo de computação valioso.
• Escalabilidade vs. Estabilidade: À medida que o número de chips aumenta, falhas que antes eram raras tornam-se eventos rotineiros.
• Limitações das Soluções Atuais: Métodos recentes como o DiLoCo reduziram a largura de banda necessária através de sincronização intermitente, mas permaneceram fundamentalmente síncronos (lock-step), tornando-se vulneráveis a falhas locais que ainda param o sistema global.
• O Dilema CAP: O artigo argumenta que o pré-treinamento atual prioriza excessivamente a Consistência (todos os nós têm a mesma visão dos pesos) em detrimento da Disponibilidade e Tolerância a Partições, o que é ineficiente em ambientes de grande escala e propensos a falhas.

2. Metodologia: Decoupled DiLoCo

Os autores propõem o Decoupled DiLoCo, uma evolução do framework DiLoCo que rompe a barreira de sincronização "lock-step" para maximizar o goodput (taxa de trabalho útil) em ambientes caóticos.

Arquitetura Principal

O sistema divide o cluster global em M "learners" (aprendizes) independentes e assíncronos, coordenados por um sincronizador central (syncer).

1. Learners (Aprendizes):

   • Cada learner executa um loop de otimização local independente em seu próprio shard de dados.
   • Eles realizam passos de otimização interna (ex: AdamW) sem esperar pelos pares.
   • Os modelos são particionados em fragmentos (subconjuntos de pesos).
   • Os learners enviam metadados de progresso e fragmentos de pesos ao syncer assincronamente.
   • Se um learner falha, os outros continuam treinando sem interrupção.

2. Syncer (Sincronizador):

   • Responsável pela reconciliação global dos estados divergentes dos learners.
   • Executa um passo de otimização externa (outer optimization) para agregar os fragmentos recebidos.
   • Quórum Mínimo (K): Em vez de esperar por todos os M learners, o syncer aguarda apenas um subconjunto mínimo (K) de learners ativos para realizar a agregação. Isso garante alta disponibilidade.
   • Janela de Graça Adaptativa (Adaptive Grace Window): O sistema utiliza o tempo ocioso (slack) entre passos de computação para esperar por mais learners, melhorando a eficiência da amostragem sem parar o sistema.
   • Mergers Pesados por Token: A agregação utiliza pesos dinâmicos baseados no número de tokens processados por cada learner desde a última atualização, garantindo que learners mais rápidos ou com mais dados tenham contribuição proporcional.

3. Técnicas de Agregação Avançadas:

   • Radial-Directional Averaging (RDA): Um método de fusão que separa a média das normas e das direções dos gradientes externos. Isso estabiliza a hiperparametrização do otimizador externo e melhora a performance ao escalar o número de learners.
   • Fragmentação Balanceada: Uso de um algoritmo greedy de bin-packing para dividir os tensores do modelo em fragmentos de tamanho equilibrado, minimizando o uso de largura de banda de pico.

4. Engenharia de Caos e Recuperação:

   • O sistema foi projetado para suportar falhas contínuas simuladas (Chaos Engineering).
   • Recuperação de Learner: Se um learner falha e reinicia, ele pode se reconectar, buscar o estado mais recente de um par saudável e continuar, sem exigir reinicialização global.
   • Scavenging (Varredura de Recursos): O sistema pode adicionar dinamicamente novos learners (recursos computacionais "sobrando" ou de diferentes gerações de hardware) durante o treinamento, aumentando o goodput sem reconfiguração complexa.

3. Principais Contribuições

• Decoupled DiLoCo: Um framework de treinamento distribuído que prioriza disponibilidade e tolerância a partições sobre consistência estrita, decompondo clusters monolíticos em aprendizes assíncronos.
• Arquitetura de Sistema Escalável: Design com um syncer central leve (rodando em CPUs) que permite a reconciliação assíncrona de parâmetros, operando eficientemente em escala de pré-treinamento.
• Validação Empírica Robusta: Evidências de que o método atinge desempenho downstream comparável ao treinamento Data-Parallel (DP) padrão em tarefas de texto e visão, para arquiteturas densas e Mixture-of-Experts (MoE), mesmo sob falhas agressivas.
• Resiliência via Engenharia de Caos: Demonstração de que o framework mantém alta disponibilidade e qualidade do modelo em ambientes simulados com milhões de chips e taxas de falha elevadas, alcançando 100% de tempo de atividade (uptime) em cenários onde o DP falharia.

4. Resultados Experimentais

Os experimentos foram realizados com modelos como Gemma (2B, 5B, 9B) e MoE, treinados com até 1,3 trilhões de tokens.

• Goodput e Uptime:

  • Em simulações com 2,4 milhões de chips e falhas frequentes (MTBI de 1 ano por chip), o Data-Parallel (DP) com elasticidade caiu para 40-58% de goodput.
  • O Decoupled DiLoCo (M=8) manteve 80-88% de goodput e 100% de uptime (o sistema nunca parou de "dar passos").
  • Regra geral: Um sistema com M learners equivale a um sistema single-learner usando M vezes menos chips em termos de resiliência.

• Desempenho do Modelo (ML Performance):

  • O Decoupled DiLoCo alcançou desempenho downstream comparável ou superior ao DP padrão em benchmarks de texto (ARC, HellaSwag, etc.) e visão (MMMU, DocVQA).
  • A qualidade do modelo não foi degradada pelo treinamento assíncrono ou pelas falhas de hardware simuladas.
  • O método funciona bem tanto para modelos densos quanto para MoE, apesar da complexidade de balanceamento de carga local nos MoE.

• Heterogeneidade e Scavenging:

  • O sistema suportou clusters heterogêneos (TPUv5e e TPUv5p misturados) sem gargalos causados pelos chips mais lentos.
  • Permitiu o uso de recursos computacionais "scavenged" (adicionados dinamicamente), acelerando o treinamento sem degradar a qualidade final.

• Eficiência de Largura de Banda:

  • Redução de ordens de magnitude na necessidade de largura de banda em comparação ao DP, permitindo treinamento em redes de datacenter menos robustas ou geograficamente distribuídas.

5. Significado e Impacto

O Decoupled DiLoCo representa uma mudança de paradigma no pré-treinamento de LLMs:

1. Mudança de Prioridade: Move o foco da consistência estrita (SPMD) para a disponibilidade, reconhecendo que em escala massiva, falhas são inevitáveis e o custo de sincronização global é proibitivo.
2. Viabilidade de Infraestrutura Distribuída: Torna viável o treinamento em clusters geograficamente distribuídos ou utilizando hardware heterogêneo e preemptível, onde a latência e a confiabilidade variam drasticamente.
3. Futuro do Treinamento: Sugere que, ao relaxar os requisitos de sincronização, é possível extrair mais computação útil de hardware propenso a falhas, alinhando-se com a visão de sistemas de fluxo de dados assíncronos (como o Pathways do Google).

Em resumo, o trabalho demonstra que é possível treinar modelos de ponta com a mesma qualidade dos métodos atuais, mas com uma resiliência sistêmica drasticamente superior, permitindo o uso eficiente de recursos computacionais em escala global e em ambientes hostis.