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

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)

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",

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),
+    }

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

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)