Resumo Técnico: Fusão de Modelos para Agregação em DiLoCo

Declaração do Problema

O treinamento de Baixa Comunicação Distribuída (DiLoCo) surgiu como um método para treinar grandes modelos de linguagem em clusters de computação mal conectados ("ilhas"), minimizando a sincronização global. No DiLoço, os trabalhadores realizam múltiplas etapas de otimização local (etapas internas) antes de agregar suas atualizações de parâmetros (pseudo-gradientes) por meio de um otimizador global. Embora seja eficaz para reduzir os custos de comunicação, o DiLoCo sofre com a degradação de desempenho conforme o número de trabalhadores ($R$) ou o número de etapas locais ($H$) aumenta. Essa degradação ocorre porque as réplicas locais divergem significamente, fazendo com que a média dos pseudo-gradientes se desalinhe com as atualizações produzidas pelo treinamento de paralelismo de dados estritamente sincronizado.

Concomitantemente, técnicas de Fusão de Modelos (Model Merging) foram desenvolvidas para combinar modelos ajustados (fine-tuned) de forma independente em um único modelo. Um desafio central na fusão é a "interferência de tarefas", onde as atualizações de parâmetros de diferentes tarefas conflitam, levando a quedas de desempenho, especialmente conforme o número de modelos a serem fundidos aumenta.

Os autores identificam uma analogia estrutural entre esses dois campos: a etapa de agregação de pseudo-gradiente no DiLoCo (com um otimizador externo SGD) é matematicamente equivalente à Aritmética de Tarefas (Task Arithmetic) de fusão de modelos. Nesta analogia:

• O modelo global compartilhado no início de uma etapa externa corresponde à inicialização pré-treinada.
• As réplicas locais dos trabalhadores correspondem a modelos "especialistas" ajustados.
• As atualizações locais (pseudo-gradientes) correspondem a "vetores de tarefa".
• O número de trabalhadores corresponde ao número de tarefas.

O problema central abordado é se técnicas avançadas de fusão de modelos, projetadas especificamente para mitigar a interferência de tarefas, podem ser reaproveitadas para melhorar a etapa de agregação no DiLoCo, restaurando assim o desempenho conforme o número de trabalhadores ou etapas locais aumenta.

Metodologia

Ponte Teórica

O artigo estabelece que a atualização externa do DiLoCo com SGD é uma aplicação direta da aritmética de tarefas. Ao visualizar a agregação de pseudo-gradientes como um problema de fusão de modelos, os autores propõem substituir a operação de média padrão por algoritmos de fusão sofisticados projetados para reduzir a interferência negativa.

Avaliação Experimental de Métodos de Fusão

Os autores testaram sistematicamente vários métodos de fusão de última geração como substitutos diretos para a etapa de agregação do DiLoCo. Os métodos avaliados incluíram:

• Baseados em Esparsidade: TIES (TrIm, Elect Sign) e DARE (Drop And REscale).
• Baseados em Ortogonalização: TSV (Task Singular Vector), Iso-C (Isotropic merging in Common subspace) e Iso-CTS.

Descobertas sobre os Métodos de Fusão:

• Métodos de esparsidade (TIES, DARE): Desempenharam-se de forma semelhante ao DiLoCo SGD padrão apenas em níveis de esparsidade muito baixos. O desempenho degradou conforme a esparsidade aumentava, sugerindo que não ofereceram uma vantagem distinta sobre a média simples neste contexto.
• TSV: Desempenhou-se mal (maior perda de validação) e foi computacionalmente proibitivo, exigindo múltiplas Decomposições de Valor Singular (SVD) por etapa.
• Iso-CTS: Superou o baseline, mas foi computacionalmente caro e não ofereceu vantagem significativa sobre o Iso-C no cenário do DiLoCo (onde as tarefas locais são curtas e similares).
• Iso-C: Emergiu como o candidato mais promissor. Ele melhora o desempenho ao substituir o espectro de valores singulares do pseudo-gradiente agregado por sua média, efetivamente achatando as direções de alta energia para gerar um subespaço comum equilibrado. Requer apenas uma SVD por matriz de parâmetros, oferecendo uma relação custo-benefência favorável.

Algoritmo Proposto: IsoLoCo

Com base no sucesso do Iso-C, os autores propõem o IsoLoCo (Isotropic Distributed Low-Communication). O IsoLoCo adapta o Iso-C para o treinamento distribuído ao:

