mirror of
https://github.com/huggingface/lerobot.git
synced 2026-08-08 17:39:44 +00:00
Haoming Song
ef88d4e52b
* 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.
71 lines
2.7 KiB
Python
71 lines
2.7 KiB
Python
#!/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)
|