DeepSpeed ZeRO

DeepSpeed ZeRO (Zero Redundancy Optimizer) elimina a redundância de memória no treinamento distribuído ao dividir (sharding) os estados do otimizador, os gradientes e os parâmetros entre as GPUs. O ZeRO possui três estágios, cada um dividindo mais estado do que o anterior. O DeepSpeed também oferece suporte ao offload para memória de CPU ou NVMe para economias adicionais. Cada estágio e nível de offload adicional reduz a memória de pico, ao custo de mais comunicação entre GPUs.

                  params        grads       opt states
                ┌──────────┐ ┌──────────┐ ┌──────────┐
ZeRO-1          │██████████│ │██████████│ │███░░░░░░░│  GPU 0
                │██████████│ │██████████│ │░░░███░░░░│  GPU 1
                │██████████│ │██████████│ │░░░░░░████│  GPU 2
                └──────────┘ └──────────┘ └──────────┘
                ┌──────────┐ ┌──────────┐ ┌──────────┐
ZeRO-2          │██████████│ │███░░░░░░░│ │███░░░░░░░│  GPU 0
                │██████████│ │░░░███░░░░│ │░░░███░░░░│  GPU 1
                │██████████│ │░░░░░░████│ │░░░░░░████│  GPU 2
                └──────────┘ └──────────┘ └──────────┘
                ┌──────────┐ ┌──────────┐ ┌──────────┐
ZeRO-3          │███░░░░░░░│ │███░░░░░░░│ │███░░░░░░░│  GPU 0
                │░░░███░░░░│ │░░░███░░░░│ │░░░███░░░░│  GPU 1
                │░░░░░░████│ │░░░░░░████│ │░░░░░░████│  GPU 2
                └──────────┘ └──────────┘ └──────────┘
  █ resident    ░ held on another GPU

ZeRO-2 faz o sharding de gradientes e estados do otimizador com menor sobrecarga de comunicação do que o ZeRO-3. Use o ZeRO-3 apenas quando o seu modelo não couber em várias GPUs com o ZeRO-2.

Instalação

Instale o DeepSpeed a partir do PyPI ou instale o Transformers com o extra deepspeed.

pip install deepspeed
pip install transformers[deepspeed]

Se você encontrar erros de instalação relacionados ao CUDA, consulte a documentação de DeepSpeed CUDA. Instalar a partir do código-fonte é a opção mais confiável, pois corresponde exatamente ao seu hardware e inclui recursos que ainda não estão disponíveis na versão do PyPI.

Configurar

O Trainer integra o DeepSpeed por meio do argumento deepspeed, que aceita um arquivo de configuração JSON. Alternativamente, use um arquivo de configuração do Accelerate em vez de TrainingArguments.

TrainingArguments

Arquivo de configuração do Accelerate

Use "auto" na sua configuração para os valores que você deseja que o DeepSpeed preencha a partir de TrainingArguments. Se você quiser especificar explicitamente um valor, certifique-se de usar o mesmo valor para o argumento do DeepSpeed e para TrainingArguments.

Consulte a referência DeepSpeed Configuration JSON para obter a lista completa de opções de configuração do DeepSpeed.

{
    "train_micro_batch_size_per_gpu": "auto",
    "gradient_accumulation_steps": "auto",
    "optimizer": {
        "type": "AdamW",
        "params": { "lr": "auto" }
    },
    "fp16": { "enabled": "auto" }
}

Configuração do DeepSpeed	TrainingArguments
`train_micro_batch_size_per_gpu`	`per_device_train_batch_size`
`gradient_accumulation_steps`	`gradient_accumulation_steps`
`optimizer.params.lr`	`learning_rate`
`fp16.enabled`	`fp16`

Passe a config para o argumento deepspeed.

from transformers import TrainingArguments

args = TrainingArguments(
    deepspeed="path/to/deepspeed_config.json",
    ...
)

# DeepSpeed launcher
deepspeed --num_gpus 4 train.py

# torchrun
torchrun --nproc_per_node 4 train.py

# Accelerate
accelerate launch --num_processes 4 train.py

Etapas do ZeRO

Selecione uma configuração de estágio ZeRO para usar como ponto de partida.

ZeRO-1

ZeRO-2

ZeRO-3

Os campos a seguir são importantes para personalizar o treinamento.

