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/distributed/test_checkpoint_wrappers.py

@@ -0,0 +1,70 @@
+#!/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.
+"""The DCP wrappers must hand accelerate exact shard directories.
+
+accelerate 1.14 resolves the load directory with a substring check ("optimizer" /
+"pytorch_model_fsdp" in the path -> use as-is) while the save side joins the shard name
+unconditionally, so a run path like `--job_name=optimizer_sweep` would save to
+`training_state/optimizer_0/` but load from `training_state/` itself. Passing the exact
+shard dir makes the containment check deterministically a no-op.
+"""
+
+from pathlib import Path
+from types import SimpleNamespace
+
+import pytest
+
+pytest.importorskip("accelerate", reason="accelerate is required (install lerobot[training])")
+
+
+def fake_accelerator() -> SimpleNamespace:
+    return SimpleNamespace(state=SimpleNamespace(fsdp_plugin=object()))
+
+
+# A parent path that trips both of accelerate's substring checks at once.
+POISONED_PARENT = Path("/outputs/train/optimizer_sweep_pytorch_model_fsdp_repro/training_state")
+
+
+def test_load_sharded_optimizer_passes_exact_shard_dir(monkeypatch):
+    import accelerate.utils
+
+    from lerobot.distributed.checkpoint import load_sharded_optimizer
+
+    seen = {}
+    monkeypatch.setattr(
+        accelerate.utils,
+        "load_fsdp_optimizer",
+        lambda plugin, accelerator, optimizer, model, input_dir: seen.update(path=input_dir),
+    )
+    load_sharded_optimizer(fake_accelerator(), optimizer=object(), model=object(), input_dir=POISONED_PARENT)
+    assert seen["path"] == str(POISONED_PARENT / "optimizer_0")
+    assert isinstance(seen["path"], str)  # str, never Path (accelerate does string checks)
+
+
+def test_load_sharded_model_passes_exact_shard_dir(monkeypatch):
+    import accelerate.utils
+
+    from lerobot.distributed.checkpoint import load_sharded_model
+
+    seen = {}
+    monkeypatch.setattr(
+        accelerate.utils,
+        "load_fsdp_model",
+        lambda plugin, accelerator, model, input_dir: seen.update(path=input_dir),
+    )
+    load_sharded_model(fake_accelerator(), model=object(), input_dir=POISONED_PARENT)
+    assert seen["path"] == str(POISONED_PARENT / "pytorch_model_fsdp_0")
+    assert isinstance(seen["path"], str)