1. Correção Isotrópica: Para matrizes de parâmetros 2D (camadas de peso), ele calcula a SVD do pseudo-gradiente médio, substitui os valores singulares pela sua média e reconstrói a atualização. Isso regulariza o espectro enquanto preserva a estrutura do subespaço.
2. Momento de Nesterov: Diferente do Iso-C padrão, o IsoLoCo incorpora o momento de Nesterov no otimizador externo. Isso é motivado por trabalhos anteriores que mostram o benefício do momento em configurações FedOpt/DiLoCo. O pseudo-gradiente corrigido é alimentado em um buffer de momento para gerar a atualização final.

Otimização de Implementação

Para lidar com o custo computacional da SVD, os autores desenvolveram uma implementação rápida do IsoLoCo usando iterações de Newton–Schulz. Isso substitui a SVD padrão por uma ortogonalização iterativa e utiliza um proxy RMS para a média dos valores singulares. Esta otimização gera uma ortogonalização até duas ordens de magnitude mais rápida com perda de degradação negligenciável.

Resultados Principais

Desempenho Através de Contagens de Trabalhadores

Experimentos foram conduzidos em modelos do tipo Llama (178M, 512M e 1B de parâmetros) usando o dataset DCLM.

• Escalabilidade com Trabalhadores: O IsoLoCo supera significativamente o DiLoCo padrão conforme o número de trabalhadores aumenta.
  • Na escala de 178M com $R=128$, o IsoLoCo reduziu a lacuna de perda de validação em relação ao baseline de paralelismo de dados em 5,8% comparado ao DiLoCo.
  • Em $R=128$, o IsoLoCo mostrou uma degradação de 9,7% em relação ao treinamento de paralelismo de dados, enquanto o DiLoCo mostrou uma degradação de 16,5%.
  • O IsoLoCo até superou o baseline de paralelismo de dados em contagens baixas de trabalhadores ($R \le 4$).
• Tamanho do Modelo: A vantagem do IsoLoCo sobre o DiLoCo persiste através dos tamanhos de modelo (178M a 1B), embora a melhoria relativa seja ligeiramente menor para modelos maiores (1,2% em 1B vs. ~2,4% em 178M), provavelmente porque modelos maiores escalam naturalmente melhor com contagens de trabalhadores.

Impacto das Etapas Internas

O IsoLoCo foi avaliado com intervalos de sincronização crescentes ($H = 30$ a $240$). Ele consistentemente superou o DiLoCo em todos os intervalos, demonstrando robustez a períodos mais longos de divergência local.

Comparação com DiLoCo Baseado em Momento

Notavelmente, a combinação do DiLoCo SGD com a agregação Iso-C (sem momento) já superou o DiLoCo padrão (que usa momento de Nesterov). Adicionar o momento de Nesterov ao Iso-C (criando o IsoLoCo) ampliou ainda mais a diferença de desempenho, indicando que a agregação baseada em fusão e os mecanismos de momento são complementares.

Compatibilidade com Muon

O artigo observa que o IsoLoCo pode ser composto com o otimizador interno Muon (MuLoCo). Nesta configuração, o IsoLoCo alcançou a menor perda de validação entre todos os métodos para $R \ge 32$, enquanto o IsoLoCo com um otimizador interno AdamW igualou ou superou o MuLoCo em contagens altas de trabalhadores com custos de ortogonalização significativamente menores.

Significância e Alegações

O artigo afirma que visualizar o treinamento distribuído de baixa comunicação através da lente da fusão de modelos fornece um poderoso framework para melhorar a escalabilidade. As principais contribuições são:

1. Unificação Teórica: Estabelecer uma analogia direta entre a agregação de pseudo-gradiente do DiLoCo e a fusão de modelos de aritmética de tarefas.
2. Validação Empírica: Demonstrar que técnicas de fusão de modelos, especificamente o Iso-C, podem servir como regras de agregação superiores para o treinamento distribuído, superando a média simples e até o DiLoCo baseado em momento.
3. IsoLoCo: Propor um algoritmo prático, potencializado por momento, que reduz significativamente a lacuna de desempenho entre o treinamento de baixa comunicação e os padrões de ouro de paralelismo de dados, particularmente em regimes com muitos trabalhadores ou longos intervalos de sincronização.
4. Eficiência: Fornecer uma implementação baseada em Newton-Schulz que torna o processo de fusão pesado em SVD computacionalmente viável para o treinamento distribuído iterativo.

Os autores concluem que a agregação inspirada em fusão é uma estratégia eficaz para o treinamento de baixa comunicação distribuída, oferecendo um caminho para escalar o treinamento através de clusters geograficamente distribuídos ou mal conectados sem sacrificar a qualidade do modelo.