diff --git a/src/lerobot/configs/accelerator.py b/src/lerobot/configs/accelerator.py index aebc50d18aa..fde1f24e6cb 100644 --- a/src/lerobot/configs/accelerator.py +++ b/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): 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 b4ece6ebac9..aada6010b97 100644 --- a/src/lerobot/distributed/__init__.py +++ b/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", diff --git a/src/lerobot/distributed/utils.py b/src/lerobot/distributed/utils.py index 33616fa2614..a3729ef42d7 100644 --- a/src/lerobot/distributed/utils.py +++ b/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) 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) diff --git a/src/lerobot/policies/act/modeling_act.py b/src/lerobot/policies/act/modeling_act.py index 2298db39761..950fb347a13 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"] @@ -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 diff --git a/src/lerobot/scripts/lerobot_train.py b/src/lerobot/scripts/lerobot_train.py index 5d55dda59a7..8f01479f931 100644 --- a/src/lerobot/scripts/lerobot_train.py +++ b/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) 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: 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 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(")