Resumo Técnico: Particionamento Adaptativo de Dados para Aprendizado Federado e Distribuído Eficiente em Energia em Sistemas Heterogêneos

1. Declaração do Problema

O treinamento de modelos distribuídos em sistemas heterogêneos é significativamente dificultado pelas variações na capacidade de processamento, comportamento de memória e condições operacionais entre os dispositivos participantes. Quando uma divisão fixa e igual de dados de treinamento é aplicada, dispositivos mais rápidos terminam seu trabalho precocemente e permanecem ociosos enquanto esperam por dispositivos "atrasados" (stragglers) mais lentos nos pontos de sincronização. Esse efeito de atraso não apenas aumenta o makespan total do treinamento, mas também eleva o consumo de energia dos dispositivos, uma vez que os participantes sincronizados devem permanecer ativos até que o progresso coletivo possa continuar.

Abordagens existentes para mitigar esse problema frequentemente dependem da seleção de participantes (excluindo dispositivos lentos), agregação assíncrona, mudanças nos métodos de otimização ou adaptação do modelo. Embora sistemas recentes tenham explorado a redistribuição de dados de treinamento, muitos dependem de controladores aprendidos, busca prévia, configurações específicas de hardware ou alterações na semântica de sincronização e nos conjuntos de participação. Permanece uma lacuna para um mecanismo de malha fechada comum que possa operar através de diferentes fronteiras de sincronização e propriedade de dados (como Aprendizado Federado e Paralelismo de Dados Distribuídos), mantendo a configuração original, o modelo e as regras de sincronização.

2. Metodologia

Os autores propõem um controlador de alocação leve, baseado em medições, que realoca um orçamento fixo de amostras nos limites das rodadas com base no tempo de treinamento observado e na energia medida por amostra. O controlador opera sem modificar o modelo, a função de perda, o otimizador ou o conjunto de dispositivos participantes.

Regimes de Execução

O estudo avalia o controlador através de três regimes de execução distintos:

1. Média de Modelo Periódica (PMA - Periodic Model Averaging): Trabalhadores treinam réplicas de modelos independentes em dados compartilhados dentro de um domínio confiável e sincronizam apenas nos limites das rodadas.
2. Paralelismo de Dados Distribuídos (DDP - Distributed Data Parallelism): Os ranks sincronizam gradientes após cada passo do otimizador. Este regime requer uma estratégia específica de dimensionamento de lote (batch-sizing) para lidar com cargas de trabalho desiguais por rank sem introduzir preenchimentos (padding) ou lotes fictícios.
3. Aprendizado Federado Hierárquico com DDP intra-silo (HFL-DDP): Os dispositivos são agrupados em "silos" fixos (domínios de propriedade de dados). Os ranks dentro de um silo usam DDP, enquanto os modelos dos silos são agregados no limite da rodada federada. A reassignação é restrita aos ranks dentro do mesmo silo para respeitar a propriedade dos dados.

O Controlador Adaptativo

O controlador funciona da seguinte forma:

• Medição: Ao final de cada rodada $r$, cada participante $i$ reporta seu tempo de processamento ($t_i^{(r)}$), consumo de energia ($e_i^{(r)}$) e contagem de amostras ($n_i^{(r)}$).
• Cálculo de Custo: Custos por amostra são derivados como $\tau_i^{(r)} = t_i^{(r)} / n_i^{(r)}$ e $\epsilon_i^{(r)} = e_i^{(r)} / n_i^{(r)}$. Estes são suavizados usando uma Média Móvel Exponencial (EMA) para reduzir o ruído.
• Divisão Proporcional Ponderada: O controlador calcula um custo combinado adimensional $c_i^{(r)}$ usando pesos selecionados pelo usuário $w_t$ (tempo) e $w_e$ (energia).
  • $c_i^{(r)} = w_t \frac{\hat{\tau}_i^{(r)}}{\text{mediana}(\hat{\tau})} + w_e \frac{\hat{\epsilon}_i^{(r)}}{\text{mediana}(\hat{\epsilon})}$
• Alocação: Os participantes são atribuídos amostras inversamente proporcionais ao seu custo. Um mecanismo de "banda morta" (deadband) evita a realocação desnecessária se as diferenças de custo forem insignificantes.
• Especificidades de DDP: Para DDP, o controlador garante que todos os ranks executem o mesmo número de passos do otimizador por rodada, construindo cronogramas de tamanho de lote por passo que redistribuem as amostras locais, preservando o objetivo da média do gradiente da amostra.

Configuração Experimental

