Etapas do DeepSpeed ZeRO: que estado do treino é repartido?
Tradução automática
Este artigo foi traduzido automaticamente a partir da versão original em inglês.
O ZeRO-1 reparte o estado do otimizador pelas GPUs em paralelismo de dados; o ZeRO-2 reparte também os gradientes, e o ZeRO-3, além disso, os parâmetros do modelo. Teste a etapa mais baixa que permita à carga de trabalho caber na memória com um débito aceitável. Uma etapa superior poupa mais memória do estado do modelo, mas pode exigir mais comunicação.
Estas etapas alteram o local onde o estado do treino é armazenado. Não reduzem automaticamente a memória das ativações.
Um exemplo de memória com pressupostos explícitos
Suponha que o treino usa pesos e gradientes em FP16, juntamente com pesos mestres do Adam em FP32 e dois momentos do otimizador. Isto atribui 2 bytes aos pesos, 2 aos gradientes e 12 ao estado do otimizador por parâmetro: 16 bytes no total.
Para 7.5 mil milhões de parâmetros em oito GPUs, o cálculo do estado do modelo é:
| Etapa | Estado repartido pelas GPUs | GB do estado do modelo por GPU |
|---|---|---|
| Paralelismo de dados convencional | Nenhum | 120 |
| ZeRO-1 | Estado do otimizador | 41.25 |
| ZeRO-2 | Estado do otimizador e gradientes | 28.125 |
| ZeRO-3 | Os três componentes | 15 |
São GB decimais. Excluem as ativações, os parâmetros reunidos temporariamente, a sobrecarga do alocador e os buffers do runtime. Uma stack de treino que armazene os dados com precisões diferentes exige outros cálculos.
O artigo sobre ZeRO explica as partições do estado e analisa a sua comunicação. O ZeRO-3 reconstrói os parâmetros necessários através de operações all-gather; a análise da repartição completa no artigo chega a cerca de 1.5 vezes o volume de comunicação do paralelismo de dados convencional. Essa proporção resulta de uma análise sob os seus pressupostos, e não é uma garantia para todas as redes e configurações.
Escolher com base no perfil de memória
Se os parâmetros couberem na memória e os estados do otimizador causarem a falha, teste a etapa 1 ou 2. Se os parâmetros totalmente replicados não couberem, teste a etapa 3. Se as ativações dominarem, combine a etapa escolhida com checkpointing de ativações ou com um lote de treino mais pequeno.
O DeepSpeed também permite transferir o estado para CPU e NVMe. Consulte a configuração atual do ZeRO e meça o custo de transferência no sistema e no armazenamento reais. A capacidade adicional pode ajudar, mas as transferências lentas podem reduzir o débito.
A secção sobre ZeRO do LLM Engineering Guide apresenta um cálculo executável para outras quantidades de parâmetros e grupos de GPU.