Conceitos
ZeRO: particionar estados do treino entre GPUs
Intuição
Imagine uma equipe de 64 pessoas em que cada uma carrega uma cópia completa da mesma enciclopédia, embora cada uma só consulte uma página por vez. Em vez de todos carregarem tudo, cada pessoa pode guardar só um pedaço e pedir aos colegas a página que precisar na hora de usá-la. Treinar um modelo grande com várias GPUs sofre do mesmo desperdício: cada GPU guarda uma cópia inteira de coisas que ela só usa em parte, e o espaço acaba antes de o modelo caber.
Como funciona
No treino com o otimizador Adam em precisão mista, o paper contabiliza, para Ψ parâmetros, 2Ψ bytes de parâmetros fp16, 2Ψ bytes de gradientes fp16 e KΨ bytes de estados do otimizador (cópia fp32 dos parâmetros, momento e variância, com K = 12), num total de 16Ψ bytes; os estados do otimizador costumam ser a maior parcela. No data parallelism (DP) básico, cada dispositivo guarda uma réplica completa de tudo isso, de modo que o DP é eficiente em computação e comunicação, mas não reduz a memória por GPU. O ZeRO-DP elimina essa redundância particionando os model states entre os Nd processos de dados, em três estágios cumulativos. No primeiro (Pos), os estados do otimizador são divididos em Nd partes, cada processo atualiza só a sua e um all-gather distribui os parâmetros atualizados, o que reduz a memória para 4Ψ + 12Ψ/Nd (até 4x). No segundo (Pos+g), cada gradiente é reduzido só no processo dono da partição, via reduce-scatter, e depois liberado, chegando a cerca de 2Ψ (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, de modo que a memória cai para 16Ψ/Nd, linear no grau de DP. Quanto à comunicação, o DP padrão move 2Ψ por passo (reduce-scatter mais all-gather, cada um com Ψ), e Pos+g mantém os mesmos 2Ψ; já Pos+g+p precisa de dois all-gathers de parâmetros (forward e backward) mais o reduce-scatter de gradientes, somando 3Ψ, ou seja, 1,5x o baseline. Os residual states (ativações, buffers, fragmentação) são tratados à parte pelo ZeRO-R, que, por exemplo, particiona os checkpoints de ativação (usados no checkpointing de ativações) entre as GPUs do model parallelism.
Na prática
No exemplo do paper, um modelo de 7,5B parâmetros com Nd = 64 ocupa 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. A implementação avaliada, o ZeRO-100B (Pos+g mais ZeRO-R, em PyTorch), treinou em 400 GPUs V100 modelos de até 170B parâmetros combinados com model parallelism, com mais de 38 Tflops por GPU em 100B; sem model parallelism, treinou modelos de até 13B, enquanto o DDP do PyTorch esgota a memória com 1,4B. O paper 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, e apresenta o ZeRO como parte da biblioteca de código aberto DeepSpeed. O estágio Pos+g+p ainda era trabalho planejado na época, e os autores reconhecem que treinar 1 trilhão de parâmetros esbarraria em falta de computação.
Use em um desenho
Pratique na escola de System Design os casos que dependem deste conceito.