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/common/test_checkpoint_legacy_contracts.py

@@ -0,0 +1,221 @@
+#!/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.
+"""Legacy-checkpoint contracts.
+
+Two contracts are pinned here so they are documented behavior, not accidents:
+
+- **The v0.6.0 hard break.** The v0.6.0 #3810 FSDP checkpoint layout
+  (full gathered ``model.safetensors`` + full ``optimizer_state.safetensors``, no DCP dirs,
+  no ``checkpoint_format`` in ``train_config.json``) is a hard break with ZERO v0.6.0-aware
+  runtime code — not even layout detection. A sharded resume pointed at such a checkpoint
+  must fail through the ORDINARY missing-artifact path (torch DCP erroring on the absent
+  ``training_state/optimizer_0/``), while the model weights remain loadable forever via
+  ``from_pretrained`` and the old ``num_processes`` key keeps feeding the topology reader.
+- **Converter equivalence.** ``dcp_to_safetensors`` (real ``merge_fsdp_weights``, no mocks)
+  on accelerate's ``save_fsdp_model`` DCP layout reproduces exactly the tensors that the
+  direct-gather ``save_pretrained`` artifact contains.
+"""
+
+import json
+from pathlib import Path
+from types import SimpleNamespace
+
+import pytest
+
+pytest.importorskip("accelerate", reason="accelerate is required (install lerobot[training])")
+
+import torch
+import torch.distributed.checkpoint as dist_cp
+from accelerate.utils.constants import FSDP_MODEL_NAME, OPTIMIZER_NAME
+from safetensors.torch import load_file
+from torch.distributed.checkpoint.api import CheckpointException
+from torch.distributed.fsdp import FSDPModule
+
+from lerobot.common.train_utils import (
+    load_training_metadata,
+    resume_after_prepare,
+    resume_before_prepare,
+)
+from lerobot.configs.accelerator import FSDPConfig
+from lerobot.configs.default import DatasetConfig
+from lerobot.configs.train import TRAIN_CONFIG_NAME, CheckpointFormat, TrainPipelineConfig
+from lerobot.distributed.checkpoint import dcp_to_safetensors, is_sharded_module
+from lerobot.optim.optimizers import save_optimizer_state
+from lerobot.utils.constants import PRETRAINED_MODEL_DIR, TRAINING_STATE_DIR, TRAINING_STEP
+from lerobot.utils.io_utils import write_json
+from lerobot.utils.random_utils import save_rng_state
+from tests.fixtures.dummy_checkpoint_policy import DummyCheckpointPolicy, make_dummy_policy
+
+
+@pytest.fixture
+def accelerate_state():
+    """accelerate's process state, as the trainer's `Accelerator()` would have initialized it.
+
+    `load_fsdp_optimizer` and `merge_fsdp_weights` both consult `PartialState` internals
+    (logging and main-process gating). Single-process CPU state; reset on teardown so no
+    global accelerate state leaks into other tests.
+    """
+    from accelerate.state import AcceleratorState, PartialState
+
+    PartialState()
+    yield
+    AcceleratorState._reset_state(reset_partial_state=True)
+
+
+def make_v060_fsdp_checkpoint(checkpoint_dir: Path) -> dict[str, torch.Tensor]:
+    """Reproduce the v0.6.0 #3810 FSDP checkpoint layout with real artifacts.
+
+    - ``pretrained_model/``: ``config.json`` + full gathered ``model.safetensors`` (real
+      ``save_pretrained`` outputs) and a ``train_config.json`` predating the v0.7 fields
+      (``checkpoint_format``/``parallelism``/``accelerator`` stripped from the draccus dump);
+    - ``training_state/``: old-style ``training_step.json`` (``{"step", "num_processes"}``,
+      no ``dp_world_size``), ``rng_state.safetensors``, and the gathered full optimizer
+      channel (``optimizer_state.safetensors`` + ``optimizer_param_groups.json``) — and,
+      crucially, NO ``optimizer_0/`` DCP directory.
+
+    Returns the saved model weights for later comparison.
+    """
+    policy = make_dummy_policy()
+    optimizer = torch.optim.Adam(policy.parameters())
+    policy.forward({"observation.state": torch.randn(2, 4)})[0].backward()
+    optimizer.step()  # real optimizer state, applied before the weights are saved
+
+    pretrained_dir = checkpoint_dir / PRETRAINED_MODEL_DIR
+    policy.save_pretrained(pretrained_dir)
+    cfg = TrainPipelineConfig(dataset=DatasetConfig(repo_id="lerobot/dummy"), batch_size=3)
+    cfg._save_pretrained(pretrained_dir)
+    config_path = pretrained_dir / TRAIN_CONFIG_NAME
+    raw = json.loads(config_path.read_text())
+    assert "checkpoint_format" in raw  # draccus dumps defaults; a v0.6.0 config predates the key
+    for key in ("checkpoint_format", "parallelism", "accelerator"):
+        raw.pop(key, None)
+    config_path.write_text(json.dumps(raw, indent=4))
+
+    training_state_dir = checkpoint_dir / TRAINING_STATE_DIR
+    training_state_dir.mkdir()
+    write_json({"step": 5000, "num_processes": 4}, training_state_dir / TRAINING_STEP)
+    save_rng_state(training_state_dir)
+    save_optimizer_state(optimizer, training_state_dir)
+    return {key: tensor.clone() for key, tensor in policy.state_dict().items()}
+
+
+def as_fsdp2_module(policy: DummyCheckpointPolicy) -> DummyCheckpointPolicy:
+    """Give the policy FSDP2's runtime identity via the in-place class swap `fully_shard` performs.
+
+    torch's `fully_shard` swaps ``module.__class__`` to a ``(FSDPModule, type(module))``
+    subclass; mirroring that swap is what makes `is_sharded_module` (and thus the sharded
+    branch of `resume_after_prepare`) see a sharded model on a CPU-only single process. The
+    parameters stay plain tensors — sufficient here, because the resume must fail at the DCP
+    read before any sharded state is touched.
+    """
+    policy.__class__ = type(f"FSDP{type(policy).__name__}", (FSDPModule, type(policy)), {})
+    assert is_sharded_module(policy)
+    return policy
+
+
+def sharded_passthrough_accelerator() -> SimpleNamespace:
+    """The accelerator surface the sharded resume touches, carrying the trainer's real plugin.
+
+    `FSDPConfig.build_plugin()` is the exact FSDP2 plugin construction `make_accelerator`
+    hands to accelerate (state_dict_type stays at the FSDP2 default, SHARDED_STATE_DICT).
+    """
+    return SimpleNamespace(
+        unwrap_model=lambda m: m,
+        wait_for_everyone=lambda: None,
+        state=SimpleNamespace(fsdp_plugin=FSDPConfig().build_plugin()),
+    )
+
+
+class TestV060HardBreak:
+    """Pin the v0.6.0 hard break as a contract.
+
+    Zero v0.6.0-aware code ships — not even layout detection — so every assertion here must
+    hold through ORDINARY code paths only: the recorded config parses with plain defaults,
+    phase-1 resume and the weights stay loadable, and the sharded phase-2 resume fails with
+    torch DCP's own missing-artifact error, never a bespoke v0.6.0 message.
+    """
+
+    def test_sharded_resume_fails_with_ordinary_missing_artifact_error(self, tmp_path, accelerate_state):
+        make_v060_fsdp_checkpoint(tmp_path)
+
+        # No checkpoint_format recorded -> plain draccus default, no layout detection anywhere.
+        cfg = TrainPipelineConfig.from_pretrained(tmp_path / PRETRAINED_MODEL_DIR / TRAIN_CONFIG_NAME)
+        assert cfg.checkpoint_format is CheckpointFormat.SAFETENSORS
+        cfg.checkpoint_path = tmp_path
+
+        # Phase 1 (RNG + step counter) is format-independent and still succeeds.
+        assert resume_before_prepare(cfg) == 5000
+
+        # Phase 2 under sharding: the recorded format skips the DCP model preflight (the
+        # weights were already loaded by from_pretrained), then the sharded optimizer load
+        # hits the absent optimizer_0/ and fails inside torch DCP — the ordinary error path.
+        assert not (tmp_path / TRAINING_STATE_DIR / f"{OPTIMIZER_NAME}_0").exists()
+        policy = as_fsdp2_module(make_dummy_policy())
+        optimizer = torch.optim.Adam(policy.parameters())
+        with pytest.raises(CheckpointException) as excinfo:
+            resume_after_prepare(cfg, sharded_passthrough_accelerator(), policy, optimizer, None)
+        message = str(excinfo.value)
+        assert "lerobot-convert-dcp" not in message  # the converter hint belongs to recorded-format=DCP
+        assert "v0.6" not in message  # no bespoke wording: the explanation lives in the migration docs
+
+    def test_weights_remain_loadable_via_from_pretrained(self, tmp_path):
+        saved_weights = make_v060_fsdp_checkpoint(tmp_path)
+        policy = DummyCheckpointPolicy.from_pretrained(tmp_path / PRETRAINED_MODEL_DIR)
+        for key, tensor in policy.state_dict().items():
+            assert torch.equal(tensor, saved_weights[key]), key
+
+    def test_topology_reader_falls_back_to_legacy_num_processes(self, tmp_path):
+        make_v060_fsdp_checkpoint(tmp_path)
+        assert load_training_metadata(tmp_path / TRAINING_STATE_DIR)["dp_world_size"] == 4
+
+
+class TestConverterEquivalence:
+    def test_dcp_to_safetensors_output_equals_direct_gather(self, tmp_path, accelerate_state):
+        """DCP -> safetensors conversion is exactly the direct-gather artifact.
+
+        The DCP checkpoint is written with torch's real `dist_cp.save` (single process, no
+        process group), replicating accelerate's `save_fsdp_model` SHARDED_STATE_DICT branch
+        byte for byte: the ``{"model": state_dict}`` nesting and the ``pytorch_model_fsdp_0``
+        directory name. The conversion runs the real `merge_fsdp_weights` — no mocks.
+        """
+        policy = make_dummy_policy()
+        with torch.no_grad():
+            for param in policy.parameters():
+                param.add_(torch.randn_like(param))  # make every tensor distinct from init
+        reference = {key: tensor.clone() for key, tensor in policy.state_dict().items()}
+
+        # The direct-gather artifact (on a single process the gather is state_dict itself).
+        direct_dir = tmp_path / "direct"
+        policy.save_pretrained(direct_dir)
+
+        # The DCP artifact, laid out exactly as accelerate's save_fsdp_model writes it.
+        pretrained_dir = tmp_path / "checkpoint" / PRETRAINED_MODEL_DIR
+        dcp_dir = pretrained_dir / f"{FSDP_MODEL_NAME}_0"
+        dcp_dir.mkdir(parents=True)
+        dist_cp.save(
+            state_dict={"model": policy.state_dict()},
+            storage_writer=dist_cp.FileSystemWriter(str(dcp_dir)),
+        )
+
+        merged_file = dcp_to_safetensors(dcp_dir, pretrained_dir)
+        assert merged_file == pretrained_dir / "model.safetensors"
+        merged = load_file(merged_file)
+        direct = load_file(direct_dir / "model.safetensors")
+        assert set(merged) == set(direct) == set(reference)
+        for key, tensor in reference.items():
+            assert torch.equal(merged[key], tensor), key
+            assert torch.equal(direct[key], tensor), key
+            assert merged[key].dtype == tensor.dtype, key

