From 978b14859f06a65b76fe61bd4a9cb6858bd8ab9d Mon Sep 17 00:00:00 2001 From: Hexu Zhao Date: Wed, 9 Sep 2026 09:56:36 +0000 Subject: [PATCH 01/11] optim: fused AdamW when all params are CUDA float tensors Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_013AWNsZ45BMx5moJogT9xxT --- src/lerobot/optim/optimizers.py | 22 ++++++++++++++++++++++ 1 file changed, 22 insertions(+) diff --git a/src/lerobot/optim/optimizers.py b/src/lerobot/optim/optimizers.py index 820579c061c..005381d3223 100644 --- a/src/lerobot/optim/optimizers.py +++ b/src/lerobot/optim/optimizers.py @@ -86,6 +86,21 @@ def build(self, params: OptimizerParams) -> torch.optim.Optimizer | dict[str, to raise NotImplementedError +def _all_cuda_float(params: OptimizerParams) -> bool: + """True when every parameter (flat list or param groups) is a floating CUDA tensor.""" + if isinstance(params, dict): + return False + tensors: list[torch.Tensor] = [] + for p in params: + if isinstance(p, dict): + tensors.extend(p["params"]) + else: + tensors.append(p) + return len(tensors) > 0 and all( + isinstance(t, torch.Tensor) and t.is_cuda and t.is_floating_point() for t in tensors + ) + + @OptimizerConfig.register_subclass("adam") @dataclass class AdamConfig(OptimizerConfig): @@ -113,6 +128,13 @@ class AdamWConfig(OptimizerConfig): def build(self, params: OptimizerParams) -> torch.optim.Optimizer: kwargs = asdict(self) kwargs.pop("grad_clip_norm") + # One fused kernel per parameter group instead of the per-parameter python loop + # (which reads every step counter back with .item()). Only when every parameter + # is a CUDA float tensor, which is what the fused kernel supports. + if not isinstance(params, dict): + params = list(params) # may be a generator; it is read twice below + if _all_cuda_float(params): + kwargs["fused"] = True return torch.optim.AdamW(params, **kwargs) From 73bb32320d6ed5ad1256bdf3cc320980334b417c Mon Sep 17 00:00:00 2001 From: Hexu Zhao Date: Wed, 9 Sep 2026 10:06:03 +0000 Subject: [PATCH 02/11] ddp: expose broadcast_buffers; skip the per-forward buffer broadcast when no BatchNorm trains The coalesced buffer broadcast blocks the host on every rank at each forward start, so it is a per-step rank barrier plus a 286 KB broadcast for buffers that never change (ACT: FrozenBatchNorm2d, positional embeddings). exp0914 split of exp0908 commit 8cf80f71: the gradient_as_bucket_view default flip is a flag that already exists at the base (--accelerator.ddp.gradient_as_bucket_view=true) and moves to the tuned baseline; the default stays False here. The code part (broadcast_buffers field + disable_buffer_broadcast_if_static) is what this commit keeps. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_012cdPcxsNZ5nE3qys9t14sj --- src/lerobot/configs/accelerator.py | 6 ++++++ src/lerobot/distributed/__init__.py | 2 +- src/lerobot/distributed/utils.py | 28 ++++++++++++++++++++++++++++ src/lerobot/scripts/lerobot_train.py | 2 ++ 4 files changed, 37 insertions(+), 1 deletion(-) diff --git a/src/lerobot/configs/accelerator.py b/src/lerobot/configs/accelerator.py index aebc50d18aa..3f346e45d03 100644 --- a/src/lerobot/configs/accelerator.py +++ b/src/lerobot/configs/accelerator.py @@ -101,6 +101,11 @@ class DDPConfig: find_unused_parameters: bool = True gradient_as_bucket_view: bool = False static_graph: bool = False + # Broadcast module buffers from rank 0 at every forward. Only needed when a buffer changes + # during training (BatchNorm running stats). `lerobot_train` turns it off after `prepare` + # when the policy has no BatchNorm module, because the call blocks every rank until all + # of them reach the forward; see `disable_buffer_broadcast_if_static`. + broadcast_buffers: bool = True def build_kwargs_handler(self) -> "DistributedDataParallelKwargs": """Build the DDP kwargs handler for `Accelerator(kwargs_handlers=[...])`. @@ -115,6 +120,7 @@ def build_kwargs_handler(self) -> "DistributedDataParallelKwargs": find_unused_parameters=self.find_unused_parameters, gradient_as_bucket_view=self.gradient_as_bucket_view, static_graph=self.static_graph, + broadcast_buffers=self.broadcast_buffers, ) diff --git a/src/lerobot/distributed/__init__.py b/src/lerobot/distributed/__init__.py index b4ece6ebac9..517d4341300 100644 --- a/src/lerobot/distributed/__init__.py +++ b/src/lerobot/distributed/__init__.py @@ -30,7 +30,7 @@ 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 +from .utils import disable_buffer_broadcast_if_static, finalize_sharded_policy, is_main_process, strip_accelerate_cp_hooks __all__ = [ "ParallelDims", diff --git a/src/lerobot/distributed/utils.py b/src/lerobot/distributed/utils.py index 33616fa2614..91defbf3dc3 100644 --- a/src/lerobot/distributed/utils.py +++ b/src/lerobot/distributed/utils.py @@ -92,3 +92,31 @@ def finalize_sharded_policy(policy: nn.Module, parallel_dims: "ParallelDims") -> for method_name in getattr(type(policy), "_fsdp_forward_methods", ()): if callable(getattr(policy, method_name, None)): register_fsdp_forward_method(policy, method_name) + + +def disable_buffer_broadcast_if_static(policy: nn.Module) -> bool: + """Turn off DDP's per-forward buffer broadcast when no buffer can change during training. + + `DistributedDataParallel` with `broadcast_buffers=True` broadcasts every buffer from rank 0 + at the start of each forward, through a coalesced broadcast that blocks the host until all + ranks have joined. That is only needed for buffers that training mutates, which in practice + means BatchNorm running statistics. A policy with no BatchNorm module (ACT uses + FrozenBatchNorm2d, whose buffers are constants) keeps identical buffers on every rank + without the broadcast, so the call is pure overhead: one rank barrier per step. + + Args: + policy: The policy as returned by `accelerator.prepare()`. + + Returns: + True when the broadcast was switched off. + """ + from torch.nn.modules.batchnorm import _BatchNorm + from torch.nn.parallel import DistributedDataParallel + + if not isinstance(policy, DistributedDataParallel) or not policy.broadcast_buffers: + return False + if any(isinstance(m, _BatchNorm) for m in policy.module.modules()): + return False + policy.broadcast_buffers = False + logging.info("DDP buffer broadcast disabled: the policy has no BatchNorm module, so no buffer changes.") + return True diff --git a/src/lerobot/scripts/lerobot_train.py b/src/lerobot/scripts/lerobot_train.py index 5d55dda59a7..b3cf327597c 100644 --- a/src/lerobot/scripts/lerobot_train.py +++ b/src/lerobot/scripts/lerobot_train.py @@ -64,6 +64,7 @@ from lerobot.datasets.factory import make_train_eval_datasets from lerobot.distributed import ( ParallelDims, + disable_buffer_broadcast_if_static, finalize_sharded_policy, is_main_process, make_accelerator, @@ -574,6 +575,7 @@ def train(cfg: TrainPipelineConfig): policy, optimizer, dataloader, lr_scheduler ) finalize_sharded_policy(policy, parallel_dims) + disable_buffer_broadcast_if_static(policy) if cfg.resume: resume_after_prepare(cfg, accelerator, policy, optimizer, lr_scheduler) From 663535e8e5ea670054ae6b96f7feddfd90f14667 Mon Sep 17 00:00:00 2001 From: Hexu Zhao Date: Wed, 9 Sep 2026 10:12:28 +0000 Subject: [PATCH 03/11] compile: wire CompileConfig; ACT compiles policy.model regionally (auto when the policy declares regions) fallback_random=True keeps the eager RNG stream (dropout, VAE randn_like) so a compiled run reproduces eager to fp32 rounding at the same seed. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_013AWNsZ45BMx5moJogT9xxT --- src/lerobot/configs/accelerator.py | 6 ++++- src/lerobot/configs/train.py | 4 ++-- src/lerobot/distributed/__init__.py | 4 +++- src/lerobot/distributed/utils.py | 30 ++++++++++++++++++++++++ src/lerobot/policies/act/modeling_act.py | 2 ++ src/lerobot/scripts/lerobot_train.py | 3 +++ 6 files changed, 45 insertions(+), 4 deletions(-) diff --git a/src/lerobot/configs/accelerator.py b/src/lerobot/configs/accelerator.py index 3f346e45d03..f402e337f6c 100644 --- a/src/lerobot/configs/accelerator.py +++ b/src/lerobot/configs/accelerator.py @@ -170,10 +170,14 @@ class CompileConfig: (per wrap unit) — the only combination proven with FSDP2. """ - enabled: bool = False + # None = auto: on when the policy declares `_compile_regions` and the run is not sharded. + enabled: bool | None = None backend: str = "inductor" mode: str | None = None regional: bool = True + # Keep eager RNG semantics inside compiled regions (dropout, the VAE's randn_like), so a + # compiled run reproduces the eager one to fp32 rounding at the same seed. + fallback_random: bool = True class ActivationCheckpointingMode(str, Enum): diff --git a/src/lerobot/configs/train.py b/src/lerobot/configs/train.py index 9381a03703b..7500d13001b 100644 --- a/src/lerobot/configs/train.py +++ b/src/lerobot/configs/train.py @@ -354,8 +354,8 @@ def _validate_distributed(self) -> None: "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.compile.enabled and self.parallelism.is_sharded: + raise ValueError("--accelerator.compile is wired for DDP/single-process runs only.") 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: diff --git a/src/lerobot/distributed/__init__.py b/src/lerobot/distributed/__init__.py index 517d4341300..a3c391de040 100644 --- a/src/lerobot/distributed/__init__.py +++ b/src/lerobot/distributed/__init__.py @@ -30,10 +30,12 @@ from .factory import guard_against_env_interference, make_accelerator, set_fsdp_wrap_modules from .parallel_dims import ParallelDims -from .utils import disable_buffer_broadcast_if_static, finalize_sharded_policy, is_main_process, strip_accelerate_cp_hooks +from .utils import apply_torch_compile, disable_buffer_broadcast_if_static, finalize_sharded_policy, is_main_process, strip_accelerate_cp_hooks __all__ = [ "ParallelDims", + "apply_torch_compile", + "disable_buffer_broadcast_if_static", "finalize_sharded_policy", "guard_against_env_interference", "is_main_process", diff --git a/src/lerobot/distributed/utils.py b/src/lerobot/distributed/utils.py index 91defbf3dc3..414a2ebd9fe 100644 --- a/src/lerobot/distributed/utils.py +++ b/src/lerobot/distributed/utils.py @@ -18,6 +18,7 @@ import logging from typing import TYPE_CHECKING +import torch import torch.distributed as dist from torch import nn @@ -120,3 +121,32 @@ def disable_buffer_broadcast_if_static(policy: nn.Module) -> bool: policy.broadcast_buffers = False logging.info("DDP buffer broadcast disabled: the policy has no BatchNorm module, so no buffer changes.") return True + + + +def apply_torch_compile(policy: nn.Module, compile_cfg) -> nn.Module: + """Compile the policy per `CompileConfig`, before `accelerator.prepare()`. + + Regional (default): the policy names its compute core in `_compile_regions` (ACT: `model`); + each named submodule is replaced by its `torch.compile` wrapper, and the loss glue around it + (which reads scalars back with `.item()`) stays eager. Non-regional: the whole policy. + """ + regions = getattr(policy, "_compile_regions", ()) + enabled = compile_cfg.enabled + if enabled is None: + enabled = bool(regions) + if not enabled: + return policy + import torch._inductor.config as inductor_config + + inductor_config.fallback_random = compile_cfg.fallback_random + kwargs = {"backend": compile_cfg.backend, "mode": compile_cfg.mode} + if not compile_cfg.regional: + regions = () + if regions: + for name in regions: + setattr(policy, name, torch.compile(getattr(policy, name), **kwargs)) + logging.info("torch.compile applied to %s (%s)", ", ".join(regions), kwargs) + return policy + logging.info("torch.compile applied to the whole policy (%s)", kwargs) + return torch.compile(policy, **kwargs) diff --git a/src/lerobot/policies/act/modeling_act.py b/src/lerobot/policies/act/modeling_act.py index 2298db39761..a5622e0ea81 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" + # Submodules `--accelerator.compile` wraps; the loss glue in `forward` stays eager. + _compile_regions = ("model",) # FSDP2 wrap units: one unit per transformer layer of both stacks. _fsdp_wrap_modules = ["ACTEncoderLayer", "ACTDecoderLayer"] diff --git a/src/lerobot/scripts/lerobot_train.py b/src/lerobot/scripts/lerobot_train.py index b3cf327597c..2e6f4b1dbd3 100644 --- a/src/lerobot/scripts/lerobot_train.py +++ b/src/lerobot/scripts/lerobot_train.py @@ -64,6 +64,7 @@ from lerobot.datasets.factory import make_train_eval_datasets from lerobot.distributed import ( ParallelDims, + apply_torch_compile, disable_buffer_broadcast_if_static, finalize_sharded_policy, is_main_process, @@ -550,6 +551,8 @@ def train(cfg: TrainPipelineConfig): # 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. + policy = apply_torch_compile(policy, cfg.accelerator.compile) + if is_main_process(): logging.info("Creating optimizer and scheduler") optimizer, lr_scheduler = make_optimizer_and_scheduler(cfg, policy) From 2668051949421a16363e513dc5a4f4055c6c9649 Mon Sep 17 00:00:00 2001 From: Hexu Zhao Date: Wed, 9 Sep 2026 10:21:05 +0000 Subject: [PATCH 04/11] compile: default mode reduce-overhead (CUDA graphs) Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_013AWNsZ45BMx5moJogT9xxT --- src/lerobot/configs/accelerator.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/src/lerobot/configs/accelerator.py b/src/lerobot/configs/accelerator.py index f402e337f6c..94e858f6308 100644 --- a/src/lerobot/configs/accelerator.py +++ b/src/lerobot/configs/accelerator.py @@ -173,7 +173,9 @@ class CompileConfig: # None = auto: on when the policy declares `_compile_regions` and the run is not sharded. enabled: bool | None = None backend: str = "inductor" - mode: str | None = None + # CUDA graphs: the eager step is bound by ~1000 kernel launches, and inductor's default + # mode replaces them with as many Triton launches; only graph replay removes the cost. + mode: str | None = "reduce-overhead" regional: bool = True # Keep eager RNG semantics inside compiled regions (dropout, the VAE's randn_like), so a # compiled run reproduces the eager one to fp32 rounding at the same seed. From d37994a18358cb44d1bf85cedc9aba8b95dc0a5d Mon Sep 17 00:00:00 2001 From: Hexu Zhao Date: Wed, 9 Sep 2026 10:30:42 +0000 Subject: [PATCH 05/11] act: keep sub-losses on device; MetricsTracker reads 0-d tensors back after the optimizer step Removes the two host syncs (l1_loss.item(), mean_kld.item()) between forward and backward, so the host can enqueue the backward while the forward runs. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_013AWNsZ45BMx5moJogT9xxT --- src/lerobot/policies/act/modeling_act.py | 7 +++++-- src/lerobot/utils/logging_utils.py | 2 ++ 2 files changed, 7 insertions(+), 2 deletions(-) diff --git a/src/lerobot/policies/act/modeling_act.py b/src/lerobot/policies/act/modeling_act.py index a5622e0ea81..950fb347a13 100644 --- a/src/lerobot/policies/act/modeling_act.py +++ b/src/lerobot/policies/act/modeling_act.py @@ -151,7 +151,10 @@ def forward(self, batch: dict[str, Tensor]) -> tuple[Tensor, dict]: num_valid = valid_mask.sum() * abs_err.shape[-1] l1_loss = (abs_err * valid_mask).sum() / num_valid.clamp_min(1) - loss_dict = {"l1_loss": l1_loss.item()} + # Sub-losses stay on the device as detached scalars: reading them back here would block + # the host between forward and backward. MetricsTracker.update_metrics converts them + # after the optimizer step, where the loss is read back anyway. + loss_dict = {"l1_loss": l1_loss.detach()} if self.config.use_vae and log_sigma_x2_hat is not None: # Calculate Dₖₗ(latent_pdf || standard_normal). Note: After computing the KL-divergence for # each dimension independently, we sum over the latent dimension to get the total @@ -160,7 +163,7 @@ def forward(self, batch: dict[str, Tensor]) -> tuple[Tensor, dict]: mean_kld = ( (-0.5 * (1 + log_sigma_x2_hat - mu_hat.pow(2) - (log_sigma_x2_hat).exp())).sum(-1).mean() ) - loss_dict["kld_loss"] = mean_kld.item() + loss_dict["kld_loss"] = mean_kld.detach() loss = l1_loss + mean_kld * self.config.kl_weight else: loss = l1_loss diff --git a/src/lerobot/utils/logging_utils.py b/src/lerobot/utils/logging_utils.py index 94d644cf784..06229599906 100644 --- a/src/lerobot/utils/logging_utils.py +++ b/src/lerobot/utils/logging_utils.py @@ -188,6 +188,8 @@ def update_metrics(self, values: dict[str, Any]) -> None: Caller-registered metrics (those passed to the constructor) are never overridden. """ for name, value in values.items(): + if isinstance(value, torch.Tensor) and value.numel() == 1: + value = value.item() # a 0-d tensor a policy left on the device, read back here if isinstance(value, bool) or not isinstance(value, (int, float)): continue if name in self._caller_metrics: From 5cc82b53e6ee0c6f70580358213ff6a0993bb733 Mon Sep 17 00:00:00 2001 From: Hexu Zhao Date: Wed, 9 Sep 2026 10:38:03 +0000 Subject: [PATCH 06/11] ddp: one allreduce bucket (bucket_cap_mb=1024) With a compiled backward every gradient is ready at once; 25 MB buckets only add a collective launch and a rank sync per bucket. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_013AWNsZ45BMx5moJogT9xxT --- src/lerobot/configs/accelerator.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/src/lerobot/configs/accelerator.py b/src/lerobot/configs/accelerator.py index 94e858f6308..f32270a85ba 100644 --- a/src/lerobot/configs/accelerator.py +++ b/src/lerobot/configs/accelerator.py @@ -106,6 +106,10 @@ class DDPConfig: # when the policy has no BatchNorm module, because the call blocks every rank until all # of them reach the forward; see `disable_buffer_broadcast_if_static`. broadcast_buffers: bool = True + # Allreduce bucket size. With a compiled backward every gradient is ready at once, so + # 25 MB buckets only add one collective launch and one rank-sync per bucket; a bucket + # larger than the model gives one allreduce per step. + bucket_cap_mb: int = 1024 def build_kwargs_handler(self) -> "DistributedDataParallelKwargs": """Build the DDP kwargs handler for `Accelerator(kwargs_handlers=[...])`. @@ -121,6 +125,7 @@ def build_kwargs_handler(self) -> "DistributedDataParallelKwargs": gradient_as_bucket_view=self.gradient_as_bucket_view, static_graph=self.static_graph, broadcast_buffers=self.broadcast_buffers, + bucket_cap_mb=self.bucket_cap_mb, ) From e5621af6e50faca5a38ad598d9a7dc1da999bb77 Mon Sep 17 00:00:00 2001 From: Hexu Zhao Date: Thu, 10 Sep 2026 22:52:32 +0000 Subject: [PATCH 07/11] distributed: type apply_torch_compile's compile_cfg (mypy) pyproject enables disallow_untyped_defs for lerobot.distributed.*, so the pre-commit mypy hook failed on the untyped `compile_cfg` parameter. Annotate it as CompileConfig and drop the extra blank line ruff-format flagged. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_017uuiVQa3BLp3KwiKPfXb5U --- src/lerobot/distributed/utils.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/src/lerobot/distributed/utils.py b/src/lerobot/distributed/utils.py index 414a2ebd9fe..a3729ef42d7 100644 --- a/src/lerobot/distributed/utils.py +++ b/src/lerobot/distributed/utils.py @@ -22,6 +22,8 @@ import torch.distributed as dist from torch import nn +from lerobot.configs.accelerator import CompileConfig + if TYPE_CHECKING: from lerobot.distributed.parallel_dims import ParallelDims @@ -123,8 +125,7 @@ def disable_buffer_broadcast_if_static(policy: nn.Module) -> bool: return True - -def apply_torch_compile(policy: nn.Module, compile_cfg) -> nn.Module: +def apply_torch_compile(policy: nn.Module, compile_cfg: CompileConfig) -> nn.Module: """Compile the policy per `CompileConfig`, before `accelerator.prepare()`. Regional (default): the policy names its compute core in `_compile_regions` (ACT: `model`); From ae4af846908b2b5e4c1bb1b1bcbf670cbc5fc9b6 Mon Sep 17 00:00:00 2001 From: Hexu Zhao Date: Tue, 15 Sep 2026 03:47:43 +0000 Subject: [PATCH 08/11] style: wrap the distributed re-export list (ruff format) Formatting only. The import line in src/lerobot/distributed/__init__.py exceeded the line length once the NUMA re-export was no longer on it. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_012cdPcxsNZ5nE3qys9t14sj --- src/lerobot/distributed/__init__.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/src/lerobot/distributed/__init__.py b/src/lerobot/distributed/__init__.py index a3c391de040..aada6010b97 100644 --- a/src/lerobot/distributed/__init__.py +++ b/src/lerobot/distributed/__init__.py @@ -30,7 +30,13 @@ from .factory import guard_against_env_interference, make_accelerator, set_fsdp_wrap_modules from .parallel_dims import ParallelDims -from .utils import apply_torch_compile, disable_buffer_broadcast_if_static, finalize_sharded_policy, is_main_process, strip_accelerate_cp_hooks +from .utils import ( + apply_torch_compile, + disable_buffer_broadcast_if_static, + finalize_sharded_policy, + is_main_process, + strip_accelerate_cp_hooks, +) __all__ = [ "ParallelDims", From bc51b2d0069a5c01680a6c8455ccca0de5097e86 Mon Sep 17 00:00:00 2001 From: Hexu Zhao Date: Wed, 16 Sep 2026 04:39:08 +0000 Subject: [PATCH 09/11] train: build the optimizer before compiling the policy torch.compile replaces `policy.model` with a wrapper, so every parameter under it is renamed to `model._orig_mod.*`. ACT selects its backbone parameter group by the `model.backbone` prefix, so with the optimizer built after the compile the backbone group came out empty and `optimizer_lr_backbone` was silently ignored. The two defaults are equal (1e-5), so a default run is unaffected, but any run that sets a different backbone learning rate lost it without a warning. Building the optimizer first keeps the policy's own parameter names, and it still happens before `prepare()`, which is what accelerate's FSDP2 path requires. Tests cover apply_torch_compile's enable logic and regional wrapping, that the wrapper keeps the same parameter objects, that a group selected by name is lost once the wrapper is in place, and that the train script keeps the two calls in this order. Co-Authored-By: Claude Opus 5 (1M context) --- src/lerobot/scripts/lerobot_train.py | 11 ++-- tests/distributed/test_policy_surface.py | 83 +++++++++++++++++++++++- 2 files changed, 88 insertions(+), 6 deletions(-) diff --git a/src/lerobot/scripts/lerobot_train.py b/src/lerobot/scripts/lerobot_train.py index 2e6f4b1dbd3..8f01479f931 100644 --- a/src/lerobot/scripts/lerobot_train.py +++ b/src/lerobot/scripts/lerobot_train.py @@ -549,14 +549,17 @@ def train(cfg: TrainPipelineConfig): **processor_kwargs, ) - # 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. - policy = apply_torch_compile(policy, cfg.accelerator.compile) - if is_main_process(): logging.info("Creating optimizer and scheduler") + # Created BEFORE compile, so that a policy selecting a parameter group by name still matches: + # torch.compile replaces `policy.model` with a wrapper, and ACT picks its backbone group by the + # `model.backbone` prefix, which reads `model._orig_mod.backbone` once wrapped. Also 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. optimizer, lr_scheduler = make_optimizer_and_scheduler(cfg, policy) + policy = apply_torch_compile(policy, cfg.accelerator.compile) + # --- resume phase 1 + dataloaders ---------------------------------------------------------- step = 0 # number of loop steps (= micro-batches consumed per data-parallel worker) if cfg.resume: diff --git a/tests/distributed/test_policy_surface.py b/tests/distributed/test_policy_surface.py index 6f20d871276..380e66259b6 100644 --- a/tests/distributed/test_policy_surface.py +++ b/tests/distributed/test_policy_surface.py @@ -21,8 +21,8 @@ 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.configs.accelerator import CompileConfig, FSDPConfig +from lerobot.distributed import apply_torch_compile, set_fsdp_wrap_modules, strip_accelerate_cp_hooks from lerobot.policies.pretrained import PreTrainedPolicy @@ -124,3 +124,82 @@ def test_size_based_policy_needs_no_names(self): def test_non_sharded_run_is_noop(self): set_fsdp_wrap_modules(_accelerator_with(None), _UndeclaredPolicy()) + + +class _RegionPolicy(nn.Module): + """Declares a compile region and selects a parameter group by name, the way ACT does.""" + + _compile_regions = ("model",) + + def __init__(self): + super().__init__() + self.model = nn.Sequential() + self.model.backbone = nn.Linear(2, 2) + self.model.trunk = nn.Linear(2, 2) + self.head = nn.Linear(2, 2) + + def backbone_group(self) -> list[nn.Parameter]: + return [p for n, p in self.named_parameters() if n.startswith("model.backbone")] + + +class _NoRegionPolicy(nn.Module): + def __init__(self): + super().__init__() + self.model = nn.Linear(2, 2) + + +class TestApplyTorchCompile: + def test_act_declares_a_region_it_owns(self): + """The declared region names must track the modeling code — this test pins the drift.""" + from lerobot.policies.act.modeling_act import ACTPolicy + + assert ACTPolicy._compile_regions == ("model",) + assert "model" in ACTPolicy.__init__.__code__.co_names + + def test_disabled_returns_the_policy_untouched(self): + policy = _RegionPolicy() + model = policy.model + assert apply_torch_compile(policy, CompileConfig(enabled=False)) is policy + assert policy.model is model + + def test_auto_follows_the_declaration(self): + declared = apply_torch_compile(_RegionPolicy(), CompileConfig()) + assert isinstance(declared.model, torch._dynamo.eval_frame.OptimizedModule) + + plain = _NoRegionPolicy() + model = plain.model + assert apply_torch_compile(plain, CompileConfig()) is plain + assert plain.model is model + + def test_regional_wraps_only_the_declared_region(self): + policy = _RegionPolicy() + head = policy.head + apply_torch_compile(policy, CompileConfig(enabled=True)) + assert isinstance(policy.model, torch._dynamo.eval_frame.OptimizedModule) + assert policy.head is head + + def test_the_wrapper_keeps_the_same_parameter_objects(self): + policy = _RegionPolicy() + before = dict(policy.named_parameters()) + apply_torch_compile(policy, CompileConfig(enabled=True)) + after = {n.replace("_orig_mod.", ""): p for n, p in policy.named_parameters()} + assert after.keys() == before.keys() + assert all(after[n] is p for n, p in before.items()) + + def test_a_group_selected_by_name_is_lost_after_compile(self): + """Why lerobot_train builds the optimizer before compiling: the wrapper renames parameters.""" + policy = _RegionPolicy() + assert len(policy.backbone_group()) == 2 + + apply_torch_compile(policy, CompileConfig(enabled=True)) + assert policy.backbone_group() == [] + assert any("model._orig_mod.backbone" in n for n, _ in policy.named_parameters()) + + def test_the_train_script_builds_the_optimizer_first(self): + """Guards the order the test above makes necessary.""" + import inspect + + from lerobot.scripts import lerobot_train + + source = inspect.getsource(lerobot_train.train) + assert source.index("make_optimizer_and_scheduler(") < source.index("apply_torch_compile(") From 710ae0f767f513a44f7aec4e942c805a6285b5c4 Mon Sep 17 00:00:00 2001 From: Hexu Zhao Date: Wed, 16 Sep 2026 04:39:08 +0000 Subject: [PATCH 10/11] tests: the compile gate replaces the placeholder rejection `test_compile_placeholder` asserted that `--accelerator.compile` is always rejected, which this branch replaces: compile is wired for DDP and single-process runs and still rejected when the run is sharded. Two tests now cover both sides of that gate. Co-Authored-By: Claude Opus 5 (1M context) --- tests/configs/test_train_config_distributed.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/tests/configs/test_train_config_distributed.py b/tests/configs/test_train_config_distributed.py index 8a227ee07ef..79e111fb4aa 100644 --- a/tests/configs/test_train_config_distributed.py +++ b/tests/configs/test_train_config_distributed.py @@ -50,12 +50,17 @@ def test_cfg_parallel_training_rejected(self): with pytest.raises(ValueError, match="inference-only"): cfg._validate_distributed() - def test_compile_placeholder(self): - cfg = make_cfg() + def test_compile_rejected_when_sharded(self): + cfg = make_cfg(parallelism=sharded()) cfg.accelerator.compile.enabled = True with pytest.raises(ValueError, match="compile"): cfg._validate_distributed() + def test_compile_accepted_without_sharding(self): + cfg = make_cfg() + cfg.accelerator.compile.enabled = True + cfg._validate_distributed() + def test_activation_checkpointing_placeholder(self): cfg = make_cfg() cfg.accelerator.activation_checkpointing.mode = ActivationCheckpointingMode.FULL From 23a91bd2c71c67399662280b986da1ed42bce3e2 Mon Sep 17 00:00:00 2001 From: Hexu Zhao Date: Wed, 16 Sep 2026 04:55:03 +0000 Subject: [PATCH 11/11] config: say what fallback_random actually guarantees The comment claimed a compiled run reproduces an eager one to fp32 rounding at the same seed. It does not. fallback_random only picks eager's RNG implementation over inductor's; fusion still moves the point at which each dropout mask is drawn, so with ACT's dropout on the two runs part at step 1. The class docstring still described compile as an unwired placeholder, which this branch changes. It now says what is wired and keeps the setup-order contract for the sharded case, which is still rejected. Co-Authored-By: Claude Opus 5 (1M context) --- src/lerobot/configs/accelerator.py | 14 ++++++++------ 1 file changed, 8 insertions(+), 6 deletions(-) diff --git a/src/lerobot/configs/accelerator.py b/src/lerobot/configs/accelerator.py index f32270a85ba..fde1f24e6cb 100644 --- a/src/lerobot/configs/accelerator.py +++ b/src/lerobot/configs/accelerator.py @@ -168,11 +168,11 @@ def build_plugin(self) -> "GradientAccumulationPlugin": @dataclass class CompileConfig: - """torch.compile knobs — a configured placeholder: wiring lands in a later round. + """torch.compile knobs, wired for DDP and single-process runs. - 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. + A sharded run is still rejected. The setup-order contract it will follow there is unchanged: + compile applies after CP dispatch install and activation checkpointing, before `fully_shard`, + regionally (per wrap unit), the only combination proven with FSDP2. """ # None = auto: on when the policy declares `_compile_regions` and the run is not sharded. @@ -182,8 +182,10 @@ class CompileConfig: # mode replaces them with as many Triton launches; only graph replay removes the cost. mode: str | None = "reduce-overhead" regional: bool = True - # Keep eager RNG semantics inside compiled regions (dropout, the VAE's randn_like), so a - # compiled run reproduces the eager one to fp32 rounding at the same seed. + # Draw the random numbers of a compiled region (dropout, the VAE's randn_like) with eager's + # implementation instead of inductor's. This does not make a compiled run reproduce an eager + # one: fusion still moves the point at which each mask is drawn, so with ACT's dropout on the + # two part at step 1. fallback_random: bool = True