mirror of
https://github.com/huggingface/lerobot.git
synced 2026-08-08 17:39:44 +00:00
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.
This commit is contained in:
@@ -20,7 +20,6 @@ from lerobot.optim.optimizers import (
|
||||
MultiAdamConfig,
|
||||
SGDConfig,
|
||||
load_optimizer_state,
|
||||
load_optimizer_state_dict,
|
||||
save_optimizer_state,
|
||||
)
|
||||
from lerobot.utils.constants import (
|
||||
@@ -66,44 +65,6 @@ def test_save_and_load_optimizer_state(model_params, optimizer, tmp_path):
|
||||
torch.testing.assert_close(optimizer.state_dict(), loaded_optimizer.state_dict())
|
||||
|
||||
|
||||
def test_save_and_load_fsdp_optimizer_state_dict_roundtrip(tmp_path):
|
||||
"""The FSDP full optimizer state dict is keyed by parameter FQNs (dotted strings), not the
|
||||
integer indices of the single-GPU path. Verify it survives the safetensors save -> read
|
||||
round-trip used by the FSDP save/resume path (save_optimizer_state(optim_state_dict=...) then
|
||||
load_optimizer_state_dict), which the flatten/unflatten "/" separator must not corrupt."""
|
||||
full_osd = {
|
||||
"state": {
|
||||
"model.layers.0.weight": {
|
||||
"step": torch.tensor(3.0),
|
||||
"exp_avg": torch.randn(4, 4),
|
||||
"exp_avg_sq": torch.randn(4, 4),
|
||||
},
|
||||
"model.layers.0.bias": {
|
||||
"step": torch.tensor(3.0),
|
||||
"exp_avg": torch.randn(4),
|
||||
"exp_avg_sq": torch.randn(4),
|
||||
},
|
||||
},
|
||||
"param_groups": [
|
||||
{"lr": 1e-4, "betas": [0.9, 0.999], "eps": 1e-8, "weight_decay": 0.0, "params": [0, 1]}
|
||||
],
|
||||
}
|
||||
|
||||
save_optimizer_state(
|
||||
torch.optim.Adam([torch.nn.Parameter(torch.randn(1))]), tmp_path, optim_state_dict=full_osd
|
||||
)
|
||||
assert (tmp_path / OPTIMIZER_STATE).is_file()
|
||||
assert (tmp_path / OPTIMIZER_PARAM_GROUPS).is_file()
|
||||
|
||||
loaded = load_optimizer_state_dict(tmp_path)
|
||||
# FQN keys must be preserved verbatim (not int-cast, not split on their dots).
|
||||
assert set(loaded["state"].keys()) == set(full_osd["state"].keys())
|
||||
for fqn, sub in full_osd["state"].items():
|
||||
for k, v in sub.items():
|
||||
torch.testing.assert_close(loaded["state"][fqn][k], v)
|
||||
assert loaded["param_groups"] == full_osd["param_groups"]
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def base_params_dict():
|
||||
return {
|
||||
|
||||
Reference in New Issue
Block a user