diff --git a/docs/source/bring_your_own_policies.mdx b/docs/source/bring_your_own_policies.mdx index 697a07691..332ab73eb 100644 --- a/docs/source/bring_your_own_policies.mdx +++ b/docs/source/bring_your_own_policies.mdx @@ -161,6 +161,16 @@ The methods called by the train/eval loops: Batches are flat dictionaries keyed by the constants in [`lerobot.utils.constants`](https://github.com/huggingface/lerobot/blob/main/src/lerobot/utils/constants.py): `OBS_STATE` (`observation.state.`), `OBS_IMAGES` (`observation.images.`), `OBS_LANGUAGE`, `ACTION`, etc. Reuse the constants — don't invent new prefixes. +If your model is large enough to warrant [sharded multi-GPU training](./multi_gpu_training#sharded-training-fsdp), also declare its FSDP wrap units — the repeated block classes sharding operates on: + +```python +class MyPolicy(PreTrainedPolicy): + ... + _fsdp_wrap_modules = ["MyTransformerBlock"] +``` + +With this one declaration, `--parallelism.dp_shard=N` works out of the box for your policy (users can still override it with `--accelerator.fsdp.wrap_modules`). Without any wrap source, sharded runs fail at startup by design. + ### Processor functions LeRobot uses `PolicyProcessorPipeline`s to normalize inputs and de-normalize outputs around your policy. For a concrete reference, see [`processor_act.py`](https://github.com/huggingface/lerobot/blob/main/src/lerobot/policies/act/processor_act.py) or [`processor_diffusion.py`](https://github.com/huggingface/lerobot/blob/main/src/lerobot/policies/diffusion/processor_diffusion.py). @@ -300,7 +310,7 @@ The file names are load-bearing: the factory does lazy imports by name, and the Two places need to know about your policy. All by name. 1. **`policies/__init__.py`** — re-export `MyPolicyConfig` and add it to `__all__`. This import is what registers your policy: `@PreTrainedConfig.register_subclass("my_policy")` runs, and from then on the factory resolves everything by convention. **Don't** re-export the modeling class; it loads lazily through the factory (so `import lerobot` stays fast). -2. **`templates/lerobot_modelcard_template.md` and the root `README.md`** — the template is what `push_model_to_hub` renders into the model card of every checkpoint trained with your policy: add a one-line description of your policy in the `model_name` branches, map it in `policy_docs` so cards link to your MDX guide, and optionally add an architecture image to `diagrams`. Then add your policy to the models table in the root `README.md`, under the right category, linking to your doc page. +2. **`templates/lerobot_modelcard_template.md` and the root `README.md`** — the template is what the end-of-training publisher renders into the model card of every checkpoint trained with your policy: add a one-line description of your policy in the `model_name` branches, map it in `policy_docs` so cards link to your MDX guide, and optionally add an architecture image to `diagrams`. Then add your policy to the models table in the root `README.md`, under the right category, linking to your doc page. Mirror an existing policy that's structurally similar to yours; the diff is small. @@ -344,7 +354,7 @@ A new policy is much easier to review — and far more useful — when it ships **Pick at least one in-tree benchmark.** LeRobot ships sim benchmarks with per-benchmark Docker images (LIBERO, LIBERO-plus, Meta-World, RoboTwin 2.0, RoboCasa365, RoboCerebra, RoboMME, VLABench and more). Pick the one that matches your policy's modality — VLAs usually go to LIBERO or VLABench; image-only BC to LIBERO or Meta-World. The full list lives under [Benchmarks](./libero) in the docs sidebar. -**Push the checkpoint & processors** to the Hub under `lerobot/_` (or your namespace if you don't have write access; a maintainer can mirror it). Use `PreTrainedPolicy.push_model_to_hub` so the repo gets `config.json`, `model.safetensors`, and a model card. +**Push the checkpoint & processors** to the Hub under `lerobot/_` (or your namespace if you don't have write access; a maintainer can mirror it). The easiest way is training with `--policy.repo_id=/` and `--policy.push_to_hub=true`: `lerobot-train` publishes the model, both processors, and a model card at the end of the run. To publish an existing checkpoint after the fact, upload its `pretrained_model/` directory (e.g. `huggingface-cli upload`), or use `lerobot-convert-dcp --push_to_hub=...` for sharded-format checkpoints. **Report results in your policy's MDX**, with the exact `lerobot-eval` command and hardware so anyone can re-run: diff --git a/docs/source/multi_gpu_training.mdx b/docs/source/multi_gpu_training.mdx index 7c212364e..82c8596f8 100644 --- a/docs/source/multi_gpu_training.mdx +++ b/docs/source/multi_gpu_training.mdx @@ -1,28 +1,29 @@ # Multi-GPU Training -This guide shows you how to train policies on multiple GPUs using [Hugging Face Accelerate](https://huggingface.co/docs/accelerate). +LeRobot trains on multiple GPUs through [Hugging Face Accelerate](https://huggingface.co/docs/accelerate). Three data-parallel layouts are supported: + +| Layout | What it does | Config | +| -------- | ------------------------------------------------------------- | ------------------------------------------------------- | +| **DDP** | Replicates the full model on every GPU | default on any multi-GPU launch | +| **FSDP** | Shards parameters, gradients, and optimizer state across GPUs | `--parallelism.dp_shard=N` | +| **HSDP** | Shards within groups of GPUs, replicates across groups | `--parallelism.dp_replicate=R --parallelism.dp_shard=S` | ## Installation -`accelerate` is included in the `training` extra. Install it with: +`accelerate` is included in the `training` extra: ```bash pip install 'lerobot[training]' ``` -## Training with Multiple GPUs +## Launching -You can launch training in two ways: +Distributed training can be launched through both `torchrun` and `accelerate launch`. Accelerate is used as a plain launcher: it does not manage the training configuration, and every distributed training setting lives in LeRobot's own config system. -### Option 1: Without config (specify parameters directly) - -You can specify all parameters directly in the command without running `accelerate config`: +With `torchrun`: ```bash -accelerate launch \ - --multi_gpu \ - --num_processes=2 \ - $(which lerobot-train) \ +torchrun --nproc-per-node=2 $(which lerobot-train) \ --dataset.repo_id=${HF_USER}/my_dataset \ --policy.type=act \ --policy.repo_id=${HF_USER}/my_trained_policy \ @@ -31,32 +32,10 @@ accelerate launch \ --wandb.enable=true ``` -**Key accelerate parameters:** - -- `--multi_gpu`: Enable multi-GPU training -- `--num_processes=2`: Number of GPUs to use -- `--mixed_precision=fp16`: Use fp16 mixed precision (or `bf16` if supported) - -### Option 2: Using accelerate config - -If you prefer to save your configuration, you can optionally configure accelerate for your hardware setup by running: +With `accelerate launch` (as a plain launcher): ```bash -accelerate config -``` - -This interactive setup will ask you questions about your training environment (number of GPUs, mixed precision settings, etc.) and saves the configuration for future use. For a simple multi-GPU setup on a single machine, you can use these recommended settings: - -- Compute environment: This machine -- Number of machines: 1 -- Number of processes: (number of GPUs you want to use) -- GPU ids to use: (leave empty to use all) -- Mixed precision: fp16 or bf16 (recommended for faster training) - -Then launch training with: - -```bash -accelerate launch $(which lerobot-train) \ +accelerate launch --num_processes=2 $(which lerobot-train) \ --dataset.repo_id=${HF_USER}/my_dataset \ --policy.type=act \ --policy.repo_id=${HF_USER}/my_trained_policy \ @@ -65,116 +44,133 @@ accelerate launch $(which lerobot-train) \ --wandb.enable=true ``` -## How It Works +With no `--parallelism.*` flags, a multi-process launch runs plain DDP. Multi-node runs use the standard `torchrun --nnodes/--node-rank/--rdzv-endpoint` flags (or `accelerate launch --num_machines/--machine_rank/--main_process_ip`). -When you launch training with accelerate: +> [!WARNING] +> Accelerate's YAML config files (`accelerate launch --config_file some.yaml`, `accelerate config`) are not supported. They configure the engine through environment variables, bypassing LeRobot's configuration system, so `train_config.json` would no longer describe the settings a run actually used. `lerobot-train` therefore refuses to start when [accelerate environment variables](https://huggingface.co/docs/accelerate/usage_guides/fsdp) are set. Put the settings in `--parallelism.*` / `--accelerator.*` flags instead, or set `LEROBOT_ALLOW_ACCELERATE_ENV=1` to acknowledge the override and proceed anyway. -1. **Automatic detection**: LeRobot automatically detects if it's running under accelerate -2. **Data distribution**: Your batch is automatically split across GPUs -3. **Gradient synchronization**: Gradients are synchronized across GPUs during backpropagation -4. **Single process logging**: Only the main process logs to wandb and saves checkpoints +## Batch semantics, learning rate, and steps -## Learning Rate and Training Steps Scaling +Each of the `dp_replicate × dp_shard` data-parallel workers loads its own `--batch_size` micro-batch every step, so one training step consumes `batch_size × dp_world_size` samples, and `× gradient_accumulation_steps` of those go into each optimizer update: -**Important:** LeRobot does **NOT** automatically scale learning rates or training steps based on the number of GPUs. This gives you full control over your training hyperparameters. - -### Why No Automatic Scaling? - -Many distributed training frameworks automatically scale the learning rate by the number of GPUs (e.g., `lr = base_lr × num_gpus`). -However, LeRobot keeps the learning rate exactly as you specify it. - -### When and How to Scale - -If you want to scale your hyperparameters when using multiple GPUs, you should do it manually: - -**Learning Rate Scaling:** - -```bash -# Example: 2 GPUs with linear LR scaling -# Base LR: 1e-4, with 2 GPUs -> 2e-4 -accelerate launch --num_processes=2 $(which lerobot-train) \ - --optimizer.lr=2e-4 \ - --dataset.repo_id=lerobot/pusht \ - --policy.type=act +``` +effective_batch_size = batch_size × dp_world_size × gradient_accumulation_steps ``` -**Training Steps Scaling:** +The training banner prints this factorization at startup. `--steps` counts loop steps (micro-batches per worker), not optimizer updates. -Since the effective batch size `bs` increases with multiple GPUs (batch_size × num_gpus), you may want to reduce the number of training steps proportionally: +Gradient accumulation is a first-class flag: ```bash -# Example: 2 GPUs with effective batch size 2x larger -# Original: batch_size=8, steps=100000 -# With 2 GPUs: batch_size=8 (16 in total), steps=50000 -accelerate launch --num_processes=2 $(which lerobot-train) \ - --batch_size=8 \ - --steps=50000 \ - --dataset.repo_id=lerobot/pusht \ - --policy.type=act +torchrun --nproc-per-node=2 $(which lerobot-train) \ + --batch_size=8 --accelerator.gradient_accumulation.steps=4 ... ``` -## Training Large Models with FSDP +**LeRobot does not auto-scale the learning rate or the number of steps** when the effective batch size grows. If you scale out and want equivalent training, please adjust manually, e.g. with 2 GPUs: double `--optimizer.lr` (linear scaling), or halve `--steps`. -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. +## Sharded training (FSDP) -An example on how to launch LeRobot training with FSDP across 4 GPUs (1 machine): +If a model is too large to train with DDP, shard it with FSDP2: ```bash -accelerate launch --config_file fsdp.yaml --num_processes=4 $(which lerobot-train) \ +torchrun --nproc-per-node=4 $(which lerobot-train) \ --dataset.repo_id=${HF_USER}/my_dataset \ --policy.type= \ + --parallelism.dp_shard=4 \ + --accelerator.mixed_precision=bf16 \ --output_dir=outputs/train/my_policy_fsdp ``` -A minimal `fsdp.yaml` (FSDP1; shards params/grads/optimizer — ZeRO-3-equivalent): +`--parallelism.dp_shard=-1` shards over however many processes the launcher started. -```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 +### Wrap units + +FSDP shards the model in units (typically the repeated transformer block) and gathers one unit at a time during forward/backward. Policies declare their wrap units via `_fsdp_wrap_modules` on the policy class. For example, ACT declares `["ACTEncoderLayer", "ACTDecoderLayer"]` and FastWAM declares `["MoTLayer"]`. For a policy without a `_fsdp_wrap_modules` declaration, pass one of the flags below. You can specify the module class name explicitly, or use a size-based policy instead: + +```bash +--accelerator.fsdp.wrap_modules='["MyTransformerBlock"]' # explicit class names +--accelerator.fsdp.min_num_params=1000000 # or: wrap every submodule above 1M params ``` -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()`. +If a policy doesn't declare `_fsdp_wrap_modules` and no `--accelerator.fsdp.wrap_modules` or `--accelerator.fsdp.min_num_params` is passed, the run fails at startup rather than silently wrapping only the root module (which would forfeit all sharding memory savings). -### FSDP checkpoints +Other sharding settings: -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: +- `--accelerator.fsdp.reshard_after_forward`: whether to keep each unit's parameters resident after forward. +- `--accelerator.fsdp.cpu_offload`: keeps parameters, gradients and optimizer states on CPU. +- `--accelerator.fsdp.ignored_modules`: a regex of module paths to keep unsharded. -- **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. +### HSDP + +Hybrid Sharded Data Parallel: parameters, gradients and optimizer states are sharded across `dp_shard` ranks, and that sharding is replicated `dp_replicate` times. Parameter all-gathers and gradient reduce-scatters stay inside a shard group; only the all-reduce that synchronizes the replicas crosses between groups. The two degrees must multiply to the world size: + +```bash +# 16 GPUs = 2 nodes × 8: shard within each node, replicate across nodes +torchrun --nnodes=2 --nproc-per-node=8 ... $(which lerobot-train) \ + --parallelism.dp_replicate=2 --parallelism.dp_shard=8 ... +``` + +## Checkpoints + +Every checkpoint contains a `pretrained_model/` directory and a `training_state/` directory: + +```text +005000/ # the training step at that checkpoint +├── pretrained_model/ +│ ├── config.json # policy config +│ ├── train_config.json # the full training config +│ ├── model.safetensors # full weights (checkpoint_format ∈ {safetensors, safetensors_dcp}, or any non-sharded run) +│ ├── pytorch_model_fsdp_0/ # DCP weight shards (checkpoint_format ∈ {dcp, safetensors_dcp}) +│ ├── policy_preprocessor.json # preprocessor config (when the run has a preprocessor) +│ ├── policy_preprocessor_step_*.safetensors # state of the stateful preprocessor steps +│ ├── policy_postprocessor.json # postprocessor config (when the run has a postprocessor) +│ └── policy_postprocessor_step_*.safetensors # state of the stateful postprocessor steps +└── training_state/ + ├── training_step.json # step counter, topology, and batch semantics + ├── rng_state.safetensors # rng states + ├── scheduler_state.json # scheduler state (when the run has a scheduler) + ├── optimizer_state.safetensors # full optimizer state (non-sharded runs) + ├── optimizer_param_groups.json # optimizer param groups (non-sharded runs) + └── optimizer_0/ # DCP optimizer shards (sharded runs) +``` + +During single-GPU or DDP training, the pipeline serializes each state dict into a single file: `model.safetensors` for the model and `optimizer_state.safetensors` for the optimizer. + +During sharded training, the optimizer state is saved as DCP shards under `training_state/optimizer_0/`, and the layout of the model under `pretrained_model/` can be configured through `--checkpoint_format`: + +| `--checkpoint_format` | Weights artifact | Use when | +| ------------------------- | -------------------------------------------- | --------------------------------------------------------------------- | +| `safetensors` _(default)_ | single `model.safetensors` only | you want every checkpoint immediately loadable with `from_pretrained` | +| `dcp` | `pytorch_model_fsdp_0/` shard directory only | gathering the full weights makes saves and resumes too slow | +| `safetensors_dcp` | both | you want fast resume _and_ immediately loadable checkpoints | + +Two things to know about gathered (`safetensors`) checkpoints from sharded runs: + +- **They store fp32 weights.** Under mixed precision training, FSDP keeps an fp32 master copy, and the checkpoint saves the master copy to make sure training resumes consistently. +- The gather is collective (all ranks participate) but only the main process writes. + +### Converting DCP checkpoints + +`lerobot-convert-dcp` merges a DCP shard directory into a regular `model.safetensors`, offline and without GPUs: + +```bash +lerobot-convert-dcp --checkpoint_dir=outputs/train/run/checkpoints/005000 +lerobot-convert-dcp --checkpoint_dir=... --delete_dcp=true --push_to_hub=${HF_USER}/my_policy +``` + +`--push_to_hub` publishes the converted directory as a model repo. + +### Resuming + +Resume with `--resume=true --config_path=.../checkpoints/last/pretrained_model/train_config.json`. Resuming from a DCP checkpoint supports resharding the model and optimizer state to the _current_ topology, which means you can resume with a different `dp_replicate/dp_shard` split. The data sampler can always resume at the right epoch and offset, but is only _sample-exact_ when the world size and batch size match the original run (a warning is logged otherwise). + +> [!NOTE] +> FSDP checkpoints written by LeRobot 0.6.x and earlier used a different on-disk layout (a gathered full optimizer state) and **cannot be resumed**. ## 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. -- Training logs, checkpoints, and hub uploads are only done by the main process to avoid conflicts. Non-main processes have console logging disabled to prevent duplicate output. -- The effective batch size is `batch_size × num_gpus`. If you use 4 GPUs with `--batch_size=8`, your effective batch size is 32. -- Learning rate scheduling is handled correctly across multiple processes—LeRobot sets `step_scheduler_with_optimizer=False` to prevent accelerate from adjusting scheduler steps based on the number of processes. -- When saving or pushing models, LeRobot automatically unwraps the model from accelerate's distributed wrapper to ensure compatibility. -- WandB integration automatically initializes only on the main process, preventing multiple runs from being created. +- Checkpoint saves and end-of-training publishes are collective (every rank enters them). Gathered weights, sidecar files and Hub uploads are written by the main process alone. +- Metrics are reduced across ranks before logging: losses are averaged, and `samples/s` reports cluster-wide throughput. +- Learning-rate scheduling is stepped once per training step regardless of the number of processes (`step_scheduler_with_optimizer=False` is baked in). -For more advanced configurations and troubleshooting, see the [Accelerate documentation](https://huggingface.co/docs/accelerate). If you want to learn more about how to train on a large number of GPUs, checkout this awesome guide: [Ultrascale Playbook](https://huggingface.co/spaces/nanotron/ultrascale-playbook). +For background on the underlying machinery, see the [Accelerate FSDP guide](https://huggingface.co/docs/accelerate/usage_guides/fsdp). To go deeper on large-scale training, check out the [Ultrascale Playbook](https://huggingface.co/spaces/nanotron/ultrascale-playbook). diff --git a/docs/source/torch_accelerators.mdx b/docs/source/torch_accelerators.mdx index 6dfbfbad5..8af72766a 100644 --- a/docs/source/torch_accelerators.mdx +++ b/docs/source/torch_accelerators.mdx @@ -40,3 +40,15 @@ lerobot-eval \ ``` However, in most cases, presence of an accelerator is detected automatically and `policy.device` parameter can be omitted from CLI commands. + +## Mixed precision + +Training precision is owned by `--accelerator.mixed_precision`, which accepts `no` (default) and `bf16`: + +```bash +lerobot-train \ + --policy.type=act \ + --accelerator.mixed_precision=bf16 ... +``` + +`bf16` requires an accelerator that supports it. diff --git a/pyproject.toml b/pyproject.toml index 17bab326e..4a9ac404d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -346,6 +346,7 @@ lerobot-record="lerobot.scripts.lerobot_record:main" lerobot-replay="lerobot.scripts.lerobot_replay:main" lerobot-setup-motors="lerobot.scripts.lerobot_setup_motors:main" lerobot-teleoperate="lerobot.scripts.lerobot_teleoperate:main" +lerobot-convert-dcp="lerobot.scripts.lerobot_convert_dcp:main" lerobot-eval="lerobot.scripts.lerobot_eval:main" lerobot-train="lerobot.scripts.lerobot_train:main" lerobot-train-tokenizer="lerobot.scripts.lerobot_train_tokenizer:main" @@ -475,6 +476,12 @@ default.extend-ignore-identifiers-re = [ # TODO: Enable mypy gradually module by module across multiple PRs # Uncomment [tool.mypy] first, then uncomment individual module overrides as they get proper type annotations +[tool.pytest.ini_options] +markers = [ + "multigpu: distributed tests needing 2-4 GPUs (CI: docker_publish.yml lane)", + "multigpu_heavy: 8-GPU sweeps and soak tests; never run in CI", +] + [tool.mypy] python_version = "3.12" ignore_missing_imports = true @@ -521,6 +528,15 @@ disallow_untyped_defs = true disallow_incomplete_defs = true check_untyped_defs = true +[[tool.mypy.overrides]] +module = "lerobot.distributed.*" +ignore_errors = false + +# extra strictness for the distributed engine +disallow_untyped_defs = true +disallow_incomplete_defs = true +check_untyped_defs = true + [[tool.mypy.overrides]] module = "lerobot.optim.*" ignore_errors = false diff --git a/src/lerobot/common/train_utils.py b/src/lerobot/common/train_utils.py index a300cc76c..ed3814d4f 100644 --- a/src/lerobot/common/train_utils.py +++ b/src/lerobot/common/train_utils.py @@ -13,16 +13,41 @@ # 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. -from pathlib import Path +"""Training-output persistence: checkpoints, two-phase resume, and hub publishing. -from huggingface_hub import HfApi, snapshot_download +Rank discipline: every function here that can +contain a collective is documented as such and must run on ALL ranks; rank-0-only file writes +sit under one grouped ``is_main_process()`` gate per contiguous region, placed below all +collectives. The leaf save/load helpers carry no rank gates of their own — the exception is +``PreTrainedPolicy._save_pretrained``, whose gate is internal because its collective gather and +its writes live in the same method. +""" + +import logging +from importlib.resources import files +from pathlib import Path +from tempfile import TemporaryDirectory +from typing import TYPE_CHECKING, Any + +import torch.distributed as dist +from huggingface_hub import HfApi, ModelCard, ModelCardData, snapshot_download from torch.optim import Optimizer from torch.optim.lr_scheduler import LRScheduler +from lerobot.__version__ import __version__ +from lerobot.configs.policies import PreTrainedConfig +from lerobot.configs.rewards import RewardModelConfig from lerobot.configs.train import TrainPipelineConfig +from lerobot.distributed.checkpoint import ( + is_sharded_module, + load_sharded_model, + load_sharded_optimizer, + save_sharded_model, + save_sharded_optimizer, +) +from lerobot.distributed.utils import is_main_process from lerobot.optim import ( load_optimizer_state, - load_optimizer_state_dict, load_scheduler_state, save_optimizer_state, save_scheduler_state, @@ -40,14 +65,39 @@ from lerobot.utils.hub import find_latest_hub_checkpoint from lerobot.utils.io_utils import load_json, write_json from lerobot.utils.random_utils import load_rng_state, save_rng_state +if TYPE_CHECKING: + from accelerate import Accelerator + + from lerobot.datasets.dataset_metadata import LeRobotDatasetMetadata + from lerobot.rewards.pretrained import PreTrainedRewardModel + def get_step_identifier(step: int, total_steps: int) -> str: + """Format a step number as the zero-padded identifier used for checkpoint directory names. + + Args: + step (int): The training step to format. + total_steps (int): The total number of training steps; sets the padding width + (minimum 6 digits). + + Returns: + str: The zero-padded step identifier, e.g. `"005000"`. + """ num_digits = max(6, len(str(total_steps))) return f"{step:0{num_digits}d}" def get_step_checkpoint_dir(output_dir: Path, total_steps: int, step: int) -> Path: - """Returns the checkpoint sub-directory corresponding to the step number.""" + """Returns the checkpoint sub-directory corresponding to the step number. + + Args: + output_dir (Path): The training run's output directory. + total_steps (int): The total number of training steps; sets the identifier padding. + step (int): The training step of the checkpoint. + + Returns: + Path: The checkpoint step directory, `output_dir/checkpoints/`. + """ step_identifier = get_step_identifier(step, total_steps) return output_dir / CHECKPOINTS_DIR / step_identifier @@ -63,37 +113,15 @@ def should_save_checkpoint(step: int, save_freq: int, total_steps: int) -> bool: return (save_freq > 0 and step % save_freq == 0) or step == total_steps -def save_training_step( - step: int, save_dir: Path, num_processes: int | None = None, batch_size: int | None = None -) -> None: - state: dict = {"step": step} - # num_processes and batch_size are recorded so a resumed run can detect a changed world size or - # batch size: the sampler's resume offset is computed from the (num_processes, batch_size) that - # produced `step`, since both scale how many sampler positions a step consumes (see - # compute_sampler_state). - if num_processes is not None: - state["num_processes"] = num_processes - if batch_size is not None: - state["batch_size"] = batch_size - write_json(state, save_dir / TRAINING_STEP) +def update_last_checkpoint(checkpoint_dir: Path) -> None: + """Point the `last` symlink in the checkpoints directory at the given checkpoint. + Any existing `last` symlink is replaced. The link target is relative to the checkpoints + directory, so the tree stays valid when the run directory is moved. -def load_training_step(save_dir: Path) -> int: - training_step = load_json(save_dir / TRAINING_STEP) - return training_step["step"] - - -def load_training_num_processes(checkpoint_dir: Path) -> int | None: - """World size recorded at checkpoint time, or None for checkpoints written before it was stored.""" - return load_json(checkpoint_dir / TRAINING_STATE_DIR / TRAINING_STEP).get("num_processes") - - -def load_training_batch_size(checkpoint_dir: Path) -> int | None: - """Per-process batch size recorded at checkpoint time, or None for older checkpoints.""" - return load_json(checkpoint_dir / TRAINING_STATE_DIR / TRAINING_STEP).get("batch_size") - - -def update_last_checkpoint(checkpoint_dir: Path) -> Path: + Args: + checkpoint_dir (Path): The checkpoint step directory the `last` link should target. + """ last_checkpoint_dir = checkpoint_dir.parent / LAST_CHECKPOINT_LINK if last_checkpoint_dir.is_symlink(): last_checkpoint_dir.unlink() @@ -101,6 +129,68 @@ def update_last_checkpoint(checkpoint_dir: Path) -> Path: last_checkpoint_dir.symlink_to(relative_target) +# --------------------------------------------------------------------------------------------- +# training_step.json +# --------------------------------------------------------------------------------------------- + + +def save_training_metadata(step: int, save_dir: Path, cfg: TrainPipelineConfig) -> None: + """Record the step counter plus everything a resume needs to reason about topology changes. + + `step` counts loop iterations (= micro-batches), so + the sampler resume offset is `step x batch_size x dp_world_size` with no grad-accum factor. + `grad_accum_steps` and the parallelism snapshot are recorded so a resume can warn precisely + when the optimizer-update cadence or the sharding topology changed. + + Args: + step (int): The training step (micro-batch counter) to record. + save_dir (Path): The `training_state/` directory to write `training_step.json` into. + cfg (TrainPipelineConfig): The training config whose batch size, gradient-accumulation, + and parallelism settings are snapshotted alongside the step. + """ + state: dict[str, Any] = { + "step": step, + "dp_world_size": cfg.parallelism.dp_world_size, + "batch_size": cfg.batch_size, + "grad_accum_steps": cfg.accelerator.gradient_accumulation.steps, + "parallelism": { + "dp_replicate": cfg.parallelism.dp_replicate, + "dp_shard": cfg.parallelism.dp_shard, + "ring_degree": cfg.parallelism.context_parallel.ring_degree, + "ulysses_degree": cfg.parallelism.context_parallel.ulysses_degree, + }, + } + write_json(state, save_dir / TRAINING_STEP) + + +def load_training_metadata(training_state_dir: Path) -> dict[str, Any]: + """Read everything `save_training_metadata` recorded, in a single pass. + + Every key is always present: fields a checkpoint predates come back as None, so a caller + reading `metadata["batch_size"]` gets a KeyError on a typo rather than a silent None. + + Args: + training_state_dir (Path): The checkpoint's `training_state/` directory. + + Returns: + dict[str, Any]: `step` plus the `dp_world_size`, `batch_size`, `grad_accum_steps` and + `parallelism` snapshot recorded alongside it (None where not recorded). + """ + state = load_json(training_state_dir / TRAINING_STEP) + return { + "step": int(state["step"]), + "dp_world_size": state.get("dp_world_size", state.get("num_processes")), + "batch_size": state.get("batch_size"), + "grad_accum_steps": state.get("grad_accum_steps"), + "parallelism": state.get("parallelism"), + } + + +# --------------------------------------------------------------------------------------------- +# Checkpoint save +# --------------------------------------------------------------------------------------------- + + def save_checkpoint( checkpoint_dir: Path, step: int, @@ -110,192 +200,301 @@ def save_checkpoint( scheduler: LRScheduler | None = None, preprocessor: PolicyProcessorPipeline | None = None, 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, + accelerator: "Accelerator | None" = None, ) -> None: """This function creates the following directory structure: 005000/ # training step at checkpoint ├── pretrained_model/ │ ├── config.json # policy config - │ ├── model.safetensors # policy weights + │ ├── model.safetensors # policy weights (checkpoint_format ∈ {safetensors, safetensors_dcp}, or any non-sharded run) + │ ├── pytorch_model_fsdp_0/ # DCP model shards (checkpoint_format ∈ {dcp, safetensors_dcp}) │ ├── train_config.json # train config - │ ├── processor.json # processor config (if preprocessor provided) - │ └── step_*.safetensors # processor state files (if any) + │ ├── policy_preprocessor.json # preprocessor config (if preprocessor provided) + │ ├── policy_preprocessor_step_*.safetensors # state of the stateful preprocessor steps + │ ├── policy_postprocessor.json # postprocessor config (if postprocessor provided) + │ └── policy_postprocessor_step_*.safetensors # state of the stateful postprocessor steps └── training_state/ - ├── optimizer_param_groups.json # optimizer param groups - ├── optimizer_state.safetensors # optimizer state + ├── optimizer_param_groups.json # optimizer param groups (non-sharded runs) + ├── optimizer_state.safetensors # optimizer state (non-sharded runs) + ├── optimizer_0/ # DCP optimizer shards (sharded runs) ├── rng_state.safetensors # rng states - ├── scheduler_state.json # scheduler state - └── training_step.json # training step + ├── scheduler_state.json # scheduler state (if scheduler provided) + └── training_step.json # training step + dp_world_size/batch_size/grad_accum + topology + + Collective: MUST be called on every rank. Rank-0-only writes are gated internally, so the + call site needs no rank branches. Args: - cfg (TrainPipelineConfig): The training config used for this run. + checkpoint_dir (Path): The checkpoint step directory to write (e.g. `.../checkpoints/005000`). step (int): The training step at that checkpoint. + cfg (TrainPipelineConfig): The training config used for this run. policy (PreTrainedPolicy): The policy to save. - optimizer (Optimizer | None, optional): The optimizer to save the state from. Defaults to None. + optimizer (Optimizer): The optimizer to save the state from. scheduler (LRScheduler | None, optional): The scheduler to save the state from. Defaults to None. - preprocessor: The preprocessor/pipeline to save. Defaults to None. - postprocessor: The postprocessor/pipeline to save. Defaults to None. - num_processes (int | None, optional): Distributed world size to record for sample-exact - 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. + preprocessor (PolicyProcessorPipeline | None, optional): The preprocessor/pipeline to save. Defaults to None. + postprocessor (PolicyProcessorPipeline | None, optional): The postprocessor/pipeline to save. + Defaults to None. + accelerator (Accelerator | None, optional): The accelerator the policy was prepared with; + used to unwrap the model and required on sharded runs, where it owns the DCP save + channels. Defaults to None (plain single-process saves). """ pretrained_dir = checkpoint_dir / PRETRAINED_MODEL_DIR - policy.save_pretrained(pretrained_dir, state_dict=model_state_dict) - cfg.save_pretrained(pretrained_dir) + fmt = cfg.checkpoint_format + policy_to_save = accelerator.unwrap_model(policy) if accelerator is not None else policy + sharded = is_sharded_module(policy_to_save) + + # -- model artifact(s): the two collective-capable calls ---------------------------------- if cfg.peft is not None: - # When using PEFT, policy.save_pretrained will only write the adapter weights + config, not the - # policy config which we need for loading the model. In this case we'll write it ourselves. - policy.config.save_pretrained(pretrained_dir) - if preprocessor is not None: - preprocessor.save_pretrained(pretrained_dir) - if postprocessor is not None: - postprocessor.save_pretrained(pretrained_dir) + # PeftModel.save_pretrained is an external API with no internal rank gate, and the + # adapters are replicated (PEFT x sharded is rejected at validation): main rank writes. + if is_main_process(): + policy_to_save.save_pretrained(pretrained_dir) + elif fmt.wants_safetensors or not sharded: + # Collective when sharded (full gather); writes happen on the main process only in all + # multi-rank layouts (the gate lives inside _save_pretrained, next to its collective gather). + policy_to_save.save_pretrained(pretrained_dir) + if fmt.wants_dcp and sharded: + save_sharded_model(accelerator, policy_to_save, pretrained_dir) + + # -- sidecar configs: ONE gate for the whole contiguous rank-0-only region ---------------- + if is_main_process(): + if fmt.wants_dcp and not fmt.wants_safetensors: + # save_pretrained did not run: keep the DCP-only checkpoint self-describing. + policy_to_save.config.save_pretrained(pretrained_dir) + cfg.save_pretrained(pretrained_dir) + if cfg.peft is not None: + # PEFT's save_pretrained writes only adapter weights + config; the policy config + # needed to reload the base model is written explicitly. + policy_to_save.config.save_pretrained(pretrained_dir) + if preprocessor is not None: + preprocessor.save_pretrained(pretrained_dir) + 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, - optim_state_dict=optim_state_dict, + checkpoint_dir, step, cfg, optimizer, scheduler, accelerator, sharded=sharded, model=policy_to_save ) + if accelerator is not None: + accelerator.wait_for_everyone() def save_training_state( checkpoint_dir: Path, - train_step: int, - optimizer: Optimizer | None = None, + step: int, + cfg: TrainPipelineConfig, + optimizer: Optimizer | dict[str, Optimizer] | None = None, scheduler: LRScheduler | None = None, - num_processes: int | None = None, - batch_size: int | None = None, - optim_state_dict: dict | None = None, + accelerator: "Accelerator | None" = None, + *, + sharded: bool = False, + model: PreTrainedPolicy | None = None, ) -> None: - """ - Saves the training step, optimizer state, scheduler state, and rng state. + """Write training_state/. Collective under sharding: call on every rank. Args: - save_dir (Path): The directory to save artifacts to. - train_step (int): Current training step. - optimizer (Optimizer | None, optional): The optimizer from which to save the state_dict. + checkpoint_dir (Path): The checkpoint step directory; `training_state/` is created inside it. + step (int): The training step at that checkpoint. + cfg (TrainPipelineConfig): The training config used for this run (its topology and + accumulation settings are recorded in `training_step.json`). + optimizer (Optimizer | dict[str, Optimizer] | None, optional): The optimizer(s) to save + the state from. Defaults to None. + scheduler (LRScheduler | None, optional): The scheduler to save the state from. Defaults to None. - scheduler (LRScheduler | None, optional): The scheduler from which to save the state_dict. - 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. + accelerator (Accelerator | None, optional): Required when `sharded` is True — it owns + the DCP optimizer save channel. Defaults to None. + sharded (bool): The model's sharding state, computed once in `save_checkpoint` and + threaded here so the two sites cannot disagree. Defaults to False. + model (PreTrainedPolicy | None, optional): Required only for the sharded optimizer + channel: torch's optimizer DCP APIs are model-coupled (the state dict is keyed by + model FQNs), so accelerate's `save_fsdp_optimizer` needs the sharded module + alongside the optimizer. Defaults to None. """ save_dir = checkpoint_dir / TRAINING_STATE_DIR + # All ranks: the directory must exist before the DCP optimizer collective writes into it + # (exist_ok makes the concurrent mkdir race-free on shared filesystems). 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, optim_state_dict=optim_state_dict) - if scheduler is not None: - save_scheduler_state(scheduler, save_dir) + + if optimizer is not None and sharded: + if accelerator is None or model is None: + raise ValueError("Saving a sharded optimizer state requires the accelerator and model.") + # Collective — all ranks write their DCP shards into optimizer_0/. + save_sharded_optimizer(accelerator, optimizer, model, save_dir) + + if is_main_process(): # ONE grouped gate for the whole rank-0-only region + save_training_metadata(step, save_dir, cfg) + save_rng_state(save_dir) + if scheduler is not None: + save_scheduler_state(scheduler, save_dir) + if optimizer is not None and not sharded: + save_optimizer_state(optimizer, save_dir) -def load_training_state( - 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. - This is used to resume a training run. +# --------------------------------------------------------------------------------------------- +# Two-phase resume +# --------------------------------------------------------------------------------------------- + + +def resume_before_prepare(cfg: TrainPipelineConfig) -> int: + """Phase 1 — before `accelerator.prepare()`: restore RNG and return the step counter. + + Pure loaders only. The sampler resume offset is *derived* from the returned step inside the + dataloader factory, and everything bound to sharded objects (model DCP shards, optimizer, + scheduler) loads in `resume_after_prepare`. Args: - 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). + cfg (TrainPipelineConfig): The resumed training config; `cfg.checkpoint_path` locates + the checkpoint to restore from. + + Returns: + int: The training step recorded in the checkpoint (micro-batch counter). Raises: - NotADirectoryError: If 'checkpoint_dir' doesn't contain a 'training_state' dir - - Returns: - tuple[int, Optimizer, LRScheduler | None]: training step, optimizer and scheduler with their - state_dict loaded. + NotADirectoryError: If the checkpoint has no `training_state/` directory. + ValueError: If the resumed topology crosses the sharded/non-sharded boundary relative + to the one recorded in the checkpoint. """ - training_state_dir = checkpoint_dir / TRAINING_STATE_DIR + training_state_dir = cfg.checkpoint_path / TRAINING_STATE_DIR if not training_state_dir.is_dir(): raise NotADirectoryError(training_state_dir) - + metadata = load_training_metadata(training_state_dir) + _guard_resume_changes(cfg, metadata) load_rng_state(training_state_dir) - step = load_training_step(training_state_dir) - if load_optimizer: - optimizer = load_optimizer_state(optimizer, training_state_dir) + return metadata["step"] + + +def _guard_resume_changes(cfg: TrainPipelineConfig, metadata: dict[str, Any]) -> None: + """Check the resumed run settings against the ones recorded in the checkpoint. + + Two tiers, both driven by the checkpoint's recorded parallelism snapshot: + + - **Hard error** when the resume crosses the sharded/non-sharded boundary in either + direction: the checkpoint's training-state artifacts only support resuming on the same + kind of topology (resharding works across sizes, not across kinds). Checkpoints without + a recorded snapshot skip this check. + - **One warning** naming every other recorded setting that differs — those changes are + legal (DCP reshards weights and optimizer state across topologies and the sampler offset + adapts), but a changed ``grad_accum_steps`` shifts the optimizer-update cadence, so the + resume says precisely what differs. The sampler-exactness warnings + (``dp_world_size``/``batch_size``) live with the sampler math in the dataloader factory. + + Args: + cfg (TrainPipelineConfig): The resumed training config, compared against the settings + recorded in the checkpoint. + metadata (dict[str, Any]): The checkpoint's recorded training metadata, as returned by + `load_training_metadata`. + + Raises: + ValueError: If the checkpoint records a sharded topology and the resumed run is + non-sharded, or vice versa. + """ + snapshot = metadata["parallelism"] + + if snapshot is not None: + recorded_sharded = ( + snapshot.get("dp_shard", 1) != 1 + or snapshot.get("ring_degree", 1) * snapshot.get("ulysses_degree", 1) > 1 + ) + if recorded_sharded != cfg.parallelism.is_sharded: + raise ValueError( + f"Cannot resume: the checkpoint was written with a " + f"{'sharded' if recorded_sharded else 'non-sharded'} topology " + f"(dp_replicate={snapshot.get('dp_replicate')}, dp_shard={snapshot.get('dp_shard')}) " + f"but this run is {'sharded' if cfg.parallelism.is_sharded else 'non-sharded'} " + f"(dp_replicate={cfg.parallelism.dp_replicate}, dp_shard={cfg.parallelism.dp_shard})." + ) + + recorded = { + "grad_accum_steps": ( + metadata["grad_accum_steps"], + cfg.accelerator.gradient_accumulation.steps, + ), + } + if snapshot is not None: + recorded.update( + { + "dp_replicate": (snapshot.get("dp_replicate"), cfg.parallelism.dp_replicate), + "dp_shard": (snapshot.get("dp_shard"), cfg.parallelism.dp_shard), + "ring_degree": ( + snapshot.get("ring_degree"), + cfg.parallelism.context_parallel.ring_degree, + ), + "ulysses_degree": ( + snapshot.get("ulysses_degree"), + cfg.parallelism.context_parallel.ulysses_degree, + ), + } + ) + changed = [f"{key}: {was} -> {now}" for key, (was, now) in recorded.items() if was not in (None, now)] + if changed and is_main_process(): + logging.warning( + "Resuming with settings that differ from the checkpoint: " + "; ".join(changed) + ". " + "Topology changes reshard safely via DCP; a changed grad_accum_steps shifts the " + "optimizer-update cadence (the step counter keeps counting micro-batches)." + ) + + +def resume_after_prepare( + cfg: TrainPipelineConfig, + accelerator: "Accelerator", + policy: PreTrainedPolicy, + optimizer: Optimizer | dict[str, Optimizer], + scheduler: LRScheduler | None, +) -> None: + """Phase 2 — after `accelerator.prepare()`: model (DCP) -> optimizer -> scheduler. + + Collective under sharding: call on every rank. The model-weight source follows the + checkpoint's own recorded `checkpoint_format` (on resume, `cfg` was parsed from the + checkpoint's train_config.json): DCP-bearing formats load shards here into the prepared + model (whose construction skipped the safetensors load); the safetensors format was already + loaded by `from_pretrained` before sharding — no model step here. + + Args: + cfg (TrainPipelineConfig): The resumed training config; `cfg.checkpoint_path` locates + the checkpoint and `cfg.checkpoint_format` selects the model-weight source. + accelerator (Accelerator): The accelerator the policy was prepared with; it unwraps the + model and owns the DCP load channels. + policy (PreTrainedPolicy): The prepared (possibly sharded) policy to load weights into. + optimizer (Optimizer | dict[str, Optimizer]): The prepared optimizer(s) to restore. + scheduler (LRScheduler | None): The scheduler to restore, or None if the run has none. + + Raises: + FileNotFoundError: If the checkpoint format declares DCP model shards but the shard + directory is missing (e.g. it was pruned before upload). + """ + checkpoint_dir = cfg.checkpoint_path + pretrained_dir = checkpoint_dir / PRETRAINED_MODEL_DIR + training_state_dir = checkpoint_dir / TRAINING_STATE_DIR + unwrapped = accelerator.unwrap_model(policy) + sharded = is_sharded_module(unwrapped) + + if cfg.checkpoint_format.wants_dcp: + from accelerate.utils.constants import FSDP_MODEL_NAME + + dcp_dir = pretrained_dir / f"{FSDP_MODEL_NAME}_0" + if not dcp_dir.is_dir(): + raise FileNotFoundError( + f"checkpoint_format={cfg.checkpoint_format.value} declares DCP model shards, " + f"but {dcp_dir} is missing. If the shards were pruned, convert what remains " + "with `lerobot-convert-dcp` or resume from a safetensors checkpoint." + ) + load_sharded_model(accelerator, unwrapped, pretrained_dir) + + if sharded: + # Requires the prepared optimizer: FSDP2's prepare rebinds param groups to DTensors but + # never migrates optimizer.state — DCP reshards it here (works across topology changes). + load_sharded_optimizer(accelerator, optimizer, unwrapped, training_state_dir) + else: + load_optimizer_state(optimizer, training_state_dir) + if scheduler is not None: - scheduler = load_scheduler_state(scheduler, training_state_dir) - - return step, optimizer, scheduler + load_scheduler_state(scheduler, training_state_dir) -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) +# --------------------------------------------------------------------------------------------- +# Hub: checkpoint push (resume artifact) and publishing (distribution artifact) +# --------------------------------------------------------------------------------------------- def push_checkpoint_to_hub( @@ -311,6 +510,16 @@ def push_checkpoint_to_hub( The model repo is created idempotently, and the commit is tagged with the checkpoint step so a checkpoint can be recovered with --policy.pretrained_revision= instead of a commit sha. + + The directory is uploaded verbatim — including DCP shards under the DCP formats: this tree + exists for *resume*, not distribution, and `resolve_resume_checkpoint` downloads it back + symmetrically. + + Args: + checkpoint_dir (Path): The local checkpoint step directory to upload. + repo_id (str): The Hub model repo to push to (created idempotently if missing). + private (bool | None): Whether a newly created repo should be private. Defaults to + None (public unless the organization's default is private). """ api = HfApi() api.create_repo(repo_id=repo_id, repo_type="model", private=private, exist_ok=True) @@ -338,6 +547,16 @@ def resolve_resume_checkpoint(repo_id: str, output_dir: Path) -> Path: into `output_dir/checkpoints//`, recreate the local `last` symlink, and return that local checkpoint dir. Used to resume training from the Hub on a machine (or HF Jobs pod) that does not have the original local run dir. + + Args: + repo_id (str): The Hub model repo holding `checkpoints//` subtrees. + output_dir (Path): The local run directory to download the checkpoint into. + + Returns: + Path: The local checkpoint step directory, `output_dir/checkpoints/`. + + Raises: + FileNotFoundError: If the repo contains no checkpoints under `checkpoints/`. """ latest = find_latest_hub_checkpoint(repo_id) if latest is None: @@ -354,3 +573,214 @@ def resolve_resume_checkpoint(repo_id: str, output_dir: Path) -> Path: checkpoint_dir = output_dir / latest update_last_checkpoint(checkpoint_dir) return checkpoint_dir + + +def publish_trained_model( + cfg: TrainPipelineConfig, + model: "PreTrainedPolicy | PreTrainedRewardModel", + preprocessor: PolicyProcessorPipeline | None, + postprocessor: PolicyProcessorPipeline | None, + dataset_meta: "LeRobotDatasetMetadata | None", + *, + peft_model: Any | None = None, +) -> None: + """Publish the complete training bundle as a distributable model repo. + + Collective-safe: call on ALL ranks — the model commit gathers sharded weights through + `save_pretrained`; uploads happen on the main process only (gated inside + `HubMixin.push_to_hub` and here). Commits, in order: (1) the model (skipped for PEFT — + adapters replace full weights), (2) the preprocessor, (3) the postprocessor, (4) the bundle + sidecar: README.md model card + train_config.json (+ adapter weights and the wrapped + policy's config in the PEFT case). Every commit uploads a freshly assembled directory, so + a published repo carries only the distributable artifacts. + + Args: + cfg (TrainPipelineConfig): The training config; saved as `train_config.json` and used + to render the model card. + model (PreTrainedPolicy | PreTrainedRewardModel): The trained model to publish; its + config supplies the target repo id, visibility, license, and tags. + preprocessor (PolicyProcessorPipeline | None): The preprocessor pipeline to publish + alongside the model, if any. + postprocessor (PolicyProcessorPipeline | None): The postprocessor pipeline to publish + alongside the model, if any. + dataset_meta (LeRobotDatasetMetadata | None): Dataset metadata for the model card, if + available. + peft_model (Any | None): The PEFT wrapper when training adapters; its adapter weights + replace the full model weights in the published repo. Defaults to None. + + Raises: + ValueError: If the model config carries no repo id (`--policy.repo_id`). + """ + model_cfg = model.config + repo_id = model_cfg.repo_id + if not repo_id: + raise ValueError("Publishing requires a repo id (--policy.repo_id).") + ignore = ["*.tmp", "*.log"] + + if peft_model is None: + # Calls are made on the exact objects that own each method (never through PEFT's + # attribute forwarding), so the peft branch below never touches this path. + model.push_to_hub(repo_id, private=model_cfg.private, ignore_patterns=ignore) + if preprocessor is not None: + preprocessor.push_to_hub(repo_id, private=model_cfg.private) + if postprocessor is not None: + postprocessor.push_to_hub(repo_id, private=model_cfg.private) + + if is_main_process(): + api = HfApi() + repo_id = api.create_repo(repo_id=repo_id, private=model_cfg.private, exist_ok=True).repo_id + with TemporaryDirectory(ignore_cleanup_errors=True) as tmp: + saved_path = Path(tmp) / repo_id + saved_path.mkdir(parents=True, exist_ok=True) + if peft_model is not None: + peft_model.save_pretrained(saved_path) # adapter weights + adapter config + model.config.save_pretrained(saved_path) # PEFT cannot write the policy config + card = generate_model_card(model_cfg, cfg=cfg, dataset_meta=dataset_meta) + card.save(str(saved_path / "README.md")) + cfg.save_pretrained(saved_path) # train_config.json + commit_info = api.upload_folder( + repo_id=repo_id, + repo_type="model", + folder_path=saved_path, + commit_message="Upload model card and train config", + allow_patterns=["*.safetensors", "*.json", "*.yaml", "*.md"], + ignore_patterns=ignore, + ) + # Contract: lerobot.jobs.hf.submit_to_hf watches for this exact "Model pushed to " + # line to end a remote run early. Keep the wording and URL format in sync. + logging.info(f"Model pushed to {commit_info.repo_url.url}") + + if dist.is_initialized(): + dist.barrier() + + +# --------------------------------------------------------------------------------------------- +# Model card +# --------------------------------------------------------------------------------------------- + +_BASE_MODEL_MAPPING = { + "smolvla": "lerobot/smolvla_base", + "pi0": "lerobot/pi0_base", + "pi05": "lerobot/pi05_base", + "pi0_fast": "lerobot/pi0fast-base", + "xvla": "lerobot/xvla-base", +} + + +def build_card_context( + cfg: TrainPipelineConfig | None, + dataset_meta: "LeRobotDatasetMetadata | None", + input_features: dict | None, + output_features: dict | None, +) -> dict: + """Collect optional data for the model-card template. + + Returns plain values only (no Markdown) — the template in + ``lerobot/templates/lerobot_modelcard_template.md`` decides how and whether to show + each one. Everything is best-effort: anything unavailable is left empty/None and the + template simply skips that section, so this never breaks a Hub push. + + Args: + cfg (TrainPipelineConfig | None): The training config supplying the training section, + if available. + dataset_meta (LeRobotDatasetMetadata | None): Dataset metadata supplying the dataset, + robot-type, and camera sections, if available. + input_features (dict | None): The policy's input feature declarations, if any. + output_features (dict | None): The policy's output feature declarations, if any. + + Returns: + dict: Template context with `training`, `input_features`, `output_features`, + `dataset`, `robot_type`, and `cameras` entries; unavailable pieces stay + empty/None. + """ + context = { + "training": None, + "input_features": input_features or {}, + "output_features": output_features or {}, + "dataset": None, + "robot_type": None, + "cameras": [], + } + + if cfg is not None: + optimizer = getattr(cfg, "optimizer", None) + context["training"] = { + "steps": cfg.steps, + "batch_size": cfg.batch_size, + "seed": cfg.seed, + "optimizer": getattr(optimizer, "type", None) if optimizer else None, + "lr": getattr(optimizer, "lr", None) if optimizer else None, + "lerobot_version": __version__, + } + + if dataset_meta is not None: + context["dataset"] = { + "repo_id": dataset_meta.repo_id, + "episodes": dataset_meta.total_episodes, + "frames": dataset_meta.total_frames, + "fps": dataset_meta.fps, + "tasks": [str(task) for task in dataset_meta.tasks.index], + } + context["robot_type"] = dataset_meta.robot_type + context["cameras"] = [key.split(".")[-1] for key in dataset_meta.camera_keys] + + return context + + +def generate_model_card( + model_cfg: PreTrainedConfig | RewardModelConfig, + cfg: TrainPipelineConfig | None = None, + dataset_meta: "LeRobotDatasetMetadata | None" = None, +) -> ModelCard: + """Render the LeRobot model card for a trained policy or reward model. + + A free function on purpose: every template variable comes from arguments — the model + config, the training config, and the dataset metadata — none from a live model, so a card + can also be rendered from a checkpoint's `config.json` alone (see `lerobot-convert-dcp`). + The config type selects the template: reward models get the reward-model card, policies the + policy card with the training/dataset sections. + + Args: + model_cfg (PreTrainedConfig | RewardModelConfig): The model config providing type, + license, tags, repo id, and — for policies — the feature declarations. + cfg (TrainPipelineConfig | None, optional): The training config for the training and + dataset card sections. Defaults to None. + dataset_meta (LeRobotDatasetMetadata | None, optional): Dataset metadata for the + dataset card sections. Defaults to None. + + Returns: + ModelCard: The rendered and validated LeRobot model card. + """ + model_type = model_cfg.type + base_model = _BASE_MODEL_MAPPING.get(model_type) + + if isinstance(model_cfg, RewardModelConfig): + tags = {"robotics", "lerobot", "reward-model", model_type} + template_card = ( + files("lerobot.templates") + .joinpath("lerobot_rewardmodel_modelcard_template.md") + .read_text("utf-8") + ) + context: dict[str, Any] = {} # the reward template renders from card_data alone + else: + tags = {"robotics", "lerobot", model_type} + template_card = ( + files("lerobot.templates").joinpath("lerobot_modelcard_template.md").read_text("utf-8") + ) + context = build_card_context(cfg, dataset_meta, model_cfg.input_features, model_cfg.output_features) + # Used by the template to pre-fill commands and the "Fine-tuned from" line. + context["policy_repo_id"] = model_cfg.repo_id + context["base_model"] = base_model + + card_data = ModelCardData( + license=model_cfg.license or "apache-2.0", + library_name="lerobot", + pipeline_tag="robotics", + tags=list(tags.union(model_cfg.tags or [])), + model_name=model_type, + datasets=cfg.dataset.repo_id if cfg is not None else None, + base_model=base_model, + ) + card = ModelCard.from_template(card_data, template_str=template_card, **context) + card.validate() + return card diff --git a/src/lerobot/configs/accelerator.py b/src/lerobot/configs/accelerator.py new file mode 100644 index 000000000..aebc50d18 --- /dev/null +++ b/src/lerobot/configs/accelerator.py @@ -0,0 +1,273 @@ +#!/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. +"""Execution-runtime configuration: everything handed to (or applied by) the `Accelerator`. + +Each sub-config mirrors the plain-typed subset of the corresponding accelerate object and +builds it at runtime (the way ``OptimizerConfig.build()`` constructs a ``torch.optim.Optimizer``), +so the whole tree round-trips through the CLI and ``train_config.json`` and parsing a config +never imports accelerate. +""" + +from dataclasses import dataclass, field +from enum import Enum +from typing import TYPE_CHECKING + +from lerobot.configs.parallelism import ParallelismConfig + +if TYPE_CHECKING: + from accelerate import Accelerator + from accelerate.utils import ( + DistributedDataParallelKwargs, + FullyShardedDataParallelPlugin, + GradientAccumulationPlugin, + ) + + +@dataclass +class FSDPConfig: + """Mirror of the `FullyShardedDataParallelPlugin` subset LeRobot supports (FSDP2 only). + + Exactly one wrap policy applies: `wrap_modules` (module *class names* forming the FSDP + units — and, later, the activation-checkpointing units) or `min_num_params` (size-based). + When both are None, the policy's own `_fsdp_wrap_modules` declaration is used; a run where + no wrap source exists at all fails loudly rather than silently wrapping only the root. + """ + + reshard_after_forward: bool = True + wrap_modules: list[str] | None = None + min_num_params: int | None = None + cpu_offload: bool = False + # Regex matched against module FQNs to exclude their parameters from sharding. + ignored_modules: str | None = None + + def __post_init__(self) -> None: + """Validate the wrap-policy fields. + + Raises: + ValueError: If both ``wrap_modules`` and ``min_num_params`` are set (they are + mutually exclusive wrap policies), or if ``min_num_params`` is < 1. + """ + if self.wrap_modules is not None and self.min_num_params is not None: + raise ValueError( + "fsdp.wrap_modules and fsdp.min_num_params are mutually exclusive wrap policies." + ) + if self.min_num_params is not None and self.min_num_params < 1: + raise ValueError(f"fsdp.min_num_params must be >= 1, got {self.min_num_params}.") + + def build_plugin(self) -> "FullyShardedDataParallelPlugin": + """Build the FSDP2 plugin for `Accelerator(fsdp_plugin=...)`. + + Returns: + FullyShardedDataParallelPlugin: FSDP2 (`fsdp_version=2`) plugin carrying the + mirrored wrap policy, resharding, CPU-offload, and ignored-modules settings. + """ + from accelerate.utils import FullyShardedDataParallelPlugin + + use_size_policy = self.min_num_params is not None + return FullyShardedDataParallelPlugin( + fsdp_version=2, + reshard_after_forward=self.reshard_after_forward, + auto_wrap_policy="size_based_wrap" if use_size_policy else "transformer_based_wrap", + # May legitimately still be None here: the policy-declared default is applied right + # before `accelerator.prepare()` (see lerobot.distributed.factory.set_fsdp_wrap_modules). + transformer_cls_names_to_wrap=list(self.wrap_modules) if self.wrap_modules else None, + min_num_params=self.min_num_params, + cpu_offload=self.cpu_offload, + ignored_modules=self.ignored_modules, + # state_dict_type stays at the FSDP2 default (SHARDED_STATE_DICT) and is never + # switched: full gathers go through torch's state-dict API, which does not consult + # the plugin. activation_checkpointing stays False: AC is LeRobot-owned. + ) + + +@dataclass +class DDPConfig: + """Mirror of the `DistributedDataParallelKwargs` subset LeRobot exposes.""" + + # Today's in-script default, kept for models with conditional computation. + find_unused_parameters: bool = True + gradient_as_bucket_view: bool = False + static_graph: bool = False + + def build_kwargs_handler(self) -> "DistributedDataParallelKwargs": + """Build the DDP kwargs handler for `Accelerator(kwargs_handlers=[...])`. + + Returns: + DistributedDataParallelKwargs: Handler carrying the mirrored DDP fields, applied + by accelerate when it wraps the model in `DistributedDataParallel`. + """ + from accelerate.utils import DistributedDataParallelKwargs + + return DistributedDataParallelKwargs( + find_unused_parameters=self.find_unused_parameters, + gradient_as_bucket_view=self.gradient_as_bucket_view, + static_graph=self.static_graph, + ) + + +@dataclass +class GradientAccumulationConfig: + """Mirror of the `GradientAccumulationPlugin` subset LeRobot supports. + + Only the step count is a knob. ``sync_with_dataloader`` is pinned to False by + :meth:`build_plugin`: the training loop cycles a finite dataloader, so accelerate's default + of syncing at every dataloader end would force an optimizer step at every dataset epoch + boundary instead of every ``steps`` micro-batches. + """ + + steps: int = 1 + + def __post_init__(self) -> None: + """Validate the accumulation step count. + + Raises: + ValueError: If ``steps`` is < 1. + """ + if self.steps < 1: + raise ValueError(f"gradient_accumulation.steps must be >= 1, got {self.steps}.") + + def build_plugin(self) -> "GradientAccumulationPlugin": + """Build the plugin for `Accelerator(gradient_accumulation_plugin=...)`. + + A named plugin argument, not a `kwargs_handlers` entry: accelerate consumes this object + through its dedicated constructor parameter — the `KwargsHandler` base class only lends + it `to_kwargs()`, so the consumption site, not the inheritance, decides its role. + + Returns: + GradientAccumulationPlugin: Carrying the mirrored step count, with + ``sync_with_dataloader=False`` pinned (see the class docstring). + """ + from accelerate.utils import GradientAccumulationPlugin + + return GradientAccumulationPlugin(num_steps=self.steps, sync_with_dataloader=False) + + +@dataclass +class CompileConfig: + """torch.compile knobs — a configured placeholder: wiring lands in a later round. + + The setup-order contract it will follow is already fixed: compile applies + after CP dispatch install and activation checkpointing, before `fully_shard`, regionally + (per wrap unit) — the only combination proven with FSDP2. + """ + + enabled: bool = False + backend: str = "inductor" + mode: str | None = None + regional: bool = True + + +class ActivationCheckpointingMode(str, Enum): + NONE = "none" + FULL = "full" + + +@dataclass +class ActivationCheckpointingConfig: + """Activation-checkpointing knobs — a configured placeholder: wiring lands in a later round. + + AC units will coincide with the FSDP wrap units (one declaration drives both), applied + before torch.compile and `fully_shard` (the same ordering contract as CompileConfig). + """ + + mode: ActivationCheckpointingMode = ActivationCheckpointingMode.NONE + + +@dataclass +class AcceleratorConfig: + """Builds the `Accelerator` — the runtime counterpart of the `parallelism` topology. + + `mixed_precision` selects accelerate-native AMP for DDP/single-GPU runs and the FSDP2 + `MixedPrecisionPolicy` for sharded runs (accelerate derives it). Sharded runs support + "no" and "bf16" only; fp16's GradScaler-over-DTensor path is unverified and fails fast + at config validation. + """ + + mixed_precision: str = "no" + gradient_accumulation: GradientAccumulationConfig = field(default_factory=GradientAccumulationConfig) + fsdp: FSDPConfig = field(default_factory=FSDPConfig) + ddp: DDPConfig = field(default_factory=DDPConfig) + compile: CompileConfig = field(default_factory=CompileConfig) + activation_checkpointing: ActivationCheckpointingConfig = field( + default_factory=ActivationCheckpointingConfig + ) + + def __post_init__(self) -> None: + """Validate the accelerate-facing scalar fields. + + Raises: + ValueError: If ``mixed_precision`` is not one of ``"no"``, ``"fp16"``, ``"bf16"``. + """ + if self.mixed_precision not in ("no", "fp16", "bf16"): + raise ValueError( + f"mixed_precision must be one of 'no', 'fp16', 'bf16', got {self.mixed_precision!r}." + ) + + def build(self, parallelism: ParallelismConfig, *, cpu: bool = False) -> "Accelerator": + """Translate the mirrored fields into a ready `Accelerator` (call once per process). + + `parallelism` must already be resolved against the world size. The degradation matrix + is encoded here and nowhere else: sharded -> FSDP2 (+HSDP via the accelerate + `ParallelismConfig` mesh), replicated-only -> DDP kwargs, single process -> plain. + + Args: + parallelism (ParallelismConfig): The resolved process topology; selects which + accelerate path (FSDP2 mesh, DDP kwargs handler, or plain) is configured. + cpu (bool): Force CPU execution even when CUDA is available. Defaults to False. + + Returns: + Accelerator: The configured accelerate entry point for this process. + """ + from accelerate import Accelerator + + kwargs: dict = { + # LeRobot steps its scheduler manually once per training step; accelerate must not + # rescale scheduler stepping by num_processes. + "step_scheduler_with_optimizer": False, + "gradient_accumulation_plugin": self.gradient_accumulation.build_plugin(), + "mixed_precision": self.mixed_precision, + "cpu": cpu, + } + if parallelism.is_sharded: + kwargs["fsdp_plugin"] = self.fsdp.build_plugin() + kwargs["parallelism_config"] = _accelerate_parallelism_config(parallelism) + elif parallelism.is_replicated_only: + kwargs["kwargs_handlers"] = [self.ddp.build_kwargs_handler()] + return Accelerator(**kwargs) + + +def _accelerate_parallelism_config(parallelism: ParallelismConfig) -> object: + """LeRobot topology -> accelerate `ParallelismConfig`. + + CP is declared honestly (`cp_size = ring x ulysses`) so accelerate builds the canonical + mesh, folds CP into the FSDP shard group (`dp_shard_cp`), and duplicates batches within CP + groups. The ring/ulysses sub-structure stays private to `lerobot.distributed.ParallelDims`. + + Args: + parallelism (ParallelismConfig): The resolved LeRobot topology to translate. + + Returns: + object: The accelerate `ParallelismConfig` mirroring `dp_replicate`, `dp_shard`, and + the collapsed `cp_size` (annotated as `object` so importing this module never + imports accelerate). + """ + from accelerate.parallelism_config import ParallelismConfig as AccelerateParallelismConfig + + return AccelerateParallelismConfig( + dp_replicate_size=parallelism.dp_replicate, + dp_shard_size=parallelism.dp_shard, + cp_size=parallelism.cp_size, + ) diff --git a/src/lerobot/configs/parallelism.py b/src/lerobot/configs/parallelism.py new file mode 100644 index 000000000..cd8ec6cad --- /dev/null +++ b/src/lerobot/configs/parallelism.py @@ -0,0 +1,190 @@ +#!/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. +"""Declarative process topology for distributed training and inference. + +The mesh convention (canonical row-major rank layout, outermost first):: + + (dp_replicate, dp_shard, ring, ulysses) + +- ``dp_replicate x dp_shard`` is the data-parallel world: HSDP replicates over + ``dp_replicate`` and shards parameters over ``dp_shard``. FSDP2's actual shard + group folds context parallelism in (``dp_shard x ring x ulysses``), matching + accelerate's ``dp_shard_cp`` flattening and torchtitan's ``fsdp`` axis. +- ``ring`` is the outer and ``ulysses`` the inner context-parallel dim + (diffusers convention: ulysses all-to-all exchanges run over adjacent, typically + NVLink-connected ranks). +- ``cfg_parallel`` (classifier-free-guidance parallelism) is a branch-parallel, + inference-only dim that sits between dp and the sequence dims. It never + affects weight sharding or checkpoints. + +This module is pure configuration: plain-typed dataclasses that draccus can +round-trip through the CLI and ``train_config.json``. Runtime objects (device +meshes, process groups) live in :mod:`lerobot.distributed`. +""" + +import os +from dataclasses import dataclass, field + + +@dataclass +class ContextParallelConfig: + """Ring x Ulysses context parallelism (sequence parallelism for attention). + + Both degrees are configured placeholders in this release: the CP engine is not implemented + yet, and enabling either degree > 1 fails fast at config validation. The fields exist now so + that the CLI surface, checkpoint metadata, and mesh math are stable when the engine lands. + """ + + ring_degree: int = 1 + ulysses_degree: int = 1 + + def __post_init__(self) -> None: + """Validate the declared context-parallel degrees. + + Raises: + ValueError: If ``ring_degree`` or ``ulysses_degree`` is < 1. + """ + if self.ring_degree < 1 or self.ulysses_degree < 1: + raise ValueError( + f"Context-parallel degrees must be >= 1, got ring_degree={self.ring_degree}, " + f"ulysses_degree={self.ulysses_degree}." + ) + + @property + def size(self) -> int: + """Total number of ranks a full sequence is sharded across.""" + return self.ring_degree * self.ulysses_degree + + +@dataclass +class ParallelismConfig: + """Degrees of every parallelism dim. Invariant: their product equals the world size. + + Degradations are expressed purely through the degrees (no mode flags): + + - single process: all degrees 1; + - DDP: ``dp_replicate == world_size`` (auto-filled when every sharding field is left at its + default — plain ``torchrun`` keeps today's out-of-the-box behavior); + - FSDP: ``dp_shard > 1`` (or ``-1`` to fill the remaining world into the shard dim); + - HSDP: ``dp_replicate > 1`` and ``dp_shard > 1``. + + ``resolve()`` turns the declared degrees into concrete ones once the world size is known and + is the single place the world-size equation is enforced. It is called by + :func:`lerobot.distributed.factory.make_accelerator`; the config is inert until then. + """ + + dp_replicate: int = 1 + # -1 is an explicit opt-in sentinel: shard over world_size // (dp_replicate * cp). + dp_shard: int = 1 + context_parallel: ContextParallelConfig = field(default_factory=ContextParallelConfig) + # Classifier-free-guidance parallelism — inference-only (cosmos/vllm-omni precedent: + # cond/uncond branches on different ranks). Reserved for the serving round; training + # validates it to 1. Meaningful values are 1 or 2 (Cosmos3 has two CFG branches). + cfg_parallel: int = 1 + + def __post_init__(self) -> None: + """Validate the declared degrees (world-size-independent checks only). + + Raises: + ValueError: If ``dp_replicate`` is < 1, ``dp_shard`` is neither >= 1 nor the + ``-1`` infer sentinel, or ``cfg_parallel`` is not 1 or 2. + """ + if self.dp_replicate < 1: + raise ValueError(f"dp_replicate must be >= 1, got {self.dp_replicate}.") + if self.dp_shard < 1 and self.dp_shard != -1: + raise ValueError(f"dp_shard must be >= 1, or -1 to infer, got {self.dp_shard}.") + if self.cfg_parallel not in (1, 2): + raise ValueError(f"cfg_parallel must be 1 or 2, got {self.cfg_parallel}.") + + @property + def cp_size(self) -> int: + """Total context-parallel size (``ring_degree * ulysses_degree``).""" + return self.context_parallel.size + + @property + def is_sharded(self) -> bool: + """True when the run uses FSDP2 (parameters sharded); selects the sharded engine path.""" + return self.dp_shard != 1 or self.cp_size > 1 + + @property + def is_replicated_only(self) -> bool: + """True for plain DDP (weights replicated, no sharding).""" + return not self.is_sharded and self.dp_replicate > 1 + + @property + def dp_world_size(self) -> int: + """Number of distinct data-parallel workers (batches are sharded this many ways). + + Returns: + int: ``dp_replicate * dp_shard``. + + Raises: + RuntimeError: If accessed while ``dp_shard`` is still the ``-1`` sentinel, i.e. + before :meth:`resolve` has bound the degrees to a world size. + """ + if self.dp_shard == -1: + raise RuntimeError("dp_world_size is undefined before resolve() fills dp_shard=-1.") + return self.dp_replicate * self.dp_shard + + def resolve(self, world_size: int) -> None: + """Bind the declared degrees to a concrete world size (idempotent). + + Fills the ``dp_shard=-1`` sentinel, auto-fills ``dp_replicate`` for the DDP degradation, + and enforces ``dp_replicate * dp_shard * cp == world_size`` with every degree echoed on + failure. + + Args: + world_size (int): Total number of launched processes (torchrun's ``WORLD_SIZE``). + + Raises: + ValueError: If a context-parallel degree is > 1 (the CP engine is not implemented + yet), if ``dp_shard=-1`` cannot be inferred because ``world_size`` is not + divisible by ``dp_replicate * cp``, or if the resolved degrees do not multiply + to ``world_size``. + """ + if self.cp_size > 1: + raise ValueError( + "Context parallelism is not implemented yet: ring_degree and ulysses_degree " + "must be 1. The fields are reserved for the CP engine round." + ) + if self.is_sharded: + if self.dp_shard == -1: + self.dp_shard, remainder = divmod(world_size, self.dp_replicate * self.cp_size) + if remainder or self.dp_shard < 1: + raise ValueError( + f"Cannot infer dp_shard: world_size={world_size} is not divisible by " + f"dp_replicate={self.dp_replicate} * cp={self.cp_size}." + ) + elif self.dp_replicate == 1: + # Untouched config on a multi-process launch: fill the DDP degradation. + self.dp_replicate = world_size + total = self.dp_replicate * self.dp_shard * self.cp_size + if total != world_size: + raise ValueError( + f"Parallelism degrees do not multiply to the world size: dp_replicate=" + f"{self.dp_replicate} * dp_shard={self.dp_shard} * ring=" + f"{self.context_parallel.ring_degree} * ulysses=" + f"{self.context_parallel.ulysses_degree} = {total} != WORLD_SIZE={world_size}." + ) + + +def world_size_from_env() -> int: + """World size as set by torchrun (or 1 outside distributed launches). + + Returns: + int: The ``WORLD_SIZE`` environment variable, or 1 when unset. + """ + return int(os.environ.get("WORLD_SIZE", "1")) diff --git a/src/lerobot/configs/train.py b/src/lerobot/configs/train.py index 1240b749c..cf9b7817e 100644 --- a/src/lerobot/configs/train.py +++ b/src/lerobot/configs/train.py @@ -18,6 +18,7 @@ import multiprocessing import os import tempfile from dataclasses import dataclass, field +from enum import Enum from pathlib import Path from typing import Any @@ -26,6 +27,8 @@ from huggingface_hub import hf_hub_download from huggingface_hub.errors import HfHubHTTPError from lerobot import envs +from lerobot.configs.accelerator import AcceleratorConfig, ActivationCheckpointingMode +from lerobot.configs.parallelism import ParallelismConfig from lerobot.optim import LRSchedulerConfig, OptimizerConfig from lerobot.utils.constants import PRETRAINED_MODEL_DIR from lerobot.utils.hub import HubMixin, find_latest_hub_checkpoint @@ -39,6 +42,34 @@ from .rewards import RewardModelConfig TRAIN_CONFIG_NAME = "train_config.json" +class CheckpointFormat(str, Enum): + """Model-artifact format inside training checkpoints. + + Selects only the *model* artifact; the training_state layout is format-independent (the + optimizer channel is always DCP under sharded runs, safetensors+json otherwise). + + - SAFETENSORS (default): a full `model.safetensors` — maximum compatibility, one gather per + save under sharding. + - DCP: sharded `pytorch_model_fsdp_0/*.distcp` only — fastest save/resume; convert with + `lerobot-convert-dcp` before distributing. + - SAFETENSORS_AND_DCP: both artifacts, written independently. + """ + + SAFETENSORS = "safetensors" + DCP = "dcp" + SAFETENSORS_AND_DCP = "safetensors_dcp" + + @property + def wants_safetensors(self) -> bool: + """True when a full `model.safetensors` artifact should be written.""" + return self in (CheckpointFormat.SAFETENSORS, CheckpointFormat.SAFETENSORS_AND_DCP) + + @property + def wants_dcp(self) -> bool: + """True when sharded DCP model shards (`pytorch_model_fsdp_0/`) should be written.""" + return self in (CheckpointFormat.DCP, CheckpointFormat.SAFETENSORS_AND_DCP) + + def _migrate_legacy_rabc_fields(config: dict[str, Any]) -> dict[str, Any] | None: """Return migrated payload for legacy RA-BC fields, or None when no migration is needed.""" legacy_fields = ( @@ -121,9 +152,16 @@ class TrainPipelineConfig(HubMixin): # Checkpoint is saved every `save_freq` training iterations and after the last training step. # A non-positive value disables periodic saving, keeping only the final checkpoint. save_freq: int = 20_000 + # Model-artifact format inside checkpoints; non-default values require a sharded run. + checkpoint_format: CheckpointFormat = CheckpointFormat.SAFETENSORS use_policy_training_preset: bool = True optimizer: OptimizerConfig | None = None scheduler: LRSchedulerConfig | None = None + # Process topology: dp_replicate / dp_shard (HSDP) and context-parallel degree placeholders. + parallelism: ParallelismConfig = field(default_factory=ParallelismConfig) + # Execution runtime handed to the Accelerator: mixed precision, gradient accumulation, + # FSDP/DDP tuning knobs, compile & activation-checkpointing placeholders. + accelerator: AcceleratorConfig = field(default_factory=AcceleratorConfig) eval: EvalConfig = field(default_factory=EvalConfig) wandb: WandBConfig = field(default_factory=WandBConfig) peft: PeftConfig | None = None @@ -291,6 +329,60 @@ class TrainPipelineConfig(HubMixin): if self.save_checkpoint_to_hub and not (self.policy is not None and self.policy.repo_id): raise ValueError("save_checkpoint_to_hub requires --policy.repo_id.") + self._validate_distributed() + + def _validate_distributed(self) -> None: + """Fail-fasts for the distributed-training scope. + + Raises: + ValueError: If the config requests anything outside the verified scope: context + parallelism or CFG parallelism (reserved placeholders), the compile or + activation-checkpointing placeholders, a DCP checkpoint format on a + non-sharded run, or — under sharded training — fp16 mixed precision, PEFT, + reward-model training, in-training environment evaluation, or multi-optimizer + configs. + """ + if self.parallelism.cp_size > 1: + raise ValueError( + "Context parallelism is not implemented yet: --parallelism.context_parallel.* " + "degrees must be 1 (reserved for the CP engine round)." + ) + if self.parallelism.cfg_parallel != 1: + raise ValueError( + "CFG parallelism is inference-only and must be 1 for training " + "(cfg_parallel is reserved for the serving round)." + ) + if self.accelerator.compile.enabled: + raise ValueError("--accelerator.compile is a placeholder and not wired yet.") + if self.accelerator.activation_checkpointing.mode is not ActivationCheckpointingMode.NONE: + raise ValueError("--accelerator.activation_checkpointing is a placeholder and not wired yet.") + if self.checkpoint_format is not CheckpointFormat.SAFETENSORS and not self.parallelism.is_sharded: + raise ValueError( + f"checkpoint_format={self.checkpoint_format.value} requires a sharded run " + "(--parallelism.dp_shard != 1); non-sharded checkpoints are always safetensors." + ) + if self.parallelism.is_sharded: + if self.accelerator.mixed_precision == "fp16": + raise ValueError( + "fp16 is not supported under sharded training (GradScaler over DTensor " + "gradients is unverified); use bf16 or full precision." + ) + if self.peft is not None: + raise ValueError("PEFT is not supported under sharded training yet.") + if self.is_reward_model_training: + raise ValueError( + "Reward-model training is not supported under sharded training yet " + "(reward models declare no FSDP wrap units and have no sharded save path)." + ) + if self.env is not None and self.env_eval_freq > 0: + raise ValueError( + "In-training environment evaluation is not supported under sharded training " + "(a rank-0-only rollout of a sharded model deadlocks on collectives); set " + "--env_eval_freq=0 and evaluate with lerobot-eval on saved checkpoints." + ) + if self.optimizer is not None and self.optimizer.builds_multiple_optimizers: + raise ValueError("Multi-optimizer configs are not supported under sharded training.") + @classmethod def __get_path_fields__(cls) -> list[str]: """Keys for draccus pretrained-path loading.""" diff --git a/src/lerobot/distributed/__init__.py b/src/lerobot/distributed/__init__.py new file mode 100644 index 000000000..b4ece6eba --- /dev/null +++ b/src/lerobot/distributed/__init__.py @@ -0,0 +1,43 @@ +#!/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. +"""Distributed-training runtime for LeRobot. + +This package owns everything that turns the declarative topology in +:class:`lerobot.configs.parallelism.ParallelismConfig` into a running engine: +mesh math (:class:`~lerobot.distributed.parallel_dims.ParallelDims`), the +`Accelerator` factory (:func:`~lerobot.distributed.factory.make_accelerator`), +sharding-aware checkpoint helpers, and small rank utilities. + +Setup-order contract (normative): +CP dispatch install -> activation checkpointing -> torch.compile -> +``fully_shard``/DDP (via ``accelerator.prepare``) -> optimizer rebind. +Only the last two steps are active today; CP/AC/compile are configured +placeholders wired in later rounds. +""" + +from .factory import guard_against_env_interference, make_accelerator, set_fsdp_wrap_modules +from .parallel_dims import ParallelDims +from .utils import finalize_sharded_policy, is_main_process, strip_accelerate_cp_hooks + +__all__ = [ + "ParallelDims", + "finalize_sharded_policy", + "guard_against_env_interference", + "is_main_process", + "make_accelerator", + "set_fsdp_wrap_modules", + "strip_accelerate_cp_hooks", +] diff --git a/src/lerobot/distributed/checkpoint.py b/src/lerobot/distributed/checkpoint.py new file mode 100644 index 000000000..05d91f2d9 --- /dev/null +++ b/src/lerobot/distributed/checkpoint.py @@ -0,0 +1,195 @@ +#!/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. +"""Sharding-aware checkpoint primitives. + +Two artifact channels with distinct owners: + +- the **distributable** ``model.safetensors``: produced by ``PreTrainedPolicy.save_pretrained`` + through :func:`full_model_state_dict` — a collective full gather when the model is sharded; +- the **resume** channel (sharded runs): torch DCP directories written/read through accelerate's + ``save/load_fsdp_model`` and ``save/load_fsdp_optimizer`` (``pytorch_model_fsdp_0/`` and + ``optimizer_0/``, names imported from accelerate constants), which reshard on load across + topology changes. + +Every function that touches sharded state is a collective and must run on ALL ranks. +""" + +from pathlib import Path +from typing import TYPE_CHECKING + +import torch +from torch import nn + +if TYPE_CHECKING: + from accelerate import Accelerator + + +def is_sharded_module(module: nn.Module) -> bool: + """True when `fully_shard` owns this module's parameters (FSDP2's in-place class swap). + + Args: + module (nn.Module): The module to inspect (a torch.compile wrapper is looked through + via `_orig_mod`). + + Returns: + bool: True when the module (or its compiled `_orig_mod`) is an `FSDPModule`. + """ + from torch.distributed.fsdp import FSDPModule + + if isinstance(module, FSDPModule): + return True + # torch.compile wraps the sharded module; mirror accelerate's `_orig_mod` check. + orig_mod = getattr(module, "_orig_mod", None) + return orig_mod is not None and isinstance(orig_mod, FSDPModule) + + +def full_model_state_dict(module: nn.Module) -> dict[str, torch.Tensor]: + """The module's full (unsharded) state dict, however its parameters are laid out. + + Sharded modules gather through torch's DCP state-dict API: a COLLECTIVE that must run on + every rank; with ``cpu_offload=True`` the full dict materializes on the main rank only and + every other rank receives a literal ``{}`` (runtime-verified — a + rank-0-gated call deadlocks). Plain modules return ``module.state_dict()`` on every rank. + + Args: + module (nn.Module): The (possibly sharded) module to read the state dict from. + + Returns: + dict[str, torch.Tensor]: The full state dict — on the main rank only (``{}`` + elsewhere) when the module is sharded, on every rank otherwise. + """ + if not is_sharded_module(module): + return module.state_dict() + + from torch.distributed.checkpoint.state_dict import StateDictOptions, get_model_state_dict + + return get_model_state_dict(module, options=StateDictOptions(full_state_dict=True, cpu_offload=True)) + + +def _fsdp_plugin(accelerator: "Accelerator") -> object: + """The accelerator's FSDP plugin, required by every DCP save/load helper below. + + Args: + accelerator (Accelerator): The accelerator that prepared the sharded model. + + Returns: + object: The FSDP plugin held by `accelerator.state`. + + Raises: + RuntimeError: If the accelerator was not configured with an FSDP plugin. + """ + plugin = getattr(accelerator.state, "fsdp_plugin", None) + if plugin is None: + raise RuntimeError("Sharded checkpointing requires an FSDP-prepared Accelerator.") + return plugin + + +def save_sharded_model(accelerator: "Accelerator", model: nn.Module, output_dir: Path) -> None: + """Write the DCP model shards (`pytorch_model_fsdp_0/`). Collective: call on all ranks. + + Args: + accelerator (Accelerator): The accelerator that prepared the sharded model. + model (nn.Module): The prepared (sharded) model to save. + output_dir (Path): The directory the shard subdirectory is created in. + """ + from accelerate.utils import save_fsdp_model + + # accelerate 1.14's DCP helpers do string containment checks on the path: + # always hand them str, never Path. + save_fsdp_model(_fsdp_plugin(accelerator), accelerator, model, str(output_dir)) + + +def load_sharded_model(accelerator: "Accelerator", model: nn.Module, input_dir: Path) -> None: + """Load DCP model shards into the prepared (sharded) model. Collective: call on all ranks. + + Args: + accelerator (Accelerator): The accelerator that prepared the sharded model. + model (nn.Module): The prepared (sharded) model to load into. + input_dir (Path): The directory containing the `pytorch_model_fsdp_0/` shard + subdirectory. + """ + from accelerate.utils import load_fsdp_model + from accelerate.utils.constants import FSDP_MODEL_NAME + + # Pass the exact shard directory: accelerate's load resolves it with a substring check + # ("pytorch_model_fsdp" in the path -> use as-is), which misfires on run paths that happen + # to contain the marker; the exact dir makes the check deterministic. + load_fsdp_model(_fsdp_plugin(accelerator), accelerator, model, str(input_dir / f"{FSDP_MODEL_NAME}_0")) + + +def save_sharded_optimizer( + accelerator: "Accelerator", optimizer: torch.optim.Optimizer, model: nn.Module, output_dir: Path +) -> None: + """Write the DCP optimizer shards (`optimizer_0/`). Collective: call on all ranks. + + Args: + accelerator (Accelerator): The accelerator that prepared the model and optimizer. + optimizer (torch.optim.Optimizer): The prepared optimizer to save the state from. + model (nn.Module): The prepared (sharded) model the optimizer state is keyed by. + output_dir (Path): The directory the shard subdirectory is created in. + """ + from accelerate.utils import save_fsdp_optimizer + + save_fsdp_optimizer(_fsdp_plugin(accelerator), accelerator, optimizer, model, str(output_dir)) + + +def load_sharded_optimizer( + accelerator: "Accelerator", optimizer: torch.optim.Optimizer, model: nn.Module, input_dir: Path +) -> None: + """Load DCP optimizer shards into the prepared optimizer. Collective: call on all ranks. + + Must run AFTER ``accelerator.prepare()``: FSDP2's prepare rebinds the optimizer's param + groups to sharded DTensors but never migrates ``optimizer.state`` — the resharding load is + the only correct way to restore it. + + Args: + accelerator (Accelerator): The accelerator that prepared the model and optimizer. + optimizer (torch.optim.Optimizer): The prepared optimizer to restore the state into. + model (nn.Module): The prepared (sharded) model the optimizer state is keyed by. + input_dir (Path): The directory containing the `optimizer_0/` shard subdirectory. + """ + from accelerate.utils import load_fsdp_optimizer + from accelerate.utils.constants import OPTIMIZER_NAME + + # Exact shard directory for the same reason as load_sharded_model: accelerate's substring + # check ("optimizer" in the path) would misread e.g. --job_name=optimizer_sweep run paths. + load_fsdp_optimizer( + _fsdp_plugin(accelerator), accelerator, optimizer, model, str(input_dir / f"{OPTIMIZER_NAME}_0") + ) + + +def dcp_to_safetensors(dcp_dir: Path, output_dir: Path, *, delete_dcp: bool = False) -> Path: + """Merge a DCP shard directory into a single `model.safetensors` (offline, single process). + + Thin wrapper over `accelerate.utils.merge_fsdp_weights`, which loads the shards without a + process group, writes safetensors directly, and — when asked — removes the merged shard + directory itself, only on the main process and only once the merge has succeeded. + + Args: + dcp_dir (Path): The DCP shard directory to merge (e.g. `.../pytorch_model_fsdp_0`). + output_dir (Path): The directory the merged `model.safetensors` is written into. + delete_dcp (bool): Whether to remove the shard directory once it has been merged. + Defaults to False. + + Returns: + Path: The written `model.safetensors` file's path. + """ + from accelerate.utils import merge_fsdp_weights + + merge_fsdp_weights( + str(dcp_dir), str(output_dir), safe_serialization=True, remove_checkpoint_dir=delete_dcp + ) + return output_dir / "model.safetensors" diff --git a/src/lerobot/distributed/factory.py b/src/lerobot/distributed/factory.py new file mode 100644 index 000000000..b1003c224 --- /dev/null +++ b/src/lerobot/distributed/factory.py @@ -0,0 +1,147 @@ +#!/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. +"""The `Accelerator` factory — the only place accelerate gets configured. + +`torchrun` is the launcher; every accelerate parameter comes from `TrainPipelineConfig` +(`cfg.parallelism` + `cfg.accelerator`) so a run is reproducible from its `train_config.json` +alone. `accelerate launch` without a `--config_file` remains equivalent (it only sets rendezvous +env vars in that mode); the yaml flow is superseded. +""" + +import os +from typing import TYPE_CHECKING + +from lerobot.configs.parallelism import world_size_from_env +from lerobot.configs.train import TrainPipelineConfig + +if TYPE_CHECKING: + from accelerate import Accelerator + + from lerobot.policies.pretrained import PreTrainedPolicy + +# Env vars through which `accelerate launch --config_file` (or a stray shell) would configure +# accelerate behind the config system's back. Plugin `__post_init__`s read these silently as +# field fallbacks (ACCELERATE_DYNAMO_* enables torch.compile through the default +# TorchDynamoPlugin; ACCELERATE_GRADIENT_ACCUMULATION_STEPS overrides the explicitly passed +# value inside Accelerator.__init__), which would make train_config.json lie about what ran. +_ACCELERATE_ENV_PREFIXES = ("FSDP_", "PARALLELISM_CONFIG_", "ACCELERATE_DYNAMO_") +_ACCELERATE_ENV_VARS = ( + "ACCELERATE_USE_FSDP", + "ACCELERATE_USE_PARALLELISM_CONFIG", + "ACCELERATE_GRADIENT_ACCUMULATION_STEPS", +) +_ENV_OVERRIDE = "LEROBOT_ALLOW_ACCELERATE_ENV" + + +def guard_against_env_interference() -> None: + """Hard-error when accelerate-configuring env vars are set. + + A silently env-overridden "reproducible" config is worse than a stop: users migrating from + the old `accelerate launch --config_file fsdp.yaml` flow get a precise error instead of a + config that lies. Set LEROBOT_ALLOW_ACCELERATE_ENV=1 to acknowledge and proceed. + + Raises: + RuntimeError: If any accelerate-configuring environment variable is set and the + LEROBOT_ALLOW_ACCELERATE_ENV override is not. + """ + if os.environ.get(_ENV_OVERRIDE): + return + offending = sorted( + name + for name in os.environ + if name in _ACCELERATE_ENV_VARS or name.startswith(_ACCELERATE_ENV_PREFIXES) + ) + if offending: + raise RuntimeError( + f"Accelerate-configuring environment variables are set: {', '.join(offending)}. " + "LeRobot manages accelerate exclusively through TrainPipelineConfig " + "(--parallelism.* / --accelerator.*); launch with plain torchrun and remove these " + "variables (the `accelerate launch --config_file` flow is superseded), or set " + f"{_ENV_OVERRIDE}=1 to acknowledge that they may override your config." + ) + + +def make_accelerator(cfg: TrainPipelineConfig) -> "Accelerator": + """Resolve the topology against the launched world and build the `Accelerator`. + + Must run once per process, before any other component needs the device or the process + group (`Accelerator.__init__` initializes both and builds the device mesh). + + Args: + cfg (TrainPipelineConfig): The full training config; `cfg.parallelism` is resolved in + place against the launched world size and `cfg.accelerator` builds the result. + + Returns: + Accelerator: The configured accelerator, with device and process group initialized. + + Raises: + ValueError: If `cfg.checkpoint_format` requires DCP but the topology resolved to a + non-sharded run. + """ + guard_against_env_interference() + cfg.parallelism.resolve(world_size_from_env()) + # The parse-time format check ran against the declared degrees, where the dp_shard=-1 + # sentinel counts as sharded; it may resolve to an unsharded run (e.g. -1 at world size 1). + # Re-check against the concrete degrees so the recorded format never lies about the + # artifacts a checkpoint will actually contain. + if cfg.checkpoint_format.wants_dcp and not cfg.parallelism.is_sharded: + raise ValueError( + f"checkpoint_format={cfg.checkpoint_format.value} requires a sharded run, but the " + f"topology resolved to a non-sharded one (dp_replicate={cfg.parallelism.dp_replicate}, " + f"dp_shard={cfg.parallelism.dp_shard}); non-sharded checkpoints are always safetensors." + ) + return cfg.accelerator.build( + cfg.parallelism, + cpu=cfg.trainable_config.device == "cpu", + ) + + +def set_fsdp_wrap_modules(accelerator: "Accelerator", policy: "PreTrainedPolicy") -> None: + """Resolve the FSDP wrap-unit class names onto the plugin before `accelerator.prepare()`. + + Resolution order: user override (`--accelerator.fsdp.wrap_modules`, already on the plugin) + -> the policy's `_fsdp_wrap_modules` declaration -> hard error. Root-only wrapping — the + silent default when no wrap source exists — is never accepted: it quietly forfeits all + sharding memory savings. + + No-op for the size-based policy (`--accelerator.fsdp.min_num_params`), which needs no class + names, and for non-sharded runs (no fsdp plugin). + + Args: + accelerator (Accelerator): The accelerator whose FSDP plugin receives the wrap-unit + class names. + policy (PreTrainedPolicy): The trainable whose class may declare `_fsdp_wrap_modules`. + + Raises: + ValueError: If sharded class-based wrapping is configured but neither a user override + nor a policy declaration supplies wrap-unit class names. + """ + plugin = getattr(accelerator.state, "fsdp_plugin", None) + if plugin is None or plugin.min_num_params: + return + if plugin.transformer_cls_names_to_wrap: # user override, set at build time + return + # getattr, not attribute access: non-policy trainables (no `_fsdp_wrap_modules` attribute) + # must reach the actionable error below, not an AttributeError. + declared = getattr(type(policy), "_fsdp_wrap_modules", None) + if not declared: + raise ValueError( + f"Policy '{type(policy).__name__}' declares no FSDP wrap units. Sharded training " + "requires wrap-unit class names: set --accelerator.fsdp.wrap_modules='[\"MyBlock\"]' " + "(or --accelerator.fsdp.min_num_params for a size-based policy), or declare " + "`_fsdp_wrap_modules` on the policy class." + ) + plugin.transformer_cls_names_to_wrap = list(declared) diff --git a/src/lerobot/distributed/parallel_dims.py b/src/lerobot/distributed/parallel_dims.py new file mode 100644 index 000000000..7a9c98661 --- /dev/null +++ b/src/lerobot/distributed/parallel_dims.py @@ -0,0 +1,112 @@ +#!/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. +"""Runtime mesh math derived from the declarative :class:`ParallelismConfig`. + +`ParallelDims` is the training script's single source of truth for topology-derived numbers +(data-parallel world size and rank, sample accounting inputs) and — once the CP engine lands — +the owner of LeRobot's private ``(dp_replicate, dp_shard, ring, ulysses)`` mesh. It is a runtime +object and is never serialized (the config it derives from is what lands in +``train_config.json``). +""" + +from dataclasses import dataclass + +import torch.distributed as dist + +from lerobot.configs.parallelism import ParallelismConfig + + +@dataclass(frozen=True) +class ParallelDims: + """Concrete parallelism degrees bound to a world size (canonical row-major rank layout).""" + + dp_replicate: int + dp_shard: int + ring: int + ulysses: int + world_size: int + device_type: str + + @classmethod + def from_config(cls, cfg: ParallelismConfig, world_size: int, device_type: str) -> "ParallelDims": + """Bind a *resolved* config to the actual runtime world size (cross-checked here). + + Args: + cfg (ParallelismConfig): The declarative topology, already resolved via + `ParallelismConfig.resolve(world_size)`. + world_size (int): The launched world size the declared degrees must multiply to. + device_type (str): The accelerator device type backing the mesh (e.g. "cuda"). + + Returns: + ParallelDims: The concrete parallelism degrees bound to this world. + + Raises: + ValueError: If the config is unresolved (`dp_shard == -1`) or its degrees do not + multiply to `world_size`. + """ + total = cfg.dp_replicate * cfg.dp_shard * cfg.cp_size + if cfg.dp_shard == -1 or total != world_size: + raise ValueError( + f"ParallelismConfig is not resolved against this world: dp_replicate=" + f"{cfg.dp_replicate} * dp_shard={cfg.dp_shard} * cp={cfg.cp_size} != " + f"world_size={world_size}. Call ParallelismConfig.resolve(world_size) first " + "(make_accelerator does this)." + ) + return cls( + dp_replicate=cfg.dp_replicate, + dp_shard=cfg.dp_shard, + ring=cfg.context_parallel.ring_degree, + ulysses=cfg.context_parallel.ulysses_degree, + world_size=world_size, + device_type=device_type, + ) + + @property + def cp_size(self) -> int: + """Total context-parallel degree (`ring * ulysses`).""" + return self.ring * self.ulysses + + @property + def is_sharded(self) -> bool: + """Whether parameters are sharded (`dp_shard > 1` or any context parallelism).""" + return self.dp_shard > 1 or self.cp_size > 1 + + @property + def dp_world_size(self) -> int: + """Number of distinct data-parallel workers — the divisor for all sample accounting.""" + return self.dp_replicate * self.dp_shard + + @property + def dp_rank(self) -> int: + """This process's data-parallel coordinate (CP peers share one dp_rank). + + With the canonical row-major layout and (ring, ulysses) innermost, CP peers are + contiguous global ranks, so the dp coordinate is the integer quotient by cp_size — + the same arithmetic accelerate's mesh-aware dataloader applies. + """ + global_rank = dist.get_rank() if dist.is_initialized() else 0 + return global_rank // self.cp_size + + def cp_mesh(self) -> None: + """Private (ring, ulysses) mesh for the CP engine — reserved for the CP round. + + Raises: + NotImplementedError: Always — context parallelism is not implemented yet. + """ + raise NotImplementedError( + "Context parallelism is not implemented yet; ParallelDims.cp_mesh is reserved for " + "the CP engine round (a private mesh aligned with accelerate's cp block)." + ) diff --git a/src/lerobot/distributed/utils.py b/src/lerobot/distributed/utils.py new file mode 100644 index 000000000..33616fa26 --- /dev/null +++ b/src/lerobot/distributed/utils.py @@ -0,0 +1,94 @@ +#!/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. +"""Rank utilities and post-`prepare()` sharding finalization.""" + +import logging +from typing import TYPE_CHECKING + +import torch.distributed as dist +from torch import nn + +if TYPE_CHECKING: + from lerobot.distributed.parallel_dims import ParallelDims + + +def is_main_process() -> bool: + """True on the process that owns rank-0-only side effects (file writes, uploads, logging). + + Torch-native on purpose: persistence code must not depend on an `Accelerator` handle — + `_save_pretrained` and the hub publishers run in contexts that have none. Outside + distributed runs every process is the main process. + + Returns: + bool: True when this process is rank 0 or no process group is initialized. + """ + return not dist.is_initialized() or dist.get_rank() == 0 + + +def strip_accelerate_cp_hooks(model: nn.Module) -> int: + """Remove accelerate's context-parallel forward-pre-hooks from every module. + + When `cp_size > 1` is declared, `accelerator.prepare()` unconditionally attaches hooks that + silently replace any `attention_mask` kwarg of `*self_attn` modules with `is_causal=True` + (`accelerate.big_modeling._attach_context_parallel_hooks`) — mask corruption for policies + with non-causal attention. LeRobot implements CP itself and never enters accelerate's CP + context, so these hooks are pure hazard. Deterministically identified by their defining + module; a version canary pins that identity. + + Args: + model (nn.Module): The prepared model to strip the hooks from (all submodules are + visited). + + Returns: + int: The number of hooks removed. + """ + removed = 0 + for module in model.modules(): + for hook_id, hook in list(module._forward_pre_hooks.items()): + if getattr(hook, "__module__", None) == "accelerate.big_modeling": + del module._forward_pre_hooks[hook_id] + module._forward_pre_hooks_with_kwargs.pop(hook_id, None) + removed += 1 + return removed + + +def finalize_sharded_policy(policy: nn.Module, parallel_dims: "ParallelDims") -> None: + """Sharding correctness protocol, applied once, immediately after `accelerator.prepare()`. + + 1. Strip accelerate's CP mask hooks (only attached when cp > 1 was declared). + 2. Register the policy's non-`forward` entry points (`_fsdp_forward_methods`) so FSDP2 + unshards parameters around `select_action` & co. — without this, any inference-style + call on a sharded policy crashes on mixed Tensor/DTensor. + + No-op for DDP/single-process runs. + + Args: + policy (nn.Module): The policy as returned by `accelerator.prepare()`. + parallel_dims (ParallelDims): The run's resolved topology; decides whether the protocol + applies. + """ + if not parallel_dims.is_sharded: + return + if parallel_dims.cp_size > 1: + removed = strip_accelerate_cp_hooks(policy) + logging.info("Stripped %d accelerate context-parallel attention-mask hooks.", removed) + + from torch.distributed.fsdp import FSDPModule, register_fsdp_forward_method + + if isinstance(policy, FSDPModule): + for method_name in getattr(type(policy), "_fsdp_forward_methods", ()): + if callable(getattr(policy, method_name, None)): + register_fsdp_forward_method(policy, method_name) diff --git a/src/lerobot/jobs/hf.py b/src/lerobot/jobs/hf.py index 80d40ad62..61370b2be 100644 --- a/src/lerobot/jobs/hf.py +++ b/src/lerobot/jobs/hf.py @@ -432,7 +432,7 @@ def submit_to_hf(cfg: TrainPipelineConfig) -> None: # Finish as soon as the model is pushed, rather than waiting out the platform's # post-run finalization before the job stage flips to COMPLETED. This matches the - # exact log line emitted by PreTrainedPolicy.push_model_to_hub — the two must stay + # exact log line emitted by lerobot.common.train_utils.publish_trained_model — the two must stay # in sync. If it ever stops matching we just fall back to stage-based completion # (~30s slower), so the contract is an optimization, not a correctness requirement. success_marker = f"Model pushed to https://huggingface.co/{repo_id}" diff --git a/src/lerobot/optim/__init__.py b/src/lerobot/optim/__init__.py index 2d564c25f..46676027b 100644 --- a/src/lerobot/optim/__init__.py +++ b/src/lerobot/optim/__init__.py @@ -20,7 +20,6 @@ from .optimizers import ( SGDConfig as SGDConfig, XVLAAdamWConfig as XVLAAdamWConfig, load_optimizer_state, - load_optimizer_state_dict, save_optimizer_state, ) from .schedulers import ( @@ -51,7 +50,6 @@ __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 0a462e1aa..820579c06 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, load_json, write_json +from lerobot.utils.io_utils import deserialize_json_into_object, write_json from lerobot.utils.utils import flatten_dict, unflatten_dict # Type alias for parameters accepted by optimizer build() methods. @@ -52,6 +52,11 @@ class OptimizerConfig(draccus.ChoiceRegistry, abc.ABC): def type(self) -> str: return self.get_choice_name(self.__class__) + @property + def builds_multiple_optimizers(self) -> bool: + """True when build() returns a dict of optimizers (unsupported under sharded training).""" + return False + @classmethod def default_choice_name(cls) -> str | None: return "adam" @@ -245,6 +250,10 @@ class MultiAdamConfig(OptimizerConfig): grad_clip_norm: float = 10.0 optimizer_groups: dict[str, dict[str, Any]] = field(default_factory=dict) + @property + def builds_multiple_optimizers(self) -> bool: + return True + def build(self, params: OptimizerParams) -> dict[str, torch.optim.Optimizer]: """Build multiple Adam optimizers. @@ -283,35 +292,27 @@ class MultiAdamConfig(OptimizerConfig): def save_optimizer_state( optimizer: torch.optim.Optimizer | dict[str, torch.optim.Optimizer], save_dir: Path, - optim_state_dict: dict | None = None, ) -> None: - """Save optimizer state to disk. + """Save optimizer state to disk (non-sharded runs; sharded runs use the DCP channel). 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, optim_state_dict=optim_state_dict) + _save_single_optimizer_state(optimizer, save_dir) -def _save_single_optimizer_state( - optimizer: torch.optim.Optimizer, save_dir: Path, optim_state_dict: dict | None = None -) -> None: +def _save_single_optimizer_state(optimizer: torch.optim.Optimizer, save_dir: Path) -> None: """Save a single optimizer's state to disk.""" - state = dict(optim_state_dict) if optim_state_dict is not None else optimizer.state_dict() + state = optimizer.state_dict() param_groups = state.pop("param_groups") flat_state = flatten_dict(state) save_file(flat_state, save_dir / OPTIMIZER_STATE) @@ -365,19 +366,3 @@ 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/act/modeling_act.py b/src/lerobot/policies/act/modeling_act.py index 1432b68a5..2298db397 100644 --- a/src/lerobot/policies/act/modeling_act.py +++ b/src/lerobot/policies/act/modeling_act.py @@ -47,6 +47,8 @@ class ACTPolicy(PreTrainedPolicy): config_class = ACTConfig name = "act" + # FSDP2 wrap units: one unit per transformer layer of both stacks. + _fsdp_wrap_modules = ["ACTEncoderLayer", "ACTDecoderLayer"] def __init__( self, diff --git a/src/lerobot/policies/factory.py b/src/lerobot/policies/factory.py index 3527a1821..ae6869d91 100644 --- a/src/lerobot/policies/factory.py +++ b/src/lerobot/policies/factory.py @@ -242,6 +242,7 @@ def make_policy( ds_meta: LeRobotDatasetMetadata | None = None, env_cfg: EnvConfig | None = None, rename_map: dict[str, str] | None = None, + defer_weight_load: bool = False, ) -> PreTrainedPolicy: """ Instantiate a policy model. @@ -252,22 +253,27 @@ def make_policy( can either initialize a new policy from scratch or load a pretrained one. Args: - cfg: The configuration for the policy to be created. If `cfg.pretrained_path` is - set, the policy will be loaded with weights from that path. - ds_meta: Dataset metadata used to infer feature shapes and types. Also provides - statistics for normalization layers. - env_cfg: Environment configuration used to infer feature shapes and types. - One of `ds_meta` or `env_cfg` must be provided. - rename_map: Optional mapping of dataset or environment feature keys to match - expected policy feature names (e.g., `"left"` → `"camera1"`). + cfg (PreTrainedConfig): The configuration for the policy to be created. If + `cfg.pretrained_path` is set, the policy will be loaded with weights from that path. + ds_meta (LeRobotDatasetMetadata | None): Dataset metadata used to infer feature shapes and + types. Also provides statistics for normalization layers. + env_cfg (EnvConfig | None): Environment configuration used to infer feature shapes and + types. One of `ds_meta` or `env_cfg` must be provided. + rename_map (dict[str, str] | None): Optional mapping of dataset or environment feature + keys to match expected policy feature names (e.g., `"left"` → `"camera1"`). + defer_weight_load (bool): Build the exact policy `from_pretrained` would build — same + config resolution, same stats-derived buffers, same device placement and eval mode — + but skip the safetensors weight load. Used when resuming from a DCP checkpoint, whose + sharded weights stream in after `accelerator.prepare()` (the distributed checkpoint + engine overwrites the random init). Returns: - An instantiated and device-placed policy model. + PreTrainedPolicy: An instantiated and device-placed policy model. Raises: ValueError: If both or neither of `ds_meta` and `env_cfg` are provided. - NotImplementedError: If attempting to use an unsupported policy-backend - combination (e.g., VQBeT with 'mps'). + NotImplementedError: If attempting to use an unsupported policy-backend combination + (e.g., VQBeT with 'mps'). """ if bool(ds_meta) == bool(env_cfg): raise ValueError("Either one of a dataset metadata or a sim env must be provided.") @@ -332,11 +338,18 @@ def make_policy( ) if cfg.pretrained_path and not cfg.use_peft: - # Load a pretrained policy and override the config if needed (for example, if there are inference-time - # hyperparameters that we want to vary). - kwargs["pretrained_name_or_path"] = cfg.pretrained_path - kwargs["revision"] = cfg.pretrained_revision - policy = policy_cls.from_pretrained(**kwargs) + if defer_weight_load: + # Same construction path as from_pretrained (config already resolved from the + # checkpoint by the caller; dataset_stats/dataset_meta kwargs identical), minus the + # weight load — parity by construction. + policy = policy_cls(**kwargs) + policy.eval() + else: + # Load a pretrained policy and override the config if needed (for example, if there + # are inference-time hyperparameters that we want to vary). + kwargs["pretrained_name_or_path"] = cfg.pretrained_path + kwargs["revision"] = cfg.pretrained_revision + policy = policy_cls.from_pretrained(**kwargs) elif cfg.pretrained_path and cfg.use_peft: # Load a pretrained PEFT model on top of the policy. The pretrained path points to the folder/repo # of the adapter and the adapter's config contains the path to the base policy. So we need the diff --git a/src/lerobot/policies/fastwam/modeling_fastwam.py b/src/lerobot/policies/fastwam/modeling_fastwam.py index 10671e717..46fb8241b 100644 --- a/src/lerobot/policies/fastwam/modeling_fastwam.py +++ b/src/lerobot/policies/fastwam/modeling_fastwam.py @@ -54,6 +54,9 @@ class FastWAMPolicy(PreTrainedPolicy): config_class = FastWAMConfig name = "fastwam" + # FSDP2 wrap units: MoTLayer is the single FSDP owner of each layer's expert blocks + # (the blocks are re-parented onto it precisely so sharding has one boundary to hook). + _fsdp_wrap_modules = ["MoTLayer"] def __init__( self, diff --git a/src/lerobot/policies/pretrained.py b/src/lerobot/policies/pretrained.py index 8823e5808..00d656aed 100644 --- a/src/lerobot/policies/pretrained.py +++ b/src/lerobot/policies/pretrained.py @@ -18,20 +18,17 @@ import builtins import dataclasses import logging import os -from importlib.resources import files +import warnings from pathlib import Path -from tempfile import TemporaryDirectory -from typing import TYPE_CHECKING, TypedDict, TypeVar, Unpack +from typing import TYPE_CHECKING, Any, ClassVar, TypedDict, TypeVar, Unpack -from huggingface_hub import HfApi, ModelCard, ModelCardData, hf_hub_download, save_torch_state_dict +from huggingface_hub import 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 +from safetensors.torch import load_model as load_model_as_safetensor from torch import Tensor, nn -from lerobot.__version__ import __version__ from lerobot.configs import PreTrainedConfig -from lerobot.configs.train import TrainPipelineConfig from lerobot.utils.device_utils import resolve_safetensors_device from lerobot.utils.hub import HubMixin from lerobot.utils.import_utils import _peft_available, require_package @@ -46,56 +43,14 @@ else: get_peft_model = None if TYPE_CHECKING: + from lerobot.configs.train import TrainPipelineConfig from lerobot.datasets.dataset_metadata import LeRobotDatasetMetadata T = TypeVar("T", bound="PreTrainedPolicy") - -def _build_card_context( - cfg: TrainPipelineConfig | None, - dataset_meta: LeRobotDatasetMetadata | None, - input_features: dict | None, - output_features: dict | None, -) -> dict: - """Collect optional data for the model-card template. - - Returns plain values only (no Markdown) — the template in - ``lerobot/templates/lerobot_modelcard_template.md`` decides how and whether to show - each one. Everything is best-effort: anything unavailable is left empty/None and the - template simply skips that section, so this never breaks a Hub push. - """ - context = { - "training": None, - "input_features": input_features or {}, - "output_features": output_features or {}, - "dataset": None, - "robot_type": None, - "cameras": [], - } - - if cfg is not None: - optimizer = getattr(cfg, "optimizer", None) - context["training"] = { - "steps": cfg.steps, - "batch_size": cfg.batch_size, - "seed": cfg.seed, - "optimizer": getattr(optimizer, "type", None) if optimizer else None, - "lr": getattr(optimizer, "lr", None) if optimizer else None, - "lerobot_version": __version__, - } - - if dataset_meta is not None: - context["dataset"] = { - "repo_id": dataset_meta.repo_id, - "episodes": dataset_meta.total_episodes, - "frames": dataset_meta.total_frames, - "fps": dataset_meta.fps, - "tasks": [str(task) for task in dataset_meta.tasks.index], - } - context["robot_type"] = dataset_meta.robot_type - context["cameras"] = [key.split(".")[-1] for key in dataset_meta.camera_keys] - - return context +# Pinned far above any policy's total size so save_torch_state_dict always emits exactly one +# `model.safetensors` (no shards, no index) — a constant, not a computed byte count. +_SINGLE_FILE_SHARD_SIZE = "1TB" class ActionSelectKwargs(TypedDict, total=False): @@ -110,6 +65,22 @@ class PreTrainedPolicy(nn.Module, HubMixin, abc.ABC): config_class: None name: None + # --- declarative parallelism/acceleration surface ---------------------------------------- + # Module CLASS names forming the FSDP2 wrap units (and, once wired, the activation- + # checkpointing units). Resolved onto the accelerate plugin right before + # `accelerator.prepare()` by `lerobot.distributed.set_fsdp_wrap_modules`; sharded training + # with no wrap source anywhere fails loudly instead of silently wrapping only the root. + _fsdp_wrap_modules: ClassVar[list[str] | None] = None + # Non-`forward` entry points that must trigger FSDP2 unshard/reshard hooks when called on a + # sharded policy (registered post-prepare via `torch.distributed.fsdp + # .register_fsdp_forward_method`); calling them unregistered crashes on mixed Tensor/DTensor. + _fsdp_forward_methods: ClassVar[tuple[str, ...]] = ("select_action", "predict_action_chunk") + # Capability gate for the (future) activation-checkpointing wiring. + supports_gradient_checkpointing: ClassVar[bool] = False + # Declarative context-parallel plan (diffusers `ContextParallelModelPlan` semantics: + # module FQN -> sequence split/gather spec). Reserved for the CP engine round. + _cp_plan: ClassVar[dict[str, Any] | None] = None + def __init__(self, config: PreTrainedConfig, *inputs, **kwargs): super().__init__() if not isinstance(config, PreTrainedConfig): @@ -127,43 +98,33 @@ 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: 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). + def _save_pretrained(self, save_directory: Path) -> None: + """Serialize this policy's parameters (and config) into `save_directory`. - 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. + Sharding is handled internally: under FSDP2 the full state dict is gathered through a + COLLECTIVE, so when the policy is sharded this method (via `save_pretrained`) must be + called on EVERY rank — a rank-0-gated call deadlocks. File writes happen on the main + process only, in all layouts (single, DDP, sharded). + + Args: + save_directory (Path): Target directory for the policy config (`config.json`) and the + safetensors weight file(s). """ - 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 + # Lazy imports: the persistence layer pulls in lerobot.distributed only when saving. + from lerobot.distributed.checkpoint import full_model_state_dict, is_sharded_module + from lerobot.distributed.utils import is_main_process - 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 - if state_dict is None: - save_model_as_safetensor(model_to_save, str(save_directory / SAFETENSORS_SINGLE_FILE)) + if is_sharded_module(model_to_save): + logging.info("Gathering the full state dict from all ranks (sharded policy).") + state_dict = full_model_state_dict(model_to_save) # collective when sharded; {} off-main + if not state_dict or not is_main_process(): + # Sharded: the gather materializes on the main rank only (emptiness check). + # Non-sharded multi-rank (DDP): every rank holds a full dict — the explicit rank + # gate prevents N ranks racing on the same files. Single process: never taken. 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)) + self.config._save_pretrained(save_directory) + save_torch_state_dict(state_dict, str(save_directory), max_shard_size=_SINGLE_FILE_SHARD_SIZE) @classmethod def from_pretrained( @@ -291,92 +252,39 @@ class PreTrainedPolicy(nn.Module, HubMixin, abc.ABC): peft_model=None, state_dict: dict[str, Tensor] | None = None, dataset_meta: LeRobotDatasetMetadata | None = None, - ): - api = HfApi() - repo_id = api.create_repo( - repo_id=self.config.repo_id, private=self.config.private, exist_ok=True - ).repo_id + ) -> None: + """Publish this policy to the Hub. - # Push the files to the repo in a single commit - with TemporaryDirectory(ignore_cleanup_errors=True) as tmp: - saved_path = Path(tmp) / repo_id + Deprecated: use :func:`lerobot.common.train_utils.publish_trained_model` instead, which + also publishes the pre/post-processors alongside the model. - if peft_model is not None: - # Since PEFT just forwards calls to `push_model_to_hub`, `self` is not the PeftModel wrapper - # but the actual policy which is why we need the PEFT model passed to us to save the adapter. - # That also means that we need to store the policy config ourselves since PEFT can't. - peft_model.save_pretrained(saved_path) - self.config.save_pretrained(saved_path) - else: - # Calls _save_pretrained and stores model tensors - self.save_pretrained(saved_path, state_dict=state_dict) + Args: + cfg (TrainPipelineConfig): The training config; saved as `train_config.json` and + used to render the model card. + peft_model: The PEFT wrapper when training adapters, whose weights replace the full + model weights in the published repo. Defaults to None. + state_dict (dict[str, Tensor] | None): Ignored; weights are now gathered internally + when the policy is sharded. Defaults to None. + dataset_meta (LeRobotDatasetMetadata | None): Dataset metadata for the model card, + if available. Defaults to None. + """ + from lerobot.common.train_utils import publish_trained_model - card = self.generate_model_card( - cfg.dataset.repo_id, - self.config.type, - self.config.license, - self.config.tags, - cfg=cfg, - dataset_meta=dataset_meta, + warnings.warn( + "PreTrainedPolicy.push_model_to_hub is deprecated and will be removed in a future " + "version. Use lerobot.common.train_utils.publish_trained_model(cfg, model, " + "preprocessor, postprocessor, dataset_meta) instead.", + FutureWarning, + stacklevel=2, + ) + if state_dict is not None: + warnings.warn( + "The `state_dict` argument is ignored: sharded weights are gathered internally " + "when the policy is saved.", + FutureWarning, + stacklevel=2, ) - card.save(str(saved_path / "README.md")) - - cfg.save_pretrained(saved_path) # Calls _save_pretrained and stores train config - - commit_info = api.upload_folder( - repo_id=repo_id, - repo_type="model", - folder_path=saved_path, - commit_message="Upload policy weights, train config and readme", - allow_patterns=["*.safetensors", "*.json", "*.yaml", "*.md"], - ignore_patterns=["*.tmp", "*.log"], - ) - - # Contract: lerobot.jobs.hf.submit_to_hf watches for this exact - # "Model pushed to " line to end a remote run early. Keep the wording - # and URL format in sync (it falls back to status polling if they drift). - logging.info(f"Model pushed to {commit_info.repo_url.url}") - - def generate_model_card( - self, - dataset_repo_id: str, - model_type: str, - license: str | None, - tags: list[str] | None, - cfg: TrainPipelineConfig | None = None, - dataset_meta: LeRobotDatasetMetadata | None = None, - ) -> ModelCard: - base_model_mapping = { - "smolvla": "lerobot/smolvla_base", - "pi0": "lerobot/pi0_base", - "pi05": "lerobot/pi05_base", - "pi0_fast": "lerobot/pi0fast-base", - "xvla": "lerobot/xvla-base", - } - - card_data = ModelCardData( - license=license or "apache-2.0", - library_name="lerobot", - pipeline_tag="robotics", - tags=list(set(tags or []).union({"robotics", "lerobot", model_type})), - model_name=model_type, - datasets=dataset_repo_id, - base_model=base_model_mapping.get(model_type), - ) - - context = _build_card_context( - cfg, dataset_meta, self.config.input_features, self.config.output_features - ) - # Used by the template to pre-fill commands and the "Fine-tuned from" line. - context["policy_repo_id"] = getattr(self.config, "repo_id", None) - context["base_model"] = base_model_mapping.get(model_type) - - template_card = ( - files("lerobot.templates").joinpath("lerobot_modelcard_template.md").read_text(encoding="utf-8") - ) - card = ModelCard.from_template(card_data, template_str=template_card, **context) - card.validate() - return card + publish_trained_model(cfg, self, None, None, dataset_meta, peft_model=peft_model) def wrap_with_peft( self, diff --git a/src/lerobot/processor/migrate_policy_normalization.py b/src/lerobot/processor/migrate_policy_normalization.py index 37df4be41..055327bfe 100644 --- a/src/lerobot/processor/migrate_policy_normalization.py +++ b/src/lerobot/processor/migrate_policy_normalization.py @@ -647,10 +647,15 @@ def main(): tags = set(tags).union({"robotics", "lerobot", policy_type}) tags = list(tags) - # Generate model card - card = policy.generate_model_card( - dataset_repo_id=dataset_repo_id, model_type=policy_type, license=license, tags=tags - ) + # Generate model card through the free helper (PreTrainedPolicy.generate_model_card was + # removed with the publisher redesign), then apply the metadata recovered above — the + # migrated policy config does not carry the original repo's card fields. + from lerobot.common.train_utils import generate_model_card + + card = generate_model_card(policy.config) + card.data.datasets = dataset_repo_id + card.data.license = license + card.data.tags = sorted(tags) # Save model card locally card.save(str(output_dir / "README.md")) diff --git a/src/lerobot/rewards/pretrained.py b/src/lerobot/rewards/pretrained.py index 185e66a3b..e162fa300 100644 --- a/src/lerobot/rewards/pretrained.py +++ b/src/lerobot/rewards/pretrained.py @@ -16,12 +16,11 @@ import abc import builtins import logging import os -from importlib.resources import files +import warnings from pathlib import Path -from tempfile import TemporaryDirectory from typing import TYPE_CHECKING, Any, TypeVar -from huggingface_hub import HfApi, ModelCard, ModelCardData, hf_hub_download +from huggingface_hub import hf_hub_download 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 @@ -61,6 +60,22 @@ class PreTrainedRewardModel(nn.Module, HubMixin, abc.ABC): raise TypeError(f"Class {cls.__name__} must define 'name'") def _save_pretrained(self, save_directory: Path) -> None: + """Serialize this reward model's parameters (and config) into `save_directory`. + + Safe to call on every rank: replicas carry identical weights, so only the main process + writes (sharded reward models are rejected at config validation — no collective gather). + + Args: + save_directory (Path): Target directory for the reward model config (`config.json`) + and `model.safetensors`. + """ + from lerobot.distributed.utils import is_main_process + + # save_checkpoint calls this on every rank; replicas carry identical + # weights, so the main process is the only writer. Sharded reward models are rejected + # at config validation, so no collective gather is needed here. + if not is_main_process(): + return 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)) @@ -175,53 +190,22 @@ class PreTrainedRewardModel(nn.Module, HubMixin, abc.ABC): """ return type(self).forward is not PreTrainedRewardModel.forward - def push_model_to_hub(self, cfg: "TrainPipelineConfig"): - api = HfApi() - repo_id = api.create_repo( - repo_id=self.config.repo_id, private=self.config.private, exist_ok=True - ).repo_id + def push_model_to_hub(self, cfg: "TrainPipelineConfig") -> None: + """Publish this reward model to the Hub. - # Push the files to the repo in a single commit - with TemporaryDirectory(ignore_cleanup_errors=True) as tmp: - saved_path = Path(tmp) / repo_id + Deprecated: use :func:`lerobot.common.train_utils.publish_trained_model` instead. - self.save_pretrained(saved_path) # Calls _save_pretrained and stores model tensors + Args: + cfg (TrainPipelineConfig): The training config; saved as `train_config.json` and + used to render the model card. + """ + from lerobot.common.train_utils import publish_trained_model - card = self.generate_model_card( - cfg.dataset.repo_id, self.config.type, self.config.license, self.config.tags - ) - card.save(str(saved_path / "README.md")) - - cfg.save_pretrained(saved_path) # Calls _save_pretrained and stores train config - - commit_info = api.upload_folder( - repo_id=repo_id, - repo_type="model", - folder_path=saved_path, - commit_message="Upload reward model weights, train config and readme", - allow_patterns=["*.safetensors", "*.json", "*.yaml", "*.md"], - ignore_patterns=["*.tmp", "*.log"], - ) - - logging.info(f"Model pushed to {commit_info.repo_url.url}") - - def generate_model_card( - self, dataset_repo_id: str, model_type: str, license: str | None, tags: list[str] | None - ) -> ModelCard: - card_data = ModelCardData( - license=license or "apache-2.0", - library_name="lerobot", - pipeline_tag="robotics", - tags=list(set(tags or []).union({"robotics", "lerobot", "reward-model", model_type})), - model_name=model_type, - datasets=dataset_repo_id, + warnings.warn( + "PreTrainedRewardModel.push_model_to_hub is deprecated and will be removed in a " + "future version. Use lerobot.common.train_utils.publish_trained_model(cfg, model, " + "preprocessor, postprocessor, dataset_meta) instead.", + FutureWarning, + stacklevel=2, ) - - template_card = ( - files("lerobot.templates") - .joinpath("lerobot_rewardmodel_modelcard_template.md") - .read_text(encoding="utf-8") - ) - card = ModelCard.from_template(card_data, template_str=template_card) - card.validate() - return card + publish_trained_model(cfg, self, None, None, None) diff --git a/src/lerobot/rewards/topreward/modeling_topreward.py b/src/lerobot/rewards/topreward/modeling_topreward.py index 4958d5449..5fee8f378 100644 --- a/src/lerobot/rewards/topreward/modeling_topreward.py +++ b/src/lerobot/rewards/topreward/modeling_topreward.py @@ -58,12 +58,11 @@ import builtins import logging import os from pathlib import Path -from tempfile import TemporaryDirectory from typing import TYPE_CHECKING, Any, TypeVar import numpy as np import torch -from huggingface_hub import HfApi, hf_hub_download +from huggingface_hub import hf_hub_download from huggingface_hub.constants import CONFIG_NAME from huggingface_hub.errors import HfHubHTTPError from torch import Tensor @@ -75,9 +74,6 @@ from lerobot.rewards.topreward.configuration_topreward import TOPRewardConfig from lerobot.rewards.topreward.processor_topreward import TOPREWARD_FEATURE_PREFIX, TOPREWARD_INPUT_KEYS from lerobot.utils.import_utils import _transformers_available, require_package -if TYPE_CHECKING: - from lerobot.configs.train import TrainPipelineConfig - if TYPE_CHECKING or _transformers_available: from transformers import Qwen3VLForConditionalGeneration else: @@ -205,34 +201,3 @@ class TOPRewardModel(PreTrainedRewardModel): instance.to(config.device) instance.eval() return instance - - def push_model_to_hub(self, cfg: TrainPipelineConfig): - """Push the TOPReward ``config.json`` + model card to the Hub.""" - api = HfApi() - repo_id = api.create_repo( - repo_id=self.config.repo_id, private=self.config.private, exist_ok=True - ).repo_id - - with TemporaryDirectory(ignore_cleanup_errors=True) as tmp: - saved_path = Path(tmp) / repo_id - saved_path.mkdir(parents=True, exist_ok=True) - - self.config._save_pretrained(saved_path) - - card = self.generate_model_card( - cfg.dataset.repo_id, self.config.type, self.config.license, self.config.tags - ) - card.save(str(saved_path / "README.md")) - - cfg.save_pretrained(saved_path) - - commit_info = api.upload_folder( - repo_id=repo_id, - repo_type="model", - folder_path=saved_path, - commit_message="Upload TOPReward config and readme", - allow_patterns=["*.json", "*.yaml", "*.md"], - ignore_patterns=["*.tmp", "*.log", "*.safetensors"], - ) - - logger.info(f"Model pushed to {commit_info.repo_url.url}") diff --git a/src/lerobot/rl/learner.py b/src/lerobot/rl/learner.py index ee95d831e..f52b720ef 100644 --- a/src/lerobot/rl/learner.py +++ b/src/lerobot/rl/learner.py @@ -74,13 +74,14 @@ from torch.optim.optimizer import Optimizer from lerobot.cameras import opencv # noqa: F401 from lerobot.common.train_utils import ( get_step_checkpoint_dir, - load_training_state as utils_load_training_state, + load_training_metadata, save_checkpoint, update_last_checkpoint, ) from lerobot.common.wandb_utils import WandBLogger from lerobot.configs import parser from lerobot.datasets import LeRobotDataset, make_dataset +from lerobot.optim import load_optimizer_state from lerobot.policies import make_policy, make_pre_post_processors from lerobot.robots import so_follower # noqa: F401 from lerobot.teleoperators import gamepad, so_leader # noqa: F401 @@ -103,7 +104,7 @@ from lerobot.utils.constants import ( from lerobot.utils.device_utils import get_safe_torch_device from lerobot.utils.io_utils import load_json, write_json from lerobot.utils.process import ProcessSignalHandler, ensure_multiprocessing_start_method -from lerobot.utils.random_utils import set_seed +from lerobot.utils.random_utils import load_rng_state, set_seed from lerobot.utils.utils import ( format_big_number, init_logging, @@ -716,15 +717,18 @@ def load_training_state( algorithm-owned tensors) from the most recent checkpoint. Args: - cfg: Training configuration. - optimizers: Optimizers to load state into. - algorithm: Algorithm whose state dict should be restored. - Required for full main-equivalent resume; - the policy itself is restored separately via ``make_policy``. - device: Device on which to place loaded algorithm tensors. + cfg (TrainRLServerPipelineConfig): Training configuration; `cfg.resume` gates the load and + `cfg.output_dir` locates the last checkpoint. + optimizers (Optimizer | dict[str, Optimizer]): Optimizers to load state into. + algorithm (RLAlgorithm | None, optional): Algorithm whose state dict should be restored. + Required for full main-equivalent resume; the policy itself is restored separately via + `make_policy`. Defaults to None. + device (str | torch.device, optional): Device on which to place loaded algorithm tensors. + Defaults to "cpu". Returns: - tuple: (optimization_step, interaction_step) or (None, None) if not resuming + tuple[int | None, int | None]: `(optimization_step, interaction_step)`, or `(None, None)` + when not resuming or when loading the training state fails. """ if not cfg.resume: return None, None @@ -736,7 +740,10 @@ def load_training_state( try: # Restore optimizers + RNG + step from the standard `training_state/` folder - step, optimizers, _ = utils_load_training_state(checkpoint_dir, optimizers, None) + training_state_dir = checkpoint_dir / TRAINING_STATE_DIR + load_rng_state(training_state_dir) + step = load_training_metadata(training_state_dir)["step"] + optimizers = load_optimizer_state(optimizers, training_state_dir) # Restore algorithm-owned tensors if algorithm is not None: diff --git a/src/lerobot/scripts/lerobot_convert_dcp.py b/src/lerobot/scripts/lerobot_convert_dcp.py new file mode 100644 index 000000000..4ca99ceb4 --- /dev/null +++ b/src/lerobot/scripts/lerobot_convert_dcp.py @@ -0,0 +1,174 @@ +#!/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. +"""Convert a DCP-format checkpoint into a distributable safetensors model, offline. + +Runs single-process (no GPUs, no process group). Example: + +```bash +lerobot-convert-dcp --checkpoint_dir=outputs/train/run/checkpoints/005000 +lerobot-convert-dcp --checkpoint_dir=... --delete_dcp=true --push_to_hub=user/my-policy +``` + +`--push_to_hub` publishes the converted directory as a model repo, degrading gracefully: the +core artifacts (model.safetensors, config.json, processor files) always upload; the README +model card is enriched with training/dataset metadata only when `train_config.json` (and the +dataset it names) are reachable, with a WARNING naming exactly what was skipped otherwise. +DCP shard artifacts are never uploaded — published repos carry safetensors only. +""" + +import logging +from dataclasses import dataclass +from pathlib import Path + +from huggingface_hub import HfApi + +from lerobot.configs import parser +from lerobot.distributed.checkpoint import dcp_to_safetensors +from lerobot.utils.constants import PRETRAINED_MODEL_DIR +from lerobot.utils.utils import init_logging + + +@dataclass +class ConvertDcpConfig: + """CLI config for the offline DCP-to-safetensors checkpoint conversion.""" + + # A checkpoint step directory (containing pretrained_model/) or a pretrained_model + # directory itself. + checkpoint_dir: Path + # Remove the DCP shard directory after a successful conversion. + delete_dcp: bool = False + # Publish the converted directory to this Hub repo id (e.g. "user/my-policy"). + push_to_hub: str | None = None + private: bool | None = None + + +def _locate_pretrained_dir(checkpoint_dir: Path) -> Path: + """Resolve the pretrained_model/ directory from a user-supplied checkpoint path. + + Args: + checkpoint_dir (Path): A checkpoint step directory (containing `pretrained_model/`) or a + `pretrained_model` directory itself. + + Returns: + Path: The nested `pretrained_model/` directory when present, otherwise `checkpoint_dir` + unchanged. + """ + nested = checkpoint_dir / PRETRAINED_MODEL_DIR + return nested if nested.is_dir() else checkpoint_dir + + +def _publish_converted(pretrained_dir: Path, repo_id: str, private: bool | None) -> None: + """Best-effort publish of a converted checkpoint dir, degrading gracefully. + + The core artifacts (model.safetensors, config.json, processor files) always upload; the README + model card gains training/dataset metadata only when `train_config.json` (and the dataset it + names) are reachable, with a WARNING naming what was skipped otherwise. DCP shard artifacts are + excluded from the upload. + + Args: + pretrained_dir (Path): The converted `pretrained_model/` directory to upload. + repo_id (str): Target Hub model repo id (e.g. "user/my-policy"); created if missing. + private (bool | None): Repo visibility passed to `create_repo`; None keeps the Hub (or + existing repo's) default. + """ + from lerobot.common.train_utils import generate_model_card + from lerobot.configs.policies import PreTrainedConfig + from lerobot.configs.train import TRAIN_CONFIG_NAME, TrainPipelineConfig + + train_cfg = None + dataset_meta = None + if (pretrained_dir / TRAIN_CONFIG_NAME).is_file(): + try: + train_cfg = TrainPipelineConfig.from_pretrained(pretrained_dir) + except Exception as e: # noqa: BLE001 — degrade, never block the upload + logging.warning(f"Could not parse {TRAIN_CONFIG_NAME} ({e}); README will lack training metadata.") + else: + logging.warning(f"{TRAIN_CONFIG_NAME} missing; README will lack training metadata.") + if train_cfg is not None: + try: + from lerobot.datasets.dataset_metadata import LeRobotDatasetMetadata + + dataset_meta = LeRobotDatasetMetadata( + repo_id=train_cfg.dataset.repo_id, + root=train_cfg.dataset.root, + revision=train_cfg.dataset.revision, + ) + except Exception as e: # noqa: BLE001 + logging.warning( + f"Dataset '{train_cfg.dataset.repo_id}' unreachable ({e}); README will lack dataset metadata." + ) + try: + model_cfg = PreTrainedConfig.from_pretrained(pretrained_dir) + card = generate_model_card(model_cfg, cfg=train_cfg, dataset_meta=dataset_meta) + card.save(str(pretrained_dir / "README.md")) + except Exception as e: # noqa: BLE001 + logging.warning(f"Could not build the model card ({e}); publishing without README.") + + api = HfApi() + repo_id = api.create_repo(repo_id=repo_id, private=private, exist_ok=True).repo_id + commit_info = api.upload_folder( + repo_id=repo_id, + repo_type="model", + folder_path=str(pretrained_dir), + commit_message="Upload converted policy (DCP -> safetensors)", + allow_patterns=["*.safetensors", "*.json", "*.yaml", "*.md"], + # The checkpoint keeps its DCP shard directory unless --delete_dcp was passed; the + # allow list above admits neither `.distcp` shards nor their `.metadata` sidecar. + ignore_patterns=["*.tmp", "*.log"], + ) + logging.info(f"Model pushed to {commit_info.repo_url.url}") + + +@parser.wrap() +def convert_checkpoint(cfg: ConvertDcpConfig) -> Path: + """Merge a checkpoint's DCP shards into `model.safetensors`, then optionally publish it. + + Args: + cfg (ConvertDcpConfig): Conversion options — the checkpoint directory to convert, whether + to delete the DCP shards after a successful merge, and the optional Hub repo id (and + visibility) to publish the converted directory to. + + Returns: + Path: The path to the merged `model.safetensors` file. + + Raises: + FileNotFoundError: If the checkpoint has no DCP shard directory, i.e. it was not saved + with `checkpoint_format=dcp` (or `safetensors_dcp`). + """ + from accelerate.utils.constants import FSDP_MODEL_NAME + + pretrained_dir = _locate_pretrained_dir(cfg.checkpoint_dir) + dcp_dir = pretrained_dir / f"{FSDP_MODEL_NAME}_0" + if not dcp_dir.is_dir(): + raise FileNotFoundError( + f"No DCP shard directory at {dcp_dir}. Point --checkpoint_dir at a checkpoint " + "saved with checkpoint_format=dcp (or safetensors_dcp)." + ) + logging.info(f"Merging {dcp_dir} -> {pretrained_dir / 'model.safetensors'}") + safetensors_path = dcp_to_safetensors(dcp_dir, pretrained_dir, delete_dcp=cfg.delete_dcp) + if cfg.push_to_hub: + _publish_converted(pretrained_dir, cfg.push_to_hub, cfg.private) + return safetensors_path + + +def main() -> None: + """`lerobot-convert-dcp` console entry point: set up logging and run the conversion.""" + init_logging() + convert_checkpoint() + + +if __name__ == "__main__": + main() diff --git a/src/lerobot/scripts/lerobot_train.py b/src/lerobot/scripts/lerobot_train.py index 8a3222df1..69aba44b1 100644 --- a/src/lerobot/scripts/lerobot_train.py +++ b/src/lerobot/scripts/lerobot_train.py @@ -16,6 +16,16 @@ """Train a policy. Requires: pip install 'lerobot[training]' (includes dataset + accelerate + wandb extras) + +Launch with torchrun for distributed runs; every parallelism/acceleration knob lives on the +config (`--parallelism.*`, `--accelerator.*`) so a run is reproducible from its +train_config.json alone: + +```bash +torchrun --nproc-per-node=8 $(which lerobot-train) \ + --dataset.repo_id=... --policy.type=act \ + --parallelism.dp_shard=8 --accelerator.mixed_precision=bf16 +``` """ import dataclasses @@ -36,14 +46,13 @@ 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, + load_training_metadata, + publish_trained_model, push_checkpoint_to_hub, + resume_after_prepare, + resume_before_prepare, save_checkpoint, should_save_checkpoint, update_last_checkpoint, @@ -53,12 +62,21 @@ from lerobot.configs import JobConfig, parser from lerobot.configs.train import TrainPipelineConfig from lerobot.datasets import EpisodeAwareSampler, compute_sampler_state from lerobot.datasets.factory import make_train_eval_datasets +from lerobot.distributed import ( + ParallelDims, + finalize_sharded_policy, + is_main_process, + make_accelerator, + set_fsdp_wrap_modules, +) from lerobot.envs import close_envs, make_env, make_env_pre_post_processors from lerobot.jobs import submit_to_hf from lerobot.optim.factory import make_optimizer_and_scheduler from lerobot.policies import PreTrainedPolicy, make_policy, make_pre_post_processors +from lerobot.policies.factory import ProcessorConfigKwargs from lerobot.rewards import make_reward_pre_post_processors from lerobot.utils.collate import lerobot_collate_fn +from lerobot.utils.constants import TRAINING_STATE_DIR from lerobot.utils.import_utils import _peft_available, register_third_party_plugins, require_package from lerobot.utils.logging_utils import AverageMeter, MetricsTracker from lerobot.utils.random_utils import set_seed @@ -92,16 +110,6 @@ def _make_eval_envs(cfg: TrainPipelineConfig) -> Iterator[dict[str, dict[int, An close_envs(envs) -def _dataloader_worker_kwargs(cfg: TrainPipelineConfig) -> dict[str, Any]: - """Return worker-only DataLoader options, disabling them for single-process loading.""" - workers_enabled = cfg.num_workers > 0 - return { - "prefetch_factor": cfg.prefetch_factor if workers_enabled else None, - "persistent_workers": cfg.persistent_workers and workers_enabled, - "multiprocessing_context": cfg.dataloader_multiprocessing_context if workers_enabled else None, - } - - def update_policy( train_metrics: MetricsTracker, policy: PreTrainedPolicy, @@ -117,23 +125,26 @@ def update_policy( Performs a single training step to update the policy's weights. This function executes the forward and backward passes, clips gradients, and steps the optimizer and - learning rate scheduler. Accelerator handles mixed-precision training automatically. + learning rate scheduler. Accelerator handles mixed-precision training automatically, and — under + gradient accumulation — suppresses gradient sync on non-final micro-batches and rescales the loss. Args: - train_metrics: A MetricsTracker instance to record training statistics. - policy: The policy model to be trained. - batch: A batch of training data. - optimizer: The optimizer used to update the policy's parameters. - grad_clip_norm: The maximum norm for gradient clipping. - accelerator: The Accelerator instance for distributed training and mixed precision. - lr_scheduler: An optional learning rate scheduler. - lock: An optional lock for thread-safe optimizer updates. - sample_weighter: Optional SampleWeighter instance for per-sample loss weighting. + train_metrics (MetricsTracker): A MetricsTracker instance to record training statistics. + policy (PreTrainedPolicy): The policy model to be trained (as returned by `accelerator.prepare`). + batch (Any): A batch of training data. + optimizer (Optimizer): The optimizer used to update the policy's parameters. + grad_clip_norm (float): The maximum norm for gradient clipping (no clipping when <= 0). + accelerator (Accelerator): The Accelerator instance for distributed training and mixed precision. + lr_scheduler (LRScheduler | None, optional): An optional learning rate scheduler, stepped once + per micro-batch. Defaults to None. + lock (Lock | None, optional): An optional lock for thread-safe optimizer updates. + Defaults to None. + sample_weighter (SampleWeighter | None, optional): Optional SampleWeighter instance for + per-sample loss weighting. Defaults to None. Returns: - A tuple containing: - - The updated MetricsTracker with new statistics for this step. - - A dictionary of outputs from the policy's forward pass, for logging purposes. + tuple[MetricsTracker, dict | None]: The updated MetricsTracker with new statistics for this + step, and the dictionary of outputs from the policy's forward pass, for logging purposes. """ start_time = time.perf_counter() policy.train() @@ -147,57 +158,65 @@ def update_policy( if sample_weighter is not None: sample_weights, weight_stats = sample_weighter.compute_batch_weights(batch) - # Let accelerator handle mixed precision - with accelerator.autocast(): - if sample_weights is not None: - # Use per-sample loss for weighted training - # Note: Policies supporting sample weighting must implement forward(batch, reduction="none") - per_sample_loss, output_dict = policy.forward(batch, reduction="none") + # Under gradient accumulation this context suppresses gradient sync (FSDP2: + # set_requires_gradient_sync) on non-final micro-batches and divides the loss; + # with gradient_accumulation_steps == 1 it is a transparent no-op. + with accelerator.accumulate(policy): + # Let accelerator handle mixed precision + with accelerator.autocast(): + # `policy(...)`, never `policy.forward(...)`: FSDP2 all-gathers parameters through + # nn.Module forward hooks, which only run via __call__. + if sample_weights is not None: + # Use per-sample loss for weighted training + # Note: Policies supporting sample weighting must implement forward(batch, reduction="none") + per_sample_loss, output_dict = policy(batch, reduction="none") - # Weighted loss: each sample's contribution is scaled by its weight. - # We divide by weight sum (not batch size) so that if some weights are zero, - # the remaining samples contribute proportionally more, preserving gradient scale. - # Weights are pre-normalized to sum to batch_size for stable training dynamics. - epsilon = 1e-6 - loss = (per_sample_loss * sample_weights).sum() / (sample_weights.sum() + epsilon) + # Weighted loss: each sample's contribution is scaled by its weight. + # We divide by weight sum (not batch size) so that if some weights are zero, + # the remaining samples contribute proportionally more, preserving gradient scale. + # Weights are pre-normalized to sum to batch_size for stable training dynamics. + epsilon = 1e-6 + loss = (per_sample_loss * sample_weights).sum() / (sample_weights.sum() + epsilon) - # Log weighting statistics - if output_dict is None: - output_dict = {} - for key, value in weight_stats.items(): - output_dict[f"sample_weight_{key}"] = value - else: - loss, output_dict = policy.forward(batch) + # Log weighting statistics + if output_dict is None: + output_dict = {} + for key, value in weight_stats.items(): + output_dict[f"sample_weight_{key}"] = value + else: + loss, output_dict = policy(batch) - # TODO(rcadene): policy.unnormalize_outputs(out_dict) + # TODO(rcadene): policy.unnormalize_outputs(out_dict) - # Use accelerator's backward method - accelerator.backward(loss) + # Use accelerator's backward method + accelerator.backward(loss) - # Clip gradients if specified - if grad_clip_norm > 0: - grad_norm = accelerator.clip_grad_norm_(policy.parameters(), grad_clip_norm) - else: - grad_norm = torch.nn.utils.clip_grad_norm_( - policy.parameters(), float("inf"), error_if_nonfinite=False - ) + # Gradients are complete only on sync micro-batches; clipping partial gradients would + # be meaningless. Always pass the full parameter list: accelerate's FSDP2 path requires + # an exact match with the prepared model's parameters for a globally correct norm. + grad_norm = None + if accelerator.sync_gradients and grad_clip_norm > 0: + grad_norm = accelerator.clip_grad_norm_(policy.parameters(), grad_clip_norm) - # Optimizer step - with lock if lock is not None else nullcontext(): - optimizer.step() + # Optimizer step (a no-op on non-final micro-batches under gradient accumulation) + with lock if lock is not None else nullcontext(): + optimizer.step() + optimizer.zero_grad() - optimizer.zero_grad() + # Step through pytorch scheduler at every batch instead of epoch + if lr_scheduler is not None: + lr_scheduler.step() - # Step through pytorch scheduler at every batch instead of epoch - if lr_scheduler is not None: - lr_scheduler.step() - - # Update internal buffers if policy has update method - if has_method(accelerator.unwrap_model(policy, keep_fp32_wrapper=True), "update"): + # Update internal buffers if policy has update method. These track optimizer updates + # (EMA, target networks), not micro-batches: gate on the sync step under accumulation. + if accelerator.sync_gradients and has_method( + accelerator.unwrap_model(policy, keep_fp32_wrapper=True), "update" + ): accelerator.unwrap_model(policy, keep_fp32_wrapper=True).update() train_metrics.loss = loss.item() - train_metrics.grad_norm = grad_norm.item() + if grad_norm is not None: + train_metrics.grad_norm = grad_norm.item() train_metrics.lr = optimizer.param_groups[0]["lr"] train_metrics.update_s = time.perf_counter() - start_time if torch.cuda.is_available(): @@ -208,81 +227,180 @@ def update_policy( return train_metrics, output_dict +def make_dataloaders( + cfg: TrainPipelineConfig, + dataset, + eval_dataset, + step: int, + parallel_dims: ParallelDims, +) -> tuple[torch.utils.data.DataLoader, torch.utils.data.DataLoader | None]: + """Build the train (and optional eval) dataloader, including the sampler resume offset. + + The sampler offset is *derived* from `step` (`resume_before_prepare` loads step + RNG only): + each loop step consumes `batch_size` samples on each of the `dp_world_size` distinct + data-parallel workers — no grad-accumulation factor, since `step` counts micro-batches. + + Args: + cfg (TrainPipelineConfig): The training config (batch size, workers, streaming, resume, seed). + dataset (LeRobotDataset | MultiLeRobotDataset): The training dataset. + eval_dataset (LeRobotDataset | None): Optional held-out split; when provided, an eval + dataloader is built (subsampled per task when `cfg.max_eval_samples > 0`). + step (int): The loop step to resume the sampler from (0 for a fresh run). + parallel_dims (ParallelDims): The resolved parallelism topology; provides the device type + and the fallback dp world size for the resume offset. + + Returns: + tuple[torch.utils.data.DataLoader, torch.utils.data.DataLoader | None]: The train + dataloader and the eval dataloader (None when no eval split exists). + """ + active_cfg = cfg.trainable_config + if not cfg.dataset.streaming: + # All non-streaming (map-style) datasets use EpisodeAwareSampler. + # The order is a pure function of (seed, epoch), so every rank independently produces the + # same permutation. accelerate then shards it disjointly across data-parallel ranks via + # BatchSamplerShard without needing a `generator` attribute to synchronize an RNG, and + # resume is sample-exact. + shuffle = False + sampler = EpisodeAwareSampler( + dataset.meta.episodes["dataset_from_index"], + dataset.meta.episodes["dataset_to_index"], + episode_indices_to_use=dataset.episodes, + drop_n_last_frames=getattr(active_cfg, "drop_n_last_frames", 0), + shuffle=True, + seed=cfg.seed if cfg.seed is not None else 0, + absolute_to_relative_idx=dataset.absolute_to_relative_idx, + ) + if cfg.resume and step > 0: + # The resume offset depends on the (dp_world_size, batch_size) that produced `step`, + # so use the values recorded in the checkpoint (falling back to the current ones for + # older checkpoints that did not store them). + metadata = load_training_metadata(cfg.checkpoint_path / TRAINING_STATE_DIR) + saved_dp_world = metadata["dp_world_size"] + saved_batch_size = metadata["batch_size"] + ckpt_dp_world = saved_dp_world or parallel_dims.dp_world_size + ckpt_batch_size = saved_batch_size or cfg.batch_size + if is_main_process() and saved_dp_world not in (None, parallel_dims.dp_world_size): + logging.warning( + f"Resuming with dp_world_size={parallel_dims.dp_world_size} but the " + f"checkpoint was written with dp_world_size={saved_dp_world}. The data order " + "resumes at the right epoch/offset, but per-rank sample-exactness requires " + "the same data-parallel world size." + ) + if is_main_process() and saved_batch_size not in (None, cfg.batch_size): + logging.warning( + f"Resuming with batch_size={cfg.batch_size} but the checkpoint was written " + f"with batch_size={saved_batch_size}. The data order resumes at the right " + "epoch/offset, but per-rank sample-exactness requires the same batch size." + ) + sampler_state = compute_sampler_state(step, len(sampler), ckpt_batch_size, ckpt_dp_world) + sampler.load_state_dict(sampler_state) + if is_main_process(): + logging.info( + f"Resuming data order at epoch {sampler_state['epoch']}, " + f"sample {sampler_state['start_index']}" + ) + else: + shuffle = True + sampler = None + + device_type = parallel_dims.device_type + # Only swap in the language-aware collate when the dataset actually + # declares language columns; otherwise stay on PyTorch's default + # collate so non-language training runs are unaffected. + collate_fn = lerobot_collate_fn if dataset.meta.has_language_columns else None + dataloader = torch.utils.data.DataLoader( + dataset, + num_workers=cfg.num_workers, + batch_size=cfg.batch_size, + shuffle=shuffle and not cfg.dataset.streaming, + sampler=sampler, + pin_memory=device_type == "cuda", + drop_last=False, + collate_fn=collate_fn, + prefetch_factor=cfg.prefetch_factor if cfg.num_workers > 0 else None, + persistent_workers=cfg.persistent_workers and cfg.num_workers > 0, + multiprocessing_context=cfg.dataloader_multiprocessing_context if cfg.num_workers > 0 else None, + ) + + # Build eval dataloader if a held-out split exists + eval_dataloader = None + if eval_dataset is not None: + eval_ds = eval_dataset + if cfg.max_eval_samples > 0 and hasattr(eval_dataset, "hf_dataset"): + task_arr = eval_dataset.hf_dataset.data.column("task_index").to_numpy() + unique_tasks = sorted(set(task_arr.tolist())) + per_task = max(1, cfg.max_eval_samples // len(unique_tasks)) + selected: list[int] = [] + for t in unique_tasks: + frames = (task_arr == t).nonzero()[0][:per_task] + selected.extend(frames.tolist()) + eval_ds = torch.utils.data.Subset(eval_dataset, selected) + + eval_collate_fn = lerobot_collate_fn if dataset.meta.has_language_columns else None + eval_dataloader = torch.utils.data.DataLoader( + eval_ds, + batch_size=cfg.batch_size, + shuffle=False, + num_workers=cfg.num_workers, + pin_memory=device_type == "cuda", + drop_last=False, + collate_fn=eval_collate_fn, + prefetch_factor=cfg.prefetch_factor if cfg.num_workers > 0 else None, + persistent_workers=cfg.persistent_workers and cfg.num_workers > 0, + multiprocessing_context=cfg.dataloader_multiprocessing_context if cfg.num_workers > 0 else None, + ) + return dataloader, eval_dataloader + + @parser.wrap() -def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None): +def train(cfg: TrainPipelineConfig): """ Main function to train a policy. This function orchestrates the entire training pipeline, including: - - Setting up logging, seeding, and device configuration. + - Setting up logging, seeding, and the distributed engine. - Creating the dataset, evaluation environment (if applicable), policy, and optimizer. - - Handling resumption from a checkpoint. + - Handling resumption from a checkpoint (two-phase, around `accelerator.prepare`). - Running the main training loop, which involves fetching data batches and calling `update_policy`. - Periodically logging metrics, saving model checkpoints, and evaluating the policy. - - Pushing the final trained model to the Hugging Face Hub if configured. + - Publishing the trained model to the Hugging Face Hub if configured. Args: - cfg: A `TrainPipelineConfig` object containing all training configurations. - accelerator: Optional Accelerator instance. If None, one will be created automatically. + cfg (TrainPipelineConfig): A `TrainPipelineConfig` object containing all training + configurations, parsed from the CLI by `parser.wrap()`. On `--resume`, it is the config + recorded in the checkpoint's `train_config.json`; when `cfg.job.is_remote`, the run is + dispatched to HF Jobs instead of executing locally. """ if cfg.job.is_remote: return submit_to_hf(cfg) require_package("accelerate", extra="training") - from accelerate import Accelerator - from accelerate.utils import DistributedDataParallelKwargs, DistributedType - cfg.validate() - - # Create Accelerator if not provided - # It will automatically detect if running in distributed mode or single-process mode - # 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: - 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). - force_cpu = cfg.trainable_config.device == "cpu" - # Drive Accelerate's autocast from policy.dtype (bf16/fp16 activate it; float32 -> full precision). - has_policy_dtype = hasattr(cfg.trainable_config, "dtype") - policy_dtype = getattr(cfg.trainable_config, "dtype", None) - mixed_precision = {"bfloat16": "bf16", "float16": "fp16", "float32": "no"}.get(policy_dtype) - # Policies without a `dtype` field fall back to `use_amp`, which would otherwise be - # silently ignored here while lerobot-eval honors it. Follow torch.autocast's default - # for the configured device so training and evaluation use the same precision. - if not has_policy_dtype and getattr(cfg.trainable_config, "use_amp", False): - device_type = torch.device(cfg.trainable_config.device).type - autocast_dtype = torch.get_autocast_dtype(device_type) - mixed_precision = {torch.bfloat16: "bf16", torch.float16: "fp16"}[autocast_dtype] - accelerator = Accelerator( - step_scheduler_with_optimizer=False, - mixed_precision=mixed_precision, - kwargs_handlers=[ddp_kwargs], - cpu=force_cpu, - ) + cfg.validate() # all fail-fasts fire here, before any distributed init + # --- engine & topology -------------------------------------------------------------------- + # The factory is the ONLY accelerate configuration site: it guards against env-var + # interference, resolves the declared parallelism degrees against the launched world, and + # builds the Accelerator from the config mirrors. + accelerator = make_accelerator(cfg) + parallel_dims = ParallelDims.from_config( + cfg.parallelism, accelerator.num_processes, accelerator.device.type + ) init_logging(accelerator=accelerator) - # Determine if this is the main process (for logging and checkpointing) - # When using accelerate, only the main process should log to avoid duplicate outputs - is_main_process = accelerator.is_main_process - - # Only log on main process - if is_main_process: + if is_main_process(): logging.info(pformat(cfg.to_dict())) - # Initialize wandb only on main process - if cfg.wandb.enable and cfg.wandb.project and is_main_process: + if cfg.wandb.enable and cfg.wandb.project and is_main_process(): wandb_logger = WandBLogger(cfg) else: wandb_logger = None - if is_main_process: + if is_main_process(): logging.info(colored("Logs will be saved locally.", "yellow", attrs=["bold"])) if cfg.seed is not None: set_seed(cfg.seed, accelerator=accelerator) - # Use accelerator's device device = accelerator.device if cfg.cudnn_deterministic: torch.backends.cudnn.deterministic = True @@ -291,21 +409,21 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None): torch.backends.cudnn.benchmark = True torch.backends.cuda.matmul.allow_tf32 = True - # Dataset loading synchronization: the global main process downloads once to the shared - # dataset root, then a barrier lets every other rank read the already-populated copy. - # LeRobotDataset skips its snapshot_download when try_load() succeeds, so no rank re-downloads. - if is_main_process: + # --- data (the main process downloads once; peers read the populated cache) ---------------- + if is_main_process(): logging.info("Creating dataset") dataset, eval_dataset = make_train_eval_datasets(cfg) - accelerator.wait_for_everyone() - - # Other ranks read from the shared copy populated by the main process. - if not is_main_process: + if not is_main_process(): dataset, eval_dataset = make_train_eval_datasets(cfg) + # --- policy (weight source decided by the resume rule) ------------------------------------- + # On resume, cfg was parsed FROM the checkpoint's train_config.json, so cfg.checkpoint_format + # IS the recorded value: DCP-bearing formats skip the safetensors load here and stream the + # sharded weights in after prepare (resume_after_prepare). + defer_weight_load = cfg.resume and cfg.checkpoint_format.wants_dcp if cfg.is_reward_model_training: - if is_main_process: + if is_main_process(): logging.info("Creating reward model") from lerobot.rewards import make_reward_model @@ -320,14 +438,16 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None): "Use it directly for inference via compute_reward() (e.g. offline precompute)." ) else: - if is_main_process: + if is_main_process(): logging.info("Creating policy") policy = make_policy( cfg=cfg.policy, ds_meta=dataset.meta, rename_map=cfg.rename_map, + defer_weight_load=defer_weight_load, ) + peft_model = None if cfg.peft is not None: if cfg.is_reward_model_training: raise ValueError("PEFT is only supported for policy training. ") @@ -339,20 +459,19 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None): logging.info("Using PEFT! Wrapping model.") peft_cli_overrides = dataclasses.asdict(cfg.peft) policy = policy.wrap_with_peft(peft_cli_overrides=peft_cli_overrides) + peft_model = policy - # Wait for all processes to finish model creation before continuing accelerator.wait_for_everyone() + # --- processors (overrides built once, as one typed mapping) ------------------------------- active_cfg = cfg.trainable_config processor_pretrained_path = active_cfg.pretrained_path - processor_kwargs = {} + processor_kwargs = ProcessorConfigKwargs() if (processor_pretrained_path and not cfg.resume) or not processor_pretrained_path: processor_kwargs["dataset_stats"] = dataset.meta.stats - if cfg.is_reward_model_training: processor_kwargs["dataset_meta"] = dataset.meta - if not cfg.is_reward_model_training and processor_pretrained_path is not None: preprocessor_overrides = { "device_processor": {"device": device.type}, @@ -398,16 +517,42 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None): **processor_kwargs, ) - if is_main_process: + # Created BEFORE prepare on the unsharded parameters — accelerate's FSDP2 path requires the + # model and optimizer in one prepare() call and rebinds the param groups itself. + if is_main_process(): logging.info("Creating optimizer and scheduler") optimizer, lr_scheduler = make_optimizer_and_scheduler(cfg, policy) - # Create sample weighter if configured (e.g., for RA-BC training) + # --- resume phase 1 + dataloaders ---------------------------------------------------------- + step = 0 # number of loop steps (= micro-batches consumed per data-parallel worker) + if cfg.resume: + step = resume_before_prepare(cfg) # step + RNG only; sharded state loads after prepare + + dataloader, eval_dataloader = make_dataloaders(cfg, dataset, eval_dataset, step, parallel_dims) + + # --- prepare & resume phase 2 --------------------------------------------------------------- + # The FSDP wrap-unit class names resolve right before prepare: user override, else the + # policy's _fsdp_wrap_modules declaration — root-only wrapping is never silently accepted. + set_fsdp_wrap_modules(accelerator, accelerator.unwrap_model(policy) if peft_model else policy) + accelerator.wait_for_everyone() + if eval_dataloader is not None: + policy, optimizer, dataloader, lr_scheduler, eval_dataloader = accelerator.prepare( + policy, optimizer, dataloader, lr_scheduler, eval_dataloader + ) + else: + policy, optimizer, dataloader, lr_scheduler = accelerator.prepare( + policy, optimizer, dataloader, lr_scheduler + ) + finalize_sharded_policy(policy, parallel_dims) + if cfg.resume: + resume_after_prepare(cfg, accelerator, policy, optimizer, lr_scheduler) + + # --- auxiliaries (after the core assembly, per the construction-order contract) ------------- sample_weighter = None if cfg.sample_weighting is not None: from lerobot.utils.sample_weighting import make_sample_weighter - if is_main_process: + if is_main_process(): logging.info(f"Creating sample weighter: {cfg.sample_weighting.type}") sample_weighter = make_sample_weighter( cfg.sample_weighting, @@ -417,20 +562,15 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None): dataset_repo_id=cfg.dataset.repo_id, ) - step = 0 # number of policy updates (forward + backward + optim) - - if cfg.resume: - # 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()) - - if is_main_process: + # --- banner (main process only; numel() reads metadata — on DTensors it is the GLOBAL shape, + # so the totals are correct even after sharding) --------------------------------------------- + # One loop step consumes one micro-batch on every dp worker; the optimizer sees + # `samples_per_step x gradient_accumulation_steps` samples per update. + samples_per_step = cfg.batch_size * parallel_dims.dp_world_size + effective_batch_size = samples_per_step * cfg.accelerator.gradient_accumulation.steps + if is_main_process(): + 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()) logging.info(colored("Output dir:", "yellow", attrs=["bold"]) + f" {cfg.output_dir}") if cfg.env is not None: logging.info(f"{cfg.env.task=}") @@ -441,125 +581,22 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None): logging.info(f"{cfg.steps=} ({format_big_number(cfg.steps)})") logging.info(f"{dataset.num_frames=} ({format_big_number(dataset.num_frames)})") logging.info(f"{dataset.num_episodes=}") - num_processes = accelerator.num_processes - effective_bs = cfg.batch_size * num_processes - logging.info(f"Effective batch size: {cfg.batch_size} x {num_processes} = {effective_bs}") + logging.info( + f"Effective batch size: {cfg.batch_size} x {parallel_dims.dp_world_size} dp workers " + f"x {cfg.accelerator.gradient_accumulation.steps} grad accum = {effective_batch_size} " + f"(topology: dp_replicate={parallel_dims.dp_replicate}, dp_shard={parallel_dims.dp_shard})" + ) logging.info(f"{num_learnable_params=} ({format_big_number(num_learnable_params)})") logging.info(f"{num_total_params=} ({format_big_number(num_total_params)})") - # create dataloader for offline training - if not cfg.dataset.streaming: - # All non-streaming (map-style) datasets use EpisodeAwareSampler. - # The order is a pure function of (seed, epoch), so every rank independently produces the - # same permutation. accelerate then shards it disjointly across ranks via BatchSamplerShard - # without needing a `generator` attribute to synchronize an RNG, and resume is sample-exact. - shuffle = False - sampler = EpisodeAwareSampler( - dataset.meta.episodes["dataset_from_index"], - dataset.meta.episodes["dataset_to_index"], - episode_indices_to_use=dataset.episodes, - drop_n_last_frames=getattr(active_cfg, "drop_n_last_frames", 0), - shuffle=True, - seed=cfg.seed if cfg.seed is not None else 0, - absolute_to_relative_idx=dataset.absolute_to_relative_idx, - ) - if cfg.resume and step > 0: - # The resume offset depends on the (num_processes, batch_size) that produced `step`, so - # use the values recorded in the checkpoint (falling back to the current ones for older - # ckpts that did not store them). - saved_num_processes = load_training_num_processes(cfg.checkpoint_path) - saved_batch_size = load_training_batch_size(cfg.checkpoint_path) - ckpt_num_processes = saved_num_processes or accelerator.num_processes - ckpt_batch_size = saved_batch_size or cfg.batch_size - if is_main_process and saved_num_processes not in (None, accelerator.num_processes): - logging.warning( - f"Resuming with num_processes={accelerator.num_processes} but the checkpoint was " - f"written with num_processes={saved_num_processes}. The data order resumes at the " - "right epoch/offset, but per-rank sample-exactness requires the same world size." - ) - if is_main_process and saved_batch_size not in (None, cfg.batch_size): - logging.warning( - f"Resuming with batch_size={cfg.batch_size} but the checkpoint was written with " - f"batch_size={saved_batch_size}. The data order resumes at the right epoch/offset, " - "but per-rank sample-exactness requires the same batch size." - ) - sampler_state = compute_sampler_state(step, len(sampler), ckpt_batch_size, ckpt_num_processes) - sampler.load_state_dict(sampler_state) - if is_main_process: - logging.info( - f"Resuming data order at epoch {sampler_state['epoch']}, " - f"sample {sampler_state['start_index']}" - ) - else: - shuffle = True - sampler = None - - # Only swap in the language-aware collate when the dataset actually - # declares language columns; otherwise stay on PyTorch's default - # collate so non-language training runs are unaffected. - collate_fn = lerobot_collate_fn if dataset.meta.has_language_columns else None - dataloader = torch.utils.data.DataLoader( - dataset, - num_workers=cfg.num_workers, - batch_size=cfg.batch_size, - shuffle=shuffle and not cfg.dataset.streaming, - sampler=sampler, - pin_memory=device.type == "cuda", - drop_last=False, - collate_fn=collate_fn, - **_dataloader_worker_kwargs(cfg), - ) - - # Build eval dataloader if a held-out split exists - eval_dataloader = None - if eval_dataset is not None: - eval_ds = eval_dataset - if cfg.max_eval_samples > 0 and hasattr(eval_dataset, "hf_dataset"): - task_arr = eval_dataset.hf_dataset.data.column("task_index").to_numpy() - unique_tasks = sorted(set(task_arr.tolist())) - per_task = max(1, cfg.max_eval_samples // len(unique_tasks)) - selected: list[int] = [] - for t in unique_tasks: - frames = (task_arr == t).nonzero()[0][:per_task] - selected.extend(frames.tolist()) - eval_ds = torch.utils.data.Subset(eval_dataset, selected) - - eval_collate_fn = lerobot_collate_fn if dataset.meta.has_language_columns else None - eval_dataloader = torch.utils.data.DataLoader( - eval_ds, - batch_size=cfg.batch_size, - shuffle=False, - num_workers=cfg.num_workers, - pin_memory=device.type == "cuda", - drop_last=False, - collate_fn=eval_collate_fn, - **_dataloader_worker_kwargs(cfg), - ) - - # Prepare everything with accelerator - accelerator.wait_for_everyone() - if eval_dataloader is not None: - policy, optimizer, dataloader, lr_scheduler, eval_dataloader = accelerator.prepare( - policy, optimizer, dataloader, lr_scheduler, eval_dataloader - ) - else: - 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() train_metrics = { - # Per-rank loss reflects only one shard of the global batch; mean recovers the loss DDP - # is actually optimizing. grad_norm and lr are already identical on every rank (post - # gradient sync / deterministic scheduler) so reducing them would be a no-op collective. + # Per-rank loss reflects only one shard of the global batch; mean recovers the loss the + # data-parallel group is actually optimizing. grad_norm and lr are already identical on + # every rank (post gradient sync / deterministic scheduler) so reducing them would be a + # no-op collective. "loss": AverageMeter("loss", ":.3f", reduction="mean"), "grad_norm": AverageMeter("grdn", ":.3f"), "lr": AverageMeter("lr", ":0.1e"), @@ -574,18 +611,16 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None): # max() because headroom is gated by the worst-case rank. train_metrics["gpu_mem_gb"] = AverageMeter("mem_gb", ":.2f", reduction="max") - # Keep global batch size for logging; MetricsTracker handles world size internally. - effective_batch_size = cfg.batch_size * accelerator.num_processes train_tracker = MetricsTracker( cfg.batch_size, dataset.num_frames, dataset.num_episodes, train_metrics, initial_step=step, - accelerator=accelerator, + dp_world_size=parallel_dims.dp_world_size, ) - if is_main_process: + if is_main_process(): progbar = tqdm( total=cfg.steps - step, desc="Training", @@ -621,7 +656,7 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None): # Note: eval and checkpoint happens *after* the `step`th training update has completed, so we # increment `step` here. step += 1 - if is_main_process: + if is_main_process(): progbar.update(1) train_tracker.step() is_log_step = cfg.log_freq > 0 and step % cfg.log_freq == 0 @@ -632,12 +667,12 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None): if is_log_step: # Collective reduce must run on every rank, before the main-process gate below. train_tracker.reduce_across_ranks() - if is_main_process: + if is_main_process(): # Cluster-wide throughput, derived from the already-reduced (max) step time so it # reflects the slowest rank — which is what actually gates the next iteration. step_time = train_tracker.update_s.avg + train_tracker.dataloading_s.avg if step_time > 0: - train_tracker.samples_per_s = effective_batch_size / step_time + train_tracker.samples_per_s = samples_per_step / step_time logging.info(train_tracker) if wandb_logger: # Policy sub-losses (latent_loss, action_loss, ...) are aggregated into the @@ -661,7 +696,7 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None): if cam_key in eval_batch and eval_batch[cam_key].dtype == torch.uint8: eval_batch[cam_key] = eval_batch[cam_key].to(dtype=torch.float32) / 255.0 eval_batch = preprocessor(eval_batch) - loss, _ = policy.forward(eval_batch) + loss, _ = policy(eval_batch) # __call__, so FSDP2 forward hooks run eval_loss_sum += loss.item() n_eval_batches += 1 eval_loss = eval_loss_sum / max(n_eval_batches, 1) @@ -669,37 +704,29 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None): eval_loss = accelerator.reduce(eval_loss, reduction="mean").item() policy.train() - if is_main_process: + if is_main_process(): logging.info(f"step {step}: eval_loss={eval_loss:.4f}") if wandb_logger: wandb_logger.log_dict({"eval_loss": eval_loss}, step=step, mode="eval") 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: + # Collective: every rank participates (gathers / DCP shard writes); rank-0-only file + # writes are gated inside save_checkpoint — no rank branches at the call site. + if is_main_process(): logging.info(f"Checkpoint policy after step {step}") - checkpoint_dir = get_step_checkpoint_dir(cfg.output_dir, cfg.steps, step) - save_checkpoint( - checkpoint_dir=checkpoint_dir, - step=step, - cfg=cfg, - policy=accelerator.unwrap_model(policy), - 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, - ) + checkpoint_dir = get_step_checkpoint_dir(cfg.output_dir, cfg.steps, step) + save_checkpoint( + checkpoint_dir=checkpoint_dir, + step=step, + cfg=cfg, + policy=policy, + optimizer=optimizer, + scheduler=lr_scheduler, + preprocessor=preprocessor, + postprocessor=postprocessor, + accelerator=accelerator, + ) + if is_main_process(): update_last_checkpoint(checkpoint_dir) if cfg.save_checkpoint_to_hub: push_checkpoint_to_hub( @@ -709,11 +736,10 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None): ) if wandb_logger: wandb_logger.log_policy(checkpoint_dir) - accelerator.wait_for_everyone() if cfg.env and is_env_eval_step: - if is_main_process: + if is_main_process(): step_id = get_step_identifier(step, cfg.steps) logging.info(f"Eval policy at step {step}") with _make_eval_envs(cfg) as eval_env, torch.no_grad(), accelerator.autocast(): @@ -749,7 +775,7 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None): dataset.num_episodes, eval_metrics, initial_step=step, - accelerator=accelerator, + dp_world_size=parallel_dims.dp_world_size, ) eval_tracker.eval_s = aggregated.pop("eval_s") eval_tracker.avg_sum_reward = aggregated.pop("avg_sum_reward") @@ -761,23 +787,22 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None): accelerator.wait_for_everyone() - if is_main_process: + if is_main_process(): progbar.close() - - 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") - if getattr(active_cfg, "push_to_hub", False): - unwrapped_model = accelerator.unwrap_model(policy) - # PEFT only applies when training a policy — reward models use the plain path. - if not cfg.is_reward_model_training and cfg.policy.use_peft: - unwrapped_model.push_model_to_hub(cfg, peft_model=unwrapped_model, dataset_meta=dataset.meta) - else: - unwrapped_model.push_model_to_hub(cfg, state_dict=model_state_dict, dataset_meta=dataset.meta) - preprocessor.push_to_hub(active_cfg.repo_id) - postprocessor.push_to_hub(active_cfg.repo_id) + # --- publish (collective-safe: all ranks; the model commit gathers sharded weights) --------- + if getattr(active_cfg, "push_to_hub", False): + unwrapped = accelerator.unwrap_model(policy) + model_to_publish = unwrapped.get_base_model() if peft_model is not None else unwrapped + publish_trained_model( + cfg, + model_to_publish, + preprocessor, + postprocessor, + dataset.meta, + peft_model=unwrapped if peft_model is not None else None, + ) # Properly clean up the distributed process group accelerator.wait_for_everyone() @@ -785,7 +810,11 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None): def _remote_target_in_argv() -> bool: - """True when the CLI requests a remote HF Jobs run (--job.target=).""" + """Detect a remote HF Jobs run request on the raw CLI, before draccus parsing. + + Returns: + bool: True when the CLI requests a remote HF Jobs run (`--job.target=`). + """ target = None args = sys.argv[1:] for i, tok in enumerate(args): diff --git a/src/lerobot/utils/hub.py b/src/lerobot/utils/hub.py index 57eb819f8..8f7b17bbb 100644 --- a/src/lerobot/utils/hub.py +++ b/src/lerobot/utils/hub.py @@ -25,6 +25,9 @@ from .constants import CHECKPOINTS_DIR T = TypeVar("T", bound="HubMixin") +# Sharded-training resume artifacts (torch DCP shard dirs + shard files). Published model repos +# carry safetensors only, so publishing uploads exclude these — checkpoint pushes (which exist +# for resume, not distribution) deliberately do not. def find_latest_hub_checkpoint( repo_id: str, *, @@ -36,6 +39,16 @@ def find_latest_hub_checkpoint( Training runs push checkpoints to ``checkpoints//`` (see ``push_checkpoint_to_hub``). This lists those step dirs and returns ``checkpoints/``, or ``None`` if the repo has no checkpoints. + + Args: + repo_id (str): The Hub model repo to inspect. + token (str | bool | None): Hub authentication token. Defaults to None (the token + cached by `huggingface-cli login`). + revision (str | None): Repo revision to list. Defaults to None (the default branch). + + Returns: + str | None: The repo-relative path `checkpoints/`, or None if the repo + has no checkpoints. """ files = HfApi().list_repo_files(repo_id=repo_id, repo_type="model", revision=revision, token=token) prefix = f"{CHECKPOINTS_DIR}/" @@ -164,7 +177,7 @@ class HubMixin: ignore_patterns: list[str] | str | None = None, delete_patterns: list[str] | str | None = None, card_kwargs: dict[str, Any] | None = None, - ) -> str: + ) -> str | None: """ Upload model checkpoint to the Hub. @@ -172,6 +185,10 @@ class HubMixin: `delete_patterns` to delete existing remote files in the same commit. See [`upload_folder`] reference for more details. + Distributed contract: call on EVERY rank. `save_pretrained` runs on all ranks — for + sharded objects it can contain a collective gather (rank-gating it would deadlock) — + while repo creation and the upload happen on the main process only. + Args: repo_id (`str`): ID of the repository to push to (example: `"username/my-model"`). @@ -197,11 +214,17 @@ class HubMixin: Additional arguments passed to the card template to customize the card. Returns: - The url of the commit of your object in the given repository. + `str` or `None`: The url of the commit of your object in the given repository, or + `None` on non-main ranks of a distributed run (only the main process uploads). """ - api = HfApi(token=token) - repo_id = api.create_repo(repo_id=repo_id, private=private, exist_ok=True).repo_id + # Lazy import: hub code must not import the distributed package at module load + # (configs -> hub is on the import path of lerobot.distributed itself). + from lerobot.distributed.utils import is_main_process + # Distributed contract: `save_pretrained` runs on EVERY rank — for sharded policies it + # contains a collective gather (rank-gating it would deadlock) and it writes into this + # rank's private tmpdir only on the main process. Repo creation and upload are then + # main-process-only. if commit_message is None: if "Policy" in self.__class__.__name__: commit_message = "Upload policy" @@ -210,10 +233,13 @@ class HubMixin: else: commit_message = f"Upload {self.__class__.__name__}" - # Push the files to the repo in a single commit with TemporaryDirectory(ignore_cleanup_errors=True) as tmp: saved_path = Path(tmp) / repo_id self.save_pretrained(saved_path, card_kwargs=card_kwargs) + if not is_main_process(): + return None + api = HfApi(token=token) + repo_id = api.create_repo(repo_id=repo_id, private=private, exist_ok=True).repo_id return api.upload_folder( repo_id=repo_id, repo_type="model", diff --git a/src/lerobot/utils/logging_utils.py b/src/lerobot/utils/logging_utils.py index fd3021fea..94d644cf7 100644 --- a/src/lerobot/utils/logging_utils.py +++ b/src/lerobot/utils/logging_utils.py @@ -14,10 +14,10 @@ # See the License for the specific language governing permissions and # limitations under the License. from collections import defaultdict -from collections.abc import Callable from typing import Any import torch +import torch.distributed as dist from .utils import format_big_number @@ -69,12 +69,31 @@ class MetricsTracker: """ A helper class to track and log metrics over time. + Args: + batch_size (int): Per-process batch size (samples per micro-batch on each + data-parallel worker). + num_frames (int): Total number of frames in the training dataset. + num_episodes (int): Total number of episodes in the training dataset. + metrics (dict[str, AverageMeter]): The meters to track, keyed by metric name. + initial_step (int): Step counter to start from (non-zero when resuming a run). + Defaults to 0. + dp_world_size (int): Number of distinct data-parallel workers + (`dp_replicate * dp_shard`), used to scale sample accounting; context-parallel + peers consume the same batch and must not be double counted. Defaults to 1. + Usage pattern: ```python # initialize, potentially with non-zero initial step (e.g. if resuming run) metrics = {"loss": AverageMeter("loss", ":.3f")} - train_metrics = MetricsTracker(cfg, dataset, metrics, initial_step=step) + train_metrics = MetricsTracker( + batch_size, + dataset.num_frames, + dataset.num_episodes, + metrics, + initial_step=step, + dp_world_size=dp_world, + ) # update metrics derived from step (samples, episodes, epochs) at each training step train_metrics.step() @@ -98,12 +117,12 @@ class MetricsTracker: "_batch_size", "_num_frames", "_avg_samples_per_ep", + "_dp_world_size", "metrics", "steps", "samples", "episodes", "epochs", - "accelerator", "_caller_metrics", ] @@ -114,22 +133,25 @@ class MetricsTracker: num_episodes: int, metrics: dict[str, AverageMeter], initial_step: int = 0, - accelerator: Callable | None = None, + dp_world_size: int = 1, ): self.__dict__.update(dict.fromkeys(self.__keys__)) self._batch_size = batch_size self._num_frames = num_frames self._avg_samples_per_ep = num_frames / num_episodes + # Sample accounting scales by the number of DISTINCT data-parallel workers, which is + # dp_replicate * dp_shard — not the world size: context-parallel peers consume the same + # batch and must not be double counted. `step` counts micro-batches, so no + # grad-accumulation factor belongs here either. + self._dp_world_size = dp_world_size self.metrics = metrics self.steps = initial_step - world_size = accelerator.num_processes if accelerator else 1 # A sample is an (observation,action) pair, where observation and action # can be on multiple timestamps. In a batch, we have `batch_size` number of samples. - self.samples = self.steps * self._batch_size * world_size + self.samples = self.steps * self._batch_size * self._dp_world_size self.episodes = self.samples / self._avg_samples_per_ep self.epochs = self.samples / self._num_frames - self.accelerator = accelerator # Meter names the caller registered up front. update_metrics() leaves these untouched, so a # policy that echoes e.g. "loss" in its output dict can't clobber the aggregated meter. self._caller_metrics: set[str] = set(self.metrics) @@ -155,8 +177,7 @@ class MetricsTracker: Updates metrics that depend on 'step' for one step. """ self.steps += 1 - world_size = self.accelerator.num_processes if self.accelerator else 1 - self.samples += self._batch_size * world_size + self.samples += self._batch_size * self._dp_world_size self.episodes = self.samples / self._avg_samples_per_ep self.epochs = self.samples / self._num_frames @@ -181,11 +202,16 @@ class MetricsTracker: across all distributed processes (in-place). This is a collective operation and MUST be invoked on every rank — typically just before - logging. With no accelerator or in single-process runs it is a no-op. Without it, metrics - reported by the main process only reflect rank 0; for bottleneck-style timings - (``dataloading_s``, ``update_s``, ...) that means the slowest worker's stall is invisible. + logging. Outside distributed runs it is a no-op. Without it, metrics reported by the + main process only reflect rank 0; for bottleneck-style timings (``dataloading_s``, + ``update_s``, ...) that means the slowest worker's stall is invisible. + + Torch-native on purpose: metrics code carries no Accelerator dependency. + Note the reduction spans the WORLD group — correct for count-free averages (loss values + are identical within a context-parallel group, so including CP peers is a weighted + no-op). """ - if self.accelerator is None or self.accelerator.num_processes <= 1: + if not dist.is_initialized() or dist.get_world_size() <= 1: return buckets: dict[str, list[str]] = defaultdict(list) @@ -195,11 +221,20 @@ class MetricsTracker: if not buckets: return - device = self.accelerator.device + device = ( + torch.device("cuda", torch.cuda.current_device()) + if torch.cuda.is_available() + else torch.device("cpu") + ) + reduce_ops = { + "mean": dist.ReduceOp.AVG, + "sum": dist.ReduceOp.SUM, + "max": dist.ReduceOp.MAX, + } for reduction, names in buckets.items(): tensor = torch.tensor([self.metrics[n].avg for n in names], dtype=torch.float32, device=device) - reduced = self.accelerator.reduce(tensor, reduction=reduction) - for name, value in zip(names, reduced.tolist(), strict=True): + dist.all_reduce(tensor, op=reduce_ops[reduction]) + for name, value in zip(names, tensor.tolist(), strict=True): meter = self.metrics[name] # Preserve avg == sum / count so a later .update() on this meter accumulates # against the cluster view, not the stale per-rank history. diff --git a/tests/common/test_checkpoint_legacy_contracts.py b/tests/common/test_checkpoint_legacy_contracts.py new file mode 100644 index 000000000..d6f0d3e24 --- /dev/null +++ b/tests/common/test_checkpoint_legacy_contracts.py @@ -0,0 +1,221 @@ +#!/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. +"""Legacy-checkpoint contracts. + +Two contracts are pinned here so they are documented behavior, not accidents: + +- **The v0.6.0 hard break.** The v0.6.0 #3810 FSDP checkpoint layout + (full gathered ``model.safetensors`` + full ``optimizer_state.safetensors``, no DCP dirs, + no ``checkpoint_format`` in ``train_config.json``) is a hard break with ZERO v0.6.0-aware + runtime code — not even layout detection. A sharded resume pointed at such a checkpoint + must fail through the ORDINARY missing-artifact path (torch DCP erroring on the absent + ``training_state/optimizer_0/``), while the model weights remain loadable forever via + ``from_pretrained`` and the old ``num_processes`` key keeps feeding the topology reader. +- **Converter equivalence.** ``dcp_to_safetensors`` (real ``merge_fsdp_weights``, no mocks) + on accelerate's ``save_fsdp_model`` DCP layout reproduces exactly the tensors that the + direct-gather ``save_pretrained`` artifact contains. +""" + +import json +from pathlib import Path +from types import SimpleNamespace + +import pytest + +pytest.importorskip("accelerate", reason="accelerate is required (install lerobot[training])") + +import torch +import torch.distributed.checkpoint as dist_cp +from accelerate.utils.constants import FSDP_MODEL_NAME, OPTIMIZER_NAME +from safetensors.torch import load_file +from torch.distributed.checkpoint.api import CheckpointException +from torch.distributed.fsdp import FSDPModule + +from lerobot.common.train_utils import ( + load_training_metadata, + resume_after_prepare, + resume_before_prepare, +) +from lerobot.configs.accelerator import FSDPConfig +from lerobot.configs.default import DatasetConfig +from lerobot.configs.train import TRAIN_CONFIG_NAME, CheckpointFormat, TrainPipelineConfig +from lerobot.distributed.checkpoint import dcp_to_safetensors, is_sharded_module +from lerobot.optim.optimizers import save_optimizer_state +from lerobot.utils.constants import PRETRAINED_MODEL_DIR, TRAINING_STATE_DIR, TRAINING_STEP +from lerobot.utils.io_utils import write_json +from lerobot.utils.random_utils import save_rng_state +from tests.fixtures.dummy_checkpoint_policy import DummyCheckpointPolicy, make_dummy_policy + + +@pytest.fixture +def accelerate_state(): + """accelerate's process state, as the trainer's `Accelerator()` would have initialized it. + + `load_fsdp_optimizer` and `merge_fsdp_weights` both consult `PartialState` internals + (logging and main-process gating). Single-process CPU state; reset on teardown so no + global accelerate state leaks into other tests. + """ + from accelerate.state import AcceleratorState, PartialState + + PartialState() + yield + AcceleratorState._reset_state(reset_partial_state=True) + + +def make_v060_fsdp_checkpoint(checkpoint_dir: Path) -> dict[str, torch.Tensor]: + """Reproduce the v0.6.0 #3810 FSDP checkpoint layout with real artifacts. + + - ``pretrained_model/``: ``config.json`` + full gathered ``model.safetensors`` (real + ``save_pretrained`` outputs) and a ``train_config.json`` predating the v0.7 fields + (``checkpoint_format``/``parallelism``/``accelerator`` stripped from the draccus dump); + - ``training_state/``: old-style ``training_step.json`` (``{"step", "num_processes"}``, + no ``dp_world_size``), ``rng_state.safetensors``, and the gathered full optimizer + channel (``optimizer_state.safetensors`` + ``optimizer_param_groups.json``) — and, + crucially, NO ``optimizer_0/`` DCP directory. + + Returns the saved model weights for later comparison. + """ + policy = make_dummy_policy() + optimizer = torch.optim.Adam(policy.parameters()) + policy.forward({"observation.state": torch.randn(2, 4)})[0].backward() + optimizer.step() # real optimizer state, applied before the weights are saved + + pretrained_dir = checkpoint_dir / PRETRAINED_MODEL_DIR + policy.save_pretrained(pretrained_dir) + cfg = TrainPipelineConfig(dataset=DatasetConfig(repo_id="lerobot/dummy"), batch_size=3) + cfg._save_pretrained(pretrained_dir) + config_path = pretrained_dir / TRAIN_CONFIG_NAME + raw = json.loads(config_path.read_text()) + assert "checkpoint_format" in raw # draccus dumps defaults; a v0.6.0 config predates the key + for key in ("checkpoint_format", "parallelism", "accelerator"): + raw.pop(key, None) + config_path.write_text(json.dumps(raw, indent=4)) + + training_state_dir = checkpoint_dir / TRAINING_STATE_DIR + training_state_dir.mkdir() + write_json({"step": 5000, "num_processes": 4}, training_state_dir / TRAINING_STEP) + save_rng_state(training_state_dir) + save_optimizer_state(optimizer, training_state_dir) + return {key: tensor.clone() for key, tensor in policy.state_dict().items()} + + +def as_fsdp2_module(policy: DummyCheckpointPolicy) -> DummyCheckpointPolicy: + """Give the policy FSDP2's runtime identity via the in-place class swap `fully_shard` performs. + + torch's `fully_shard` swaps ``module.__class__`` to a ``(FSDPModule, type(module))`` + subclass; mirroring that swap is what makes `is_sharded_module` (and thus the sharded + branch of `resume_after_prepare`) see a sharded model on a CPU-only single process. The + parameters stay plain tensors — sufficient here, because the resume must fail at the DCP + read before any sharded state is touched. + """ + policy.__class__ = type(f"FSDP{type(policy).__name__}", (FSDPModule, type(policy)), {}) + assert is_sharded_module(policy) + return policy + + +def sharded_passthrough_accelerator() -> SimpleNamespace: + """The accelerator surface the sharded resume touches, carrying the trainer's real plugin. + + `FSDPConfig.build_plugin()` is the exact FSDP2 plugin construction `make_accelerator` + hands to accelerate (state_dict_type stays at the FSDP2 default, SHARDED_STATE_DICT). + """ + return SimpleNamespace( + unwrap_model=lambda m: m, + wait_for_everyone=lambda: None, + state=SimpleNamespace(fsdp_plugin=FSDPConfig().build_plugin()), + ) + + +class TestV060HardBreak: + """Pin the v0.6.0 hard break as a contract. + + Zero v0.6.0-aware code ships — not even layout detection — so every assertion here must + hold through ORDINARY code paths only: the recorded config parses with plain defaults, + phase-1 resume and the weights stay loadable, and the sharded phase-2 resume fails with + torch DCP's own missing-artifact error, never a bespoke v0.6.0 message. + """ + + def test_sharded_resume_fails_with_ordinary_missing_artifact_error(self, tmp_path, accelerate_state): + make_v060_fsdp_checkpoint(tmp_path) + + # No checkpoint_format recorded -> plain draccus default, no layout detection anywhere. + cfg = TrainPipelineConfig.from_pretrained(tmp_path / PRETRAINED_MODEL_DIR / TRAIN_CONFIG_NAME) + assert cfg.checkpoint_format is CheckpointFormat.SAFETENSORS + cfg.checkpoint_path = tmp_path + + # Phase 1 (RNG + step counter) is format-independent and still succeeds. + assert resume_before_prepare(cfg) == 5000 + + # Phase 2 under sharding: the recorded format skips the DCP model preflight (the + # weights were already loaded by from_pretrained), then the sharded optimizer load + # hits the absent optimizer_0/ and fails inside torch DCP — the ordinary error path. + assert not (tmp_path / TRAINING_STATE_DIR / f"{OPTIMIZER_NAME}_0").exists() + policy = as_fsdp2_module(make_dummy_policy()) + optimizer = torch.optim.Adam(policy.parameters()) + with pytest.raises(CheckpointException) as excinfo: + resume_after_prepare(cfg, sharded_passthrough_accelerator(), policy, optimizer, None) + message = str(excinfo.value) + assert "lerobot-convert-dcp" not in message # the converter hint belongs to recorded-format=DCP + assert "v0.6" not in message # no bespoke wording: the explanation lives in the migration docs + + def test_weights_remain_loadable_via_from_pretrained(self, tmp_path): + saved_weights = make_v060_fsdp_checkpoint(tmp_path) + policy = DummyCheckpointPolicy.from_pretrained(tmp_path / PRETRAINED_MODEL_DIR) + for key, tensor in policy.state_dict().items(): + assert torch.equal(tensor, saved_weights[key]), key + + def test_topology_reader_falls_back_to_legacy_num_processes(self, tmp_path): + make_v060_fsdp_checkpoint(tmp_path) + assert load_training_metadata(tmp_path / TRAINING_STATE_DIR)["dp_world_size"] == 4 + + +class TestConverterEquivalence: + def test_dcp_to_safetensors_output_equals_direct_gather(self, tmp_path, accelerate_state): + """DCP -> safetensors conversion is exactly the direct-gather artifact. + + The DCP checkpoint is written with torch's real `dist_cp.save` (single process, no + process group), replicating accelerate's `save_fsdp_model` SHARDED_STATE_DICT branch + byte for byte: the ``{"model": state_dict}`` nesting and the ``pytorch_model_fsdp_0`` + directory name. The conversion runs the real `merge_fsdp_weights` — no mocks. + """ + policy = make_dummy_policy() + with torch.no_grad(): + for param in policy.parameters(): + param.add_(torch.randn_like(param)) # make every tensor distinct from init + reference = {key: tensor.clone() for key, tensor in policy.state_dict().items()} + + # The direct-gather artifact (on a single process the gather is state_dict itself). + direct_dir = tmp_path / "direct" + policy.save_pretrained(direct_dir) + + # The DCP artifact, laid out exactly as accelerate's save_fsdp_model writes it. + pretrained_dir = tmp_path / "checkpoint" / PRETRAINED_MODEL_DIR + dcp_dir = pretrained_dir / f"{FSDP_MODEL_NAME}_0" + dcp_dir.mkdir(parents=True) + dist_cp.save( + state_dict={"model": policy.state_dict()}, + storage_writer=dist_cp.FileSystemWriter(str(dcp_dir)), + ) + + merged_file = dcp_to_safetensors(dcp_dir, pretrained_dir) + assert merged_file == pretrained_dir / "model.safetensors" + merged = load_file(merged_file) + direct = load_file(direct_dir / "model.safetensors") + assert set(merged) == set(direct) == set(reference) + for key, tensor in reference.items(): + assert torch.equal(merged[key], tensor), key + assert torch.equal(direct[key], tensor), key + assert merged[key].dtype == tensor.dtype, key diff --git a/tests/common/test_checkpoint_save_resume.py b/tests/common/test_checkpoint_save_resume.py new file mode 100644 index 000000000..290ad32e8 --- /dev/null +++ b/tests/common/test_checkpoint_save_resume.py @@ -0,0 +1,186 @@ +#!/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. +"""Checkpoint save/resume round-trips on the non-sharded paths.""" + +from types import SimpleNamespace + +import pytest +import torch +from safetensors.torch import load_file + +from lerobot.common.train_utils import ( + load_training_metadata, + resume_after_prepare, + resume_before_prepare, + save_checkpoint, +) +from lerobot.configs.default import DatasetConfig +from lerobot.configs.train import CheckpointFormat, TrainPipelineConfig +from lerobot.utils.constants import PRETRAINED_MODEL_DIR, TRAINING_STATE_DIR, TRAINING_STEP +from lerobot.utils.io_utils import load_json, write_json +from tests.fixtures.dummy_checkpoint_policy import make_dummy_policy + + +def make_cfg(**overrides) -> TrainPipelineConfig: + cfg = TrainPipelineConfig(dataset=DatasetConfig(repo_id="lerobot/dummy"), batch_size=3) + cfg.parallelism.resolve(1) + for name, value in overrides.items(): + setattr(cfg, name, value) + return cfg + + +def passthrough_accelerator() -> SimpleNamespace: + """The accelerator surface save/resume touches on non-sharded runs.""" + return SimpleNamespace(unwrap_model=lambda m: m, wait_for_everyone=lambda: None) + + +class TestSaveCheckpoint: + def test_non_sharded_layout(self, tmp_path): + policy = make_dummy_policy() + optimizer = torch.optim.Adam(policy.parameters()) + save_checkpoint( + tmp_path, + step=7, + cfg=make_cfg(), + policy=policy, + optimizer=optimizer, + accelerator=passthrough_accelerator(), + ) + pretrained = tmp_path / PRETRAINED_MODEL_DIR + state = tmp_path / TRAINING_STATE_DIR + assert (pretrained / "model.safetensors").is_file() + assert (pretrained / "config.json").is_file() + assert (pretrained / "train_config.json").is_file() + assert (state / TRAINING_STEP).is_file() + assert (state / "rng_state.safetensors").is_file() + assert (state / "optimizer_state.safetensors").is_file() + # single-file artifact, no index, weights intact + weights = load_file(pretrained / "model.safetensors") + assert torch.allclose(weights["net.weight"], torch.full_like(weights["net.weight"], 0.5)) + assert not list(pretrained.glob("*.index.json")) + + def test_training_step_records_topology(self, tmp_path): + cfg = make_cfg() + cfg.accelerator.gradient_accumulation.steps = 4 + policy = make_dummy_policy() + save_checkpoint( + tmp_path, + step=11, + cfg=cfg, + policy=policy, + optimizer=torch.optim.Adam(policy.parameters()), + accelerator=passthrough_accelerator(), + ) + metadata = load_training_metadata(tmp_path / TRAINING_STATE_DIR) + assert metadata["dp_world_size"] == 1 + assert metadata["batch_size"] == 3 + assert metadata["grad_accum_steps"] == 4 + + def test_dp_world_size_legacy_fallback(self, tmp_path): + """Pre-v0.7 checkpoints recorded num_processes; the reader falls back to it.""" + state_dir = tmp_path / TRAINING_STATE_DIR + state_dir.mkdir(parents=True) + write_json({"step": 5, "num_processes": 4}, state_dir / TRAINING_STEP) + metadata = load_training_metadata(tmp_path / TRAINING_STATE_DIR) + assert metadata["dp_world_size"] == 4 + assert metadata["batch_size"] is None + + +class TestResume: + def _checkpointed_run(self, tmp_path): + policy = make_dummy_policy() + optimizer = torch.optim.Adam(policy.parameters(), lr=0.123) + # give the optimizer real state + policy.forward({"observation.state": torch.randn(2, 4)})[0].backward() + optimizer.step() + cfg = make_cfg() + save_checkpoint( + tmp_path, + step=42, + cfg=cfg, + policy=policy, + optimizer=optimizer, + accelerator=passthrough_accelerator(), + ) + cfg.checkpoint_path = tmp_path + return cfg, policy, optimizer + + def test_two_phase_resume_round_trip(self, tmp_path): + cfg, _, optimizer = self._checkpointed_run(tmp_path) + assert resume_before_prepare(cfg) == 42 + + fresh_policy = make_dummy_policy() + fresh_optimizer = torch.optim.Adam(fresh_policy.parameters(), lr=0.999) + resume_after_prepare(cfg, passthrough_accelerator(), fresh_policy, fresh_optimizer, None) + restored = fresh_optimizer.state_dict() + original = optimizer.state_dict() + assert restored["param_groups"][0]["lr"] == original["param_groups"][0]["lr"] + for key, tensor in original["state"][0].items(): + assert torch.equal(restored["state"][0][key], tensor), key + + def test_resume_warns_on_changed_cadence_and_topology(self, tmp_path, caplog): + """The recorded grad-accum factor and parallelism snapshot must be compared on + resume, with one warning naming the diff.""" + import logging + + cfg, _, _ = self._checkpointed_run(tmp_path) + cfg.accelerator.gradient_accumulation.steps = 4 + cfg.parallelism.dp_replicate = 2 # same dp_world_size story is irrelevant here + with caplog.at_level(logging.WARNING): + assert resume_before_prepare(cfg) == 42 + warning = next(m for m in caplog.messages if "differ from the checkpoint" in m) + assert "grad_accum_steps: 1 -> 4" in warning + assert "dp_replicate: 1 -> 2" in warning + + def test_resume_unchanged_settings_stay_silent(self, tmp_path, caplog): + import logging + + cfg, _, _ = self._checkpointed_run(tmp_path) + with caplog.at_level(logging.WARNING): + resume_before_prepare(cfg) + assert not [m for m in caplog.messages if "differ from the checkpoint" in m] + + def test_resume_rejects_non_sharded_checkpoint_on_sharded_run(self, tmp_path): + """Resharding works across sizes, not across kinds: non-sharded -> sharded is rejected.""" + cfg, _, _ = self._checkpointed_run(tmp_path) + cfg.parallelism.dp_shard = 2 + with pytest.raises(ValueError, match="Cannot resume"): + resume_before_prepare(cfg) + + def test_resume_rejects_sharded_checkpoint_on_non_sharded_run(self, tmp_path): + """The symmetric direction: a checkpoint recorded sharded cannot resume non-sharded.""" + cfg, _, _ = self._checkpointed_run(tmp_path) + state_file = tmp_path / TRAINING_STATE_DIR / TRAINING_STEP + state = load_json(state_file) + state["parallelism"]["dp_shard"] = 2 + write_json(state, state_file) + with pytest.raises(ValueError, match="Cannot resume"): + resume_before_prepare(cfg) + + def test_resume_before_prepare_requires_training_state(self, tmp_path): + cfg = make_cfg() + cfg.checkpoint_path = tmp_path + with pytest.raises(NotADirectoryError): + resume_before_prepare(cfg) + + def test_dcp_format_integrity_preflight(self, tmp_path): + """A checkpoint declaring DCP shards without the shard dir fails with the converter hint.""" + pytest.importorskip("accelerate", reason="accelerate is required (install lerobot[training])") + cfg, policy, optimizer = self._checkpointed_run(tmp_path) + cfg.parallelism.dp_shard = 2 # pretend the recorded run was sharded + cfg.checkpoint_format = CheckpointFormat.DCP + with pytest.raises(FileNotFoundError, match="lerobot-convert-dcp"): + resume_after_prepare(cfg, passthrough_accelerator(), policy, optimizer, None) diff --git a/tests/common/test_publish.py b/tests/common/test_publish.py new file mode 100644 index 000000000..1eac9b0d4 --- /dev/null +++ b/tests/common/test_publish.py @@ -0,0 +1,148 @@ +#!/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. +"""publish_trained_model: commit set, card, log-line contract, PEFT branch (hub fully mocked).""" + +import logging +from pathlib import Path +from types import SimpleNamespace + +import pytest + +import lerobot.common.train_utils as train_utils +import lerobot.utils.hub as hub +from lerobot.common.train_utils import generate_model_card, publish_trained_model +from lerobot.configs.default import DatasetConfig +from lerobot.configs.train import TrainPipelineConfig +from tests.fixtures.dummy_checkpoint_policy import make_dummy_policy + + +class FakeHfApi: + """Records every repo/upload interaction; shared across both HfApi import sites.""" + + calls: list[dict] = [] + + def __init__(self, *args, **kwargs): + pass + + def create_repo(self, repo_id, private=None, exist_ok=False, **kwargs): + return SimpleNamespace(repo_id=repo_id) + + def upload_folder(self, *, repo_id, folder_path, commit_message, **kwargs): + FakeHfApi.calls.append( + { + "repo_id": repo_id, + "commit_message": commit_message, + "files": sorted(p.name for p in Path(folder_path).iterdir()), + "ignore_patterns": kwargs.get("ignore_patterns"), + } + ) + return SimpleNamespace(repo_url=SimpleNamespace(url=f"https://huggingface.co/{repo_id}")) + + +@pytest.fixture +def mocked_hub(monkeypatch): + FakeHfApi.calls = [] + monkeypatch.setattr(train_utils, "HfApi", FakeHfApi) + monkeypatch.setattr(hub, "HfApi", FakeHfApi) + # card.validate() hits the Hub; publishing must work offline in tests + monkeypatch.setattr(train_utils.ModelCard, "validate", lambda self: None) + return FakeHfApi + + +def make_cfg() -> TrainPipelineConfig: + cfg = TrainPipelineConfig(dataset=DatasetConfig(repo_id="user/dataset")) + cfg.parallelism.resolve(1) + return cfg + + +class RecordingProcessor: + def __init__(self): + self.pushed_to = None + + def push_to_hub(self, repo_id, **kwargs): + self.pushed_to = repo_id + + +class TestPublishTrainedModel: + def test_commit_set_and_log_contract(self, mocked_hub, caplog): + policy = make_dummy_policy(repo_id="user/policy") + pre, post = RecordingProcessor(), RecordingProcessor() + with caplog.at_level(logging.INFO): + publish_trained_model(make_cfg(), policy, pre, post, dataset_meta=None) + + # commit 1: the model through HubMixin (config.json + model.safetensors in a tmpdir) + model_commit = mocked_hub.calls[0] + assert {"config.json", "model.safetensors"} <= set(model_commit["files"]) + # commits 2-3: processors + assert pre.pushed_to == "user/policy" and post.pushed_to == "user/policy" + # commit 4: the bundle sidecar + bundle = mocked_hub.calls[-1] + assert {"README.md", "train_config.json"} <= set(bundle["files"]) + # the exact line lerobot.jobs.hf watches to end remote runs early + assert any( + m.startswith("Model pushed to https://huggingface.co/user/policy") for m in caplog.messages + ) + + def test_peft_branch_skips_model_commit(self, mocked_hub): + policy = make_dummy_policy(repo_id="user/policy") + + class FakePeftModel: + def save_pretrained(self, path): + (Path(path) / "adapter_model.safetensors").write_bytes(b"x") + + publish_trained_model(make_cfg(), policy, None, None, dataset_meta=None, peft_model=FakePeftModel()) + assert len(mocked_hub.calls) == 1 # only the bundle commit + bundle = mocked_hub.calls[0] + # adapter weights + the wrapped policy's config + card + train config, no full weights + assert {"README.md", "adapter_model.safetensors", "config.json", "train_config.json"} <= set( + bundle["files"] + ) + assert "model.safetensors" not in bundle["files"] + + def test_missing_repo_id_fails_loudly(self, mocked_hub): + policy = make_dummy_policy(repo_id=None) + with pytest.raises(ValueError, match="repo id"): + publish_trained_model(make_cfg(), policy, None, None, dataset_meta=None) + + +class TestGenerateModelCard: + def test_free_function_renders_from_arguments(self, monkeypatch): + monkeypatch.setattr(train_utils.ModelCard, "validate", lambda self: None) + policy = make_dummy_policy(repo_id="user/policy") + card = generate_model_card(policy.config, cfg=make_cfg(), dataset_meta=None) + assert card.data.library_name == "lerobot" + assert card.data.datasets == "user/dataset" + assert "lerobot" in card.data.tags + + +class TestDeprecatedPushModelToHub: + """`push_model_to_hub` stays callable for external scripts, delegating to the publisher.""" + + def test_policy_shim_warns_and_publishes(self, mocked_hub): + policy = make_dummy_policy(repo_id="user/policy") + with pytest.warns(FutureWarning, match="push_model_to_hub is deprecated"): + policy.push_model_to_hub(make_cfg()) + + # Same artifacts the method produced before: weights + config, then card + train config. + model_commit = mocked_hub.calls[0] + assert {"config.json", "model.safetensors"} <= set(model_commit["files"]) + bundle = mocked_hub.calls[-1] + assert {"README.md", "train_config.json"} <= set(bundle["files"]) + + def test_policy_shim_warns_that_state_dict_is_ignored(self, mocked_hub): + policy = make_dummy_policy(repo_id="user/policy") + with pytest.warns(FutureWarning, match="`state_dict` argument is ignored"): + policy.push_model_to_hub(make_cfg(), state_dict=policy.state_dict()) diff --git a/tests/configs/test_accelerator_config.py b/tests/configs/test_accelerator_config.py new file mode 100644 index 000000000..d70ef00aa --- /dev/null +++ b/tests/configs/test_accelerator_config.py @@ -0,0 +1,135 @@ +#!/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. + +import json + +import draccus +import pytest + +from lerobot.configs.accelerator import ( + AcceleratorConfig, + ActivationCheckpointingConfig, + ActivationCheckpointingMode, + CompileConfig, + DDPConfig, + FSDPConfig, + GradientAccumulationConfig, +) +from lerobot.configs.parallelism import ParallelismConfig + + +class TestFieldValidation: + def test_wrap_policies_mutually_exclusive(self): + with pytest.raises(ValueError, match="mutually exclusive"): + FSDPConfig(wrap_modules=["Block"], min_num_params=1000) + + def test_min_num_params_positive(self): + with pytest.raises(ValueError, match="min_num_params"): + FSDPConfig(min_num_params=0) + + def test_mixed_precision_choices(self): + with pytest.raises(ValueError, match="mixed_precision"): + AcceleratorConfig(mixed_precision="tf32") + + def test_gradient_accumulation_positive(self): + with pytest.raises(ValueError, match="gradient_accumulation.steps"): + GradientAccumulationConfig(steps=0) + + +class TestDraccusRoundTrip: + @pytest.mark.parametrize( + "cfg", + [ + AcceleratorConfig(), + AcceleratorConfig( + mixed_precision="bf16", + gradient_accumulation=GradientAccumulationConfig(steps=4), + fsdp=FSDPConfig( + reshard_after_forward=False, + wrap_modules=["ACTEncoderLayer", "ACTDecoderLayer"], + cpu_offload=True, + ignored_modules=r".*pos_embed.*", + ), + ddp=DDPConfig(find_unused_parameters=False, static_graph=True), + compile=CompileConfig(enabled=True, mode="max-autotune", regional=False), + activation_checkpointing=ActivationCheckpointingConfig(mode=ActivationCheckpointingMode.FULL), + ), + AcceleratorConfig(fsdp=FSDPConfig(min_num_params=1_000_000)), + ], + ) + def test_encode_json_decode_identity(self, cfg): + payload = json.loads(json.dumps(draccus.encode(cfg))) + assert draccus.decode(AcceleratorConfig, payload) == cfg + + def test_pre_existing_config_without_fields_gets_defaults(self): + assert draccus.decode(AcceleratorConfig, {}) == AcceleratorConfig() + + +class TestRuntimeBuilders: + """The mirrors must translate into real accelerate objects (plugins built lazily).""" + + @pytest.fixture(autouse=True) + def _requires_accelerate(self): + pytest.importorskip("accelerate", reason="accelerate is required (install lerobot[training])") + + def test_fsdp_plugin_translation(self): + plugin = FSDPConfig( + reshard_after_forward=False, wrap_modules=["MyBlock"], cpu_offload=True + ).build_plugin() + assert plugin.fsdp_version == 2 + assert plugin.reshard_after_forward is False + assert plugin.transformer_cls_names_to_wrap == ["MyBlock"] + # bools are normalized into torch offload policies by the plugin itself + assert type(plugin.cpu_offload).__name__ == "CPUOffloadPolicy" + # LeRobot never switches state_dict_type: FSDP2's SHARDED default must hold + assert plugin.state_dict_type.name == "SHARDED_STATE_DICT" + assert not plugin.activation_checkpointing + + def test_fsdp_plugin_size_based_policy(self): + plugin = FSDPConfig(min_num_params=1024).build_plugin() + assert plugin.min_num_params == 1024 + assert plugin.transformer_cls_names_to_wrap is None + + def test_ddp_kwargs_translation(self): + handler = DDPConfig(find_unused_parameters=False, gradient_as_bucket_view=True).build_kwargs_handler() + assert handler.find_unused_parameters is False + assert handler.gradient_as_bucket_view is True + + def test_gradient_accumulation_plugin_translation(self): + plugin = GradientAccumulationConfig(steps=4).build_plugin() + assert plugin.num_steps == 4 + assert plugin.sync_with_dataloader is False + + def test_gradient_accumulation_never_syncs_with_dataloader(self, monkeypatch): + """The loop cycles a finite dataloader, so accelerate's default + sync_with_dataloader=True would force an optimizer step at every dataset epoch + boundary instead of every num_steps micro-batches.""" + captured = {} + + class FakeAccelerator: + def __init__(self, **kwargs): + captured.update(kwargs) + + monkeypatch.setattr("accelerate.Accelerator", FakeAccelerator) + parallelism = ParallelismConfig() + parallelism.resolve(1) + AcceleratorConfig(gradient_accumulation=GradientAccumulationConfig(steps=4)).build( + parallelism, cpu=True + ) + ga_plugin = captured["gradient_accumulation_plugin"] + assert ga_plugin.num_steps == 4 + assert ga_plugin.sync_with_dataloader is False + assert "gradient_accumulation_steps" not in captured diff --git a/tests/configs/test_parallelism_config.py b/tests/configs/test_parallelism_config.py new file mode 100644 index 000000000..6afb5892a --- /dev/null +++ b/tests/configs/test_parallelism_config.py @@ -0,0 +1,118 @@ +#!/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. + +import json + +import draccus +import pytest + +from lerobot.configs.parallelism import ContextParallelConfig, ParallelismConfig + + +class TestResolve: + def test_single_process_defaults(self): + cfg = ParallelismConfig() + cfg.resolve(1) + assert (cfg.dp_replicate, cfg.dp_shard) == (1, 1) + assert not cfg.is_sharded and not cfg.is_replicated_only + assert cfg.dp_world_size == 1 + + def test_untouched_config_fills_ddp(self): + """Plain `torchrun --nproc-per-node=8` with a default config resolves to DDP.""" + cfg = ParallelismConfig() + cfg.resolve(8) + assert cfg.dp_replicate == 8 + assert cfg.is_replicated_only and not cfg.is_sharded + assert cfg.dp_world_size == 8 + + def test_full_shard_sentinel(self): + cfg = ParallelismConfig(dp_shard=-1) + assert cfg.is_sharded # sharded even before resolve: -1 is an explicit opt-in + cfg.resolve(8) + assert cfg.dp_shard == 8 and cfg.dp_replicate == 1 + + def test_hsdp_sentinel_infers_shard(self): + cfg = ParallelismConfig(dp_replicate=2, dp_shard=-1) + cfg.resolve(8) + assert (cfg.dp_replicate, cfg.dp_shard) == (2, 4) + assert cfg.dp_world_size == 8 + + def test_explicit_hsdp(self): + cfg = ParallelismConfig(dp_replicate=2, dp_shard=4) + cfg.resolve(8) + assert cfg.is_sharded and not cfg.is_replicated_only + + def test_product_mismatch_lists_all_degrees(self): + cfg = ParallelismConfig(dp_replicate=2, dp_shard=2) + with pytest.raises(ValueError, match=r"dp_replicate=2 \* dp_shard=2.*WORLD_SIZE=8"): + cfg.resolve(8) + + def test_explicit_replicate_must_match_world(self): + cfg = ParallelismConfig(dp_replicate=4) + with pytest.raises(ValueError, match="WORLD_SIZE=8"): + cfg.resolve(8) + + def test_sentinel_indivisible_world(self): + cfg = ParallelismConfig(dp_replicate=3, dp_shard=-1) + with pytest.raises(ValueError, match="not divisible"): + cfg.resolve(8) + + def test_cp_fails_fast(self): + cfg = ParallelismConfig(dp_shard=-1, context_parallel=ContextParallelConfig(ulysses_degree=2)) + with pytest.raises(ValueError, match="not implemented"): + cfg.resolve(8) + + +class TestFieldValidation: + @pytest.mark.parametrize("kwargs", [{"dp_replicate": 0}, {"dp_shard": 0}, {"dp_shard": -2}]) + def test_bad_dp_degrees(self, kwargs): + with pytest.raises(ValueError): + ParallelismConfig(**kwargs) + + def test_cfg_parallel_capped_at_two(self): + ParallelismConfig(cfg_parallel=2) # reserved but representable + with pytest.raises(ValueError, match="cfg_parallel"): + ParallelismConfig(cfg_parallel=3) + + @pytest.mark.parametrize("kwargs", [{"ring_degree": 0}, {"ulysses_degree": -1}]) + def test_bad_cp_degrees(self, kwargs): + with pytest.raises(ValueError): + ContextParallelConfig(**kwargs) + + def test_dp_world_size_undefined_before_resolve(self): + with pytest.raises(RuntimeError, match="resolve"): + _ = ParallelismConfig(dp_shard=-1).dp_world_size + + +class TestDraccusRoundTrip: + @pytest.mark.parametrize( + "cfg", + [ + ParallelismConfig(), + ParallelismConfig(dp_replicate=2, dp_shard=4, cfg_parallel=2), + ParallelismConfig( + dp_shard=-1, + context_parallel=ContextParallelConfig(ring_degree=2, ulysses_degree=4), + ), + ], + ) + def test_encode_json_decode_identity(self, cfg): + payload = json.loads(json.dumps(draccus.encode(cfg))) + assert draccus.decode(ParallelismConfig, payload) == cfg + + def test_pre_existing_config_without_fields_gets_defaults(self): + """Checkpoints written before this feature parse with default topology.""" + assert draccus.decode(ParallelismConfig, {}) == ParallelismConfig() diff --git a/tests/configs/test_train_config_distributed.py b/tests/configs/test_train_config_distributed.py new file mode 100644 index 000000000..8a227ee07 --- /dev/null +++ b/tests/configs/test_train_config_distributed.py @@ -0,0 +1,118 @@ +#!/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. +"""TrainPipelineConfig integration for the distributed fields: fail-fasts + config compat.""" + +import draccus +import pytest + +from lerobot.configs.accelerator import ActivationCheckpointingMode +from lerobot.configs.default import DatasetConfig, PeftConfig +from lerobot.configs.parallelism import ContextParallelConfig, ParallelismConfig +from lerobot.configs.train import CheckpointFormat, TrainPipelineConfig +from lerobot.optim.optimizers import AdamConfig, MultiAdamConfig + + +def make_cfg(**overrides) -> TrainPipelineConfig: + cfg = TrainPipelineConfig(dataset=DatasetConfig(repo_id="lerobot/dummy")) + for name, value in overrides.items(): + setattr(cfg, name, value) + return cfg + + +def sharded() -> ParallelismConfig: + return ParallelismConfig(dp_shard=-1) + + +class TestDistributedFailFasts: + def test_defaults_pass(self): + make_cfg()._validate_distributed() + + def test_cp_reserved(self): + cfg = make_cfg(parallelism=ParallelismConfig(context_parallel=ContextParallelConfig(ring_degree=2))) + with pytest.raises(ValueError, match="not implemented"): + cfg._validate_distributed() + + def test_cfg_parallel_training_rejected(self): + cfg = make_cfg(parallelism=ParallelismConfig(cfg_parallel=2)) + with pytest.raises(ValueError, match="inference-only"): + cfg._validate_distributed() + + def test_compile_placeholder(self): + cfg = make_cfg() + cfg.accelerator.compile.enabled = True + with pytest.raises(ValueError, match="compile"): + cfg._validate_distributed() + + def test_activation_checkpointing_placeholder(self): + cfg = make_cfg() + cfg.accelerator.activation_checkpointing.mode = ActivationCheckpointingMode.FULL + with pytest.raises(ValueError, match="activation_checkpointing"): + cfg._validate_distributed() + + def test_dcp_format_requires_sharding(self): + cfg = make_cfg(checkpoint_format=CheckpointFormat.DCP) + with pytest.raises(ValueError, match="sharded"): + cfg._validate_distributed() + cfg.parallelism = sharded() + cfg._validate_distributed() + + def test_fp16_rejected_when_sharded(self): + cfg = make_cfg(parallelism=sharded()) + cfg.accelerator.mixed_precision = "fp16" + with pytest.raises(ValueError, match="fp16"): + cfg._validate_distributed() + cfg.accelerator.mixed_precision = "bf16" + cfg._validate_distributed() + + def test_peft_rejected_when_sharded(self): + cfg = make_cfg(parallelism=sharded(), peft=PeftConfig()) + with pytest.raises(ValueError, match="PEFT"): + cfg._validate_distributed() + + def test_env_eval_rejected_when_sharded(self): + cfg = make_cfg(parallelism=sharded(), env_eval_freq=1000) + cfg.env = object() # any configured env triggers the check + with pytest.raises(ValueError, match="environment evaluation"): + cfg._validate_distributed() + + def test_multi_optimizer_rejected_when_sharded(self): + cfg = make_cfg(parallelism=sharded(), optimizer=MultiAdamConfig()) + with pytest.raises(ValueError, match="Multi-optimizer"): + cfg._validate_distributed() + cfg.optimizer = AdamConfig() + cfg._validate_distributed() + + +class TestConfigCompat: + def test_checkpoint_format_round_trip(self): + for fmt in CheckpointFormat: + assert draccus.decode(CheckpointFormat, draccus.encode(fmt)) is fmt + + def test_wants_predicates(self): + assert CheckpointFormat.SAFETENSORS.wants_safetensors + assert not CheckpointFormat.SAFETENSORS.wants_dcp + assert CheckpointFormat.DCP.wants_dcp and not CheckpointFormat.DCP.wants_safetensors + both = CheckpointFormat.SAFETENSORS_AND_DCP + assert both.wants_safetensors and both.wants_dcp + + +def test_reward_model_rejected_when_sharded(): + """Sharded reward runs previously failed late (missing wrap + units, DTensor serialization at the first checkpoint) instead of at validation.""" + cfg = make_cfg(parallelism=sharded()) + cfg.reward_model = object() # any configured reward model triggers the check + with pytest.raises(ValueError, match="Reward-model"): + cfg._validate_distributed() diff --git a/tests/distributed/test_accelerate_canaries.py b/tests/distributed/test_accelerate_canaries.py new file mode 100644 index 000000000..c4f6c9ce7 --- /dev/null +++ b/tests/distributed/test_accelerate_canaries.py @@ -0,0 +1,139 @@ +#!/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. +"""Version canaries for the accelerate/torch seams LeRobot's distributed engine relies on. + +LeRobot deliberately builds on a few accelerate internals that are not covered by a public +stability promise. These tests exist to fail LOUDLY on a +dependency upgrade — on a CPU runner, before any distributed job can be corrupted — whenever one +of those seams moves. If a canary fails, re-audit the corresponding integration seam before bumping +the pin; do not simply update the assertion. +""" + +import inspect + +import pytest + +pytest.importorskip("accelerate", reason="accelerate is required (install lerobot[training])") + + +def test_fsdp_checkpoint_name_constants(): + """Checkpoint dir names are imported from accelerate; the on-disk layout depends on them.""" + from accelerate.utils.constants import FSDP_MODEL_NAME, OPTIMIZER_NAME + + assert FSDP_MODEL_NAME == "pytorch_model_fsdp" + assert OPTIMIZER_NAME == "optimizer" + + +def test_parallelism_config_mesh_dim_contract(): + """FSDP2 shards over the flattened dp_shard_cp dim; the dataloader keys on exact root names.""" + from accelerate.parallelism_config import ParallelismConfig + + pc = ParallelismConfig(dp_replicate_size=2, dp_shard_size=2, cp_size=2) + assert pc.fsdp_dim_names == ["dp_replicate", "dp_shard_cp"] + assert pc.dp_shard_cp_dim_names == ["dp_shard", "cp"] + assert pc.dp_cp_dim_names == ["dp_replicate", "dp_shard", "cp"] + # Degenerate FSDP-only case still shards over the flattened name. + pc_fsdp = ParallelismConfig(dp_replicate_size=1, dp_shard_size=4) + assert pc_fsdp.fsdp_dim_names == ["dp_shard_cp"] + + +def test_accelerator_accepts_parallelism_config(): + from accelerate import Accelerator + + params = inspect.signature(Accelerator.__init__).parameters + assert "parallelism_config" in params + assert "fsdp_plugin" in params + assert "gradient_accumulation_plugin" in params + + +def test_dataloader_is_mesh_aware(): + """prepare_data_loader must accept the device mesh that makes CP peers share batches.""" + from accelerate.data_loader import prepare_data_loader + + assert "torch_device_mesh" in inspect.signature(prepare_data_loader).parameters + + +def test_cp_mask_stripping_hook_seam(): + """finalize_sharded_policy strips this exact hook. + + If accelerate renames or moves it, the strip becomes a silent no-op and CP training would + inherit mask-corrupting hooks — hence a canary rather than a runtime hasattr. + """ + from accelerate.big_modeling import _attach_context_parallel_hooks + + assert callable(_attach_context_parallel_hooks) + assert _attach_context_parallel_hooks.__module__ == "accelerate.big_modeling" + + +def test_fsdp_plugin_mirrored_fields_exist(): + """AcceleratorConfig mirrors a plain-typed subset of the plugin; the fields must survive.""" + from accelerate.utils import FullyShardedDataParallelPlugin + + fields = {f.name for f in FullyShardedDataParallelPlugin.__dataclass_fields__.values()} + assert { + "fsdp_version", + "reshard_after_forward", + "auto_wrap_policy", + "transformer_cls_names_to_wrap", + "min_num_params", + "cpu_offload", + "ignored_modules", + "activation_checkpointing", + "state_dict_type", + } <= fields + + +def test_merge_fsdp_weights_signature(): + """The DCP->safetensors converter is a thin wrapper over this accelerate utility.""" + from accelerate.utils import merge_fsdp_weights + + params = inspect.signature(merge_fsdp_weights).parameters + assert {"checkpoint_dir", "output_path", "safe_serialization"} <= set(params) + + +def test_fsdp_save_load_helpers_exist(): + from accelerate.utils import ( + load_fsdp_model, + load_fsdp_optimizer, + save_fsdp_model, + save_fsdp_optimizer, + ) + + for fn in (save_fsdp_model, load_fsdp_model, save_fsdp_optimizer, load_fsdp_optimizer): + assert callable(fn) + + +def test_torch_fsdp2_seams(): + """isinstance(FSDPModule) discrimination + non-forward entry registration + full gather.""" + from torch.distributed.checkpoint.state_dict import ( + StateDictOptions, + get_model_state_dict, # noqa: F401 + ) + from torch.distributed.fsdp import FSDPModule, register_fsdp_forward_method # noqa: F401 + + options = inspect.signature(StateDictOptions).parameters + assert {"full_state_dict", "cpu_offload"} <= set(options) + + +def test_accelerate_version_floor(): + import accelerate + from packaging import version + + if version.parse(accelerate.__version__) < version.parse("1.14.0"): + pytest.fail( + f"accelerate {accelerate.__version__} < 1.14.0: the FSDP2 auto-wrap fallback fix " + "(#3999) and the bf16->fp32 master-weight upcast this design relies on are absent." + ) diff --git a/tests/distributed/test_checkpoint_wrappers.py b/tests/distributed/test_checkpoint_wrappers.py new file mode 100644 index 000000000..19f474b8c --- /dev/null +++ b/tests/distributed/test_checkpoint_wrappers.py @@ -0,0 +1,70 @@ +#!/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. +"""The DCP wrappers must hand accelerate exact shard directories. + +accelerate 1.14 resolves the load directory with a substring check ("optimizer" / +"pytorch_model_fsdp" in the path -> use as-is) while the save side joins the shard name +unconditionally, so a run path like `--job_name=optimizer_sweep` would save to +`training_state/optimizer_0/` but load from `training_state/` itself. Passing the exact +shard dir makes the containment check deterministically a no-op. +""" + +from pathlib import Path +from types import SimpleNamespace + +import pytest + +pytest.importorskip("accelerate", reason="accelerate is required (install lerobot[training])") + + +def fake_accelerator() -> SimpleNamespace: + return SimpleNamespace(state=SimpleNamespace(fsdp_plugin=object())) + + +# A parent path that trips both of accelerate's substring checks at once. +POISONED_PARENT = Path("/outputs/train/optimizer_sweep_pytorch_model_fsdp_repro/training_state") + + +def test_load_sharded_optimizer_passes_exact_shard_dir(monkeypatch): + import accelerate.utils + + from lerobot.distributed.checkpoint import load_sharded_optimizer + + seen = {} + monkeypatch.setattr( + accelerate.utils, + "load_fsdp_optimizer", + lambda plugin, accelerator, optimizer, model, input_dir: seen.update(path=input_dir), + ) + load_sharded_optimizer(fake_accelerator(), optimizer=object(), model=object(), input_dir=POISONED_PARENT) + assert seen["path"] == str(POISONED_PARENT / "optimizer_0") + assert isinstance(seen["path"], str) # str, never Path (accelerate does string checks) + + +def test_load_sharded_model_passes_exact_shard_dir(monkeypatch): + import accelerate.utils + + from lerobot.distributed.checkpoint import load_sharded_model + + seen = {} + monkeypatch.setattr( + accelerate.utils, + "load_fsdp_model", + lambda plugin, accelerator, model, input_dir: seen.update(path=input_dir), + ) + load_sharded_model(fake_accelerator(), model=object(), input_dir=POISONED_PARENT) + assert seen["path"] == str(POISONED_PARENT / "pytorch_model_fsdp_0") + assert isinstance(seen["path"], str) diff --git a/tests/distributed/test_multigpu_training.py b/tests/distributed/test_multigpu_training.py new file mode 100644 index 000000000..b536ff2f1 --- /dev/null +++ b/tests/distributed/test_multigpu_training.py @@ -0,0 +1,486 @@ +#!/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. +"""End-to-end multi-GPU tests for the distributed core. + +Sized for a 4-GPU CI lane, these tests execute the sharded code paths nothing else in the tree can +reach — ``fully_shard`` via ``accelerator.prepare``, the DCP branches of ``save_checkpoint`` / +``save_training_state`` / ``resume_after_prepare``, the collective gather inside +``save_pretrained``, and HSDP/DDP gradient reduction — against the tiny +``DummyCheckpointPolicy`` fixture on synthetic data (no datasets, no network, no site paths). + +Run on a node with at least 4 GPUs:: + + pytest -m multigpu tests/distributed/test_multigpu_training.py -v + +Mechanics: + +- Plain pytest, no ``torchrun``: each test launches its own ranks with + ``torch.multiprocessing.spawn`` (spawn start method) and a per-test free TCP port; workers set + the torchrun-equivalent env (``RANK``/``LOCAL_RANK``/``WORLD_SIZE``/``MASTER_*``) that + accelerate's ``env://`` initialization consumes. +- Deadlock watchdog (:func:`_spawn`): the spawn context is polled with a deadline instead of a + blocking join, so a hung collective — the exact failure mode the all-ranks contracts guard + against — fails the test with ``TimeoutError`` (all workers SIGKILLed) rather than hanging CI. + A worker exception propagates through ``ProcessContext.join``, which tears down the survivors. +- Workers configure accelerate exclusively through the LeRobot config mirrors + (``AcceleratorConfig.build(ParallelismConfig)`` after ``resolve(world_size)``) — the same + construction path ``make_accelerator`` takes; see :func:`_build_accelerator` for why the + factory itself is not called. +- Without GPUs every test skips (``torch.cuda.device_count()`` gate), so the file is safe to + collect and run in the CPU lanes. +""" + +import json +import os +import socket +import time +from pathlib import Path + +import pytest +import torch +import torch.distributed as dist +import torch.multiprocessing as mp +from safetensors.torch import load_file + +from lerobot.common.train_utils import resume_after_prepare, resume_before_prepare, save_checkpoint +from lerobot.configs.default import DatasetConfig +from lerobot.configs.train import CheckpointFormat, TrainPipelineConfig +from lerobot.distributed.checkpoint import full_model_state_dict, is_sharded_module +from lerobot.utils.constants import PRETRAINED_MODEL_DIR, TRAINING_STATE_DIR + +# The spawned children re-import this module by name, so this import must resolve there too: +# torch.multiprocessing propagates the parent's sys.path through the spawn preparation data. +from tests.fixtures.dummy_checkpoint_policy import DummyCheckpointConfig, DummyCheckpointPolicy + +SEED = 20260712 +HIDDEN = 8 # DummyCheckpointPolicy is one Linear(hidden, hidden): 4 ranks shard dim 0 evenly +BATCH_SIZE = 2 +SAVE_STEP = 2 # optimizer steps run before saving in the round-trip workers +PARITY_STEPS = 3 +GA_UPDATES = 3 +SAMPLES_PER_UPDATE = 4 # per rank per optimizer update — the fixed effective batch of test 5 +GRAD_CLIP_NORM = 100.0 # generous: exercises the clip call without perturbing parity +# Generous headroom for cold NCCL init plus the lerobot re-import in 4 spawned children, while +# still bounding a deadlocked collective to minutes instead of a hung CI job. +WATCHDOG_TIMEOUT_S = 240.0 +_JOIN_POLL_S = 5.0 + + +# ------------------------------------------------------------------------------------------- +# Spawn infrastructure +# ------------------------------------------------------------------------------------------- + + +def _find_free_port() -> int: + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: + sock.bind(("127.0.0.1", 0)) + return sock.getsockname()[1] + + +def _spawn(world_size: int, worker, *args, timeout_s: float = WATCHDOG_TIMEOUT_S) -> None: + """Run ``worker(rank, world_size, port, *args)`` on ``world_size`` fresh processes. + + Watchdog approach: ``mp.spawn(join=False)`` returns a ``ProcessContext`` whose ``join`` is + polled under a deadline. On timeout every surviving worker is SIGKILLed and the test fails + with ``TimeoutError`` — a deadlock can never hang CI. When a worker raises, ``join`` itself + kills the remaining ranks and re-raises the worker's exception into the test. + """ + port = _find_free_port() + context = mp.spawn(worker, args=(world_size, port, *args), nprocs=world_size, join=False) + deadline = time.monotonic() + timeout_s + while not context.join(timeout=_JOIN_POLL_S): + if time.monotonic() >= deadline: + for process in context.processes: + if process.is_alive(): + process.kill() + for process in context.processes: + process.join(timeout=10) + raise TimeoutError( + f"{getattr(worker, '__name__', worker)}: {world_size} workers still running " + f"after {timeout_s}s — presumed deadlock; all workers killed." + ) + + +def _init_worker_env(rank: int, world_size: int, port: int) -> None: + """Give the worker the torchrun-equivalent env accelerate's ``env://`` init consumes.""" + # The tests configure accelerate through the config mirrors only; drop any accelerate env + # fallbacks inherited from the launching shell (what guard_against_env_interference would + # reject in production — here the env is simply owned by the test). + for name in list(os.environ): + if name.startswith(("FSDP_", "PARALLELISM_CONFIG_", "ACCELERATE_")): + del os.environ[name] + os.environ["MASTER_ADDR"] = "127.0.0.1" + os.environ["MASTER_PORT"] = str(port) + os.environ["RANK"] = str(rank) + os.environ["LOCAL_RANK"] = str(rank) + os.environ["WORLD_SIZE"] = str(world_size) + # The fp32 parity tolerances below assume true-fp32 matmuls. + torch.backends.cuda.matmul.allow_tf32 = False + torch.backends.cudnn.allow_tf32 = False + + +# ------------------------------------------------------------------------------------------- +# Shared building blocks +# ------------------------------------------------------------------------------------------- + + +def _make_cfg( + world_size: int, + *, + dp_replicate: int = 1, + dp_shard: int = 1, + checkpoint_format: CheckpointFormat = CheckpointFormat.SAFETENSORS, + grad_accum: int = 1, +) -> TrainPipelineConfig: + cfg = TrainPipelineConfig(dataset=DatasetConfig(repo_id="lerobot/dummy"), batch_size=BATCH_SIZE) + cfg.checkpoint_format = checkpoint_format + cfg.parallelism.dp_replicate = dp_replicate + cfg.parallelism.dp_shard = dp_shard + cfg.accelerator.mixed_precision = "no" # fp32 end to end: the parity tests depend on it + cfg.accelerator.gradient_accumulation.steps = grad_accum + # The dummy policy declares no _fsdp_wrap_modules; the size-based wrap policy shards its + # Linear without needing class names (the set_fsdp_wrap_modules no-op branch). + cfg.accelerator.fsdp.min_num_params = 1 + cfg.parallelism.resolve(world_size) + return cfg + + +def _build_accelerator(cfg: TrainPipelineConfig): + """``cfg.accelerator.build(cfg.parallelism)`` — make_accelerator's construction path. + + Deliberately not ``make_accelerator`` itself: the factory additionally derives ``cpu=`` from + ``cfg.trainable_config`` (no policy config is attached to these synthetic cfgs) and re-runs + the env guard — both owned explicitly by the tests (see ``_init_worker_env``). + """ + return cfg.accelerator.build(cfg.parallelism) + + +def _make_policy(seed: int) -> DummyCheckpointPolicy: + """Identically seeded on every rank, so shard/replicate starts from one common init.""" + torch.manual_seed(seed) + return DummyCheckpointPolicy(DummyCheckpointConfig(hidden=HIDDEN, device="cpu")) + + +def _batch(step: int, rank: int, device: torch.device) -> dict[str, torch.Tensor]: + """Deterministic per-(step, rank) batch: every dp worker sees distinct, reproducible data.""" + generator = torch.Generator().manual_seed(SEED + 1000 * step + rank) + return {"observation.state": torch.randn(BATCH_SIZE, HIDDEN, generator=generator).to(device)} + + +def _gather_full(model, optimizer) -> tuple[dict, dict]: + """Full (unsharded) model + optimizer state via torch's DCP state-dict API — a COLLECTIVE. + + With ``cpu_offload=True`` the dicts materialize on the main rank only; every other rank + receives a literal ``{}``. + """ + from torch.distributed.checkpoint.state_dict import ( + StateDictOptions, + get_model_state_dict, + get_optimizer_state_dict, + ) + + options = StateDictOptions(full_state_dict=True, cpu_offload=True) + return ( + get_model_state_dict(model, options=options), + get_optimizer_state_dict(model, optimizer, options=options), + ) + + +def _assert_tree_equal(reference, actual, path: str) -> None: + """Exact (bitwise for tensors) equality of nested state dicts, with a failing path.""" + if isinstance(reference, torch.Tensor): + assert isinstance(actual, torch.Tensor), f"{path}: {type(actual)} is not a tensor" + assert reference.dtype == actual.dtype, f"{path}: {reference.dtype} != {actual.dtype}" + assert reference.shape == actual.shape, f"{path}: {reference.shape} != {actual.shape}" + assert torch.equal(reference.cpu(), actual.cpu()), f"{path}: tensor values differ" + elif isinstance(reference, dict): + assert isinstance(actual, dict), f"{path}: {type(actual)} is not a dict" + assert set(reference) == set(actual), f"{path}: keys {set(reference) ^ set(actual)} differ" + for key in reference: + _assert_tree_equal(reference[key], actual[key], f"{path}.{key}") + elif isinstance(reference, list | tuple): + assert type(reference) is type(actual) and len(reference) == len(actual), path + for index, (ref_item, actual_item) in enumerate(zip(reference, actual, strict=True)): + _assert_tree_equal(ref_item, actual_item, f"{path}[{index}]") + else: + assert reference == actual, f"{path}: {reference!r} != {actual!r}" + + +# ------------------------------------------------------------------------------------------- +# Workers (module-level: torch.multiprocessing.spawn pickles them by reference) +# ------------------------------------------------------------------------------------------- + + +def _train_and_save_worker(rank: int, world_size: int, port: int, tmp_dir: str, fmt_value: str) -> None: + """FSDP2 (dp_shard=world_size): train SAVE_STEP steps, save_checkpoint, store the gathered + full model/optimizer state as the rank-0 reference for the resume workers.""" + _init_worker_env(rank, world_size, port) + tmp = Path(tmp_dir) + fmt = CheckpointFormat(fmt_value) + cfg = _make_cfg(world_size, dp_shard=world_size, checkpoint_format=fmt) + accelerator = _build_accelerator(cfg) + policy = _make_policy(SEED) + optimizer = torch.optim.Adam(policy.parameters(), lr=1e-2) + # FSDP2 requires model and optimizer in one prepare() call (accelerate rebinds param groups). + policy, optimizer = accelerator.prepare(policy, optimizer) + assert is_sharded_module(accelerator.unwrap_model(policy)), "prepare() did not shard the policy" + + for step in range(SAVE_STEP): + loss, _ = policy(_batch(step, rank, accelerator.device)) + accelerator.backward(loss) + optimizer.step() + optimizer.zero_grad() + + checkpoint_dir = tmp / "checkpoint" + save_checkpoint( + checkpoint_dir, step=SAVE_STEP, cfg=cfg, policy=policy, optimizer=optimizer, accelerator=accelerator + ) + + model_state, optimizer_state = _gather_full(policy, optimizer) + if accelerator.is_main_process: + from accelerate.utils.constants import FSDP_MODEL_NAME, OPTIMIZER_NAME + + pretrained_dir = checkpoint_dir / PRETRAINED_MODEL_DIR + assert (pretrained_dir / f"{FSDP_MODEL_NAME}_0").is_dir() == fmt.wants_dcp + assert (pretrained_dir / "model.safetensors").is_file() == fmt.wants_safetensors + assert (pretrained_dir / "config.json").is_file() + assert (pretrained_dir / "train_config.json").is_file() + # Sharded runs always use the DCP optimizer channel, never the safetensors one. + assert (checkpoint_dir / TRAINING_STATE_DIR / f"{OPTIMIZER_NAME}_0").is_dir() + assert not (checkpoint_dir / TRAINING_STATE_DIR / "optimizer_state.safetensors").exists() + torch.save({"model": model_state, "optimizer": optimizer_state}, tmp / "reference_state.pt") + accelerator.wait_for_everyone() + dist.destroy_process_group() + + +def _resume_and_verify_worker(rank: int, world_size: int, port: int, tmp_dir: str, fmt_value: str) -> None: + """Two-phase resume at dp_shard=world_size; the gathered state must match the saved + reference exactly (DCP round-trips are bit-exact).""" + _init_worker_env(rank, world_size, port) + tmp = Path(tmp_dir) + cfg = _make_cfg(world_size, dp_shard=world_size, checkpoint_format=CheckpointFormat(fmt_value)) + cfg.checkpoint_path = tmp / "checkpoint" + accelerator = _build_accelerator(cfg) + + assert resume_before_prepare(cfg) == SAVE_STEP # phase 1: RNG + step counter only + + # Deliberately different init: the DCP load must overwrite every parameter. + policy = _make_policy(SEED + 1) + optimizer = torch.optim.Adam(policy.parameters(), lr=1e-2) + policy, optimizer = accelerator.prepare(policy, optimizer) + resume_after_prepare(cfg, accelerator, policy, optimizer, None) # phase 2: DCP reshard-load + + model_state, optimizer_state = _gather_full(policy, optimizer) + if accelerator.is_main_process: + reference = torch.load(tmp / "reference_state.pt", map_location="cpu", weights_only=True) + _assert_tree_equal(reference["model"], model_state, "model") + _assert_tree_equal(reference["optimizer"], optimizer_state, "optimizer") + accelerator.wait_for_everyone() + dist.destroy_process_group() + + +def _loss_parity_worker( + rank: int, world_size: int, port: int, tmp_dir: str, dp_replicate: int, dp_shard: int, tag: str +) -> None: + """Train PARITY_STEPS fp32 steps on per-rank deterministic data; rank 0 records the + dp-mean loss of every step. Gradient averaging spans the same rank set in any (R, S) + factorization of the world, so the loss trajectory is topology-invariant.""" + _init_worker_env(rank, world_size, port) + cfg = _make_cfg(world_size, dp_replicate=dp_replicate, dp_shard=dp_shard) + accelerator = _build_accelerator(cfg) + policy = _make_policy(SEED) + optimizer = torch.optim.SGD(policy.parameters(), lr=0.05) + policy, optimizer = accelerator.prepare(policy, optimizer) + assert is_sharded_module(accelerator.unwrap_model(policy)) == (dp_shard > 1) + + per_step_losses = [] + for step in range(PARITY_STEPS): + loss, _ = policy(_batch(step, rank, accelerator.device)) + per_step_losses.append(accelerator.gather(loss.detach().reshape(1)).double().mean().item()) + accelerator.backward(loss) + optimizer.step() + optimizer.zero_grad() + + if accelerator.is_main_process: + (Path(tmp_dir) / f"losses_{tag}.json").write_text(json.dumps(per_step_losses)) + accelerator.wait_for_everyone() + dist.destroy_process_group() + + +def _save_pretrained_all_ranks_worker(rank: int, world_size: int, port: int, tmp_dir: str) -> None: + """The all-ranks contract: every rank calls save_pretrained, the + collective gather completes (watchdog proves no deadlock), and only rank 0 writes files.""" + _init_worker_env(rank, world_size, port) + cfg = _make_cfg(world_size, dp_shard=world_size) + accelerator = _build_accelerator(cfg) + policy = _make_policy(SEED) + # FSDP2 prepare requires an optimizer alongside the model even though this test never steps it. + optimizer = torch.optim.SGD(policy.parameters(), lr=0.1) + policy, optimizer = accelerator.prepare(policy, optimizer) + unwrapped = accelerator.unwrap_model(policy) + assert is_sharded_module(unwrapped) + + # Gather semantics: the full dict materializes on the main rank; every other rank + # receives the literal empty dict. + reference = full_model_state_dict(unwrapped) + if accelerator.is_main_process: + assert set(reference) == {"net.weight", "net.bias"} + else: + assert reference == {} + + # Every rank targets its own directory so writes are attributable per rank. + target = Path(tmp_dir) / f"rank_{rank}" + unwrapped.save_pretrained(target) + accelerator.wait_for_everyone() + + if accelerator.is_main_process: + weights = load_file(target / "model.safetensors") + assert set(weights) == set(reference) + for key, tensor in reference.items(): + assert torch.equal(weights[key], tensor), key + assert (target / "config.json").is_file() + else: + assert list(target.rglob("*")) == [], f"rank {rank} wrote files despite the rank-0 gate" + dist.destroy_process_group() + + +def _grad_accum_worker( + rank: int, world_size: int, port: int, tmp_dir: str, micro_batch_size: int, grad_accum: int, tag: str +) -> None: + """DDP fp32 with the exact accumulate/clip/step/zero_grad pattern of + ``lerobot_train.update_policy``; rank 0 records the final weights.""" + _init_worker_env(rank, world_size, port) + assert micro_batch_size * grad_accum == SAMPLES_PER_UPDATE # fixed effective batch + cfg = _make_cfg(world_size, dp_replicate=world_size, grad_accum=grad_accum) + accelerator = _build_accelerator(cfg) + # The GradientAccumulationPlugin wiring, un-overridden by any env fallback. + assert accelerator.gradient_accumulation_steps == grad_accum + policy = _make_policy(SEED) + optimizer = torch.optim.SGD(policy.parameters(), lr=0.05) + policy, optimizer = accelerator.prepare(policy, optimizer) + + # One fixed per-rank sample stream, consumed in order by both variants: update k always + # covers rows [k * SAMPLES_PER_UPDATE, (k + 1) * SAMPLES_PER_UPDATE). + generator = torch.Generator().manual_seed(SEED + 7919 * rank) + stream = torch.randn(GA_UPDATES * SAMPLES_PER_UPDATE, HIDDEN, generator=generator) + + updates_applied = 0 + for micro_step in range(GA_UPDATES * grad_accum): + rows = stream[micro_step * micro_batch_size : (micro_step + 1) * micro_batch_size] + batch = {"observation.state": rows.to(accelerator.device)} + # update_policy's pattern: accumulate() suppresses grad sync and rescales the loss on + # non-final micro-batches, and AcceleratedOptimizer makes step()/zero_grad() no-ops + # until sync_gradients is True. + with accelerator.accumulate(policy): + loss, _ = policy(batch) + accelerator.backward(loss) + if accelerator.sync_gradients: + accelerator.clip_grad_norm_(policy.parameters(), GRAD_CLIP_NORM) + updates_applied += 1 + optimizer.step() + optimizer.zero_grad() + assert updates_applied == GA_UPDATES # exactly one optimizer update per accumulation window + + if accelerator.is_main_process: + state = {key: value.cpu() for key, value in accelerator.unwrap_model(policy).state_dict().items()} + torch.save(state, Path(tmp_dir) / f"weights_{tag}.pt") + accelerator.wait_for_everyone() + dist.destroy_process_group() + + +# ------------------------------------------------------------------------------------------- +# Tests +# ------------------------------------------------------------------------------------------- + + +@pytest.mark.multigpu +@pytest.mark.skipif(torch.cuda.device_count() < 4, reason="requires 4 GPUs") +def test_fsdp2_train_save_resume_round_trip(tmp_path): + """FSDP2 dp_shard=4, checkpoint_format=safetensors_dcp: train -> save_checkpoint -> resume. + + A second spawn resumes through the two-phase path and its gathered model weights and Adam + state tensors must match the pre-save gathered reference exactly (DCP round-trips are + bit-exact). + """ + fmt = CheckpointFormat.SAFETENSORS_AND_DCP.value + _spawn(4, _train_and_save_worker, str(tmp_path), fmt) + _spawn(4, _resume_and_verify_worker, str(tmp_path), fmt) + + +@pytest.mark.multigpu +@pytest.mark.skipif(torch.cuda.device_count() < 4, reason="requires 4 GPUs") +def test_hsdp_loss_parity_with_ddp(tmp_path): + """Same seed and per-rank data: DDP (dp_replicate=4) vs HSDP (2x2), fp32, no AMP. + + Both topologies average gradients over the same four ranks, so per-step dp-mean losses must + match within tolerance. Exact parity is not expected: DDP all-reduces where HSDP + reduce-scatters within the shard group and all-reduces across replicas, and the different + reduction orders accumulate fp32 rounding — rtol=1e-4 leaves orders of magnitude of headroom + over that noise while still catching any real divergence (wrong averaging, wrong data). + """ + _spawn(4, _loss_parity_worker, str(tmp_path), 4, 1, "ddp") + _spawn(4, _loss_parity_worker, str(tmp_path), 2, 2, "hsdp") + ddp_losses = json.loads((tmp_path / "losses_ddp.json").read_text()) + hsdp_losses = json.loads((tmp_path / "losses_hsdp.json").read_text()) + assert len(ddp_losses) == len(hsdp_losses) == PARITY_STEPS + for step, (ddp_loss, hsdp_loss) in enumerate(zip(ddp_losses, hsdp_losses, strict=True)): + assert hsdp_loss == pytest.approx(ddp_loss, rel=1e-4, abs=1e-6), f"step {step}" + + +@pytest.mark.multigpu +@pytest.mark.skipif(torch.cuda.device_count() < 4, reason="requires 4 GPUs") +def test_changed_topology_resume(tmp_path): + """Save at dp_shard=4 (format=dcp), resume at dp_shard=2 on 2 ranks. + + The DCP load reshards both the model weights and the optimizer state across the topology + change; the post-resume gathered state must equal the pre-save gathered reference exactly + (cross-topology resharding is runtime-verified). + """ + fmt = CheckpointFormat.DCP.value + _spawn(4, _train_and_save_worker, str(tmp_path), fmt) + _spawn(2, _resume_and_verify_worker, str(tmp_path), fmt) + + +@pytest.mark.multigpu +@pytest.mark.skipif(torch.cuda.device_count() < 4, reason="requires 4 GPUs") +def test_save_pretrained_all_ranks_no_deadlock(tmp_path): + """dp_shard=4: save_pretrained on ALL ranks completes under the watchdog. + + Rank 0 writes model.safetensors (+ config.json) whose tensors equal the gathered full state; + ranks 1-3 write nothing. A rank-gated call would deadlock in the collective gather and be + killed by :func:`_spawn`'s timeout — completing at all is half of what this test asserts. + """ + _spawn(4, _save_pretrained_all_ranks_worker, str(tmp_path)) + + +@pytest.mark.multigpu +@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="requires 2 GPUs") +def test_gradient_accumulation_equivalence(tmp_path): + """Fixed effective batch on 2-rank DDP fp32: (batch=4, GA=1) vs (batch=2, GA=2). + + Both variants consume the identical per-rank sample stream in the same order for + GA_UPDATES optimizer updates, using update_policy's accumulate/clip/step pattern. The final + weights must agree: accumulate() rescales each micro-loss by 1/GA, so summed mean-of-2 + gradients equal the mean-of-4 gradient up to fp32 summation order — hence allclose with + rtol=1e-5/atol=1e-6 (roughly 100x the observed associativity noise), not bitwise equality. + """ + _spawn(2, _grad_accum_worker, str(tmp_path), 4, 1, "ga1") + _spawn(2, _grad_accum_worker, str(tmp_path), 2, 2, "ga2") + ga1 = torch.load(tmp_path / "weights_ga1.pt", weights_only=True) + ga2 = torch.load(tmp_path / "weights_ga2.pt", weights_only=True) + assert set(ga1) == set(ga2) == {"net.weight", "net.bias"} + for key in ga1: + assert torch.allclose(ga1[key], ga2[key], rtol=1e-5, atol=1e-6), key diff --git a/tests/distributed/test_parallel_dims_and_factory.py b/tests/distributed/test_parallel_dims_and_factory.py new file mode 100644 index 000000000..f63e03fdb --- /dev/null +++ b/tests/distributed/test_parallel_dims_and_factory.py @@ -0,0 +1,124 @@ +#!/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. + +import pytest + +from lerobot.configs.parallelism import ContextParallelConfig, ParallelismConfig +from lerobot.distributed import ParallelDims, guard_against_env_interference, is_main_process +from lerobot.distributed.factory import _ENV_OVERRIDE + + +class TestIsMainProcess: + def test_true_outside_distributed(self): + assert is_main_process() is True + + +class TestParallelDims: + def _resolved(self, world_size: int = 8, **kwargs) -> ParallelismConfig: + cfg = ParallelismConfig(**kwargs) + cfg.resolve(world_size) + return cfg + + def test_from_resolved_config(self): + dims = ParallelDims.from_config(self._resolved(dp_replicate=2, dp_shard=4), 8, "cpu") + assert dims.dp_world_size == 8 + assert dims.is_sharded + assert dims.cp_size == 1 + assert dims.dp_rank == 0 # no process group in unit tests + + def test_rejects_unresolved_config(self): + with pytest.raises(ValueError, match="resolve"): + ParallelDims.from_config(ParallelismConfig(dp_shard=-1), 8, "cpu") + + def test_rejects_world_mismatch(self): + with pytest.raises(ValueError, match="world_size=4"): + ParallelDims.from_config(self._resolved(8), 4, "cpu") + + def test_cp_mesh_reserved(self): + dims = ParallelDims(dp_replicate=1, dp_shard=2, ring=2, ulysses=2, world_size=8, device_type="cpu") + assert dims.dp_rank == 0 and dims.dp_world_size == 2 + with pytest.raises(NotImplementedError): + dims.cp_mesh() + + def test_cp_peers_share_dp_rank_arithmetic(self): + """Row-major layout: cp is innermost, so dp_rank = global_rank // cp_size.""" + dims = ParallelDims(dp_replicate=1, dp_shard=2, ring=1, ulysses=2, world_size=4, device_type="cpu") + # Without a process group the global rank is 0; the arithmetic contract is what matters. + assert dims.cp_size == 2 + assert dims.dp_rank == 0 // dims.cp_size + + def test_config_placeholder_degrees_flow_through(self): + cfg = ParallelismConfig( + dp_replicate=1, + dp_shard=2, + context_parallel=ContextParallelConfig(ring_degree=2, ulysses_degree=2), + ) + # resolve() rejects cp>1 this round; ParallelDims math itself is already cp-aware. + dims = ParallelDims( + dp_replicate=cfg.dp_replicate, + dp_shard=cfg.dp_shard, + ring=cfg.context_parallel.ring_degree, + ulysses=cfg.context_parallel.ulysses_degree, + world_size=8, + device_type="cpu", + ) + assert dims.dp_world_size == 2 and dims.cp_size == 4 + + +class TestEnvGuard: + # ACCELERATE_DYNAMO_*/ACCELERATE_GRADIENT_ACCUMULATION_STEPS are silent config overrides + # inside accelerate itself — the guard must catch them too. + _POISON = ( + "ACCELERATE_USE_FSDP", + "FSDP_VERSION", + "PARALLELISM_CONFIG_DP_SHARD_SIZE", + "ACCELERATE_DYNAMO_BACKEND", + "ACCELERATE_GRADIENT_ACCUMULATION_STEPS", + ) + + def test_clean_env_passes(self, monkeypatch): + for name in self._POISON + (_ENV_OVERRIDE,): + monkeypatch.delenv(name, raising=False) + guard_against_env_interference() + + @pytest.mark.parametrize("name", _POISON) + def test_accelerate_env_rejected_with_actionable_error(self, name, monkeypatch): + monkeypatch.delenv(_ENV_OVERRIDE, raising=False) + monkeypatch.setenv(name, "true") + with pytest.raises(RuntimeError, match=name): + guard_against_env_interference() + + def test_override_acknowledges(self, monkeypatch): + monkeypatch.setenv("FSDP_VERSION", "2") + monkeypatch.setenv(_ENV_OVERRIDE, "1") + guard_against_env_interference() + + +def test_make_accelerator_rejects_format_after_sentinel_resolution(monkeypatch): + """dp_shard=-1 counts as sharded at parse time but can resolve + to an unsharded run (world size 1), which would write a safetensors-only checkpoint whose + recorded checkpoint_format=dcp fails its own validation on resume.""" + from lerobot.configs.default import DatasetConfig + from lerobot.configs.train import CheckpointFormat, TrainPipelineConfig + from lerobot.distributed.factory import make_accelerator + + monkeypatch.delenv("WORLD_SIZE", raising=False) + cfg = TrainPipelineConfig(dataset=DatasetConfig(repo_id="lerobot/dummy")) + cfg.parallelism.dp_shard = -1 + cfg.checkpoint_format = CheckpointFormat.DCP + cfg._validate_distributed() # passes: the sentinel is declared as sharded + with pytest.raises(ValueError, match="resolved to a non-sharded"): + make_accelerator(cfg) diff --git a/tests/distributed/test_policy_surface.py b/tests/distributed/test_policy_surface.py new file mode 100644 index 000000000..6f20d8712 --- /dev/null +++ b/tests/distributed/test_policy_surface.py @@ -0,0 +1,126 @@ +#!/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. +"""The declarative policy surface and its distributed-side consumers.""" + +from types import SimpleNamespace + +import pytest +import torch +from torch import nn + +from lerobot.configs.accelerator import FSDPConfig +from lerobot.distributed import set_fsdp_wrap_modules, strip_accelerate_cp_hooks +from lerobot.policies.pretrained import PreTrainedPolicy + + +class TestDeclarativeAttributes: + def test_base_defaults(self): + assert PreTrainedPolicy._fsdp_wrap_modules is None + assert PreTrainedPolicy._fsdp_forward_methods == ("select_action", "predict_action_chunk") + assert PreTrainedPolicy.supports_gradient_checkpointing is False + assert PreTrainedPolicy._cp_plan is None + + def test_act_wrap_units_name_real_classes(self): + """The declared class names must track the modeling code — this test pins the drift.""" + from lerobot.policies.act import modeling_act + + for name in modeling_act.ACTPolicy._fsdp_wrap_modules: + assert isinstance(getattr(modeling_act, name), type), name + + def test_fastwam_wrap_units_name_real_classes(self): + from lerobot.policies.fastwam import modeling_fastwam + from lerobot.policies.fastwam.wan import modular + + for name in modeling_fastwam.FastWAMPolicy._fsdp_wrap_modules: + assert isinstance(getattr(modular, name), type), name + + +class _SelfAttn(nn.Module): + def forward(self, x, attention_mask=None, is_causal=False): + return x, attention_mask, is_causal + + +class _TinyModel(nn.Module): + def __init__(self): + super().__init__() + self.self_attn = _SelfAttn() + + +class TestStripAccelerateCpHooks: + def test_strips_the_real_accelerate_hook_and_restores_mask_semantics(self): + """Attach accelerate's actual mask-stripping hook, strip it, verify masks survive.""" + pytest.importorskip("accelerate", reason="accelerate is required (install lerobot[training])") + from accelerate.big_modeling import _attach_context_parallel_hooks + + model = _TinyModel() + mask = torch.ones(2, 2) + + _attach_context_parallel_hooks(model) + _, hooked_mask, hooked_causal = model.self_attn(torch.zeros(1), attention_mask=mask) + assert hooked_mask is None and hooked_causal is True # the hazard is real + + assert strip_accelerate_cp_hooks(model) == 1 + _, clean_mask, clean_causal = model.self_attn(torch.zeros(1), attention_mask=mask) + assert clean_mask is mask and clean_causal is False + assert not model.self_attn._forward_pre_hooks + assert not model.self_attn._forward_pre_hooks_with_kwargs + + def test_user_hooks_survive(self): + model = _TinyModel() + model.self_attn.register_forward_pre_hook(lambda m, args: None) + assert strip_accelerate_cp_hooks(model) == 0 + assert len(model.self_attn._forward_pre_hooks) == 1 + + +class _DeclaredPolicy: + _fsdp_wrap_modules = ["DeclaredBlock"] + + +class _UndeclaredPolicy: + _fsdp_wrap_modules = None + + +def _accelerator_with(plugin) -> SimpleNamespace: + return SimpleNamespace(state=SimpleNamespace(fsdp_plugin=plugin)) + + +class TestSetFsdpWrapModules: + @pytest.fixture(autouse=True) + def _requires_accelerate(self): + pytest.importorskip("accelerate", reason="accelerate is required (install lerobot[training])") + + def test_policy_declaration_fills_plugin(self): + plugin = FSDPConfig().build_plugin() + set_fsdp_wrap_modules(_accelerator_with(plugin), _DeclaredPolicy()) + assert plugin.transformer_cls_names_to_wrap == ["DeclaredBlock"] + + def test_user_override_wins(self): + plugin = FSDPConfig(wrap_modules=["UserBlock"]).build_plugin() + set_fsdp_wrap_modules(_accelerator_with(plugin), _DeclaredPolicy()) + assert plugin.transformer_cls_names_to_wrap == ["UserBlock"] + + def test_no_wrap_source_fails_loudly(self): + plugin = FSDPConfig().build_plugin() + with pytest.raises(ValueError, match="_fsdp_wrap_modules"): + set_fsdp_wrap_modules(_accelerator_with(plugin), _UndeclaredPolicy()) + + def test_size_based_policy_needs_no_names(self): + plugin = FSDPConfig(min_num_params=1024).build_plugin() + set_fsdp_wrap_modules(_accelerator_with(plugin), _UndeclaredPolicy()) + assert plugin.transformer_cls_names_to_wrap is None + + def test_non_sharded_run_is_noop(self): + set_fsdp_wrap_modules(_accelerator_with(None), _UndeclaredPolicy()) diff --git a/tests/fixtures/dummy_checkpoint_policy.py b/tests/fixtures/dummy_checkpoint_policy.py new file mode 100644 index 000000000..f6825ca85 --- /dev/null +++ b/tests/fixtures/dummy_checkpoint_policy.py @@ -0,0 +1,87 @@ +#!/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 diff --git a/tests/optim/test_optimizers.py b/tests/optim/test_optimizers.py index 5b480f70d..d18565562 100644 --- a/tests/optim/test_optimizers.py +++ b/tests/optim/test_optimizers.py @@ -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 { diff --git a/tests/policies/test_policies.py b/tests/policies/test_policies.py index 285b87d4c..e6e41ac1f 100644 --- a/tests/policies/test_policies.py +++ b/tests/policies/test_policies.py @@ -301,8 +301,12 @@ 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.""" +def test_save_pretrained_single_file_artifact(dummy_dataset_metadata, tmp_path): + """The distributable checkpoint is one unsharded safetensors file. + + The former `state_dict=` variant of this test died with the #3810 save override: the + kwarg would now be silently swallowed by HubMixin's **push_to_hub_kwargs. + """ policy_cls = get_policy_class("act") policy_cfg = make_policy_config("act") features = dataset_to_policy_features(dummy_dataset_metadata.features) @@ -313,8 +317,8 @@ def test_save_pretrained_with_state_dict(dummy_dataset_metadata, tmp_path): 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()) + save_dir = tmp_path / "single_file_artifact" + policy.save_pretrained(save_dir) # A single, unsharded safetensors file (no sharded set + index). assert (save_dir / SAFETENSORS_SINGLE_FILE).is_file() diff --git a/tests/rewards/test_reward_model_base.py b/tests/rewards/test_reward_model_base.py index 1c4dad642..05a28f39e 100644 --- a/tests/rewards/test_reward_model_base.py +++ b/tests/rewards/test_reward_model_base.py @@ -22,6 +22,7 @@ from types import SimpleNamespace import pytest import torch +from lerobot.common.train_utils import generate_model_card from lerobot.configs.rewards import RewardModelConfig from lerobot.optim.optimizers import AdamWConfig from lerobot.rewards.pretrained import PreTrainedRewardModel @@ -326,7 +327,7 @@ def test_train_pipeline_config_from_pretrained_strips_legacy_rabc_when_disabled( # --------------------------------------------------------------------------- -# PreTrainedRewardModel hub upload: push_model_to_hub + generate_model_card. +# PreTrainedRewardModel hub upload: publish_trained_model + generate_model_card. # We test the generation side (offline) fully, and the upload side with HfApi # mocked so nothing actually hits the network. # --------------------------------------------------------------------------- @@ -336,6 +337,13 @@ def _make_dummy_reward_model(**config_kwargs): return _DummyHubReward(_DummyHubRewardConfig(**config_kwargs)), _DummyHubRewardConfig +def _make_train_cfg(dataset_repo_id: str): + from lerobot.configs.default import DatasetConfig + from lerobot.configs.train import TrainPipelineConfig + + return TrainPipelineConfig(dataset=DatasetConfig(repo_id=dataset_repo_id)) + + @pytest.fixture def _offline_model_card(monkeypatch): """``ModelCard.validate`` does a live ``POST`` to huggingface.co — bypass it @@ -353,12 +361,7 @@ def test_reward_model_generate_model_card_renders_expected_fields(_offline_model tags=["robot", "sim"], ) - card = model.generate_model_card( - dataset_repo_id="user/my_dataset", - model_type=model.config.type, - license=model.config.license, - tags=model.config.tags, - ) + card = generate_model_card(model.config, cfg=_make_train_cfg("user/my_dataset")) # Metadata (YAML header) — ModelCardData fields. assert card.data.license == "mit" @@ -380,21 +383,16 @@ def test_reward_model_generate_model_card_uses_default_license(_offline_model_ca """When config.license is None the card falls back to apache-2.0.""" model, _ = _make_dummy_reward_model() - card = model.generate_model_card( - dataset_repo_id="user/my_dataset", - model_type=model.config.type, - license=model.config.license, - tags=None, - ) + card = generate_model_card(model.config, cfg=_make_train_cfg("user/my_dataset")) assert card.data.license == "apache-2.0" -def test_reward_model_push_model_to_hub_uploads_expected_files(monkeypatch, _offline_model_card): - """``push_model_to_hub`` must: +def test_publish_trained_model_uploads_expected_reward_files(monkeypatch, _offline_model_card): + """Publishing a reward model through ``publish_trained_model`` must: 1. create the repo, - 2. assemble a temp folder with weights + config.json + train_config.json + README.md, - 3. call ``api.upload_folder`` on that folder. + 2. push the model through ``HubMixin.push_to_hub`` (weights + config.json), + 3. upload a bundle sidecar with train_config.json + the reward-specific README.md. All network calls are mocked. """ from huggingface_hub.constants import CONFIG_NAME @@ -430,18 +428,80 @@ def test_reward_model_push_model_to_hub_uploads_expected_files(monkeypatch, _off uploaded["files"] = sorted(p.name for p in Path(folder_path).iterdir()) return fake_commit_info - from lerobot.rewards import pretrained as reward_pretrained + import lerobot.common.train_utils as train_utils + import lerobot.utils.hub as hub_module + from lerobot.common.train_utils import publish_trained_model - monkeypatch.setattr(reward_pretrained, "HfApi", lambda *a, **kw: _FakeHfApi()) + all_files: set[str] = set() - model.push_model_to_hub(train_cfg) + class _RecordingFakeHfApi(_FakeHfApi): + def __init__(self, *args, **kwargs): + pass + + def upload_folder(self, *, repo_id, repo_type, folder_path, commit_message, **_kwargs): + result = super().upload_folder( + repo_id=repo_id, + repo_type=repo_type, + folder_path=folder_path, + commit_message=commit_message, + ) + all_files.update(uploaded["files"]) + return result + + monkeypatch.setattr(train_utils, "HfApi", _RecordingFakeHfApi) + monkeypatch.setattr(hub_module, "HfApi", _RecordingFakeHfApi) + + publish_trained_model(train_cfg, model, None, None, dataset_meta=None) assert uploaded["create_repo_id"] == "user/my_reward" assert uploaded["upload_repo_id"] == "user/my_reward" assert uploaded["upload_repo_type"] == "model" - assert uploaded["commit_message"] == "Upload reward model weights, train config and readme" - # Minimum required files that must be uploaded with a reward model. - assert CONFIG_NAME in uploaded["files"] # config.json - assert TRAIN_CONFIG_NAME in uploaded["files"] # train_config.json - assert "README.md" in uploaded["files"] - assert any(name.endswith(".safetensors") for name in uploaded["files"]) + # Minimum required files across the publish commits. + assert CONFIG_NAME in all_files # config.json (model commit) + assert TRAIN_CONFIG_NAME in all_files # train_config.json (bundle commit) + assert "README.md" in all_files # reward-specific card (bundle commit) + assert any(name.endswith(".safetensors") for name in all_files) # weights (model commit) + + +def test_save_pretrained_writes_nothing_off_main_rank(tmp_path, monkeypatch): + """save_checkpoint calls save_pretrained on every rank; the + reward serializer must gate its writes so DDP replicas do not race on the same files.""" + import lerobot.distributed.utils as dist_utils + + model, _ = _make_dummy_reward_model() + monkeypatch.setattr(dist_utils, "is_main_process", lambda: False) + model.save_pretrained(tmp_path) + assert not any(tmp_path.iterdir()) + + +def test_reward_model_push_model_to_hub_shim_warns_and_publishes(monkeypatch, _offline_model_card): + """The deprecated ``push_model_to_hub`` stays callable, delegating to the publisher.""" + from huggingface_hub.constants import CONFIG_NAME + + import lerobot.common.train_utils as train_utils + import lerobot.utils.hub as hub_module + from lerobot.configs.train import TRAIN_CONFIG_NAME + + all_files: set[str] = set() + + class _FakeHfApi: + def __init__(self, *args, **kwargs): + pass + + def create_repo(self, repo_id, private=None, exist_ok=False, **kwargs): + return SimpleNamespace(repo_id=repo_id) + + def upload_folder(self, *, repo_id, folder_path, **_kwargs): + all_files.update(p.name for p in Path(folder_path).iterdir()) + return SimpleNamespace(repo_url=SimpleNamespace(url=f"https://huggingface.co/{repo_id}")) + + monkeypatch.setattr(train_utils, "HfApi", _FakeHfApi) + monkeypatch.setattr(hub_module, "HfApi", _FakeHfApi) + + model, _ = _make_dummy_reward_model(repo_id="user/my_reward") + with pytest.warns(FutureWarning, match="push_model_to_hub is deprecated"): + model.push_model_to_hub(_make_train_cfg("user/my_dataset")) + + assert CONFIG_NAME in all_files + assert TRAIN_CONFIG_NAME in all_files + assert "README.md" in all_files diff --git a/tests/scripts/test_convert_dcp.py b/tests/scripts/test_convert_dcp.py new file mode 100644 index 000000000..430369993 --- /dev/null +++ b/tests/scripts/test_convert_dcp.py @@ -0,0 +1,123 @@ +#!/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. +"""lerobot-convert-dcp: locating, converting, and graceful-degradation publishing.""" + +import logging +import shutil +from pathlib import Path +from types import SimpleNamespace + +import pytest + +pytest.importorskip("accelerate", reason="accelerate is required (install lerobot[training])") + +import lerobot.distributed.checkpoint as dist_checkpoint +from lerobot.scripts.lerobot_convert_dcp import ( + ConvertDcpConfig, + _locate_pretrained_dir, + _publish_converted, + convert_checkpoint, +) +from lerobot.utils.constants import PRETRAINED_MODEL_DIR + + +@pytest.fixture +def fake_merge(monkeypatch): + """Stand in for accelerate.utils.merge_fsdp_weights: writes a marker safetensors file.""" + import accelerate.utils + + def merge(checkpoint_dir, output_path, safe_serialization=True, remove_checkpoint_dir=False): + assert isinstance(checkpoint_dir, str) and isinstance(output_path, str) # str, not Path + (Path(output_path) / "model.safetensors").write_bytes(b"merged") + # Mirror accelerate: the shard directory is removed by the merge itself, when asked. + if remove_checkpoint_dir: + shutil.rmtree(checkpoint_dir) + + monkeypatch.setattr(accelerate.utils, "merge_fsdp_weights", merge) + + +def make_dcp_checkpoint(tmp_path: Path) -> Path: + pretrained = tmp_path / PRETRAINED_MODEL_DIR + dcp_dir = pretrained / "pytorch_model_fsdp_0" + dcp_dir.mkdir(parents=True) + (dcp_dir / "__0_0.distcp").write_bytes(b"shard") + (pretrained / "config.json").write_text("{}") + return tmp_path + + +class TestConvert: + def test_locate_accepts_step_dir_or_pretrained_dir(self, tmp_path): + step_dir = make_dcp_checkpoint(tmp_path) + pretrained = step_dir / PRETRAINED_MODEL_DIR + assert _locate_pretrained_dir(step_dir) == pretrained + assert _locate_pretrained_dir(pretrained) == pretrained + + def test_convert_keeps_dcp_by_default(self, tmp_path, fake_merge): + step_dir = make_dcp_checkpoint(tmp_path) + out = convert_checkpoint(ConvertDcpConfig(checkpoint_dir=step_dir)) + assert out.read_bytes() == b"merged" + assert (step_dir / PRETRAINED_MODEL_DIR / "pytorch_model_fsdp_0").is_dir() + + def test_convert_delete_dcp(self, tmp_path, fake_merge): + step_dir = make_dcp_checkpoint(tmp_path) + convert_checkpoint(ConvertDcpConfig(checkpoint_dir=step_dir, delete_dcp=True)) + assert not (step_dir / PRETRAINED_MODEL_DIR / "pytorch_model_fsdp_0").exists() + + def test_missing_shards_error_names_the_format(self, tmp_path): + with pytest.raises(FileNotFoundError, match="checkpoint_format=dcp"): + convert_checkpoint(ConvertDcpConfig(checkpoint_dir=tmp_path)) + + +class TestPublishGracefulDegradation: + def _mock_api(self, monkeypatch): + calls = {} + + class FakeApi: + def create_repo(self, repo_id, private=None, exist_ok=False): + return SimpleNamespace(repo_id=repo_id) + + def upload_folder(self, *, repo_id, folder_path, allow_patterns, **kwargs): + calls["repo_id"] = repo_id + calls["files"] = sorted(p.name for p in Path(folder_path).iterdir()) + calls["allow_patterns"] = allow_patterns + return SimpleNamespace(repo_url=SimpleNamespace(url=f"https://huggingface.co/{repo_id}")) + + import lerobot.scripts.lerobot_convert_dcp as mod + + monkeypatch.setattr(mod, "HfApi", FakeApi) + return calls + + def test_missing_train_config_warns_and_uploads_core(self, tmp_path, monkeypatch, caplog): + calls = self._mock_api(monkeypatch) + pretrained = make_dcp_checkpoint(tmp_path) / PRETRAINED_MODEL_DIR + (pretrained / "model.safetensors").write_bytes(b"w") + with caplog.at_level(logging.WARNING): + _publish_converted(pretrained, "user/converted", private=None) + assert any("train_config.json missing" in m for m in caplog.messages) + assert "model.safetensors" in calls["files"] + # The DCP shard directory is still on disk (--delete_dcp defaults to False) but the + # allow list admits neither `.distcp` shards nor their `.metadata` sidecar. + assert set(calls["allow_patterns"]) == {"*.safetensors", "*.json", "*.yaml", "*.md"} + # config.json is not parseable as a policy config here -> card skipped with a warning + assert any("model card" in m for m in caplog.messages) + + def test_dcp_to_safetensors_passes_str_paths(self, tmp_path, fake_merge): + """accelerate 1.14's DCP helpers do string containment checks.""" + dcp_dir = tmp_path / "pytorch_model_fsdp_0" + dcp_dir.mkdir() + out = dist_checkpoint.dcp_to_safetensors(dcp_dir, tmp_path, delete_dcp=True) + assert out == tmp_path / "model.safetensors" + assert not dcp_dir.exists() diff --git a/tests/utils/test_logging_utils.py b/tests/utils/test_logging_utils.py index 111bd0f01..dbc64bba3 100644 --- a/tests/utils/test_logging_utils.py +++ b/tests/utils/test_logging_utils.py @@ -17,6 +17,7 @@ import pytest import torch +import lerobot.utils.logging_utils as logging_utils from lerobot.utils.logging_utils import AverageMeter, MetricsTracker @@ -25,19 +26,6 @@ def mock_metrics(): return {"loss": AverageMeter("loss", ":.3f"), "accuracy": AverageMeter("accuracy", ":.2f")} -class MockAccelerator: - def __init__(self, num_processes: int, reduce_fn=None): - self.num_processes = num_processes - self.device = torch.device("cpu") - self._reduce_fn = reduce_fn - - def reduce(self, tensor, reduction="mean"): - # In single-process tests we just want a deterministic stand-in for accelerate's reduce. - if self._reduce_fn is not None: - return self._reduce_fn(tensor, reduction) - return tensor - - def test_average_meter_initialization(): meter = AverageMeter("loss", ":.2f") assert meter.name == "loss" @@ -96,14 +84,14 @@ def test_metrics_tracker_step(mock_metrics): assert tracker.epochs == tracker.samples / 1000 -def test_metrics_tracker_initialization_with_accelerator(mock_metrics): +def test_metrics_tracker_initialization_with_dp_world(mock_metrics): tracker = MetricsTracker( batch_size=32, num_frames=1000, num_episodes=50, metrics=mock_metrics, initial_step=10, - accelerator=MockAccelerator(num_processes=2), + dp_world_size=2, ) assert tracker.steps == 10 assert tracker.samples == 10 * 32 * 2 @@ -111,14 +99,14 @@ def test_metrics_tracker_initialization_with_accelerator(mock_metrics): assert tracker.epochs == tracker.samples / 1000 -def test_metrics_tracker_step_with_accelerator(mock_metrics): +def test_metrics_tracker_step_with_dp_world(mock_metrics): tracker = MetricsTracker( batch_size=32, num_frames=1000, num_episodes=50, metrics=mock_metrics, initial_step=5, - accelerator=MockAccelerator(num_processes=2), + dp_world_size=2, ) tracker.step() assert tracker.steps == 6 @@ -178,53 +166,37 @@ def test_average_meter_reduction_stored(): assert meter.reduction == "max" -def test_metrics_tracker_reduce_across_ranks_no_accelerator(): +def test_metrics_tracker_reduce_across_ranks_outside_distributed(): metrics = {"update_s": AverageMeter("update_s", reduction="max")} tracker = MetricsTracker(batch_size=32, num_frames=1000, num_episodes=50, metrics=metrics) tracker.update_s = 0.5 - tracker.reduce_across_ranks() # no-op without accelerator + tracker.reduce_across_ranks() # no-op without an initialized process group assert tracker.update_s.avg == 0.5 -def test_metrics_tracker_reduce_across_ranks_single_process(): - metrics = {"update_s": AverageMeter("update_s", reduction="max")} - tracker = MetricsTracker( - batch_size=32, - num_frames=1000, - num_episodes=50, - metrics=metrics, - accelerator=MockAccelerator(num_processes=1), - ) - tracker.update_s = 0.5 - tracker.reduce_across_ranks() # no-op when world size is 1 - assert tracker.update_s.avg == 0.5 - - -def test_metrics_tracker_reduce_across_ranks_invokes_reduce(): +def test_metrics_tracker_reduce_across_ranks_invokes_all_reduce(monkeypatch): captured = {} - def fake_reduce(tensor, reduction): - captured["reduction"] = reduction + def fake_all_reduce(tensor, op): + captured["op"] = op captured["values"] = tensor.clone() # Pretend the slowest rank reported 0.9 instead of this rank's 0.4. - return torch.tensor([0.9], dtype=tensor.dtype, device=tensor.device) + tensor.fill_(0.9) + + monkeypatch.setattr(logging_utils.dist, "is_initialized", lambda: True) + monkeypatch.setattr(logging_utils.dist, "get_world_size", lambda: 4) + monkeypatch.setattr(logging_utils.dist, "all_reduce", fake_all_reduce) metrics = { "loss": AverageMeter("loss"), # reduction="none" -> not touched "update_s": AverageMeter("update_s", reduction="max"), } - tracker = MetricsTracker( - batch_size=32, - num_frames=1000, - num_episodes=50, - metrics=metrics, - accelerator=MockAccelerator(num_processes=4, reduce_fn=fake_reduce), - ) + tracker = MetricsTracker(batch_size=32, num_frames=1000, num_episodes=50, metrics=metrics) tracker.loss = 1.0 tracker.update_s = 0.4 tracker.reduce_across_ranks() - assert captured["reduction"] == "max" + assert captured["op"] == logging_utils.dist.ReduceOp.MAX assert torch.allclose(captured["values"], torch.tensor([0.4])) assert tracker.update_s.avg == pytest.approx(0.9) # Metrics without a reduction stay untouched. diff --git a/tests/utils/test_train_utils.py b/tests/utils/test_train_utils.py index 3098c080d..3d686eec3 100644 --- a/tests/utils/test_train_utils.py +++ b/tests/utils/test_train_utils.py @@ -15,24 +15,22 @@ # limitations under the License. from pathlib import Path -from unittest.mock import MagicMock, Mock, patch +from unittest.mock import MagicMock import pytest from lerobot.common.train_utils import ( get_step_checkpoint_dir, get_step_identifier, - load_training_batch_size, - load_training_num_processes, - load_training_state, - load_training_step, + load_training_metadata, push_checkpoint_to_hub, - save_checkpoint, + save_training_metadata, save_training_state, - save_training_step, should_save_checkpoint, update_last_checkpoint, ) +from lerobot.configs.default import DatasetConfig +from lerobot.configs.train import TrainPipelineConfig from lerobot.utils.constants import ( CHECKPOINTS_DIR, LAST_CHECKPOINT_LINK, @@ -69,38 +67,23 @@ def test_get_step_checkpoint_dir(): assert step_dir == output_dir / CHECKPOINTS_DIR / "000005" -def test_save_load_training_step(tmp_path): - save_training_step(5000, tmp_path) +def make_cfg(batch_size: int = 32) -> TrainPipelineConfig: + cfg = TrainPipelineConfig(dataset=DatasetConfig(repo_id="lerobot/dummy"), batch_size=batch_size) + cfg.parallelism.resolve(1) + return cfg + + +def test_save_training_metadata_writes_the_step_file(tmp_path): + save_training_metadata(5000, tmp_path, make_cfg()) assert (tmp_path / TRAINING_STEP).is_file() -def test_load_training_step(tmp_path): - step = 5000 - save_training_step(step, tmp_path) - loaded_step = load_training_step(tmp_path) - assert loaded_step == step - - -def test_save_training_state_records_num_processes(tmp_path, optimizer, scheduler): - save_training_state(tmp_path, 10, optimizer, scheduler, num_processes=4) - assert load_training_num_processes(tmp_path) == 4 - - -def test_load_training_num_processes_absent_returns_none(tmp_path, optimizer, scheduler): - # Checkpoints written before the world size was recorded must still load (back-compat). - save_training_state(tmp_path, 10, optimizer, scheduler) - assert load_training_num_processes(tmp_path) is None - - -def test_save_training_state_records_batch_size(tmp_path, optimizer, scheduler): - save_training_state(tmp_path, 10, optimizer, scheduler, batch_size=32) - assert load_training_batch_size(tmp_path) == 32 - - -def test_load_training_batch_size_absent_returns_none(tmp_path, optimizer, scheduler): - # Checkpoints written before the batch size was recorded must still load (back-compat). - save_training_state(tmp_path, 10, optimizer, scheduler) - assert load_training_batch_size(tmp_path) is None +def test_save_training_state_records_topology(tmp_path, optimizer, scheduler): + save_training_state(tmp_path, 10, make_cfg(batch_size=32), optimizer, scheduler) + metadata = load_training_metadata(tmp_path / TRAINING_STATE_DIR) + assert metadata["step"] == 10 + assert metadata["dp_world_size"] == 1 + assert metadata["batch_size"] == 32 def test_update_last_checkpoint(tmp_path): @@ -112,32 +95,12 @@ def test_update_last_checkpoint(tmp_path): assert last_checkpoint.resolve() == checkpoint -@patch("lerobot.common.train_utils.save_training_state") -def test_save_checkpoint(mock_save_training_state, tmp_path, optimizer): - policy = Mock() - cfg = Mock() - save_checkpoint(tmp_path, 10, cfg, policy, optimizer) - policy.save_pretrained.assert_called_once() - cfg.save_pretrained.assert_called_once() - mock_save_training_state.assert_called_once() +# save_checkpoint round-trips (all formats, real policies) live in +# tests/common/test_checkpoint_save_resume.py. -@patch("lerobot.common.train_utils.save_training_state") -def test_save_checkpoint_peft(mock_save_training_state, tmp_path, optimizer): - policy = Mock() - policy.config = Mock() - policy.config.save_pretrained = Mock() - cfg = Mock() - cfg.use_peft = True - save_checkpoint(tmp_path, 10, cfg, policy, optimizer) - policy.save_pretrained.assert_called_once() - cfg.save_pretrained.assert_called_once() - policy.config.save_pretrained.assert_called_once() - mock_save_training_state.assert_called_once() - - -def test_save_training_state(tmp_path, optimizer, scheduler): - save_training_state(tmp_path, 10, optimizer, scheduler) +def test_save_training_state_layout(tmp_path, optimizer, scheduler): + save_training_state(tmp_path, 10, make_cfg(), optimizer, scheduler) assert (tmp_path / TRAINING_STATE_DIR).is_dir() assert (tmp_path / TRAINING_STATE_DIR / TRAINING_STEP).is_file() assert (tmp_path / TRAINING_STATE_DIR / RNG_STATE).is_file() @@ -146,27 +109,8 @@ def test_save_training_state(tmp_path, optimizer, scheduler): assert (tmp_path / TRAINING_STATE_DIR / SCHEDULER_STATE).is_file() -def test_save_load_training_state(tmp_path, optimizer, scheduler): - save_training_state(tmp_path, 10, optimizer, scheduler) - loaded_step, loaded_optimizer, loaded_scheduler = 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 +# The two-phase resume (resume_before_prepare / resume_after_prepare) is covered in +# tests/common/test_checkpoint_save_resume.py with real policies and optimizer state. def test_push_checkpoint_to_hub_creates_repo_and_uploads(tmp_path, monkeypatch):