perf(act): training on 8x B200, 28.0 ms to 13.0 ms per step, mostly from the eager model's kernel launches

src/lerobot/configs/accelerator.py

@@ -101,6 +101,15 @@ 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
+    # 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=[...])`.
@@ -115,6 +124,8 @@ 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,
+            bucket_cap_mb=self.bucket_cap_mb,
         )
 
 
@@ -157,17 +168,25 @@ 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.
     """
 
-    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
+    # 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
+    # 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
 
 
 class ActivationCheckpointingMode(str, Enum):

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:

src/lerobot/distributed/__init__.py

@@ -30,10 +30,18 @@
 
 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 (
+    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",

src/lerobot/distributed/utils.py

@@ -18,9 +18,12 @@
 import logging
 from typing import TYPE_CHECKING
 
+import torch
 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
 
@@ -92,3 +95,59 @@ 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
+
+
+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`);
+    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)

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)
 
 

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"]
 
@@ -149,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
@@ -158,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

src/lerobot/scripts/lerobot_train.py

@@ -64,6 +64,8 @@
 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,
     make_accelerator,
@@ -547,12 +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.
     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:
@@ -574,6 +581,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)
 

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:

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

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(")