Entrenamiento con múltiples GPU

LeRobot entrena en varias GPU a través de Hugging Face Accelerate. Se admiten tres configuraciones de paralelismo de datos:

Instalación

accelerate está incluido en el extra training:

pip install 'lerobot[training]'

Lanzamiento

El entrenamiento distribuido puede iniciarse tanto con torchrun como con accelerate launch. Accelerate se utiliza como un lanzador simple: no gestiona la configuración del entrenamiento, y cada ajuste de entrenamiento distribuido reside en el propio sistema de configuración de LeRobot.

Con 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

Con accelerate launch (como lanzador simple):

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

Sin las banderas --parallelism.*, una ejecución multiproceso ejecuta DDP simple. Las ejecuciones multinodo utilizan las banderas estándar torchrun --nnodes/--node-rank/--rdzv-endpoint (o accelerate launch --num_machines/--machine_rank/--main_process_ip).

Los archivos de configuración YAML de Accelerate (accelerate launch --config_file some.yaml, accelerate config) no son compatibles. Configuran el motor a través de variables de entorno, lo que omite el sistema de configuración de LeRobot, por lo que train_config.json ya no describiría la configuración que se utilizó realmente en una ejecución. Por lo tanto, lerobot-train se niega a iniciar cuando se establecen variables de entorno de accelerate. En su lugar, coloque la configuración en las banderas --parallelism.* / --accelerator.*, o establezca LEROBOT_ALLOW_ACCELERATE_ENV=1 para reconocer la sobrescritura y continuar de todos modos.

Semántica de lotes, tasa de aprendizaje y pasos

Cada uno de los dp_replicate × dp_shard trabajadores de paralelismo de datos carga su propio micro-lote --batch_size en cada paso, por lo que un paso de entrenamiento consume batch_size × dp_world_size muestras, y × gradient_accumulation_steps de ellas se incorporan a cada actualización del optimizador:

effective_batch_size = batch_size × dp_world_size × gradient_accumulation_steps

La insignia de entrenamiento imprime esta factorización al inicio. --steps cuenta los pasos del bucle (micro-lotes por trabajador), no las actualizaciones del optimizador.

La acumulación de gradientes es una opción de primera clase:

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

LeRobot no escala automáticamente la tasa de aprendizaje ni el número de pasos cuando aumenta el tamaño efectivo del lote. Si escalas hacia fuera y deseas un entrenamiento equivalente, ajústalo manualmente, por ejemplo, con 2 GPUs: duplica --optimizer.lr (escalado lineal) o divide por dos --steps.

Entrenamiento fragmentado (FSDP)

Si un modelo es demasiado grande para entrenarlo con DDP, divídelo en fragmentos con 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 realiza el particionamiento en tantos procesos como haya iniciado el lanzador.

--accelerator.mixed_precision admite tanto bf16 como fp16. Más detalles en precisión mixta

Unidades de envoltura

FSDP divide el modelo en unidades (típicamente el bloque de transformer repetido) y reúne una unidad a la vez durante el forward/backward. Las políticas declaran sus unidades de envoltura mediante _fsdp_wrap_modules en la clase de la política. Por ejemplo, ACT declara ["ACTEncoderLayer", "ACTDecoderLayer"] y FastWAM declara ["MoTLayer"]. Para una política sin una declaración de _fsdp_wrap_modules, pase una de las banderas siguientes. Puede especificar el nombre de la clase del módulo explícitamente, o usar una política basada en el tamaño en su lugar:

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

FSDP no puede dividir los pesos vinculados (por ejemplo, una cabeza de lenguaje vinculada a la incrustación de entrada) entre unidades de envoltura. torch 2.12 y versiones posteriores se detienen en el primer paso de entrenamiento y nombran el parámetro, mientras que las versiones anteriores los desvinculan silenciosamente. Envolver la clase de bloque repetido suele mantenerlos juntos en la unidad raíz.

Si una política no declara _fsdp_wrap_modules y no se pasa --accelerator.fsdp.wrap_modules ni --accelerator.fsdp.min_num_params, la ejecución falla al inicio en lugar de envolver silenciosamente solo el módulo raíz (lo cual haría perder todos los ahorros de memoria por particionamiento).

Otros ajustes de sharding:

• --accelerator.fsdp.reshard_after_forward: si se mantienen los parámetros de cada unidad residentes después del forward.
• --accelerator.fsdp.cpu_offload: mantiene los parámetros, los gradientes y los estados del optimizador en la CPU.
• --accelerator.fsdp.ignored_modules: una expresión regular de rutas de módulos para mantenerlos sin fragmentar.
• --accelerator.fsdp.mixed_precision_policy.{param_dtype,reduce_dtype,output_dtype,cast_forward_inputs}: los campos de MixedPrecisionPolicy de PyTorch, aplicados a cada unidad de FSDP. Una vez que se establece cualquier campo, la política se construye únicamente a partir de estos campos y reemplaza a la derivada de --accelerator.mixed_precision, por lo que también debe establecerse param_dtype. Antes de torch 2.15, un reduce_dtype no establecido sigue a param_dtype. A partir de torch 2.15, sigue el dtype propio de los parámetros (float32 con precisión mixta). float16 requiere --accelerator.mixed_precision=fp16, que construye el escalador de pérdida.