zero_optimization define o estágio do ZeRO.
```
{ "zero_optimization": { "stage": 3 } }
```
Defina os argumentos de tamanho de lote e acúmulo de gradiente como "auto". Se você definir manualmente esses valores de forma que conflitem com TrainingArguments, o treinamento continuará silenciosamente com os valores incorretos.
```
{
    "train_micro_batch_size_per_gpu": "auto",
    "train_batch_size": "auto",
    "gradient_accumulation_steps": "auto",
    "gradient_clipping": "auto"
}
```
bf16 define a precisão do treinamento. Defina-o como "auto" para que ele espelhe a flag bf16 em TrainingArguments.
```
{ "bf16": { "enabled": "auto" } }
```
stage3_gather_16bit_weights_on_model_save executa um all-gather em todos os GPUs antes de salvar, reconstruindo os tensores completos a partir de seus shards. Este é um argumento do ZeRO-3.
```
{
    "zero_optimization": {
        "stage": 3,
        "stage3_gather_16bit_weights_on_model_save": true
    }
}
```
Defina overlap_comm como true para ocultar a latência do all-reduce atrás da passada reversa. allgather_bucket_size e reduce_bucket_size trocam a velocidade de comunicação por memória de GPU. Valores menores usam menos memória, mas deixam a comunicação mais lenta.
```
{
    "zero_optimization": {
        "stage": 2,
        "overlap_comm": true,
        "allgather_bucket_size": 2e8,
        "reduce_bucket_size": 2e8,
        "contiguous_gradients": true
    }
}
```
offload_optimizer transfere o otimizador para a memória da CPU. Para economizar ainda mais memória, transfira também os parâmetros do modelo com offload_param (somente ZeRO-3). Defina pin_memory como true para acelerar as transferências entre CPU e GPU, mas isso trava a RAM, tornando-a indisponível para outros processos.
```
{
    "zero_optimization": {
        "stage": 3,
        "offload_optimizer": { "device": "cpu", "pin_memory": true },
        "offload_param":     { "device": "cpu", "pin_memory": true }
    }
}
```

optimizer e scheduler padrão para o otimizador e o agendador configurados em TrainingArguments. Defina como "auto" para que o DeepSpeed leia os valores de TrainingArguments, a menos que você precise de um otimizador nativo do DeepSpeed, como o LAMB.

{
    "optimizer": {
        "type": "AdamW",
        "params": { "lr": "auto", "betas": "auto", "eps": "auto", "weight_decay": "auto" }
    },
    "scheduler": {
        "type": "WarmupDecayLR",
        "params": { "total_num_steps": "auto", "warmup_min_lr": "auto", "warmup_max_lr": "auto", "warmup_num_steps": "auto" }
    }
}

Se você estiver transferindo o otimizador, defina zero_force_ds_cpu_optimizer como false para usar o otimizador CPU Adam do DeepSpeed.

{
    "zero_force_ds_cpu_optimizer": false
}

Checkpoints

O DeepSpeed salva checkpoints em um formato particionado que não pode ser carregado diretamente com from_pretrained(). Defina load_best_model_at_end() como True para que o Trainer acompanhe e recarregue o melhor checkpoint ao final do treinamento.

from transformers import TrainingArguments, Trainer

args = TrainingArguments(
    deepspeed="ds_config_zero3.json",
    load_best_model_at_end=True,
    ...
)
# after training, save a normal transformers checkpoint
trainer.save_model("./best-model")

Definir save_only_model=True ignora o salvamento do estado completo do otimizador, o que significa que você não pode recarregar o melhor modelo ao final do treinamento. Defina também stage3_gather_16bit_weights_on_model_save: true para reconstruir os pesos completos a partir de seus fragmentos. Isso é necessário para salvar um artefato de modelo consolidado de 16 bits ou um state dict de 16 bits com ZeRO-3. O Transformers levanta um erro quando save_only_model=True é combinado com load_best_model_at_end=True.

Para retomar com diferentes configurações de paralelismo, consulte o guia Universal Checkpointing do DeepSpeed.

Próximas etapas

Leia o capítulo Zero Redundancy Optimizer do The Ultra-Scale Playbook para saber mais sobre como o ZeRO funciona.
Leia os artigos do ZeRO: Memory Optimizations Toward Training Trillion Parameter Models, Democratizing Billion-Scale Model Training e Breaking the GPU Memory Wall for Extreme Scale Deep Learning.

Update on GitHub