FSDP ou DeepSpeed ZeRO: qual se adapta à minha configuração de treino?
Tradução automática
Este artigo foi traduzido automaticamente a partir da versão original em inglês.
Comece com FSDP quando uma infraestrutura de treino em PyTorch precisar de particionamento distribuído do estado. Teste DeepSpeed se precisar das suas etapas ZeRO selecionáveis ou das funções de transferência para CPU e NVMe. Nenhum framework oferece sempre o melhor desempenho; compare a configuração necessária no seu cluster.
Ambos podem distribuir os parâmetros, os gradientes e o estado do otimizador entre GPUs. As suas APIs, políticas e formas de gerir checkpoints são diferentes.
O particionamento completo segue o mesmo processo básico
Com o particionamento completo dos parâmetros, cada GPU armazena uma parte do modelo. Antes de calcular um grupo de camadas, o framework reúne os parâmetros de que esse grupo precisa. Depois efetua o cálculo, liberta os parâmetros reunidos segundo a sua política de reparticionamento e combina e distribui os gradientes através de reduce-scatter.
As decisões sobre agrupamento e pré-carregamento controlam a memória temporária e a sobreposição da comunicação com o cálculo. Particionar cada pequena operação separadamente pode fazer com que a comunicação domine o trabalho; reunir grupos muito grandes pode criar um pico temporário de memória.
| Requisito | Candidato a testar |
|---|---|
| APIs nativas de PyTorch e estado de tensores distribuídos | FSDP2 |
| Particionamento do estado do otimizador com parâmetros ainda replicados | DeepSpeed ZeRO-1 ou ZeRO-2 |
| Particionamento completo do estado | FSDP ou ZeRO-3 |
| Transferência do estado para NVMe | DeepSpeed ZeRO-Infinity |
O tutorial de FSDP2 usa fully_shard e parâmetros DTensor. Também disponibiliza políticas de reparticionamento: manter os parâmetros após a passagem em frente consome mais memória, mas reduz as operações de reunião posteriores. Crie o otimizador depois de aplicar as transformações de particionamento necessárias.
Compare um passo de treino completo
Mantenha constantes o modelo, os dados, a precisão, o comprimento da sequência, o tamanho efetivo do lote e a topologia das GPUs. Meça o pico de memória alocada, o tempo de cada passo e a comunicação. Verifique se um checkpoint guardado consegue restaurar o estado do modelo, o estado do otimizador e o progresso de treino pretendido.
Hugging Face Accelerate pode selecionar FSDP ou DeepSpeed através de configurações de lançamento para código preparado para essas integrações. Isso não significa que qualquer script de treino possa mudar sem ajustes, sobretudo na gestão de checkpoints, na transferência de estado e no tratamento de otimizadores personalizados.
A comparação de FSDP no LLM Engineering Guide relaciona estas escolhas com a memória usada pelo estado de treino.