HSDP

Paralelismo de datos con fragmentación híbrida: los parámetros, los gradientes y los estados del optimizador se fragmentan entre dp_shard rangos, y esa fragmentación se replica dp_replicate veces. Las operaciones all-gather de parámetros y reduce-scatter de gradientes se mantienen dentro de un grupo de fragmentos; solo el all-reduce que sincroniza las réplicas cruza entre grupos. Los dos grados deben multiplicarse para obtener el tamaño del 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

El ajuste fino con PEFT (--peft.*) funciona con FSDP2 y HSDP. El modelo base congelado se fragmenta junto con los adaptadores, de modo que cada rango solo conserva su fragmento de los 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

• Las unidades de envoltura provienen de la declaración _fsdp_wrap_modules de la política (o de --accelerator.fsdp.wrap_modules).
• Use precisión mixta cuando el modelo base esté almacenado en bf16 o fp16. PEFT mantiene sus adaptadores en fp32, mientras que las copias de --peft.full_training_modules conservan el dtype del modelo base, y FSDP2 necesita un único dtype para los parámetros entrenables que fragmenta juntos. Con la precisión mixta activada, accelerate convierte todos los parámetros entrenables a fp32.
• Los puntos de control almacenan el adaptador en lugar de los pesos completos, en cada --checkpoint_format (ver Puntos de control).
• --accelerator.fsdp.ignored_modules se compara con los nombres de los módulos del modelo PEFT, que comienzan con base_model.model.. Comience el patrón con .* para que coincidan.

Puntos de control

Cada punto de control contiene un directorio pretrained_model/ y un directorio 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 el entrenamiento con una sola GPU o con DDP, el pipeline serializa cada state dict en un único archivo: model.safetensors para el modelo y optimizer_state.safetensors para el optimizador.

Durante el entrenamiento fragmentado, el estado del optimizador se guarda como fragmentos de DCP bajo training_state/optimizer_0/, y el diseño del modelo bajo pretrained_model/ se puede configurar mediante --checkpoint_format:

Una ejecución de PEFT almacena su adaptador en lugar de los pesos completos: adapter_model.safetensors y adapter_config.json reemplazan a model.safetensors, y los fragmentos de DCP contienen solo el adaptador. El modelo base se carga desde la ruta registrada en adapter_config.json.

Dos cosas que debes saber sobre los puntos de control reunidos (safetensors) de ejecuciones fragmentadas:

• Almacenan pesos en fp32. En el entrenamiento con precisión mixta, FSDP mantiene una copia maestra en fp32, y el punto de control guarda la copia maestra para garantizar que el entrenamiento se reanude de manera consistente.
• La operación de gather es colectiva (participan todos los rangos), pero solo el proceso principal escribe.

Conversión de puntos de control DCP

lerobot-convert-dcp fusiona un directorio de fragmentos DCP en un model.safetensors regular, de forma offline y sin 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 el directorio convertido como un repositorio de modelos. Los fragmentos de una ejecución de PEFT se fusionan en adapter_model.safetensors. Al igual que cualquier punto de control de dcp, un punto de control de PEFT guardado con --checkpoint_format=dcp se carga con --policy.path solo después de esta conversión.

Reanudación

Reanude con --resume=true --config_path=.../checkpoints/last/pretrained_model/train_config.json. Reanudar desde un checkpoint de DCP admite el re-particionado (resharding) del modelo y del estado del optimizador a la topología actual, lo que significa que puede reanudar con una división diferente de dp_replicate/dp_shard. El muestreador de datos siempre puede reanudar en la época y el desplazamiento correctos, pero solo es exacto a nivel de muestra cuando el tamaño del mundo y el tamaño del lote coinciden con los de la ejecución original (de lo contrario, se registra una advertencia).

Los puntos de control de FSDP escritos por LeRobot 0.6.x y versiones anteriores utilizaban un esquema en disco diferente (un estado completo del optimizador reunido) y no se pueden reanudar.

Notas

• Las guardas de puntos de control y las publicaciones al final del entrenamiento son colectivas (todos los rangos las inician). Los pesos recopilados, los archivos secundarios y las subidas al Hub se escriben únicamente por el proceso principal.
• Las métricas se reducen entre rangos antes de registrarlas: las pérdidas se promedian y samples/s informa el rendimiento de todo el clúster.
• La programación de la tasa de aprendizaje se actualiza una vez por cada paso de entrenamiento, independientemente del número de procesos (step_scheduler_with_optimizer=False está integrado).

Para conocer el contexto de la infraestructura subyacente, consulte la guía de FSDP de Accelerate. Para profundizar en el entrenamiento a gran escala, eche un vistazo al Ultrascale Playbook.