From 73782447f2ca420f8d71c9bd0a169ece5968d2d6 Mon Sep 17 00:00:00 2001 From: Maxime Ellerbach Date: Mon, 22 Jun 2026 13:51:21 +0200 Subject: [PATCH] feat(train): FSDP checkpoint saving (#3810) * feat(train): FSDP checkpoint saving * adding docs for FSDP * adding a test for the fsdp checkpoint path * cleanup * fixing final upload to hub * refactored initial implementation to use torch fsdp api and adding new tests --- docs/source/multi_gpu_training.mdx | 55 ++++++++++++ src/lerobot/common/train_utils.py | 88 +++++++++++++++++-- src/lerobot/optim/__init__.py | 2 + src/lerobot/optim/optimizers.py | 35 ++++++-- src/lerobot/policies/pretrained.py | 43 +++++++++- src/lerobot/scripts/lerobot_train.py | 32 ++++++- tests/optim/test_optimizers.py | 39 +++++++++ tests/policies/test_policies.py | 24 ++++++ tests/training/test_multi_gpu.py | 121 ++++++++++++++++++++++++--- tests/utils/test_train_utils.py | 15 ++++ 10 files changed, 423 insertions(+), 31 deletions(-) diff --git a/docs/source/multi_gpu_training.mdx b/docs/source/multi_gpu_training.mdx index d7369e8f8..7907340c3 100644 --- a/docs/source/multi_gpu_training.mdx +++ b/docs/source/multi_gpu_training.mdx @@ -113,6 +113,61 @@ accelerate launch --num_processes=2 $(which lerobot-train) \ --policy=act ``` +## Training Large Models with FSDP + +DDP replicates the full model on every GPU, so a model that doesn't fit on one GPU won't fit under +DDP either. For large models, use **FSDP** (Fully Sharded Data Parallel), which shards parameters, +gradients, and optimizer state across GPUs. See the [accelerate FSDP guide](https://huggingface.co/docs/accelerate/usage_guides/fsdp) for background. + +An example on how to launch LeRobot training with FSDP across 4 GPUs (1 machine): + +```bash +accelerate launch --config_file fsdp.yaml --num_processes=4 $(which lerobot-train) \ + --dataset.repo_id=${HF_USER}/my_dataset \ + --policy.type= \ + --output_dir=outputs/train/my_policy_fsdp +``` + +A minimal `fsdp.yaml` (FSDP1; shards params/grads/optimizer — ZeRO-3-equivalent): + +```yaml +compute_environment: LOCAL_MACHINE +distributed_type: FSDP +mixed_precision: bf16 +num_machines: 1 +num_processes: 4 +fsdp_config: + fsdp_version: 1 + fsdp_sharding_strategy: FULL_SHARD # params + grads + optimizer (ZeRO-3) + fsdp_auto_wrap_policy: TRANSFORMER_BASED_WRAP + fsdp_transformer_layer_cls_to_wrap: # repeated block class to shard + fsdp_use_orig_params: true # required: optimizer is built pre-prepare + fsdp_state_dict_type: FULL_STATE_DICT +``` + +Set `fsdp_transformer_layer_cls_to_wrap` to your model's repeated transformer-block class so each +block is sharded as its own unit. `fsdp_use_orig_params: true` is required because LeRobot builds the +optimizer before `accelerator.prepare()`. + +### FSDP checkpoints + +LeRobot gathers the full state dict across all ranks and the main process writes it as a single +`model.safetensors`, loadable as usual with `Policy.from_pretrained(...)`. Two things to look out for: + +- **Checkpoints store fp32 weights.** Under mixed precision (`bf16`/`fp16`) FSDP keeps an fp32 master + copy, and the checkpoint saves it (~2× the bf16 size on disk) so training can resume consistently + with the fp32 optimizer state; `from_pretrained` casts back to the policy dtype on load. FSDP-specific + caveat: an fp32 checkpoint is materialized in full precision on the target device _before_ casting, + so loading it for inference on a tight GPU can OOM even when the bf16 model would fit — load on CPU + first, or cast `model.safetensors` to the deployment dtype offline. +- The sharded optimizer state is gathered into a full (world-size-independent) state dict and saved + alongside the model in the same `optimizer_state.safetensors` / `optimizer_param_groups.json` + format as single-GPU training, so **resume-from-checkpoint is supported** with `--resume=true`. + Resume reshards both the model and the optimizer state to the _current_ FSDP topology, so you can + resume an FSDP checkpoint on a different number of GPUs. Note that the data sampler is only + sample-exact when the world size and batch size match the original run (a warning is logged + otherwise); the optimizer/model state itself is unaffected. + ## Notes - The `--policy.use_amp` flag in `lerobot-train` is only used when **not** running with accelerate. When using accelerate, mixed precision is controlled by accelerate's configuration. diff --git a/src/lerobot/common/train_utils.py b/src/lerobot/common/train_utils.py index 2d23b4003..5ae593bb8 100644 --- a/src/lerobot/common/train_utils.py +++ b/src/lerobot/common/train_utils.py @@ -21,6 +21,7 @@ from torch.optim.lr_scheduler import LRScheduler from lerobot.configs.train import TrainPipelineConfig from lerobot.optim import ( load_optimizer_state, + load_optimizer_state_dict, load_scheduler_state, save_optimizer_state, save_scheduler_state, @@ -98,6 +99,8 @@ def save_checkpoint( postprocessor: PolicyProcessorPipeline | None = None, num_processes: int | None = None, batch_size: int | None = None, + model_state_dict: dict | None = None, + optim_state_dict: dict | None = None, ) -> None: """This function creates the following directory structure: @@ -127,9 +130,18 @@ def save_checkpoint( resume. Defaults to None (not recorded). batch_size (int | None, optional): Per-process batch size to record for sample-exact resume. Defaults to None (not recorded). + model_state_dict: Pre-gathered full (unsharded) model state dict. Required under FSDP, + where `policy.state_dict()` would return sharded tensors; the caller gathers it via a + cross-rank collective and passes it here so rank 0 can write it directly. It holds + FSDP's fp32 master weights and is saved as-is (the loader casts to the policy dtype on + read). When None (DDP / single-GPU), the model is saved the normal way. Defaults to None. + optim_state_dict: Pre-gathered full (unsharded) optimizer state dict. Required under FSDP + (gathered alongside `model_state_dict` via `gather_fsdp_state_dicts`); saved in the same + safetensors format as the single-GPU path. When None, `optimizer.state_dict()` is used. + Defaults to None. """ pretrained_dir = checkpoint_dir / PRETRAINED_MODEL_DIR - policy.save_pretrained(pretrained_dir) + policy.save_pretrained(pretrained_dir, state_dict=model_state_dict) cfg.save_pretrained(pretrained_dir) if cfg.peft is not None: # When using PEFT, policy.save_pretrained will only write the adapter weights + config, not the @@ -140,7 +152,13 @@ def save_checkpoint( if postprocessor is not None: postprocessor.save_pretrained(pretrained_dir) save_training_state( - checkpoint_dir, step, optimizer, scheduler, num_processes=num_processes, batch_size=batch_size + checkpoint_dir, + step, + optimizer, + scheduler, + num_processes=num_processes, + batch_size=batch_size, + optim_state_dict=optim_state_dict, ) @@ -151,6 +169,7 @@ def save_training_state( scheduler: LRScheduler | None = None, num_processes: int | None = None, batch_size: int | None = None, + optim_state_dict: dict | None = None, ) -> None: """ Saves the training step, optimizer state, scheduler state, and rng state. @@ -164,19 +183,21 @@ def save_training_state( Defaults to None. num_processes (int | None, optional): Distributed world size to record. Defaults to None. batch_size (int | None, optional): Per-process batch size to record. Defaults to None. + optim_state_dict: Pre-gathered full optimizer state dict (for FSDP). Saved instead of + `optimizer.state_dict()` when provided. Defaults to None. """ save_dir = checkpoint_dir / TRAINING_STATE_DIR save_dir.mkdir(parents=True, exist_ok=True) save_training_step(train_step, save_dir, num_processes=num_processes, batch_size=batch_size) save_rng_state(save_dir) if optimizer is not None: - save_optimizer_state(optimizer, save_dir) + save_optimizer_state(optimizer, save_dir, optim_state_dict=optim_state_dict) if scheduler is not None: save_scheduler_state(scheduler, save_dir) def load_training_state( - checkpoint_dir: Path, optimizer: Optimizer, scheduler: LRScheduler | None + checkpoint_dir: Path, optimizer: Optimizer, scheduler: LRScheduler | None, load_optimizer: bool = True ) -> tuple[int, Optimizer, LRScheduler | None]: """ Loads the training step, optimizer state, scheduler state, and rng state. @@ -186,6 +207,10 @@ def load_training_state( checkpoint_dir (Path): The checkpoint directory. Should contain a 'training_state' dir. optimizer (Optimizer): The optimizer to load the state_dict to. scheduler (LRScheduler | None): The scheduler to load the state_dict to (can be None). + load_optimizer (bool, optional): Whether to load the optimizer state from disk. Defaults to + True. Set to False under FSDP, where the sharded optimizer state must be loaded after + `accelerator.prepare()` via `load_fsdp_optimizer_state` (the optimizer is returned + untouched here). Raises: NotADirectoryError: If 'checkpoint_dir' doesn't contain a 'training_state' dir @@ -200,8 +225,61 @@ def load_training_state( load_rng_state(training_state_dir) step = load_training_step(training_state_dir) - optimizer = load_optimizer_state(optimizer, training_state_dir) + if load_optimizer: + optimizer = load_optimizer_state(optimizer, training_state_dir) if scheduler is not None: scheduler = load_scheduler_state(scheduler, training_state_dir) return step, optimizer, scheduler + + +def gather_fsdp_state_dicts(model, optimizer) -> tuple[dict, dict]: + """Gather the full (unsharded) model and optimizer state dicts under FSDP. + + `model.state_dict()` and `FSDP.optim_state_dict(...)` are cross-rank collectives, so this must be + called on *every* rank with the prepared (FSDP-wrapped) `model` and `optimizer`. With + `rank0_only=True` and `offload_to_cpu=True`, every rank runs the all-gather but only rank 0 + materializes the full dicts (the others get empty dicts) and they are kept on CPU to bound GPU + memory. The returned optimizer state dict is keyed by parameter FQNs and is world-size + independent; `load_fsdp_optimizer_state` reshards it on resume. + + Returns: + (model_state_dict, optim_state_dict): full dicts on rank 0, empty dicts on other ranks. + """ + from torch.distributed.fsdp import ( + FullOptimStateDictConfig, + FullStateDictConfig, + FullyShardedDataParallel as FSDP, # noqa F401 + StateDictType, + ) + + state_cfg = FullStateDictConfig(offload_to_cpu=True, rank0_only=True) + optim_cfg = FullOptimStateDictConfig(offload_to_cpu=True, rank0_only=True) + with FSDP.state_dict_type(model, StateDictType.FULL_STATE_DICT, state_cfg, optim_cfg): + model_state_dict = model.state_dict() + optim_state_dict = FSDP.optim_state_dict(model, optimizer) + return model_state_dict, optim_state_dict + + +def load_fsdp_optimizer_state(model, optimizer, checkpoint_dir: Path) -> None: + """Load the FSDP optimizer state (saved as safetensors) and reshard it into the optimizer. + + This is a cross-rank collective and must be called on every rank *after* `accelerator.prepare()` + with the prepared (FSDP-wrapped) `model` and `optimizer`. The saved state is the full, + world-size-independent optimizer state (keyed by parameter FQNs); `FSDP.optim_state_dict_to_load` + reshards it to the current FSDP topology, so resume on a different number of GPUs works. + """ + from torch.distributed.fsdp import ( + FullOptimStateDictConfig, + FullStateDictConfig, + FullyShardedDataParallel as FSDP, # noqa F401 + StateDictType, + ) + + # Every rank reads the same full state from the (shared) checkpoint dir, so rank0_only=False. + full_osd = load_optimizer_state_dict(checkpoint_dir / TRAINING_STATE_DIR) + state_cfg = FullStateDictConfig(rank0_only=False) + optim_cfg = FullOptimStateDictConfig(rank0_only=False) + with FSDP.state_dict_type(model, StateDictType.FULL_STATE_DICT, state_cfg, optim_cfg): + sharded_osd = FSDP.optim_state_dict_to_load(model=model, optim=optimizer, optim_state_dict=full_osd) + optimizer.load_state_dict(sharded_osd) diff --git a/src/lerobot/optim/__init__.py b/src/lerobot/optim/__init__.py index 46676027b..2d564c25f 100644 --- a/src/lerobot/optim/__init__.py +++ b/src/lerobot/optim/__init__.py @@ -20,6 +20,7 @@ from .optimizers import ( SGDConfig as SGDConfig, XVLAAdamWConfig as XVLAAdamWConfig, load_optimizer_state, + load_optimizer_state_dict, save_optimizer_state, ) from .schedulers import ( @@ -50,6 +51,7 @@ __all__ = [ "VQBeTSchedulerConfig", # State management "load_optimizer_state", + "load_optimizer_state_dict", "load_scheduler_state", "save_optimizer_state", "save_scheduler_state", diff --git a/src/lerobot/optim/optimizers.py b/src/lerobot/optim/optimizers.py index 0bdd7a37e..0a462e1aa 100644 --- a/src/lerobot/optim/optimizers.py +++ b/src/lerobot/optim/optimizers.py @@ -27,7 +27,7 @@ from lerobot.utils.constants import ( OPTIMIZER_PARAM_GROUPS, OPTIMIZER_STATE, ) -from lerobot.utils.io_utils import deserialize_json_into_object, write_json +from lerobot.utils.io_utils import deserialize_json_into_object, load_json, write_json from lerobot.utils.utils import flatten_dict, unflatten_dict # Type alias for parameters accepted by optimizer build() methods. @@ -281,28 +281,37 @@ class MultiAdamConfig(OptimizerConfig): def save_optimizer_state( - optimizer: torch.optim.Optimizer | dict[str, torch.optim.Optimizer], save_dir: Path + optimizer: torch.optim.Optimizer | dict[str, torch.optim.Optimizer], + save_dir: Path, + optim_state_dict: dict | None = None, ) -> None: """Save optimizer state to disk. Args: optimizer: Either a single optimizer or a dictionary of optimizers. save_dir: Directory to save the optimizer state. + optim_state_dict: Pre-gathered optimizer state dict (for FSDP, where the sharded state must + be gathered across ranks first). If provided, it is saved directly instead of calling + ``optimizer.state_dict()``. Only supported for a single optimizer. Defaults to None. """ if isinstance(optimizer, dict): # Handle dictionary of optimizers + if optim_state_dict is not None: + raise ValueError("optim_state_dict is not supported for a dict of optimizers") for name, opt in optimizer.items(): optimizer_dir = save_dir / name optimizer_dir.mkdir(exist_ok=True, parents=True) _save_single_optimizer_state(opt, optimizer_dir) else: # Handle single optimizer - _save_single_optimizer_state(optimizer, save_dir) + _save_single_optimizer_state(optimizer, save_dir, optim_state_dict=optim_state_dict) -def _save_single_optimizer_state(optimizer: torch.optim.Optimizer, save_dir: Path) -> None: +def _save_single_optimizer_state( + optimizer: torch.optim.Optimizer, save_dir: Path, optim_state_dict: dict | None = None +) -> None: """Save a single optimizer's state to disk.""" - state = optimizer.state_dict() + state = dict(optim_state_dict) if optim_state_dict is not None else optimizer.state_dict() param_groups = state.pop("param_groups") flat_state = flatten_dict(state) save_file(flat_state, save_dir / OPTIMIZER_STATE) @@ -356,3 +365,19 @@ def _load_single_optimizer_state(optimizer: torch.optim.Optimizer, save_dir: Pat optimizer.load_state_dict(loaded_state_dict) return optimizer + + +def load_optimizer_state_dict(save_dir: Path) -> dict: + """Read a saved optimizer state dict (safetensors + json) back into a plain dict. + + Unlike `load_optimizer_state`, this does not load into an optimizer and preserves the original + ``state`` keys verbatim (e.g. FSDP parameter FQNs, which are not integer-castable). It is used by + the FSDP resume path, where the full state must be resharded via `FSDP.optim_state_dict_to_load` + before being loaded into the (sharded) optimizer. + """ + flat_state = load_file(save_dir / OPTIMIZER_STATE) + state = unflatten_dict(flat_state) + return { + "state": state.get("state", {}), + "param_groups": load_json(save_dir / OPTIMIZER_PARAM_GROUPS), + } diff --git a/src/lerobot/policies/pretrained.py b/src/lerobot/policies/pretrained.py index a69487f3f..a7aabb3f3 100644 --- a/src/lerobot/policies/pretrained.py +++ b/src/lerobot/policies/pretrained.py @@ -23,7 +23,7 @@ from typing import TypedDict, TypeVar, Unpack import packaging import safetensors -from huggingface_hub import HfApi, ModelCard, ModelCardData, hf_hub_download +from huggingface_hub import HfApi, ModelCard, ModelCardData, hf_hub_download, save_torch_state_dict from huggingface_hub.constants import SAFETENSORS_SINGLE_FILE from huggingface_hub.errors import HfHubHTTPError from safetensors.torch import load_model as load_model_as_safetensor, save_model as save_model_as_safetensor @@ -129,10 +129,43 @@ class PreTrainedPolicy(nn.Module, HubMixin, abc.ABC): if not getattr(cls, "name", None): raise TypeError(f"Class {cls.__name__} must define 'name'") - def _save_pretrained(self, save_directory: Path) -> None: + def save_pretrained( + self, + save_directory: str | Path, + *, + state_dict: dict[str, Tensor] | None = None, + repo_id: str | None = None, + push_to_hub: bool = False, + card_kwargs: dict | None = None, + **push_to_hub_kwargs, + ) -> str | None: + """Save the policy to a directory (and optionally push to the Hub). + + Overrides `HubMixin.save_pretrained` to add a `state_dict` argument (mirroring + `transformers.PreTrainedModel.save_pretrained`). Under FSDP, `self.state_dict()` would + return sharded tensors, so the caller gathers the full state dict via a cross-rank + collective and passes it here for `_save_pretrained` to write directly. + """ + save_directory = Path(save_directory) + save_directory.mkdir(parents=True, exist_ok=True) + self._save_pretrained(save_directory, state_dict=state_dict) + if push_to_hub: + if repo_id is None: + repo_id = save_directory.name + return self.push_to_hub(repo_id=repo_id, card_kwargs=card_kwargs, **push_to_hub_kwargs) + return None + + def _save_pretrained(self, save_directory: Path, state_dict: dict[str, Tensor] | None = None) -> None: self.config._save_pretrained(save_directory) model_to_save = self.module if hasattr(self, "module") else self - save_model_as_safetensor(model_to_save, str(save_directory / SAFETENSORS_SINGLE_FILE)) + if state_dict is None: + save_model_as_safetensor(model_to_save, str(save_directory / SAFETENSORS_SINGLE_FILE)) + return + # A pre-gathered (e.g. FSDP full) state dict was supplied: write it directly. + # `save_torch_state_dict` discards shared-tensor duplicates just like `save_model` does; + # pin `max_shard_size` above the total size so the output stays a single `model.safetensors` + total_bytes = sum(t.numel() * t.element_size() for t in state_dict.values()) + save_torch_state_dict(state_dict, str(save_directory), max_shard_size=max(total_bytes, 1)) @classmethod def from_pretrained( @@ -270,6 +303,7 @@ class PreTrainedPolicy(nn.Module, HubMixin, abc.ABC): self, cfg: TrainPipelineConfig, peft_model=None, + state_dict: dict[str, Tensor] | None = None, ): api = HfApi() repo_id = api.create_repo( @@ -287,7 +321,8 @@ class PreTrainedPolicy(nn.Module, HubMixin, abc.ABC): peft_model.save_pretrained(saved_path) self.config.save_pretrained(saved_path) else: - self.save_pretrained(saved_path) # Calls _save_pretrained and stores model tensors + # Calls _save_pretrained and stores model tensors + self.save_pretrained(saved_path, state_dict=state_dict) card = self.generate_model_card( cfg.dataset.repo_id, self.config.type, self.config.license, self.config.tags, cfg=cfg diff --git a/src/lerobot/scripts/lerobot_train.py b/src/lerobot/scripts/lerobot_train.py index c94223c62..45281dac9 100644 --- a/src/lerobot/scripts/lerobot_train.py +++ b/src/lerobot/scripts/lerobot_train.py @@ -34,8 +34,10 @@ from torch.optim import Optimizer from tqdm import tqdm from lerobot.common.train_utils import ( + gather_fsdp_state_dicts, get_step_checkpoint_dir, get_step_identifier, + load_fsdp_optimizer_state, load_training_batch_size, load_training_num_processes, load_training_state, @@ -189,6 +191,7 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None): require_package("accelerate", extra="training") from accelerate import Accelerator + from accelerate.utils import DistributedDataParallelKwargs, DistributedType cfg.validate() @@ -197,8 +200,6 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None): # We set step_scheduler_with_optimizer=False to prevent accelerate from adjusting the lr_scheduler steps based on the num_processes # We set find_unused_parameters=True to handle models with conditional computation if accelerator is None: - from accelerate.utils import DistributedDataParallelKwargs - ddp_kwargs = DistributedDataParallelKwargs(find_unused_parameters=True) # Accelerate auto-detects the device based on the available hardware and ignores the policy.device setting. # Force the device to be CPU when the active config's device is set to CPU (works for both policy and reward model training). @@ -371,7 +372,12 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None): step = 0 # number of policy updates (forward + backward + optim) if cfg.resume: - step, optimizer, lr_scheduler = load_training_state(cfg.checkpoint_path, optimizer, lr_scheduler) + # Under FSDP the optimizer state is sharded and must be loaded after `accelerator.prepare()` + # (see load_fsdp_optimizer_state below), so skip the optimizer here and load it then. + is_fsdp = accelerator.distributed_type == DistributedType.FSDP + step, optimizer, lr_scheduler = load_training_state( + cfg.checkpoint_path, optimizer, lr_scheduler, load_optimizer=not is_fsdp + ) num_learnable_params = sum(p.numel() for p in policy.parameters() if p.requires_grad) num_total_params = sum(p.numel() for p in policy.parameters()) @@ -461,6 +467,12 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None): policy, optimizer, dataloader, lr_scheduler = accelerator.prepare( policy, optimizer, dataloader, lr_scheduler ) + + # FSDP optimizer state is sharded across ranks, so it can only be loaded once the optimizer and + # model are FSDP-wrapped (i.e. after `prepare`). Collective: every rank must participate. + if cfg.resume and accelerator.distributed_type == DistributedType.FSDP: + load_fsdp_optimizer_state(policy, optimizer, cfg.checkpoint_path) + dl_iter = cycle(dataloader) policy.train() @@ -559,6 +571,14 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None): train_tracker.reset_averages() if cfg.save_checkpoint and is_saving_step: + # Under FSDP, gathering the full model + optimizer state dicts is a cross-rank collective, + # so all ranks must participate; rank 0 then writes the materialized dicts. For DDP / + # single-GPU the state dicts are saved the normal way inside save_checkpoint. + is_fsdp = accelerator.distributed_type == DistributedType.FSDP + if is_fsdp: + model_state_dict, optim_state_dict = gather_fsdp_state_dicts(policy, optimizer) + else: + model_state_dict, optim_state_dict = None, None if is_main_process: logging.info(f"Checkpoint policy after step {step}") checkpoint_dir = get_step_checkpoint_dir(cfg.output_dir, cfg.steps, step) @@ -573,6 +593,8 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None): postprocessor=postprocessor, num_processes=accelerator.num_processes, batch_size=cfg.batch_size, + model_state_dict=model_state_dict, + optim_state_dict=optim_state_dict, ) update_last_checkpoint(checkpoint_dir) if wandb_logger: @@ -635,6 +657,8 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None): if eval_env: close_envs(eval_env) + is_fsdp = accelerator.distributed_type == DistributedType.FSDP + model_state_dict = accelerator.get_state_dict(policy) if is_fsdp else None if is_main_process: logging.info("End of training") @@ -644,7 +668,7 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None): if not cfg.is_reward_model_training and cfg.policy.use_peft: unwrapped_model.push_model_to_hub(cfg, peft_model=unwrapped_model) else: - unwrapped_model.push_model_to_hub(cfg) + unwrapped_model.push_model_to_hub(cfg, state_dict=model_state_dict) preprocessor.push_to_hub(active_cfg.repo_id) postprocessor.push_to_hub(active_cfg.repo_id) diff --git a/tests/optim/test_optimizers.py b/tests/optim/test_optimizers.py index d18565562..5b480f70d 100644 --- a/tests/optim/test_optimizers.py +++ b/tests/optim/test_optimizers.py @@ -20,6 +20,7 @@ from lerobot.optim.optimizers import ( MultiAdamConfig, SGDConfig, load_optimizer_state, + load_optimizer_state_dict, save_optimizer_state, ) from lerobot.utils.constants import ( @@ -65,6 +66,44 @@ 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 { diff --git a/tests/policies/test_policies.py b/tests/policies/test_policies.py index e9388b3ed..285b87d4c 100644 --- a/tests/policies/test_policies.py +++ b/tests/policies/test_policies.py @@ -23,6 +23,7 @@ import torch pytest.importorskip("datasets", reason="datasets is required (install lerobot[dataset])") +from huggingface_hub.constants import SAFETENSORS_SINGLE_FILE from packaging import version from safetensors.torch import load_file @@ -300,6 +301,29 @@ def test_save_and_load_pretrained(dummy_dataset_metadata, tmp_path, policy_name: torch.testing.assert_close(list(policy.parameters()), list(loaded_policy.parameters()), rtol=0, atol=0) +def test_save_pretrained_with_state_dict(dummy_dataset_metadata, tmp_path): + """Exercise the FSDP checkpoint path: save_pretrained with a pre-gathered state_dict.""" + policy_cls = get_policy_class("act") + policy_cfg = make_policy_config("act") + features = dataset_to_policy_features(dummy_dataset_metadata.features) + policy_cfg.output_features = {key: ft for key, ft in features.items() if ft.type is FeatureType.ACTION} + policy_cfg.input_features = { + key: ft for key, ft in features.items() if key not in policy_cfg.output_features + } + policy = policy_cls(policy_cfg) + policy.to(policy_cfg.device) + + save_dir = tmp_path / "fsdp_state_dict" + policy.save_pretrained(save_dir, state_dict=policy.state_dict()) + + # A single, unsharded safetensors file (no sharded set + index). + assert (save_dir / SAFETENSORS_SINGLE_FILE).is_file() + assert not (save_dir / f"{SAFETENSORS_SINGLE_FILE}.index.json").exists() + + loaded_policy = policy_cls.from_pretrained(save_dir, config=policy_cfg) + torch.testing.assert_close(list(policy.parameters()), list(loaded_policy.parameters()), rtol=0, atol=0) + + @pytest.mark.parametrize("multikey", [True, False]) def test_multikey_construction(multikey: bool): """ diff --git a/tests/training/test_multi_gpu.py b/tests/training/test_multi_gpu.py index 638dc3131..d37f1e35d 100644 --- a/tests/training/test_multi_gpu.py +++ b/tests/training/test_multi_gpu.py @@ -58,7 +58,46 @@ def download_dataset(repo_id, episodes): print(f"Dataset {repo_id} downloaded successfully") -def run_accelerate_training(config_args, num_processes=4, temp_dir=None): +def _write_multi_gpu_config(f, num_processes): + f.write("compute_environment: LOCAL_MACHINE\n") + f.write("distributed_type: MULTI_GPU\n") + f.write("mixed_precision: 'no'\n") + f.write(f"num_processes: {num_processes}\n") + f.write("use_cpu: false\n") + f.write("gpu_ids: all\n") + f.write("downcast_bf16: 'no'\n") + f.write("machine_rank: 0\n") + f.write("main_training_function: main\n") + f.write("num_machines: 1\n") + f.write("rdzv_backend: static\n") + f.write("same_network: true\n") + + +def _write_fsdp_config(f, num_processes): + # FSDP1 with FULL_SHARD (ZeRO-3-equivalent) and FULL_STATE_DICT, matching + # docs/source/multi_gpu_training.mdx. ACT's repeated transformer blocks are the wrap units; + # fsdp_use_orig_params is required because LeRobot builds the optimizer before prepare(). + f.write("compute_environment: LOCAL_MACHINE\n") + f.write("distributed_type: FSDP\n") + f.write("mixed_precision: 'no'\n") + f.write(f"num_processes: {num_processes}\n") + f.write("use_cpu: false\n") + f.write("gpu_ids: all\n") + f.write("machine_rank: 0\n") + f.write("main_training_function: main\n") + f.write("num_machines: 1\n") + f.write("rdzv_backend: static\n") + f.write("same_network: true\n") + f.write("fsdp_config:\n") + f.write(" fsdp_version: 1\n") + f.write(" fsdp_sharding_strategy: FULL_SHARD\n") + f.write(" fsdp_auto_wrap_policy: TRANSFORMER_BASED_WRAP\n") + f.write(" fsdp_transformer_layer_cls_to_wrap: ACTEncoderLayer,ACTDecoderLayer\n") + f.write(" fsdp_use_orig_params: true\n") + f.write(" fsdp_state_dict_type: FULL_STATE_DICT\n") + + +def run_accelerate_training(config_args, num_processes=4, temp_dir=None, distributed_type="MULTI_GPU"): """ Helper function to run training with accelerate launch. @@ -66,6 +105,7 @@ def run_accelerate_training(config_args, num_processes=4, temp_dir=None): config_args: List of config arguments to pass to lerobot_train.py num_processes: Number of processes (GPUs) to use temp_dir: Temporary directory for outputs + distributed_type: "MULTI_GPU" (DDP) or "FSDP" — selects the generated accelerate config. Returns: subprocess.CompletedProcess result @@ -75,18 +115,10 @@ def run_accelerate_training(config_args, num_processes=4, temp_dir=None): # Write YAML config with open(config_path, "w") as f: - f.write("compute_environment: LOCAL_MACHINE\n") - f.write("distributed_type: MULTI_GPU\n") - f.write("mixed_precision: 'no'\n") - f.write(f"num_processes: {num_processes}\n") - f.write("use_cpu: false\n") - f.write("gpu_ids: all\n") - f.write("downcast_bf16: 'no'\n") - f.write("machine_rank: 0\n") - f.write("main_training_function: main\n") - f.write("num_machines: 1\n") - f.write("rdzv_backend: static\n") - f.write("same_network: true\n") + if distributed_type == "FSDP": + _write_fsdp_config(f, num_processes) + else: + _write_multi_gpu_config(f, num_processes) cmd = [ "accelerate", @@ -211,3 +243,66 @@ class TestMultiGPUTraining: # Verify optimizer state exists optimizer_state = training_state_dir / "optimizer_state.safetensors" assert optimizer_state.exists(), f"No optimizer state in checkpoint {checkpoint_dir}" + + def test_fsdp_optimizer_save_and_resume(self): + """ + Test that FSDP saves the (gathered) optimizer state and can resume from it. + + Trains a few steps under FSDP, verifies the gathered optimizer state is written next to the + rest of the training state, then resumes from the checkpoint for more steps and checks it + completes without shape/key errors in the FSDP optimizer load path. + """ + # Pre-download dataset to avoid race conditions + download_dataset("lerobot/pusht", episodes=[0]) + + with tempfile.TemporaryDirectory() as temp_dir: + output_dir = Path(temp_dir) / "outputs" + + config_args = [ + "--dataset.repo_id=lerobot/pusht", + "--dataset.episodes=[0]", + "--policy.type=act", + "--policy.device=cuda", + "--policy.push_to_hub=false", + f"--output_dir={output_dir}", + "--batch_size=4", + "--steps=10", + "--eval_freq=-1", + "--log_freq=5", + "--save_freq=10", + "--seed=42", + "--num_workers=0", + ] + + result = run_accelerate_training( + config_args, num_processes=2, temp_dir=temp_dir, distributed_type="FSDP" + ) + assert result.returncode == 0, ( + f"FSDP training failed:\nSTDOUT:\n{result.stdout}\n\nSTDERR:\n{result.stderr}" + ) + + # The gathered optimizer state must be written under FSDP (proves the save collective ran), + # in the same safetensors format as single-GPU training. + training_state_dir = output_dir / "checkpoints" / "last" / "training_state" + optimizer_state = training_state_dir / "optimizer_state.safetensors" + optimizer_param_groups = training_state_dir / "optimizer_param_groups.json" + assert optimizer_state.exists(), f"FSDP optimizer state not saved in {training_state_dir}" + assert optimizer_param_groups.exists(), ( + f"FSDP optimizer param groups not saved in {training_state_dir}" + ) + + # Resume from the checkpoint for more steps. A successful run proves load_fsdp_optimizer + # accepts the saved state and reshards it without shape/key errors. + resume_config = output_dir / "checkpoints" / "last" / "pretrained_model" / "train_config.json" + resume_args = [ + f"--config_path={resume_config}", + "--resume=true", + "--steps=20", + ] + resume_result = run_accelerate_training( + resume_args, num_processes=2, temp_dir=temp_dir, distributed_type="FSDP" + ) + assert resume_result.returncode == 0, ( + f"FSDP resume failed:\nSTDOUT:\n{resume_result.stdout}\n\nSTDERR:\n{resume_result.stderr}" + ) + assert "End of training" in resume_result.stdout or "End of training" in resume_result.stderr diff --git a/tests/utils/test_train_utils.py b/tests/utils/test_train_utils.py index c171763c2..e3705409b 100644 --- a/tests/utils/test_train_utils.py +++ b/tests/utils/test_train_utils.py @@ -136,3 +136,18 @@ def test_save_load_training_state(tmp_path, optimizer, scheduler): assert loaded_step == 10 assert loaded_optimizer is optimizer assert loaded_scheduler is scheduler + + +def test_load_training_state_skip_optimizer(tmp_path, optimizer, scheduler): + # FSDP loads optimizer separately (after accelerator.prepare) + # load_training_state(load_optimizer=False) must restore step + scheduler but leave the + # optimizer untouched and never touch the on-disk optimizer state. + save_training_state(tmp_path, 10, optimizer, scheduler) + with patch("lerobot.common.train_utils.load_optimizer_state") as mock_load_optimizer_state: + loaded_step, loaded_optimizer, loaded_scheduler = load_training_state( + tmp_path, optimizer, scheduler, load_optimizer=False + ) + mock_load_optimizer_state.assert_not_called() + assert loaded_step == 10 + assert loaded_optimizer is optimizer + assert loaded_scheduler is scheduler