多 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 \
  --trackio.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 \
  --trackio.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 统计循环步数（每个工作进程的 micro-batch 数），而非优化器更新次数。

梯度累积是一级标志：

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 会在启动器启动的所有进程上进行分片。

--accelerator.mixed_precision 同时支持 bf16 和 fp16。详情请参阅 混合精度

包装单元

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

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

FSDP 无法将共享权重（例如与输入嵌入绑定的语言头）拆分到不同的包装单元中。torch 2.12 及更高版本会在第一个训练步骤停止并指明该参数，而旧版本则会静默地解除绑定。通常，将重复的块类进行包装可以确保它们保留在根单元中。

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

其他分片设置：

• --accelerator.fsdp.reshard_after_forward：是否在前向传播后保留每个单元的参数驻留。
• --accelerator.fsdp.cpu_offload：将参数、梯度和优化器状态保留在 CPU 上。
• --accelerator.fsdp.ignored_modules：用于指定保持未分片状态的模块路径的正则表达式。
• --accelerator.fsdp.mixed_precision_policy.{param_dtype,reduce_dtype,output_dtype,cast_forward_inputs}：PyTorch MixedPrecisionPolicy 的字段，应用于每个 FSDP 单元。一旦设置了任意字段，该策略将仅基于这些字段构建，并替换由 --accelerator.mixed_precision 派生的策略，因此也需设置 param_dtype。在 torch 2.15 之前，未设置的 reduce_dtype 会跟随 param_dtype。从 torch 2.15 开始，它会跟随参数自身的 dtype（在混合精度下为 float32）。float16 需要 --accelerator.mixed_precision=fp16，以构建 loss scaler。

HSDP

混合分片数据并行：参数、梯度和优化器状态在 dp_shard 个 rank 之间进行分片，并且该分片被复制 dp_replicate 次。参数的 all-gather 和梯度的 reduce-scatter 仅在分片组内部进行；只有用于同步副本的 all-reduce 会跨越组间。这两个维度的乘积必须等于 world size：

# 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

PEFT 微调（--peft.*）在 FSDP2 和 HSDP 下均可正常工作。冻结的基础模型与适配器一起进行分片，因此每个 rank 仅持有基础权重的对应分片：

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

• Wrap 单元仍然来自策略的 _fsdp_wrap_modules 声明（或 --accelerator.fsdp.wrap_modules）。
• 当基础模型以 bf16 或 fp16 存储时，请使用混合精度。PEFT 会将其适配器保持在 fp32，而 --peft.full_training_modules 的副本则保留基础模型的数据类型，并且 FSDP2 需要为一起分片的可训练参数使用单一的数据类型。启用混合精度后，accelerate 会将所有可训练参数上转为 fp32。
• 检查点存储的是适配器而非完整权重，适用于所有 --checkpoint_format（参见 检查点）。
• --accelerator.fsdp.ignored_modules 会与 PEFT 模型的模块名称进行匹配，这些名称以 base_model.model. 开头。请在模式开头添加 .* 以匹配它们。

检查点

每个检查点都包含一个 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/
    ├── 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

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

在分片训练期间，优化器状态会以 DCP 分片的形式保存在 training_state/optimizer_0/ 下，而 pretrained_model/ 下模型的布局可以通过 --checkpoint_format 进行配置：

PEFT 运行会存储其适配器而非完整权重：adapter_model.safetensors 和 adapter_config.json 取代了 model.safetensors，且 DCP 分片仅包含适配器。基础模型将从 adapter_config.json 中记录的路径加载。

关于分片运行生成的聚合（safetensors）检查点，有两点需要了解：

• 它们存储 fp32 权重。 在混合精度训练下，FSDP 会保留一份 fp32 主副本，而检查点会保存该主副本，以确保训练能够一致地恢复。
• gather 是集合操作（所有 rank 均参与），但只有主进程执行写入。

转换 DCP 检查点

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 会将转换后的目录发布为模型仓库。PEFT 运行的分片会合并到 adapter_model.safetensors 中。与任何 dcp 检查点一样，使用 --checkpoint_format=dcp 保存的 PEFT 检查点只有在完成此转换后，才能通过 --policy.path 加载。

恢复

使用 --resume=true --config_path=.../checkpoints/last/pretrained_model/train_config.json 恢复训练。从 DCP 检查点恢复支持将模型和优化器状态重新分片到 当前 拓扑，这意味着你可以使用不同的 dp_replicate/dp_shard 划分来恢复训练。数据采样器始终可以在正确的 epoch 和偏移量处恢复，但仅当 world size 和 batch size 与原始运行一致时才是 样本级精确 的（否则将记录警告）。

LeRobot 0.6.x 及更早版本写入的 FSDP 检查点采用了不同的磁盘布局（聚合的完整优化器状态），因此无法恢复。

备注

• Checkpoint 保存和训练结束时的发布操作是集合操作（每个 rank 都会进入）。聚合后的权重、sidecar 文件和 Hub 上传仅由主进程写入。
• 在记录日志之前，指标会跨 rank 进行归约：损失值取平均值，samples/s 报告整个集群的吞吐量。
• 学习率调度在每个训练步骤中步进一次，与进程数量无关（step_scheduler_with_optimizer=False 已内置）。

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