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.
88 lines
2.7 KiB
Python
88 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.
|
|
"""A minimal real PreTrainedPolicy for checkpoint/publish unit tests (CPU, tiny)."""
|
|
|
|
from dataclasses import dataclass
|
|
|
|
import torch
|
|
from torch import Tensor, nn
|
|
|
|
from lerobot.configs.policies import PreTrainedConfig
|
|
from lerobot.optim.optimizers import AdamConfig, OptimizerConfig
|
|
from lerobot.policies.pretrained import PreTrainedPolicy
|
|
|
|
|
|
@PreTrainedConfig.register_subclass("dummy_checkpoint")
|
|
@dataclass
|
|
class DummyCheckpointConfig(PreTrainedConfig):
|
|
hidden: int = 4
|
|
|
|
@property
|
|
def observation_delta_indices(self) -> list | None:
|
|
return None
|
|
|
|
@property
|
|
def action_delta_indices(self) -> list | None:
|
|
return None
|
|
|
|
@property
|
|
def reward_delta_indices(self) -> list | None:
|
|
return None
|
|
|
|
def get_optimizer_preset(self) -> OptimizerConfig:
|
|
return AdamConfig(lr=1e-3)
|
|
|
|
def get_scheduler_preset(self) -> None:
|
|
return None
|
|
|
|
def validate_features(self) -> None:
|
|
pass
|
|
|
|
|
|
class DummyCheckpointPolicy(PreTrainedPolicy):
|
|
config_class = DummyCheckpointConfig
|
|
name = "dummy_checkpoint"
|
|
|
|
def __init__(self, config: DummyCheckpointConfig, **kwargs):
|
|
super().__init__(config)
|
|
self.net = nn.Linear(config.hidden, config.hidden)
|
|
|
|
def get_optim_params(self) -> dict:
|
|
return self.parameters()
|
|
|
|
def reset(self) -> None:
|
|
pass
|
|
|
|
def forward(self, batch: dict[str, Tensor]) -> tuple[Tensor, dict | None]:
|
|
out = self.net(batch["observation.state"])
|
|
return out.mean(), None
|
|
|
|
def predict_action_chunk(self, batch: dict[str, Tensor], **kwargs) -> Tensor:
|
|
return self.net(batch["observation.state"])
|
|
|
|
def select_action(self, batch: dict[str, Tensor], **kwargs) -> Tensor:
|
|
return self.net(batch["observation.state"])
|
|
|
|
|
|
def make_dummy_policy(repo_id: str | None = None) -> DummyCheckpointPolicy:
|
|
config = DummyCheckpointConfig(device="cpu")
|
|
if repo_id is not None:
|
|
config.repo_id = repo_id
|
|
policy = DummyCheckpointPolicy(config)
|
|
with torch.no_grad():
|
|
policy.net.weight.fill_(0.5)
|
|
return policy
|