# Treinamento Multi-GPU

LeRobot treina em múltiplas GPUs por meio do [Hugging Face Accelerate](https://huggingface.co/docs/accelerate). Três layouts de paralelismo de dados são suportados:

| Layout   | O que ele faz                                                  | Config                                                  |
| -------- | ------------------------------------------------------------- | ------------------------------------------------------- |
| **DDP**  | Replica o modelo completo em cada GPU                        | padrão em qualquer inicialização multi-GPU                         |
| **FSDP** | Fragmenta parâmetros, gradientes e estado do otimizador entre GPUs | `--parallelism.dp_shard=N`                              |
| **HSDP** | Shards dentro de grupos de GPUs, replicado entre grupos        | `--parallelism.dp_replicate=R --parallelism.dp_shard=S` |

## Instalação

`accelerate` está incluído no extra `training`:

```bash
pip install 'lerobot[training]'
```

## Lançamento

O treinamento distribuído pode ser iniciado tanto por meio de `torchrun` quanto de `accelerate launch`. O Accelerate é usado como um lançador simples: ele não gerencia a configuração do treinamento, e todas as configurações de treinamento distribuído ficam no próprio sistema de configuração do LeRobot.

Com `torchrun`:

```bash
torchrun --nproc-per-node=2 $(which lerobot-train) \
  --dataset.repo_id=${HF_USER}/my_dataset \
  --policy.type=act \
  --policy.repo_id=${HF_USER}/my_trained_policy \
  --output_dir=outputs/train/act_multi_gpu \
  --job_name=act_multi_gpu \
  --trackio.enable=true
```

Com `accelerate launch` (como um lançador simples):

```bash
accelerate launch --num_processes=2 $(which lerobot-train) \
  --dataset.repo_id=${HF_USER}/my_dataset \
  --policy.type=act \
  --policy.repo_id=${HF_USER}/my_trained_policy \
  --output_dir=outputs/train/act_multi_gpu \
  --job_name=act_multi_gpu \
  --trackio.enable=true
```

Sem as flags `--parallelism.*`, uma inicialização multiprocessos executa DDP simples. Execuções multinó utilizam as flags padrão `torchrun --nnodes/--node-rank/--rdzv-endpoint` (ou `accelerate launch --num_machines/--machine_rank/--main_process_ip`).

> [!WARNING]
> Os arquivos de configuração YAML do Accelerate (`accelerate launch --config_file some.yaml`, `accelerate config`) não são suportados. Eles configuram o engine por meio de variáveis de ambiente, contornando o sistema de configuração do LeRobot, de modo que `train_config.json` deixaria de descrever as configurações efetivamente usadas em uma execução. Por isso, `lerobot-train` se recusa a iniciar quando [variáveis de ambiente do accelerate](https://huggingface.co/docs/accelerate/usage_guides/fsdp) estão definidas. Coloque as configurações nas flags `--parallelism.*` / `--accelerator.*` em vez disso, ou defina `LEROBOT_ALLOW_ACCELERATE_ENV=1` para reconhecer a substituição e prosseguir mesmo assim.

## Semântica de lote, taxa de aprendizado e passos

Cada um dos `dp_replicate × dp_shard` workers de paralelismo de dados carrega seu próprio micro-lote `--batch_size` a cada passo, de modo que um passo de treinamento consome `batch_size × dp_world_size` amostras, e `× gradient_accumulation_steps` delas entram em cada atualização do otimizador:

```
effective_batch_size = batch_size × dp_world_size × gradient_accumulation_steps
```

O banner de treinamento imprime essa fatoração na inicialização. `--steps` conta os passos do loop (micro-lotes por worker), não as atualizações do otimizador.

A acumulação de gradiente é uma flag de primeira classe:

```bash
torchrun --nproc-per-node=2 $(which lerobot-train) \
  --batch_size=8 --accelerator.gradient_accumulation.steps=4 ...
```

**O LeRobot não escala automaticamente a taxa de aprendizado ou o número de etapas** quando o tamanho efetivo do lote aumenta. Se você escalar para mais GPUs e quiser um treinamento equivalente, ajuste manualmente, por exemplo, com 2 GPUs: dobre `--optimizer.lr` (escala linear) ou reduza `--steps` pela metade.

## Treinamento fragmentado (FSDP)

Se um modelo for grande demais para treinar com DDP, faça o sharding dele com FSDP2:

```bash
torchrun --nproc-per-node=4 $(which lerobot-train) \
  --dataset.repo_id=${HF_USER}/my_dataset \
  --policy.type=<your_policy> \
  --parallelism.dp_shard=4 \
  --accelerator.mixed_precision=bf16 \
  --output_dir=outputs/train/my_policy_fsdp
```

`--parallelism.dp_shard=-1` faz o particionamento entre o número de processos iniciado pelo launcher.

`--accelerator.mixed_precision` suporta tanto `bf16` quanto `fp16`. Detalhes em [precisão mista](./torch_accelerators#mixed-precision)

### Unidades de envolvimento

O FSDP divide o modelo em unidades (tipicamente o bloco de transformer repetido) e reúne uma unidade de cada vez durante o forward/backward. As políticas declaram suas unidades de wrap via `_fsdp_wrap_modules` na classe da política. Por exemplo, o ACT declara `["ACTEncoderLayer", "ACTDecoderLayer"]` e o FastWAM declara `["MoTLayer"]`. Para uma política sem uma declaração de `_fsdp_wrap_modules`, passe um dos flags abaixo. Você pode especificar o nome da classe do módulo explicitamente, ou usar uma política baseada em tamanho:

```bash
--accelerator.fsdp.wrap_modules='["MyTransformerBlock"]'   # explicit class names
--accelerator.fsdp.min_num_params=1000000                  # or: wrap every submodule above 1M params
```

O FSDP não consegue dividir pesos vinculados (por exemplo, uma cabeça de linguagem vinculada ao embedding de entrada) entre unidades de wrap. O torch 2.12 e versões mais novas param no primeiro passo de treinamento e nomeiam o parâmetro, enquanto versões mais antigas os desvinculam silenciosamente. Envolver a classe do bloco repetido geralmente os mantém juntos na unidade raiz.

Se uma política não declarar `_fsdp_wrap_modules` e nenhum `--accelerator.fsdp.wrap_modules` ou `--accelerator.fsdp.min_num_params` for passado, a execução falha na inicialização em vez de envolver silenciosamente apenas o módulo raiz (o que resultaria na perda de todas as economias de memória de sharding).

Outras configurações de sharding:

- `--accelerator.fsdp.reshard_after_forward`: se os parâmetros de cada unidade devem permanecer residentes após o forward.
- `--accelerator.fsdp.cpu_offload`: mantém parâmetros, gradientes e estados do otimizador na CPU.
- `--accelerator.fsdp.ignored_modules`: uma expressão regular de caminhos de módulos para manter não fragmentados.
- `--accelerator.fsdp.mixed_precision_policy.{param_dtype,reduce_dtype,output_dtype,cast_forward_inputs}`: os campos do `MixedPrecisionPolicy` do PyTorch, aplicados a cada unidade FSDP. Assim que qualquer campo for definido, a política é construída apenas a partir desses campos e substitui a derivada de `--accelerator.mixed_precision`, portanto, defina também `param_dtype`. Antes do torch 2.15, um `reduce_dtype` não definido segue o `param_dtype`. A partir do torch 2.15, ele segue o dtype próprio dos parâmetros (float32 sob precisão mista). `float16` requer `--accelerator.mixed_precision=fp16`, que constrói o loss scaler.

### HSDP

Paralelismo de Dados Híbrido Fragmentado: os parâmetros, gradientes e estados do otimizador são fragmentados entre `dp_shard` ranks, e essa fragmentação é replicada `dp_replicate` vezes. As operações all-gather de parâmetros e reduce-scatter de gradientes permanecem dentro de um grupo de fragmentos; apenas o all-reduce que sincroniza as réplicas cruza entre os grupos. Os dois graus devem multiplicar para o tamanho do mundo:

```bash
# 16 GPUs = 2 nodes × 8: shard within each node, replicate across nodes
torchrun --nnodes=2 --nproc-per-node=8 ... $(which lerobot-train) \
  --parallelism.dp_replicate=2 --parallelism.dp_shard=8 ...
```

### PEFT

[Fine-tuning com PEFT](./peft_training) (`--peft.*`) funciona com FSDP2 e HSDP. O modelo base congelado é particionado junto com os adaptadores, de modo que cada rank mantém apenas sua parte dos pesos base:

```bash
torchrun --nproc-per-node=4 $(which lerobot-train) \
  --policy.path=${HF_USER}/my_base_policy \
  --peft.method_type=LORA --peft.r=16 \
  --parallelism.dp_shard=4 \
  --accelerator.mixed_precision=bf16 \
  --output_dir=outputs/train/my_policy_fsdp_lora
```

- As unidades de encapsulamento ainda vêm da declaração `_fsdp_wrap_modules` da política (ou `--accelerator.fsdp.wrap_modules`).
- Use precisão mista quando o modelo base estiver armazenado em bf16 ou fp16. O PEFT mantém seus adaptadores em fp32, enquanto as cópias de `--peft.full_training_modules` mantêm o dtype do modelo base, e o FSDP2 precisa de um único dtype para os parâmetros treináveis que ele fragmenta juntos. Com a precisão mista ativada, o accelerate faz o upcast de todos os parâmetros treináveis para fp32.
- Os checkpoints armazenam o adaptador em vez dos pesos completos, em cada `--checkpoint_format` (veja [Checkpoints](#checkpoints)).
- `--accelerator.fsdp.ignored_modules` é comparado com os nomes dos módulos do modelo PEFT, que começam com `base_model.model.`. Inicie o padrão com `.*` para correspondê-los.

## Checkpoints

Cada checkpoint contém um diretório `pretrained_model/` e um diretório `training_state/`:

```text
005000/  # the training step at that checkpoint
├── pretrained_model/
│   ├── config.json  # policy config
│   ├── train_config.json  # the full training config
│   ├── model.safetensors  # full weights (checkpoint_format ∈ {safetensors, safetensors_dcp}, or any non-sharded run)
│   ├── pytorch_model_fsdp_0/  # DCP weight shards (checkpoint_format ∈ {dcp, safetensors_dcp})
│   ├── policy_preprocessor.json  # preprocessor config (when the run has a preprocessor)
│   ├── policy_preprocessor_step_*.safetensors  # state of the stateful preprocessor steps
│   ├── policy_postprocessor.json  # postprocessor config (when the run has a postprocessor)
│   └── policy_postprocessor_step_*.safetensors  # state of the stateful postprocessor steps
└── training_state/
    ├── optimizer_param_groups.json  # optimizer param groups (non-sharded runs)
    ├── optimizer_state.safetensors  # optimizer state (non-sharded runs)
    ├── optimizer_0/  # DCP optimizer shards (sharded runs)
    ├── rng_state.safetensors  # rng states
    ├── scheduler_state.json  # scheduler state (when the run has a scheduler)
    ├── scaler_state.json  # fp16 loss-scaler state (mixed_precision=fp16 only)
    └── training_step.json  # step + dp_world_size/batch_size/grad_accum/precision + topology
```

Durante o treinamento em GPU única ou DDP, o pipeline serializa cada state dict em um único arquivo: `model.safetensors` para o modelo e `optimizer_state.safetensors` para o otimizador.

Durante o treinamento particionado, o estado do otimizador é salvo como shards DCP sob `training_state/optimizer_0/`, e o layout do modelo sob `pretrained_model/` pode ser configurado por meio de `--checkpoint_format`:

| `--checkpoint_format`     | Artefato de pesos                             | Use quando                                                              |
| ------------------------- | -------------------------------------------- | --------------------------------------------------------------------- |
| `safetensors` _(padrão)_ | apenas um único `model.safetensors`              | você deseja que cada checkpoint seja imediatamente carregável com `from_pretrained` |
| `dcp`                     | apenas o diretório de shard `pytorch_model_fsdp_0/` | reunir todos os pesos torna os salvamentos e retomas muito lentos           |
| `safetensors_dcp`         | ambos                                         | você deseja retomar rapidamente _e_ carregar imediatamente checkpoints           |

Uma execução de PEFT armazena seu adaptador em vez dos pesos completos: `adapter_model.safetensors` e `adapter_config.json` substituem `model.safetensors`, e os shards do DCP contêm apenas o adaptador. O modelo base é carregado a partir do caminho registrado em `adapter_config.json`.

Duas coisas a saber sobre checkpoints agregados (`safetensors`) de execuções fragmentadas:

- **Eles armazenam pesos em fp32.** No treinamento com precisão mista, o FSDP mantém uma cópia mestre em fp32, e o checkpoint salva a cópia mestre para garantir que o treinamento retome de forma consistente.
- O gather é coletivo (todos os ranks participam), mas apenas o processo principal grava.

### Convertendo checkpoints DCP

`lerobot-convert-dcp` mescla um diretório de shard DCP em um `model.safetensors` regular, offline e sem GPUs:

```bash
lerobot-convert-dcp --checkpoint_dir=outputs/train/run/checkpoints/005000
lerobot-convert-dcp --checkpoint_dir=... --delete_dcp=true --push_to_hub=${HF_USER}/my_policy
```

`--push_to_hub` publica o diretório convertido como um repositório de modelo. Os shards de uma execução PEFT são mesclados em `adapter_model.safetensors`. Como qualquer checkpoint `dcp`, um checkpoint PEFT salvo com `--checkpoint_format=dcp` é carregado com `--policy.path` somente após essa conversão.

### Retomando

Retome com `--resume=true --config_path=.../checkpoints/last/pretrained_model/train_config.json`. Retomar de um checkpoint DCP suporta o resharing do modelo e do estado do otimizador para a topologia _atual_, o que significa que você pode retomar com uma divisão diferente de `dp_replicate/dp_shard`. O data sampler sempre pode retomar na época e no deslocamento corretos, mas é _exato por amostra_ apenas quando o tamanho do mundo e o tamanho do lote correspondem à execução original (um aviso é registrado caso contrário).

> [!NOTE]
> Os checkpoints de FSDP gravados pelo LeRobot 0.6.x e versões anteriores utilizavam um layout diferente no disco (um estado completo do otimizador agregado) e **não podem ser retomados**.

## Observações

- Os salvamentos de checkpoint e as publicações no final do treinamento são coletivos (todos os ranks os executam). Os pesos agregados, os arquivos auxiliares e os uploads para o Hub são gravados apenas pelo processo principal.
- As métricas são reduzidas entre os ranks antes do registro: as perdas são médias e `samples/s` relata a taxa de transferência em todo o cluster.
- O agendamento da taxa de aprendizado é avançado uma vez por etapa de treinamento, independentemente do número de processos (`step_scheduler_with_optimizer=False` está embutido).

Para mais informações sobre a mecânica subjacente, consulte o [guia de FSDP do Accelerate](https://huggingface.co/docs/accelerate/usage_guides/fsdp). Para aprofundar-se no treinamento em larga escala, confira o [Ultrascale Playbook](https://huggingface.co/spaces/nanotron/ultrascale-playbook).