Papers
ZeRO: Memory Optimizations Toward Training Trillion Parameter Models
Samyam Rajbhandari, Jeff Rasley, Olatunji Ruwase, Yuxiong He · SC 2020 (arXiv 1910.02054), 2020
O problema: para onde vai a memória durante o treino
Os autores partem de um paradoxo: um GPT-2 de 1,5 bilhão de parâmetros ocupa só 3 GB de pesos em precisão de 16 bits, mas não treina em uma GPU de 32 GB com TensorFlow ou PyTorch. Eles dividem o consumo em dois grupos. O primeiro são os model states: estados do otimizador, gradientes e parâmetros. O segundo são os residual states: ativações, buffers temporários e memória fragmentada. Para o Adam em treino de precisão mista, o paper contabiliza, para Ψ parâmetros, 2Ψ bytes de parâmetros fp16 e 2Ψ bytes de gradientes fp16. Além disso, o otimizador mantém uma cópia fp32 dos parâmetros, o momento e a variância, 4Ψ bytes cada, que o paper agrupa como KΨ com K = 12. O total é 2Ψ + 2Ψ + KΨ = 16Ψ bytes, ou seja, ao menos 24 GB para o modelo de 1,5B, e os estados do otimizador costumam ser a maior parcela. Entre os resíduos, o paper cita que as ativações desse mesmo modelo, com sequência de 1K e batch 32, pedem cerca de 60 GB. O checkpointing de ativações reduz isso para cerca de 8 GB ao custo de 33% de recomputação. Um buffer fundido fp32 para 1,5B de parâmetros ocupa 6 GB, e a fragmentação, que faz um pedido de memória falhar quando não há um bloco contíguo grande o bastante, chegou a causar falta de memória com mais de 30% ainda livre em casos extremos.
Os limites de data parallel e de model parallel
No data parallelism (DP) básico, cada dispositivo guarda uma réplica completa do modelo; o mini-batch é dividido, cada processo calcula forward e backward sobre sua fatia e os gradientes são promediados. O resultado é boa eficiência de computação e comunicação, mas nenhuma redução de memória por dispositivo: segundo o paper, o DP básico esgota GPUs de 32 GB com modelos acima de 1,4B parâmetros. O model parallelism (MP) corta cada camada verticalmente entre dispositivos e, assim, particiona os estados; em troca, a computação fica fina demais e exige muita comunicação entre camadas. Funciona bem dentro de um nó, onde a banda entre GPUs é alta, e degrada fora dele. Os autores testaram um modelo de 40B com Megatron-LM em dois nós DGX-2 e observaram cerca de 5 Tflops por GPU V100, menos de 5% do pico. Mais adiante citam a queda de banda de 300 GB/s por link (NVSwitch) para 12,5 GB/s (Infiniband EDR) ao sair do nó. O pipeline parallelism, que fatia o modelo horizontalmente por camadas, também tem custos: o GPipe pede um batch proporcional ao número de partições, e o PipeDream guarda cópias defasadas de parâmetros. A partir daí o paper tira três observações. O DP escala melhor que o MP. O DP desperdiça memória por replicar tudo. E nem todo estado é necessário o tempo todo: os parâmetros de uma camada só são exigidos durante o forward e o backward dela.
ZeRO-DP: três estágios de particionamento dos model states
A ideia central do ZeRO-DP é eliminar a redundância do DP particionando os model states entre os Nd processos de dados em vez de replicá-los, mantendo a granularidade de computação do DP. Os três estágios são cumulativos. No primeiro (Pos), os estados do otimizador são divididos em Nd partes iguais; cada processo guarda e atualiza apenas a sua e, ao fim do passo, um all-gather distribui os parâmetros atualizados. O consumo cai de 16Ψ para 4Ψ + 12Ψ/Nd, perto de 4Ψ para Nd grande: redução de até 4x. No segundo (Pos+g), cada gradiente é reduzido apenas no processo dono daquela partição, via reduce-scatter com buckets, e depois liberado, chegando a cerca de 2Ψ, redução de até 8x. No terceiro (Pos+g+p), os parâmetros também são particionados e quem precisa de uma partição alheia a recebe por broadcast; a memória cai de 16Ψ para 16Ψ/Nd, linear no grau de DP. No exemplo do paper, um modelo de 7,5B com Nd = 64 precisa de 120 GB por dispositivo no DP padrão, cerca de 31,4 GB com Pos, 16,6 GB com Pos+g e cerca de 1,9 GB com Pos+g+p. Com Nd = 64, os estágios comportam modelos de até 7,5B, 14B e 128B; com Nd = 1024 e todos os estágios, o paper afirma que cabem 1 trilhão de parâmetros. Para esse trilhão, o estado total seria de cerca de 16 TB, ou 16 GB por GPU em 1024 GPUs.
O custo de comunicação: de graça até Pos+g, 1,5x com Pos+g+p
A pergunta natural é se o ZeRO troca memória por comunicação. O paper responde medindo o volume de dados movimentado por processo em cada passo. No DP padrão, o all-reduce dos gradientes é feito em duas etapas, reduce-scatter e all-gather, cada uma movendo Ψ elementos, totalizando 2Ψ por passo. Com Pos+g, o gradiente só precisa de um reduce-scatter (volume Ψ), e depois cada processo atualiza sua partição e um all-gather reúne os parâmetros atualizados (outro Ψ). O total continua 2Ψ, igual ao DP padrão, mesmo com redução de memória de até 8x. Com Pos+g+p, o all-gather dos parâmetros precisa acontecer duas vezes: uma espalhada ao longo do forward, com os parâmetros descartados após o uso, e outra no backward, em ordem inversa. Somado ao reduce-scatter dos gradientes, o total é 3Ψ, isto é, 1,5x o baseline, em troca de memória proporcional a Nd. O paper também analisa o ZeRO-R: no Megatron-LM com checkpointing de ativações, cada bloco transformer faz seis all-reduces (dois no forward, dois na recomputação e dois no backward) e move 12 × seq × hidden, e o particionamento de ativações (Pa) adiciona apenas um all-gather de seq × hidden por bloco, menos de 10% do volume original do MP. Como Pa reduz a memória de ativações na proporção do grau de MP, o batch pode crescer e o volume de comunicação do DP, inversamente proporcional ao batch, pode cair.
ZeRO-R, resultados e legado
Para os residual states, o ZeRO-R tem três peças. A primeira é o particionamento de checkpoints de ativação entre as GPUs do MP, que no MP convencional ficam replicados; no exemplo do paper (100B, batch 32, sequência 1024, MP 16), o custo de cerca de 33 GB por GPU cai para cerca de 2 GB, e esses 2 GB ainda podem ser descarregados para a CPU. A segunda são buffers de tamanho constante, que evitam que o buffer fundido cresça com o modelo (para 3B parâmetros, um buffer de 32 bits pediria 12 GB). A terceira é a desfragmentação durante o treino, com blocos contíguos pré-alocados para checkpoints e gradientes. Na implementação, chamada ZeRO-100B (Pos+g mais ZeRO-R, em PyTorch), o paper reporta, em 400 GPUs V100, modelos de até 170B parâmetros combinados com MP, acima de 15 Petaflops agregados e mais de 38 Tflops por GPU em 100B. Reporta também escala super-linear de 64 a 400 GPUs em um modelo de 60B, porque menos memória de estados permite batches maiores por GPU. Sem MP, treinam-se modelos de até 13B, enquanto o DDP do PyTorch esgota a memória com 1,4B. O Pos+g+p ainda era trabalho planejado, e os autores admitem que um trilhão de parâmetros esbarraria em falta de computação (mais de um ano de treino). O paper cita o Turing-NLG, de 17B, como modelo viabilizado pelo ZeRO, e o apresenta como parte da biblioteca de código aberto DeepSpeed.
Veja as ideias em ação
Pratique na escola de System Design os desenhos que usam este paper.