# DeepSpeed ZeRO

[DeepSpeed](https://www.deepspeed.ai/) 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.

```text
                  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 particionamento (sharding) dos gradientes e dos 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`.

```shell
pip install deepspeed
pip install transformers[deepspeed]
```

Se você encontrar erros de instalação relacionados ao CUDA, consulte a documentação de [DeepSpeed CUDA](./debugging#deepspeed-cuda). [Instalar a partir do código-fonte](https://www.deepspeed.ai/tutorials/advanced-install/#install-deepspeed-from-source) é 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](/docs/transformers/main/pt/main_classes/trainer#transformers.Trainer) integra o DeepSpeed por meio do argumento [deepspeed](/docs/transformers/main/pt/main_classes/trainer#transformers.TrainingArguments.deepspeed), que aceita um arquivo de configuração JSON. Alternativamente, use um [arquivo de configuração do Accelerate](./accelerate#accelerate-config-file) em vez de [TrainingArguments](/docs/transformers/main/pt/main_classes/trainer#transformers.TrainingArguments).

Use `"auto"` na sua configuração para os valores que você deseja que o DeepSpeed preencha a partir de [TrainingArguments](/docs/transformers/main/pt/main_classes/trainer#transformers.TrainingArguments). Se você quiser especificar explicitamente um valor, certifique-se de usar o *mesmo* valor para o argumento do DeepSpeed e para [TrainingArguments](/docs/transformers/main/pt/main_classes/trainer#transformers.TrainingArguments).

> [!NOTE]
> Consulte a referência [DeepSpeed Configuration JSON](https://www.deepspeed.ai/docs/config-json/) para obter a lista completa de opções de configuração do DeepSpeed.

```json
{
    "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](/docs/transformers/main/pt/main_classes/trainer#transformers.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`.

```py
from transformers import TrainingArguments

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

```cli
# DeepSpeed launcher
deepspeed --num_gpus 4 train.py

# torchrun
torchrun --nproc_per_node 4 train.py

# Accelerate
accelerate launch --num_processes 4 train.py
```

> [!NOTE]
> O Accelerate ignora o argumento `deepspeed` em [TrainingArguments](/docs/transformers/main/pt/main_classes/trainer#transformers.TrainingArguments).

Execute o comando [accelerate config](https://huggingface.co/docs/accelerate/en/package_reference/cli#accelerate-config) e responda às perguntas sobre seu hardware e configuração de treinamento para criar um arquivo `default_config.yaml` no seu cache.

```yaml
distributed_type: DEEPSPEED
deepspeed_config:
  deepspeed_config_file: path/to/ds_config.json
machine_rank: 0
num_machines: 1
num_processes: 4
```

Execute [accelerate launch](https://huggingface.co/docs/accelerate/en/package_reference/cli#accelerate-launch) com um script baseado em [Trainer](/docs/transformers/main/pt/main_classes/trainer#transformers.Trainer).

```shell
accelerate launch --config_file deepspeed_config.yaml train.py
```

## Etapas do ZeRO

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

```json
{
    "bf16": { "enabled": "auto" },
    "zero_optimization": { "stage": 1 },
    "gradient_clipping": "auto",
    "train_micro_batch_size_per_gpu": "auto",
    "train_batch_size": "auto",
    "gradient_accumulation_steps": "auto"
}
```

```json
{
    "bf16": { "enabled": "auto" },
    "zero_optimization": {
        "stage": 2,
        "overlap_comm": true,
        "allgather_bucket_size": 2e8,
        "reduce_bucket_size": 2e8,
        "contiguous_gradients": true
    },
    "gradient_clipping": "auto",
    "train_micro_batch_size_per_gpu": "auto",
    "train_batch_size": "auto",
    "gradient_accumulation_steps": "auto"
}
```

> [!WARNING]
> ZeRO-3 faz o particionamento dos parâmetros durante a inicialização. Você deve instanciar [TrainingArguments](/docs/transformers/main/pt/main_classes/trainer#transformers.TrainingArguments) antes de carregar seu modelo — se o modelo já estiver em cada GPU antes da configuração do DeepSpeed, nenhuma memória será economizada.

```json
{
    "bf16": { "enabled": "auto" },
    "zero_optimization": {
        "stage": 3,
        "overlap_comm": true,
        "contiguous_gradients": true,
        "reduce_bucket_size": "auto",
        "stage3_prefetch_bucket_size": "auto",
        "stage3_param_persistence_threshold": "auto",
        "stage3_gather_16bit_weights_on_model_save": true,
        "offload_optimizer": { "device": "cpu", "pin_memory": true },
        "offload_param":     { "device": "cpu", "pin_memory": true }
    },
    "gradient_clipping": "auto",
    "train_micro_batch_size_per_gpu": "auto",
    "train_batch_size": "auto",
    "gradient_accumulation_steps": "auto"
}
```

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

- `zero_optimization` define o estágio do ZeRO.

    ```json
    { "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](/docs/transformers/main/pt/main_classes/trainer#transformers.TrainingArguments), o treinamento continuará silenciosamente com os valores incorretos.

    ```json
    {
        "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](/docs/transformers/main/pt/main_classes/trainer#transformers.TrainingArguments).

    ```json
    { "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.

    ```json
    {
        "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.

    ```json
    {
        "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.

    ```json
    {
        "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](/docs/transformers/main/pt/main_classes/trainer#transformers.TrainingArguments). Defina como `"auto"` para que o DeepSpeed leia os valores de [TrainingArguments](/docs/transformers/main/pt/main_classes/trainer#transformers.TrainingArguments), a menos que você precise de um otimizador nativo do DeepSpeed, como o LAMB.

    ```json
    {
        "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.

    ```json
    {
        "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()](/docs/transformers/main/pt/main_classes/model#transformers.PreTrainedModel.from_pretrained). Defina `load_best_model_at_end()` como `True` para que o Trainer acompanhe e recarregue o melhor checkpoint ao final do treinamento.

```py
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 shards. 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`.

> [!TIP]
> Para retomar com diferentes configurações de paralelismo, consulte o guia [Universal Checkpointing](https://www.deepspeed.ai/tutorials/universal-checkpointing) do DeepSpeed.

## Próximas etapas

- Leia o capítulo [Zero Redundancy Optimizer](https://nanotron-ultrascale-playbook.static.hf.space/index.html#zero_redundancy_optimizer_(zero)) 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](https://hf.co/papers/1910.02054), [Democratizing Billion-Scale Model Training](https://hf.co/papers/2101.06840) e [Breaking the GPU Memory Wall for Extreme Scale Deep Learning](https://hf.co/papers/2104.07857).