tests/common/test_checkpoint_save_resume.py

@@ -0,0 +1,186 @@
+#!/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.
+"""Checkpoint save/resume round-trips on the non-sharded paths."""
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+from safetensors.torch import load_file
+
+from lerobot.common.train_utils import (
+    load_training_metadata,
+    resume_after_prepare,
+    resume_before_prepare,
+    save_checkpoint,
+)
+from lerobot.configs.default import DatasetConfig
+from lerobot.configs.train import CheckpointFormat, TrainPipelineConfig
+from lerobot.utils.constants import PRETRAINED_MODEL_DIR, TRAINING_STATE_DIR, TRAINING_STEP
+from lerobot.utils.io_utils import load_json, write_json
+from tests.fixtures.dummy_checkpoint_policy import make_dummy_policy
+
+
+def make_cfg(**overrides) -> TrainPipelineConfig:
+    cfg = TrainPipelineConfig(dataset=DatasetConfig(repo_id="lerobot/dummy"), batch_size=3)
+    cfg.parallelism.resolve(1)
+    for name, value in overrides.items():
+        setattr(cfg, name, value)
+    return cfg
+
+
+def passthrough_accelerator() -> SimpleNamespace:
+    """The accelerator surface save/resume touches on non-sharded runs."""
+    return SimpleNamespace(unwrap_model=lambda m: m, wait_for_everyone=lambda: None)
+
+
+class TestSaveCheckpoint:
+    def test_non_sharded_layout(self, tmp_path):
+        policy = make_dummy_policy()
+        optimizer = torch.optim.Adam(policy.parameters())
+        save_checkpoint(
+            tmp_path,
+            step=7,
+            cfg=make_cfg(),
+            policy=policy,
+            optimizer=optimizer,
+            accelerator=passthrough_accelerator(),
+        )
+        pretrained = tmp_path / PRETRAINED_MODEL_DIR
+        state = tmp_path / TRAINING_STATE_DIR
+        assert (pretrained / "model.safetensors").is_file()
+        assert (pretrained / "config.json").is_file()
+        assert (pretrained / "train_config.json").is_file()
+        assert (state / TRAINING_STEP).is_file()
+        assert (state / "rng_state.safetensors").is_file()
+        assert (state / "optimizer_state.safetensors").is_file()
+        # single-file artifact, no index, weights intact
+        weights = load_file(pretrained / "model.safetensors")
+        assert torch.allclose(weights["net.weight"], torch.full_like(weights["net.weight"], 0.5))
+        assert not list(pretrained.glob("*.index.json"))
+
+    def test_training_step_records_topology(self, tmp_path):
+        cfg = make_cfg()
+        cfg.accelerator.gradient_accumulation.steps = 4
+        policy = make_dummy_policy()
+        save_checkpoint(
+            tmp_path,
+            step=11,
+            cfg=cfg,
+            policy=policy,
+            optimizer=torch.optim.Adam(policy.parameters()),
+            accelerator=passthrough_accelerator(),
+        )
+        metadata = load_training_metadata(tmp_path / TRAINING_STATE_DIR)
+        assert metadata["dp_world_size"] == 1
+        assert metadata["batch_size"] == 3
+        assert metadata["grad_accum_steps"] == 4
+
+    def test_dp_world_size_legacy_fallback(self, tmp_path):
+        """Pre-v0.7 checkpoints recorded num_processes; the reader falls back to it."""
+        state_dir = tmp_path / TRAINING_STATE_DIR
+        state_dir.mkdir(parents=True)
+        write_json({"step": 5, "num_processes": 4}, state_dir / TRAINING_STEP)
+        metadata = load_training_metadata(tmp_path / TRAINING_STATE_DIR)
+        assert metadata["dp_world_size"] == 4
+        assert metadata["batch_size"] is None
+
+
+class TestResume:
+    def _checkpointed_run(self, tmp_path):
+        policy = make_dummy_policy()
+        optimizer = torch.optim.Adam(policy.parameters(), lr=0.123)
+        # give the optimizer real state
+        policy.forward({"observation.state": torch.randn(2, 4)})[0].backward()
+        optimizer.step()
+        cfg = make_cfg()
+        save_checkpoint(
+            tmp_path,
+            step=42,
+            cfg=cfg,
+            policy=policy,
+            optimizer=optimizer,
+            accelerator=passthrough_accelerator(),
+        )
+        cfg.checkpoint_path = tmp_path
+        return cfg, policy, optimizer
+
+    def test_two_phase_resume_round_trip(self, tmp_path):
+        cfg, _, optimizer = self._checkpointed_run(tmp_path)
+        assert resume_before_prepare(cfg) == 42
+
+        fresh_policy = make_dummy_policy()
+        fresh_optimizer = torch.optim.Adam(fresh_policy.parameters(), lr=0.999)
+        resume_after_prepare(cfg, passthrough_accelerator(), fresh_policy, fresh_optimizer, None)
+        restored = fresh_optimizer.state_dict()
+        original = optimizer.state_dict()
+        assert restored["param_groups"][0]["lr"] == original["param_groups"][0]["lr"]
+        for key, tensor in original["state"][0].items():
+            assert torch.equal(restored["state"][0][key], tensor), key
+
+    def test_resume_warns_on_changed_cadence_and_topology(self, tmp_path, caplog):
+        """The recorded grad-accum factor and parallelism snapshot must be compared on
+        resume, with one warning naming the diff."""
+        import logging
+
+        cfg, _, _ = self._checkpointed_run(tmp_path)
+        cfg.accelerator.gradient_accumulation.steps = 4
+        cfg.parallelism.dp_replicate = 2  # same dp_world_size story is irrelevant here
+        with caplog.at_level(logging.WARNING):
+            assert resume_before_prepare(cfg) == 42
+        warning = next(m for m in caplog.messages if "differ from the checkpoint" in m)
+        assert "grad_accum_steps: 1 -> 4" in warning
+        assert "dp_replicate: 1 -> 2" in warning
+
+    def test_resume_unchanged_settings_stay_silent(self, tmp_path, caplog):
+        import logging
+
+        cfg, _, _ = self._checkpointed_run(tmp_path)
+        with caplog.at_level(logging.WARNING):
+            resume_before_prepare(cfg)
+        assert not [m for m in caplog.messages if "differ from the checkpoint" in m]
+
+    def test_resume_rejects_non_sharded_checkpoint_on_sharded_run(self, tmp_path):
+        """Resharding works across sizes, not across kinds: non-sharded -> sharded is rejected."""
+        cfg, _, _ = self._checkpointed_run(tmp_path)
+        cfg.parallelism.dp_shard = 2
+        with pytest.raises(ValueError, match="Cannot resume"):
+            resume_before_prepare(cfg)
+
+    def test_resume_rejects_sharded_checkpoint_on_non_sharded_run(self, tmp_path):
+        """The symmetric direction: a checkpoint recorded sharded cannot resume non-sharded."""
+        cfg, _, _ = self._checkpointed_run(tmp_path)
+        state_file = tmp_path / TRAINING_STATE_DIR / TRAINING_STEP
+        state = load_json(state_file)
+        state["parallelism"]["dp_shard"] = 2
+        write_json(state, state_file)
+        with pytest.raises(ValueError, match="Cannot resume"):
+            resume_before_prepare(cfg)
+
+    def test_resume_before_prepare_requires_training_state(self, tmp_path):
+        cfg = make_cfg()
+        cfg.checkpoint_path = tmp_path
+        with pytest.raises(NotADirectoryError):
+            resume_before_prepare(cfg)
+
+    def test_dcp_format_integrity_preflight(self, tmp_path):
+        """A checkpoint declaring DCP shards without the shard dir fails with the converter hint."""
+        pytest.importorskip("accelerate", reason="accelerate is required (install lerobot[training])")
+        cfg, policy, optimizer = self._checkpointed_run(tmp_path)
+        cfg.parallelism.dp_shard = 2  # pretend the recorded run was sharded
+        cfg.checkpoint_format = CheckpointFormat.DCP
+        with pytest.raises(FileNotFoundError, match="lerobot-convert-dcp"):
+            resume_after_prepare(cfg, passthrough_accelerator(), policy, optimizer, None)

