多 GPU 训练

LeRobot 通过 Hugging Face Accelerate 在多个 GPU 上训练。支持三种数据并行布局：

安装

accelerate 包含在 training 附加项中：

pip install 'lerobot[training]'

启动

分布式训练可以通过 torchrun 和 accelerate launch 两种方式启动。Accelerate 仅作为普通启动器使用：它不管理训练配置，所有分布式训练设置都位于 LeRobot 自己的配置系统中。

使用 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 \
  --wandb.enable=true

使用 accelerate launch（作为普通启动器）：

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 \
  --wandb.enable=true

在没有 --parallelism.* 标志的情况下，多进程启动运行普通的 DDP。多节点运行使用标准的 torchrun --nnodes/--node-rank/--rdzv-endpoint 标志（或 accelerate launch --num_machines/--machine_rank/--main_process_ip）。

不支持 Accelerate 的 YAML 配置文件（accelerate launch --config_file some.yaml、accelerate config）。它们通过环境变量配置引擎，绕过了 LeRobot 的配置系统，因此 train_config.json 将无法再描述实际运行所使用的设置。因此，当设置了 accelerate 环境变量 时，lerobot-train 会拒绝启动。请改为将设置放入 --parallelism.* / --accelerator.* 标志，或设置 LEROBOT_ALLOW_ACCELERATE_ENV=1 以确认覆盖并继续运行。

批次语义、学习率和步数

每个 dp_replicate × dp_shard 数据并行工作进程每一步都会加载自己的 --batch_size 微批次，因此一个训练步骤消耗 batch_size × dp_world_size 个样本，其中 × gradient_accumulation_steps 会进入每次优化器更新：

effective_batch_size = batch_size × dp_world_size × gradient_accumulation_steps

训练横幅会在启动时打印此分解。--steps 统计的是循环步数（每个工作进程的微批次数），而非优化器更新次数。

梯度累积是一等标志：

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

当有效批次大小增长时，LeRobot 不会自动缩放学习率或步数。如果你扩展规模并希望获得等效训练，请手动调整，例如使用 2 个 GPU 时：将 --optimizer.lr 加倍（线性缩放），或将 --steps 减半。

分片训练（FSDP）

如果模型太大而无法使用 DDP 训练，请使用 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 会在启动器启动的任意数量进程上分片。

封装单元

FSDP 以单元（通常是重复的 transformer 块）为单位对模型进行分片，并在前向/反向传播时一次收集一个单元。policy 通过 policy 类上的 _fsdp_wrap_modules 声明其封装单元。例如，ACT 声明 ["ACTEncoderLayer", "ACTDecoderLayer"]，FastWAM 声明 ["MoTLayer"]。对于没有 _fsdp_wrap_modules 声明的 policy，请传递下面标志之一。你可以显式指定模块类名，或改用基于大小的 policy：

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

如果 policy 未声明 _fsdp_wrap_modules 且未传入 --accelerator.fsdp.wrap_modules 或 --accelerator.fsdp.min_num_params，运行会在启动时失败，而不是静默地仅封装根模块（这会丧失所有分片内存节省）。

其他分片设置：

• --accelerator.fsdp.reshard_after_forward：是否在前向传播后保持每个单元的参数驻留。
• --accelerator.fsdp.cpu_offload：将参数、梯度和优化器 state 保留在 CPU 上。
• --accelerator.fsdp.ignored_modules：要保持不分片的模块路径正则表达式。

HSDP

混合分片数据并行：参数、梯度和优化器 state 在 dp_shard 个 rank 上分片，并且该分片被复制 dp_replicate 次。参数 all-gather 和梯度 reduce-scatter 保留在分片组内；只有同步副本的 all-reduce 才跨组进行。这两个维度相乘必须等于总进程数：

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

checkpoint

每个 checkpoint 都包含一个 pretrained_model/ 目录和一个 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/
    ├── training_step.json  # step counter, topology, and batch semantics
    ├── rng_state.safetensors  # rng states
    ├── scheduler_state.json  # scheduler state (when the run has a scheduler)
    ├── optimizer_state.safetensors  # full optimizer state (non-sharded runs)
    ├── optimizer_param_groups.json  # optimizer param groups (non-sharded runs)
    └── optimizer_0/  # DCP optimizer shards (sharded runs)

在单 GPU 或 DDP 训练期间，流水线会将每个 state 字典序列化为单个文件：模型为 model.safetensors，优化器为 optimizer_state.safetensors。

在分片训练期间，优化器 state 保存为 training_state/optimizer_0/ 下的 DCP 分片，而 pretrained_model/ 下模型的布局可以通过 --checkpoint_format 配置：

关于分片运行中收集（safetensors）的 checkpoint，有两点需要了解：

• 它们存储 fp32 权重。 在混合精度训练下，FSDP 会保留一份 fp32 主副本，checkpoint 保存该主副本以确保训练一致地恢复。
• 收集是集合操作（所有 rank 都参与），但只有主进程写入。

转换 DCP checkpoint

lerobot-convert-dcp 将 DCP 分片目录合并为常规的 model.safetensors，离线进行且无需 GPU：

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 将转换后的目录发布为模型仓库。

恢复

使用 --resume=true --config_path=.../checkpoints/last/pretrained_model/train_config.json 恢复。从 DCP checkpoint 恢复支持将模型和优化器 state 重新分片到 当前 拓扑，这意味着你可以用不同的 dp_replicate/dp_shard 划分来恢复。数据采样器始终可以在正确的 epoch 和偏移处恢复，但只有在世界大小和批次大小与原始运行匹配时才是 样本精确 的（否则会记录一条警告）。

LeRobot 0.6.x 及更早版本写入的 FSDP checkpoint 使用了不同的磁盘布局（收集的完整优化器 state），无法恢复。

说明

• checkpoint 保存和训练结束时的发布是集合操作（每个 rank 都会进入这些操作）。收集的权重、附属文件和 Hub 上传仅由主进程写入。
• 指标在记录前会跨 rank 归约：损失取平均，samples/s 报告集群范围的吞吐量。
• 无论进程数量如何，学习率调度每个训练步骤步进一次（step_scheduler_with_optimizer=False 已内置）。

关于底层机制的背景，请参阅 Accelerate FSDP 指南。要深入了解大规模训练，请查看 Ultrascale Playbook。