ParaBlock: Communication-Computation Parallel Block Coordinate Federated Learning for Large Language Models

Resumo Técnico: ParaBlock: Comunicação-Computação em Paralelo para Aprendizado Federado de Coordenadas de Blocos para Grandes Modelos de Linguagem

Problema
O Aprendizado Federado (FL) emergiu como um paradigma de preservação de privacidade para o treinamento de modelos através de dados descentralizados. Recentemente, esquemas de descendente de coordenadas de blocos federados (BCD) ganharam tração para o treinamento de modelos de grande escala, permitindo que os clientes atualizem apenas um subconjunto (bloco) dos parâmetros do modelo localmente, em vez de todo o modelo. No entanto, na era dos Grandes Modelos de Linguagem (LLMs), mesmo um único bloco pode conter dezenas a centenas de milhões de parâmetros. Essa escala cria um gargalo crítico: a latência de comunicação necessária para transmitir esses grandes blocos entre os clientes e o servidor torna-se significativa, particularmente para clientes de borda com recursos limitados.

Os métodos tradicionais de federado BCD operam em um design de thread única (comunicação-computação-comunicação), onde os clientes devem esperar que a transmissão do modelo seja concluída antes de iniciar a próxima etapa de computação. Embora esse tempo de comunicação fosse negligenciável para modelos menores, agora ele prejudica severamente a escalabilidade e a eficiência do ajuste fino (fine-tuning) federado de LLMs, causando um tempo de ociosidade substancial para os clientes.

Metodologia: ParaBlock
Para abordar essa latência, os autores propõem o ParaBlock, um novo método de BCD federado que introduz uma arquitetura de comunicação-computação em paralelo. Diferente das abordagens padrão de thread única, o ParaBlock estabelece duas threads paralelas nos clientes locais:

1. Thread de Computação: Realiza as atualizações de coordenadas de blocos locais (ajuste fino do bloco atribuído) usando dados locais.
2. Thread de Comunicação: Sincroniza-se com o servidor para enviar as atualizações locais da rodada anterior e receber as atualizações globais agregadas.

Mecanismos Chave:

• Sobreposição (Overlapping): Os clientes executam a computação local enquanto comunicam simultaneamente com o servidor. Isso elimina o tempo de ociosidade tipicamente causado pela espera da transmissão de rede.
• Obsolescência e Correção: Como a thread de comunicação possui um atraso de uma rodada (os clientes recebem atualizações globais da rodada $t-1$ enquanto computam para a rodada $t$), os parâmetros do modelo local podem tornar-se inconsistentes com o modelo global. O ParaBlock aborda isso via um mecanismo de correção. Quando um cliente atualiza seu modelo local para a próxima rodada, ele não apenas aplica a nova atualização local ao bloco atual, mas também corrige os parâmetros do bloco anterior usando a atualização global recebida ($\Delta_{t-1}$) para substituir a atualização local obsoleta ($\Delta^i_{t-1}$).
• Consistência Recursiva: A análise teórica demonstra que essa correção garante que, após uma rodada global, os parâmetros do modelo local para o bloco previamente atualizado alinhem-se exatamente com os parâmetros do modelo global, mitigando efetivamente o impacto do atraso de uma rodada.

Principais Contribuições

1. Arquitetura Inovadora: A proposta do ParaBlock, que substitui o design tradicional de FL de thread única por um esquema paralelo de duas threads especificamente projetado para reduzir a latência de comunicação no ajuste fino federado de LLMs.
2. Garantias Teóricas: Os autores fornecem uma análise de convergência rigorosa para objetivos não convexos. Eles provam que o ParaBlock atinge uma taxa de convergência de $O(1/\sqrt{T})$, que é consistente com os métodos padrão de descendente de coordenadas de blocos federados. Isso confirma que os ganhos de eficiência do paralelismo não comprometem o desempenho de convergência teórica.
3. Validação Empírica: Experimentos extensos foram realizados no ajuste fino dos modelos Llama 3-8B e Llama 3.2-3B em tarefas de seguimento de instruções gerais (Alpaca-GPT4) e raciocínio matemático (MathInstruct/GSM8K).

Resultados

• Desempenho: O ParaBlock mantém um forte desempenho em tarefas de downstream, alcançando pontuações comparáveis ou superiores às linhas de base de BCD federado padrão (ex: FedBCD) e superando significativamente vários métodos de Aprendizado de Ajuste Fino de Parâmetros Eficientes (PEFT) baseados em FL (ex: FedIT, FLoRA) em tarefas de seguimento de instruções e raciocínio matemático.
• Eficiência: O método reduz significativamente o tempo de execução (wall-clock runtime). Em cenários com baixa largura de banda de rede (50M/s) e pequenos tamanhos de lote (batch sizes), o ParaBlock reduziu o tempo de execução em mais de 30% em comparação com linhas de base de coordenação de blocos como o FedCyBGD. Comparado a métodos baseados em LoRA como o FedIT, o ParaBlock alcançou uma redução de tempo de execução de aproximadamente 40% na maioria das configurações.
• Robustez: Estudos de ablação mostraram que o método permanece eficaz sob distribuições de dados heterogêneas (não-IID) e em diferentes arquiteturas de modelos (incluindo Qwen-2.5). O método também demonstrou compatibilidade com configurações de participação parcial em cenários de cross-silo.

Significância
O artigo afirma que o ParaBlock oferece uma solução simples e eficaz para o gargalo de comunicação inerente ao BCD federado para LLMs. Ao sobrepor comunicação e computação, ele permite o implante escalável do ajuste fino federado para grandes modelos em clientes com recursos limitados, sem sacrificar as garantias de convergência ou a utilidade do modelo. O trabalho destaca que a redução da latência de comunicação através do paralelismo arquitetural é uma estratégia viável e necessária para a próxima geração de treinamento de LLM federado, distinta e ortogonal a técnicas existentes como quantização ou poda (pruning).