A avaliação foi conduzida em um testbed de Sistema-on-Chip (SoC) de onze dispositivos heterogêneos apresentando arquiteturas Arm, RISC-V e x86. Três tarefas de aprendizado supervisionado foram utilizadas: MNIST (classificação de imagens), ERA5 (regressão de séries temporais) e IMDB (classificação de sentimento de texto). O estudo comparou as políticas adaptativas contra:

• Alocação igual: Divisão uniforme de amostras.
• Perfilamento estático: Uma medição única seguida de alocação fixa.
• Seleção tipo Oort: Um baseline que exclui participantes lentos.
• FedAsync: PMA assíncrono.
• Dimensionamento de lote EMA: Um baseline de vazão (throughput) por passo para DDP.

3. Principais Contribuições

O artigo apresenta quatro contribuições primárias:

1. Um Controlador Baseado em Medição: Um mecanismo leve que realoca um orçamento fixo de amostras com base no tempo e energia observados por amostra, preservando o conjunto completo de participação e as semânticas de sincronização.
2. Design Unificado: Um design de alocação comum aplicável a PMA, DDP e HFL-DDP, com restrições específicas para fronteiras de propriedade de dados em configurações hierárquicas.
3. Objetivo Consciente de Energia: Uma função objetivo com pesos ajustáveis que expõe a interação entre energia por amostra, makespan de treinamento e energia do dispositivo sincronizado, desafiando a suposição de que maximizar o peso de energia sempre minimiza a energia total.
4. Avaliação Abrangente: Um estudo replicado através de três tarefas de aprendizado e três regimes de sincronização em um testbed heterogêneo, conectando mudanças de alocação a desequilíbrio de computação, exposição à sincronização, makespan e energia medida.

4. Resultados

Experimentos no testbed heterogêneo produziram os seguintes achados:

• Redução no Makespan e na Energia: O particionamento adaptativo (especificamente as políticas "Apenas Tempo" e "Baixa Energia") reduziu consistentemente o makespan de treinamento e a energia medida dos dispositivos em comparação com a Alocação Igual. No PMA, a política "Apenas Tempo" reduziu o makespan mediano em aproximadamente 71–80% e a energia em 28–40% entre as tarefas.
• Mitigação de Stragglers: O controlador reduziu significativamente o desequilíbrio de computação. No DDP, o coeficiente de variação para o tempo de passagem direta (forward-pass) caiu de ~0,8–1,1 (Igual) para ~0,02–0,18 (adaptativo), correlacionando-se diretamente com a redução da espera de sincronização.
• Convergência: As políticas adaptativas atingiram regiões de baixo objetivo (perda/MSE de teste) mais rapidamente do que a Alocação Igual, muitas vezes alcançando uma qualidade final comparável em muito menos tempo.
• Resposta de Energia Não Monotônica: Aumentar o peso da energia ($w_e$) nem sempre reduziu a energia total. Atribuir mais trabalho a dispositivos energeticamente eficientes, porém mais lentos, pode estender a janela de execução sincronizada, aumentando a energia acumulada de todos os participantes. Um pequeno peso de energia ($w_e = 0,25$) frequentemente proporcionou um melhor compromisso do que "Apenas Energia" ($w_e = 1$).
• Overhead: O cálculo do escalonamento foi desprezível, consumindo menos de 0,02% da duração da rodada.

5. Significância e Alegações

O artigo afirma que o particionamento de dados baseado em medição fornece um método prático para mitigar stragglers em diferentes regimes de treinamento distribuído sem alterar a configuração de aprendizado ou as regras de sincronização.

• Praticidade: O método não requer perfilamento prévio, políticas aprendidas ou mudanças na arquitetura do otimizador ou do modelo. Ele opera como um controlador de malha fechada usando telemetria de tempo de execução.
• Otimização Dependente de Contexto: Os autores enfatizam que o equilíbrio ideal entre tempo e energia depende da tarefa específica e do regime de execução. Embora "Apenas Tempo" seja o padrão mais confiável para minimizar o makespan, um pequeno peso de energia pode ser benéfico em contextos específicos (ex: HFL-DDP em certas tarefas).
• Distinção de Trabalhos Anteriores: Diferente dos métodos de seleção de participantes (ex: Oort) que excluem dispositivos, esta abordagem retém todos os participantes provisionados. Diferente dos métodos assíncronos, ela preserva as semânticas de sincronização do regime alvo.
• Limitações: Os autores notam modestamente que os resultados são específicos para o testbed de onze dispositivos avaliado e condições de rede fixas. Eles não alegam superioridade universal sobre todos os métodos existentes, mas demonstram que um simples realocamento de orçamento fixo baseado em medições pode superar divisões iguais estáticas e o perfilamento de uma única vez em ambientes heterogêneos.