Treinamento Multi-GPU

LeRobot treina em múltiplas GPUs por meio do Hugging Face Accelerate. Três layouts de paralelismo de dados são suportados:

Instalação

accelerate está incluído no extra training:

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:

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):

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).

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 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:

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:

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

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:

--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:

# 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.*) 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:

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).
• --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/:

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:

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:

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).

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. Para aprofundar-se no treinamento em larga escala, confira o Ultrascale Playbook.