From 2b4c5f49e3939f0073bcf5f027de2927fb467276 Mon Sep 17 00:00:00 2001 From: pepijn Date: Mon, 18 May 2026 18:58:17 +0000 Subject: [PATCH] fix(pi05): disable foreach AdamW by default Avoid the multi-tensor AdamW temporary that can OOM full PI05/PI052 fine-tuning near GPU memory limits. Co-authored-by: Cursor --- src/lerobot/optim/optimizers.py | 1 + src/lerobot/policies/pi05/configuration_pi05.py | 2 ++ 2 files changed, 3 insertions(+) diff --git a/src/lerobot/optim/optimizers.py b/src/lerobot/optim/optimizers.py index 0bdd7a37e..3d9ec6b98 100644 --- a/src/lerobot/optim/optimizers.py +++ b/src/lerobot/optim/optimizers.py @@ -104,6 +104,7 @@ class AdamWConfig(OptimizerConfig): eps: float = 1e-8 weight_decay: float = 1e-2 grad_clip_norm: float = 10.0 + foreach: bool | None = None def build(self, params: OptimizerParams) -> torch.optim.Optimizer: kwargs = asdict(self) diff --git a/src/lerobot/policies/pi05/configuration_pi05.py b/src/lerobot/policies/pi05/configuration_pi05.py index 124e85cc9..d15cc0ee8 100644 --- a/src/lerobot/policies/pi05/configuration_pi05.py +++ b/src/lerobot/policies/pi05/configuration_pi05.py @@ -93,6 +93,7 @@ class PI05Config(PreTrainedConfig): optimizer_eps: float = 1e-8 optimizer_weight_decay: float = 0.01 optimizer_grad_clip_norm: float = 1.0 + optimizer_foreach: bool | None = False # Scheduler settings: see openpi `CosineDecaySchedule` # Note: These will auto-scale if --steps < scheduler_decay_steps @@ -152,6 +153,7 @@ class PI05Config(PreTrainedConfig): eps=self.optimizer_eps, weight_decay=self.optimizer_weight_decay, grad_clip_norm=self.optimizer_grad_clip_norm, + foreach=self.optimizer_foreach, ) def get_scheduler_preset(self):