diff --git a/docs/source/multi_gpu_training.mdx b/docs/source/multi_gpu_training.mdx index 0bec365c9..7907340c3 100644 --- a/docs/source/multi_gpu_training.mdx +++ b/docs/source/multi_gpu_training.mdx @@ -152,12 +152,21 @@ 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 thigs to look out for: +`model.safetensors`, loadable as usual with `Policy.from_pretrained(...)`. Two things to look out for: -- With mixed precision, (`bf16`/`fp16`) FSDP keeps an fp32 master copy, so the checkpoint is fp32 - (~2× the bf16 size on disk) and is cast back to the policy dtype on load. -- **Optimizer state is not saved under FSDP**, so **resume-from-checkpoint is not supported**. - Saved weights are fully usable for evaluation and fine-tuning. +- **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 diff --git a/src/lerobot/common/train_utils.py b/src/lerobot/common/train_utils.py index cd8d43381..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, @@ -99,6 +100,7 @@ def save_checkpoint( 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: @@ -133,6 +135,10 @@ def save_checkpoint( 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, state_dict=model_state_dict) @@ -146,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, ) @@ -157,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. @@ -170,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. @@ -192,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 @@ -206,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/scripts/lerobot_train.py b/src/lerobot/scripts/lerobot_train.py index a235d6248..59e12bf81 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, @@ -369,7 +371,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()) @@ -459,6 +466,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() @@ -557,31 +570,30 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None): train_tracker.reset_averages() if cfg.save_checkpoint and is_saving_step: - # All ranks must call get_state_dict; rank 0 gets the - # full state dict, others get an empty dict. + # 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 - model_state_dict = accelerator.get_state_dict(policy) + 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) - if is_fsdp: - # TODO(fsdp): sharded optimizer-state save/resume is not wired up yet. - logging.warning( - "FSDP checkpoint: saving model weights only (optimizer state skipped; " - "resume-from-checkpoint not supported under FSDP yet)." - ) save_checkpoint( checkpoint_dir=checkpoint_dir, step=step, cfg=cfg, policy=accelerator.unwrap_model(policy), - optimizer=None if is_fsdp else optimizer, + optimizer=optimizer, scheduler=lr_scheduler, preprocessor=preprocessor, 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: 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/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