feat(train): parallel training framework — FSDP2, HSDP, gradient accumulation, and DCP checkpoints (#4010)

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

tests/configs/test_accelerator_config.py

@@ -0,0 +1,135 @@
+#!/usr/bin/env python
+
+# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import json
+
+import draccus
+import pytest
+
+from lerobot.configs.accelerator import (
+    AcceleratorConfig,
+    ActivationCheckpointingConfig,
+    ActivationCheckpointingMode,
+    CompileConfig,
+    DDPConfig,
+    FSDPConfig,
+    GradientAccumulationConfig,
+)
+from lerobot.configs.parallelism import ParallelismConfig
+
+
+class TestFieldValidation:
+    def test_wrap_policies_mutually_exclusive(self):
+        with pytest.raises(ValueError, match="mutually exclusive"):
+            FSDPConfig(wrap_modules=["Block"], min_num_params=1000)
+
+    def test_min_num_params_positive(self):
+        with pytest.raises(ValueError, match="min_num_params"):
+            FSDPConfig(min_num_params=0)
+
+    def test_mixed_precision_choices(self):
+        with pytest.raises(ValueError, match="mixed_precision"):
+            AcceleratorConfig(mixed_precision="tf32")
+
+    def test_gradient_accumulation_positive(self):
+        with pytest.raises(ValueError, match="gradient_accumulation.steps"):
+            GradientAccumulationConfig(steps=0)
+
+
+class TestDraccusRoundTrip:
+    @pytest.mark.parametrize(
+        "cfg",
+        [
+            AcceleratorConfig(),
+            AcceleratorConfig(
+                mixed_precision="bf16",
+                gradient_accumulation=GradientAccumulationConfig(steps=4),
+                fsdp=FSDPConfig(
+                    reshard_after_forward=False,
+                    wrap_modules=["ACTEncoderLayer", "ACTDecoderLayer"],
+                    cpu_offload=True,
+                    ignored_modules=r".*pos_embed.*",
+                ),
+                ddp=DDPConfig(find_unused_parameters=False, static_graph=True),
+                compile=CompileConfig(enabled=True, mode="max-autotune", regional=False),
+                activation_checkpointing=ActivationCheckpointingConfig(mode=ActivationCheckpointingMode.FULL),
+            ),
+            AcceleratorConfig(fsdp=FSDPConfig(min_num_params=1_000_000)),
+        ],
+    )
+    def test_encode_json_decode_identity(self, cfg):
+        payload = json.loads(json.dumps(draccus.encode(cfg)))
+        assert draccus.decode(AcceleratorConfig, payload) == cfg
+
+    def test_pre_existing_config_without_fields_gets_defaults(self):
+        assert draccus.decode(AcceleratorConfig, {}) == AcceleratorConfig()
+
+
+class TestRuntimeBuilders:
+    """The mirrors must translate into real accelerate objects (plugins built lazily)."""
+
+    @pytest.fixture(autouse=True)
+    def _requires_accelerate(self):
+        pytest.importorskip("accelerate", reason="accelerate is required (install lerobot[training])")
+
+    def test_fsdp_plugin_translation(self):
+        plugin = FSDPConfig(
+            reshard_after_forward=False, wrap_modules=["MyBlock"], cpu_offload=True
+        ).build_plugin()
+        assert plugin.fsdp_version == 2
+        assert plugin.reshard_after_forward is False
+        assert plugin.transformer_cls_names_to_wrap == ["MyBlock"]
+        # bools are normalized into torch offload policies by the plugin itself
+        assert type(plugin.cpu_offload).__name__ == "CPUOffloadPolicy"
+        # LeRobot never switches state_dict_type: FSDP2's SHARDED default must hold
+        assert plugin.state_dict_type.name == "SHARDED_STATE_DICT"
+        assert not plugin.activation_checkpointing
+
+    def test_fsdp_plugin_size_based_policy(self):
+        plugin = FSDPConfig(min_num_params=1024).build_plugin()
+        assert plugin.min_num_params == 1024
+        assert plugin.transformer_cls_names_to_wrap is None
+
+    def test_ddp_kwargs_translation(self):
+        handler = DDPConfig(find_unused_parameters=False, gradient_as_bucket_view=True).build_kwargs_handler()
+        assert handler.find_unused_parameters is False
+        assert handler.gradient_as_bucket_view is True
+
+    def test_gradient_accumulation_plugin_translation(self):
+        plugin = GradientAccumulationConfig(steps=4).build_plugin()
+        assert plugin.num_steps == 4
+        assert plugin.sync_with_dataloader is False
+
+    def test_gradient_accumulation_never_syncs_with_dataloader(self, monkeypatch):
+        """The loop cycles a finite dataloader, so accelerate's default
+        sync_with_dataloader=True would force an optimizer step at every dataset epoch
+        boundary instead of every num_steps micro-batches."""
+        captured = {}
+
+        class FakeAccelerator:
+            def __init__(self, **kwargs):
+                captured.update(kwargs)
+
+        monkeypatch.setattr("accelerate.Accelerator", FakeAccelerator)
+        parallelism = ParallelismConfig()
+        parallelism.resolve(1)
+        AcceleratorConfig(gradient_accumulation=GradientAccumulationConfig(steps=4)).build(
+            parallelism, cpu=True
+        )
+        ga_plugin = captured["gradient_accumulation_plugin"]
+        assert ga_plugin.num_steps == 4
+        assert ga_plugin.sync_with_dataloader is False
+        assert "gradient_accumulation_steps" not in captured

tests/configs/test_parallelism_config.py

@@ -0,0 +1,118 @@
+#!/usr/bin/env python
+
+# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import json
+
+import draccus
+import pytest
+
+from lerobot.configs.parallelism import ContextParallelConfig, ParallelismConfig
+
+
+class TestResolve:
+    def test_single_process_defaults(self):
+        cfg = ParallelismConfig()
+        cfg.resolve(1)
+        assert (cfg.dp_replicate, cfg.dp_shard) == (1, 1)
+        assert not cfg.is_sharded and not cfg.is_replicated_only
+        assert cfg.dp_world_size == 1
+
+    def test_untouched_config_fills_ddp(self):
+        """Plain `torchrun --nproc-per-node=8` with a default config resolves to DDP."""
+        cfg = ParallelismConfig()
+        cfg.resolve(8)
+        assert cfg.dp_replicate == 8
+        assert cfg.is_replicated_only and not cfg.is_sharded
+        assert cfg.dp_world_size == 8
+
+    def test_full_shard_sentinel(self):
+        cfg = ParallelismConfig(dp_shard=-1)
+        assert cfg.is_sharded  # sharded even before resolve: -1 is an explicit opt-in
+        cfg.resolve(8)
+        assert cfg.dp_shard == 8 and cfg.dp_replicate == 1
+
+    def test_hsdp_sentinel_infers_shard(self):
+        cfg = ParallelismConfig(dp_replicate=2, dp_shard=-1)
+        cfg.resolve(8)
+        assert (cfg.dp_replicate, cfg.dp_shard) == (2, 4)
+        assert cfg.dp_world_size == 8
+
+    def test_explicit_hsdp(self):
+        cfg = ParallelismConfig(dp_replicate=2, dp_shard=4)
+        cfg.resolve(8)
+        assert cfg.is_sharded and not cfg.is_replicated_only
+
+    def test_product_mismatch_lists_all_degrees(self):
+        cfg = ParallelismConfig(dp_replicate=2, dp_shard=2)
+        with pytest.raises(ValueError, match=r"dp_replicate=2 \* dp_shard=2.*WORLD_SIZE=8"):
+            cfg.resolve(8)
+
+    def test_explicit_replicate_must_match_world(self):
+        cfg = ParallelismConfig(dp_replicate=4)
+        with pytest.raises(ValueError, match="WORLD_SIZE=8"):
+            cfg.resolve(8)
+
+    def test_sentinel_indivisible_world(self):
+        cfg = ParallelismConfig(dp_replicate=3, dp_shard=-1)
+        with pytest.raises(ValueError, match="not divisible"):
+            cfg.resolve(8)
+
+    def test_cp_fails_fast(self):
+        cfg = ParallelismConfig(dp_shard=-1, context_parallel=ContextParallelConfig(ulysses_degree=2))
+        with pytest.raises(ValueError, match="not implemented"):
+            cfg.resolve(8)
+
+
+class TestFieldValidation:
+    @pytest.mark.parametrize("kwargs", [{"dp_replicate": 0}, {"dp_shard": 0}, {"dp_shard": -2}])
+    def test_bad_dp_degrees(self, kwargs):
+        with pytest.raises(ValueError):
+            ParallelismConfig(**kwargs)
+
+    def test_cfg_parallel_capped_at_two(self):
+        ParallelismConfig(cfg_parallel=2)  # reserved but representable
+        with pytest.raises(ValueError, match="cfg_parallel"):
+            ParallelismConfig(cfg_parallel=3)
+
+    @pytest.mark.parametrize("kwargs", [{"ring_degree": 0}, {"ulysses_degree": -1}])
+    def test_bad_cp_degrees(self, kwargs):
+        with pytest.raises(ValueError):
+            ContextParallelConfig(**kwargs)
+
+    def test_dp_world_size_undefined_before_resolve(self):
+        with pytest.raises(RuntimeError, match="resolve"):
+            _ = ParallelismConfig(dp_shard=-1).dp_world_size
+
+
+class TestDraccusRoundTrip:
+    @pytest.mark.parametrize(
+        "cfg",
+        [
+            ParallelismConfig(),
+            ParallelismConfig(dp_replicate=2, dp_shard=4, cfg_parallel=2),
+            ParallelismConfig(
+                dp_shard=-1,
+                context_parallel=ContextParallelConfig(ring_degree=2, ulysses_degree=4),
+            ),
+        ],
+    )
+    def test_encode_json_decode_identity(self, cfg):
+        payload = json.loads(json.dumps(draccus.encode(cfg)))
+        assert draccus.decode(ParallelismConfig, payload) == cfg
+
+    def test_pre_existing_config_without_fields_gets_defaults(self):
+        """Checkpoints written before this feature parse with default topology."""
+        assert draccus.decode(ParallelismConfig, {}) == ParallelismConfig()

tests/configs/test_train_config_distributed.py

@@ -0,0 +1,118 @@
+#!/usr/bin/env python
+
+# Copyright 2026 The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""TrainPipelineConfig integration for the distributed fields: fail-fasts + config compat."""
+
+import draccus
+import pytest
+
+from lerobot.configs.accelerator import ActivationCheckpointingMode
+from lerobot.configs.default import DatasetConfig, PeftConfig
+from lerobot.configs.parallelism import ContextParallelConfig, ParallelismConfig
+from lerobot.configs.train import CheckpointFormat, TrainPipelineConfig
+from lerobot.optim.optimizers import AdamConfig, MultiAdamConfig
+
+
+def make_cfg(**overrides) -> TrainPipelineConfig:
+    cfg = TrainPipelineConfig(dataset=DatasetConfig(repo_id="lerobot/dummy"))
+    for name, value in overrides.items():
+        setattr(cfg, name, value)
+    return cfg
+
+
+def sharded() -> ParallelismConfig:
+    return ParallelismConfig(dp_shard=-1)
+
+
+class TestDistributedFailFasts:
+    def test_defaults_pass(self):
+        make_cfg()._validate_distributed()
+
+    def test_cp_reserved(self):
+        cfg = make_cfg(parallelism=ParallelismConfig(context_parallel=ContextParallelConfig(ring_degree=2)))
+        with pytest.raises(ValueError, match="not implemented"):
+            cfg._validate_distributed()
+
+    def test_cfg_parallel_training_rejected(self):
+        cfg = make_cfg(parallelism=ParallelismConfig(cfg_parallel=2))
+        with pytest.raises(ValueError, match="inference-only"):
+            cfg._validate_distributed()
+
+    def test_compile_placeholder(self):
+        cfg = make_cfg()
+        cfg.accelerator.compile.enabled = True
+        with pytest.raises(ValueError, match="compile"):
+            cfg._validate_distributed()
+
+    def test_activation_checkpointing_placeholder(self):
+        cfg = make_cfg()
+        cfg.accelerator.activation_checkpointing.mode = ActivationCheckpointingMode.FULL
+        with pytest.raises(ValueError, match="activation_checkpointing"):
+            cfg._validate_distributed()
+
+    def test_dcp_format_requires_sharding(self):
+        cfg = make_cfg(checkpoint_format=CheckpointFormat.DCP)
+        with pytest.raises(ValueError, match="sharded"):
+            cfg._validate_distributed()
+        cfg.parallelism = sharded()
+        cfg._validate_distributed()
+
+    def test_fp16_rejected_when_sharded(self):
+        cfg = make_cfg(parallelism=sharded())
+        cfg.accelerator.mixed_precision = "fp16"
+        with pytest.raises(ValueError, match="fp16"):
+            cfg._validate_distributed()
+        cfg.accelerator.mixed_precision = "bf16"
+        cfg._validate_distributed()
+
+    def test_peft_rejected_when_sharded(self):
+        cfg = make_cfg(parallelism=sharded(), peft=PeftConfig())
+        with pytest.raises(ValueError, match="PEFT"):
+            cfg._validate_distributed()
+
+    def test_env_eval_rejected_when_sharded(self):
+        cfg = make_cfg(parallelism=sharded(), env_eval_freq=1000)
+        cfg.env = object()  # any configured env triggers the check
+        with pytest.raises(ValueError, match="environment evaluation"):
+            cfg._validate_distributed()
+
+    def test_multi_optimizer_rejected_when_sharded(self):
+        cfg = make_cfg(parallelism=sharded(), optimizer=MultiAdamConfig())
+        with pytest.raises(ValueError, match="Multi-optimizer"):
+            cfg._validate_distributed()
+        cfg.optimizer = AdamConfig()
+        cfg._validate_distributed()
+
+
+class TestConfigCompat:
+    def test_checkpoint_format_round_trip(self):
+        for fmt in CheckpointFormat:
+            assert draccus.decode(CheckpointFormat, draccus.encode(fmt)) is fmt
+
+    def test_wants_predicates(self):
+        assert CheckpointFormat.SAFETENSORS.wants_safetensors
+        assert not CheckpointFormat.SAFETENSORS.wants_dcp
+        assert CheckpointFormat.DCP.wants_dcp and not CheckpointFormat.DCP.wants_safetensors
+        both = CheckpointFormat.SAFETENSORS_AND_DCP
+        assert both.wants_safetensors and both.wants_dcp
+
+
+def test_reward_model_rejected_when_sharded():
+    """Sharded reward runs previously failed late (missing wrap
+    units, DTensor serialization at the first checkpoint) instead of at validation."""
+    cfg = make_cfg(parallelism=sharded())
+    cfg.reward_model = object()  # any configured reward model triggers the check
+    with pytest.raises(ValueError, match="Reward-model"):
+        cfg._validate_distributed()