tests/common/test_publish.py

@@ -0,0 +1,148 @@
+#!/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.
+"""publish_trained_model: commit set, card, log-line contract, PEFT branch (hub fully mocked)."""
+
+import logging
+from pathlib import Path
+from types import SimpleNamespace
+
+import pytest
+
+import lerobot.common.train_utils as train_utils
+import lerobot.utils.hub as hub
+from lerobot.common.train_utils import generate_model_card, publish_trained_model
+from lerobot.configs.default import DatasetConfig
+from lerobot.configs.train import TrainPipelineConfig
+from tests.fixtures.dummy_checkpoint_policy import make_dummy_policy
+
+
+class FakeHfApi:
+    """Records every repo/upload interaction; shared across both HfApi import sites."""
+
+    calls: list[dict] = []
+
+    def __init__(self, *args, **kwargs):
+        pass
+
+    def create_repo(self, repo_id, private=None, exist_ok=False, **kwargs):
+        return SimpleNamespace(repo_id=repo_id)
+
+    def upload_folder(self, *, repo_id, folder_path, commit_message, **kwargs):
+        FakeHfApi.calls.append(
+            {
+                "repo_id": repo_id,
+                "commit_message": commit_message,
+                "files": sorted(p.name for p in Path(folder_path).iterdir()),
+                "ignore_patterns": kwargs.get("ignore_patterns"),
+            }
+        )
+        return SimpleNamespace(repo_url=SimpleNamespace(url=f"https://huggingface.co/{repo_id}"))
+
+
+@pytest.fixture
+def mocked_hub(monkeypatch):
+    FakeHfApi.calls = []
+    monkeypatch.setattr(train_utils, "HfApi", FakeHfApi)
+    monkeypatch.setattr(hub, "HfApi", FakeHfApi)
+    # card.validate() hits the Hub; publishing must work offline in tests
+    monkeypatch.setattr(train_utils.ModelCard, "validate", lambda self: None)
+    return FakeHfApi
+
+
+def make_cfg() -> TrainPipelineConfig:
+    cfg = TrainPipelineConfig(dataset=DatasetConfig(repo_id="user/dataset"))
+    cfg.parallelism.resolve(1)
+    return cfg
+
+
+class RecordingProcessor:
+    def __init__(self):
+        self.pushed_to = None
+
+    def push_to_hub(self, repo_id, **kwargs):
+        self.pushed_to = repo_id
+
+
+class TestPublishTrainedModel:
+    def test_commit_set_and_log_contract(self, mocked_hub, caplog):
+        policy = make_dummy_policy(repo_id="user/policy")
+        pre, post = RecordingProcessor(), RecordingProcessor()
+        with caplog.at_level(logging.INFO):
+            publish_trained_model(make_cfg(), policy, pre, post, dataset_meta=None)
+
+        # commit 1: the model through HubMixin (config.json + model.safetensors in a tmpdir)
+        model_commit = mocked_hub.calls[0]
+        assert {"config.json", "model.safetensors"} <= set(model_commit["files"])
+        # commits 2-3: processors
+        assert pre.pushed_to == "user/policy" and post.pushed_to == "user/policy"
+        # commit 4: the bundle sidecar
+        bundle = mocked_hub.calls[-1]
+        assert {"README.md", "train_config.json"} <= set(bundle["files"])
+        # the exact line lerobot.jobs.hf watches to end remote runs early
+        assert any(
+            m.startswith("Model pushed to https://huggingface.co/user/policy") for m in caplog.messages
+        )
+
+    def test_peft_branch_skips_model_commit(self, mocked_hub):
+        policy = make_dummy_policy(repo_id="user/policy")
+
+        class FakePeftModel:
+            def save_pretrained(self, path):
+                (Path(path) / "adapter_model.safetensors").write_bytes(b"x")
+
+        publish_trained_model(make_cfg(), policy, None, None, dataset_meta=None, peft_model=FakePeftModel())
+        assert len(mocked_hub.calls) == 1  # only the bundle commit
+        bundle = mocked_hub.calls[0]
+        # adapter weights + the wrapped policy's config + card + train config, no full weights
+        assert {"README.md", "adapter_model.safetensors", "config.json", "train_config.json"} <= set(
+            bundle["files"]
+        )
+        assert "model.safetensors" not in bundle["files"]
+
+    def test_missing_repo_id_fails_loudly(self, mocked_hub):
+        policy = make_dummy_policy(repo_id=None)
+        with pytest.raises(ValueError, match="repo id"):
+            publish_trained_model(make_cfg(), policy, None, None, dataset_meta=None)
+
+
+class TestGenerateModelCard:
+    def test_free_function_renders_from_arguments(self, monkeypatch):
+        monkeypatch.setattr(train_utils.ModelCard, "validate", lambda self: None)
+        policy = make_dummy_policy(repo_id="user/policy")
+        card = generate_model_card(policy.config, cfg=make_cfg(), dataset_meta=None)
+        assert card.data.library_name == "lerobot"
+        assert card.data.datasets == "user/dataset"
+        assert "lerobot" in card.data.tags
+
+
+class TestDeprecatedPushModelToHub:
+    """`push_model_to_hub` stays callable for external scripts, delegating to the publisher."""
+
+    def test_policy_shim_warns_and_publishes(self, mocked_hub):
+        policy = make_dummy_policy(repo_id="user/policy")
+        with pytest.warns(FutureWarning, match="push_model_to_hub is deprecated"):
+            policy.push_model_to_hub(make_cfg())
+
+        # Same artifacts the method produced before: weights + config, then card + train config.
+        model_commit = mocked_hub.calls[0]
+        assert {"config.json", "model.safetensors"} <= set(model_commit["files"])
+        bundle = mocked_hub.calls[-1]
+        assert {"README.md", "train_config.json"} <= set(bundle["files"])
+
+    def test_policy_shim_warns_that_state_dict_is_ignored(self, mocked_hub):
+        policy = make_dummy_policy(repo_id="user/policy")
+        with pytest.warns(FutureWarning, match="`state_dict` argument is ignored"):
+            policy.push_model_to_hub(make_cfg(), state_dict=policy.state_dict())