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_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()