Accelerate ND‑Parallel: Guia para Treinamento Multi‑GPU Eficiente

TL;DR

• O Accelerate adicionou ferramentas ND‑Parallel para compor múltiplas estratégias de paralelismo (DP, FSDP, TP, CP) em um único script de treinamento.
• A integração com Axolotl oferece configs de exemplo e campos de alto nível para ativar combinações de paralelismo sem reescrever o código principal.
• Principais parâmetros: ParallelismConfig com dp_replicate_size, dp_shard_size, tp_size (tensor_parallel_size no Axolotl) e sharding de contexto para sequências muito longas.

Contexto e antecedentes

Treinar modelos muito grandes em múltiplas GPUs exige escolher e combinar estratégias de paralelismo para equilibrar memória, computação e comunicação. O paralelismo por dados (DP) replica o modelo e o estado do otimizador em cada dispositivo e divide o minibatch entre réplicas; os gradientes são sincronizados antes da atualização. DP aumenta o throughput, mas exige que o modelo caiba em um único dispositivo. Quando modelos não cabem em uma única GPU, o Fully Sharded Data Parallel (FSDP) faz sharding de pesos, gradientes e estado do otimizador entre dispositivos, reunindo parâmetros de uma camada apenas quando necessário. O Tensor Parallelism (TP) faz sharding permanente dos parâmetros e divide a computação de grandes camadas lineares. O Context Parallelism (CP) faz sharding da sequência de entrada para lidar com contextos extremamente longos. O Accelerate, em colaboração com Axolotl, fornece uma abordagem integrada e configurável para compor essas estratégias, permitindo escalar do treino em um único dispositivo até setups multi‑node, controlando trade‑offs de memória e comunicação. O repositório do Accelerate contém um script de treinamento end‑to‑end e o Axolotl traz configs de exemplo testadas para começar rapidamente.

O que há de novo

• O Accelerate expõe uma forma simples de usar qualquer combinação de estratégias de paralelismo a partir de um único script de treinamento, via ParallelismConfig e plugins.
• Axolotl integra técnicas ND‑Parallel para simplificar fine‑tuning em escala e compor paralelismos com várias técnicas de ajuste fino; basta adicionar campos de paralelismo ao arquivo de config do Axolotl.
• São fornecidos configs de exemplo e scripts que demonstram como configurar dataloaders, otimizadores e loops de treino enquanto se compõem estratégias de paralelismo.

Por que isso importa (impacto para desenvolvedores/empresas)

• Iteração mais rápida: desenvolvedores podem testar combinações de DP, FSDP, TP e CP sem reescrever infraestrutura de treino.
• Escala para modelos maiores: composições de sharding e model parallelizam permitem treinar modelos que não cabem em uma única GPU ou um único nó.
• Controle de custo e recursos: ajustando réplicas, tamanhos de shard e grupos TP, equipes podem trocar memória por custo de comunicação e alinhar configurações à topologia do cluster (NVLink intra‑nó vs Infiniband inter‑nó).
• Reusabilidade: campos de config do Axolotl e ParallelismConfig do Accelerate facilitam reutilizar o mesmo script entre setups pequenos e grandes.

Detalhes técnicos ou Implementação

Estratégias de paralelismo principais

| Estratégia | O que faz | Quando usar | Principais trade‑offs |---|---|---|---| | Paralelismo por Dados (DP) | Replica modelo e estado do otimizador; divide lotes e sincroniza gradientes | Quando o modelo cabe em um dispositivo e deseja‑se maior throughput | Baixo padrão de comunicação (all‑reduce por passo), mas exige que o modelo caiba no dispositivo; dp_replicate_size controla réplicas |Fully Sharded Data Parallel (FSDP) | Faz sharding de pesos, gradientes e estado do otimizador; reúne parâmetros por camada | Quando o modelo é grande demais para um único dispositivo | Troca memória por comunicação (gather por camada); granularidade de gather afeta memória vs comunicação; dp_shard_size configura o grau |Tensor Parallelism (TP) | Faz sharding permanente de grandes camadas lineares entre dispositivos | Quando camadas únicas são grandes demais para gather em FSDP | Redução de memória linear com o tamanho do grupo TP, mas exige sincronização frequente de ativações; melhor dentro de um nó; tp_size (tensor_parallel_size no Axolotl) define o grupo |Context Parallelism (CP) | Faz sharding da entrada pela dimensão de sequência para processar contextos muito longos | Quando treina‑se sequências muito longas (centenas de milhares a milhões de tokens) | Reduz memória de atenção por dispositivo ao dividir a matriz de atenção; cada query ainda necessita dos K e V completos para softmax |

Composição de estratégias e considerações de cluster

• DP é estratégia de topo: dp_replicate_size=2 gera duas réplicas completas, cada qual pode ser novamente sharded por FSDP ou TP.
• O FSDP pode ser aplicado em múltiplos nós tratando todos os dispositivos como um único domínio de sharding (ex.: 4 nós × 8 GPUs = 32 dispositivos). Comunicação se torna mais custosa entre nós, por isso costuma‑se evitar FSDP além de um nó completo.
• TP é mais eficaz dentro dos limites de um único nó devido às necessidades de comunicação de alta largura de banda (NVLink). TP entre nós pode ser ineficiente.
• Para contextos muito longos, a atenção cresce quadraticamente com o comprimento da sequência; o post ilustra que sequências extremamente grandes (ex.: 128k tokens) podem tornar as ativações de atenção enormes, motivando o CP.

Parâmetros no Accelerate e Axolotl

• ParallelismConfig no Accelerate expõe campos como dp_replicate_size, dp_shard_size e tp_size para configurar DP, FSDP e TP e sua composição.
• No Axolotl, campos similares existem (por exemplo tensor_parallel_size, dp_shard_size) de modo que acrescentar técnicas ND‑Parallel a um config é muitas vezes apenas adicionar campos.
• O repositório do Accelerate inclui um FullyShardedDataParallelPlugin preparado para uso com dp_shard_size e exemplos end‑to‑end.

Principais conclusões

• Accelerate + Axolotl permitem compor DP, FSDP, TP e CP num único fluxo de treino.
• Use DP quando o modelo couber em um dispositivo; aumente réplicas com dp_replicate_size para escalar throughput.
• Use FSDP para sharding do estado quando o modelo não couber em um único dispositivo; ajuste a granularidade de gather.
• Use TP para sharding de camadas grandes dentro de um nó quando o gather do FSDP for inviável.
• Use CP para treinar sequências muito longas ao shardear a dimensão de sequência.

Referências

• Blog post: https://huggingface.co/blog/accelerate-nd-parallel