* feat(train): parallel training engine with FSDP2, HSDP, and DCP checkpoints
Replace the FSDP1 training path with a config-owned parallel-training
engine:
- Topology and runtime configs (--parallelism.*, --accelerator.*):
dp_replicate x dp_shard degrees select single-process, DDP (unchanged
default), FSDP2, or HSDP; mixed precision, first-class gradient
accumulation, and FSDP/DDP tuning knobs are mirrored as plain
dataclasses that build the accelerate objects at runtime, so every
run is reproducible from its train_config.json alone. Accelerate env
vars are guarded against configuring the engine behind the config
system's back.
- Declarative policy surface: policies declare FSDP2 wrap units
(_fsdp_wrap_modules) and non-forward entry points
(_fsdp_forward_methods); a shared engine resolves them around
accelerator.prepare(). Context-parallel fields are reserved and
validated to 1.
- Checkpoints: selectable --checkpoint_format (safetensors | dcp |
safetensors_dcp); the sharded optimizer channel is always DCP;
two-phase resume (step+RNG before prepare, DCP model/optimizer after)
reshards across GPU-topology changes; lerobot-convert-dcp merges DCP
shards into a distributable model.safetensors offline.
- Publishing: PreTrainedPolicy.push_model_to_hub is replaced by the
free publish_trained_model (model + processors + card + train config,
all-ranks gather with main-rank writes);
PreTrainedPolicy._save_pretrained gathers state dicts internally,
removing the state_dict= threading from save_pretrained.
- lerobot_train is restructured around the engine: optimizer built
before the single prepare() call, deferred weight load on DCP
resumes, collective save_checkpoint with no call-site rank branches,
dp-world-size-based sample accounting.
Breaking changes: FSDP checkpoints from lerobot <= 0.6.x are not
resumable (weights stay loadable via from_pretrained; pin
lerobot==0.6.x to finish old runs); the `accelerate launch
--config_file` yaml flow is superseded by the config flags; training
autocast is owned exclusively by --accelerator.mixed_precision
(policy.dtype only casts parameters).
Also fixes: reward-model hub publishing crash (TypeError on extra
kwargs).
Verified by ~200 new CPU tests (config round-trips, checkpoint
round-trips per format, two-phase resume, publisher contracts,
converter equivalence, accelerate canaries), a 5-test 4-GPU suite
(FSDP2 save/resume bit-exactness, HSDP/DDP loss parity,
changed-topology resume, all-ranks save_pretrained, grad-accum
equivalence), and end-to-end ACT (1/4/8 GPUs) + FastWAM 6B
(FSDP2 + HSDP) training runs.
* chore(agents): adding additional infos to AGENTS.md
* adding `lerobot-train` requirement inside PR checklist
* prefer using code already implemented from transformers / diffusers instead of re-implementing in tree
---------
Signed-off-by: Maxime Ellerbach <maxime.ellerbach@huggingface.co>
* chore(docs): add more guidance to bring your own policies tutorial
* removing normalization to avoid confusion with processors
* trailing whitespace
* Update docs/source/bring_your_own_policies.mdx
Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
Signed-off-by: Maxime Ellerbach <maxime@ellerbach.net>
* Update docs/source/bring_your_own_policies.mdx
Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
Signed-off-by: Maxime Ellerbach <maxime@ellerbach.net>
* adding get optim params and predict_action chunk
* removing extra quotes
---------
Signed-off-by: Maxime Ellerbach <maxime@ellerbach.net>
* feat: Register external policies
* ruff fix
* move policy util functions to policy factory
* refactor register_third_party_devices -> register_third_party_plugins
* feat: Update docs with bring your own policies
* Improve docs for new policies
* fix: Inconsistent quotation marks
* fix: Remove print statement
* fix: wrong base class name in documentation
* fix: Handle better how the models are parsed
* fix: precommit passing
* Update docs/source/bring_your_own_policies.mdx
Co-authored-by: Steven Palma <imstevenpmwork@ieee.org>
Signed-off-by: Daniel San José Pro <42489409+danielsanjosepro@users.noreply.github.com>
---------
Signed-off-by: Steven Palma <imstevenpmwork@ieee.org>
Signed-off-by: Daniel San José Pro <42489409+danielsanjosepro@users.noreply.github.com>
Co-authored-by: Steven Palma <imstevenpmwork@ieee.org>
Haoming Song