fix

src/lerobot/configs/train.py

@@ -63,10 +63,6 @@ class TrainPipelineConfig(HubMixin):
     scheduler: LRSchedulerConfig | None = None
     eval: EvalConfig = field(default_factory=EvalConfig)
     wandb: WandBConfig = field(default_factory=WandBConfig)
-    # Accelerate configuration for multi-GPU training
-    use_accelerate: bool = False
-    gradient_accumulation_steps: int = 1
-    mixed_precision: str = "no"  # Options: "no", "fp16", "bf16"
 
     def __post_init__(self):
         self.checkpoint_path = None

src/lerobot/optim/schedulers.py

@@ -92,20 +92,20 @@ class CosineDecayWithWarmupSchedulerConfig(LRSchedulerConfig):
         def lr_lambda(current_step):
             def linear_warmup_schedule(current_step):
                 if current_step <= 0:
-                    return 0.1  # Start at 10% of peak LR
+                    return 0.1  # Start at 10% of peak LR instead of 0.1%
                 if current_step >= self.num_warmup_steps:
                     return 1.0  # Reach 100% at end of warmup
-                # Linear interpolation from 0.1 to 1.0
+                # Linear interpolation from 10% to 100% of peak LR
                 return 0.1 + 0.9 * (current_step / self.num_warmup_steps)
 
             def cosine_decay_schedule(current_step):
-                # Steps since warmup ended (this was the bug!)
                 decay_step = current_step - self.num_warmup_steps
-                decay_step = min(decay_step, self.num_decay_steps)
+                decay_step = max(0, min(decay_step, self.num_decay_steps))
 
                 cosine_decay = 0.5 * (1 + math.cos(math.pi * decay_step / self.num_decay_steps))
                 alpha = self.decay_lr / self.peak_lr
-                return (1 - alpha) * cosine_decay + alpha
+                decayed = (1 - alpha) * cosine_decay + alpha
+                return decayed
 
             if current_step < self.num_warmup_steps:
                 return linear_warmup_schedule(current_step)

src/lerobot/scripts/train.py

@@ -91,49 +91,42 @@ def update_policy(
         - A dictionary of outputs from the policy's forward pass, for logging purposes.
     """
     start_time = time.perf_counter()
-    device = get_device_from_parameters(policy)
+    device = get_device_from_parameters(policy) if accelerator is None else accelerator.device
     policy.train()
 
-    # Handle mixed precision differently for accelerate vs non-accelerate
     if accelerator is not None:
-        # Accelerate handles mixed precision internally
+        # Use accelerate's autocast and backward
-        with accelerator.autocast() if use_amp else nullcontext():
+        with accelerator.autocast():
             loss, output_dict = policy.forward(batch)
-        # Use accelerator's backward method
         accelerator.backward(loss)
-    else:
-        # Original behavior for non-accelerate
-        with torch.autocast(device_type=device.type) if use_amp else nullcontext():
-            loss, output_dict = policy.forward(batch)
-            # TODO(rcadene): policy.unnormalize_outputs(out_dict)
-        grad_scaler.scale(loss).backward()
 
-    if accelerator is not None:
+        # Use accelerate's gradient clipping
-        # Accelerate handles gradient scaling internally
+        if grad_clip_norm > 0:
             grad_norm = accelerator.clip_grad_norm_(policy.parameters(), grad_clip_norm)
             if grad_norm is None:
                 grad_norm = 0.0
+        else:
+            grad_norm = torch.tensor(0.0, device=device)
+
         with lock if lock is not None else nullcontext():
             optimizer.step()
         optimizer.zero_grad()
     else:
-        # Original gradient handling for non-accelerate
+        # Original single-GPU path
-        # Unscale the gradient of the optimizer's assigned params in-place **prior to gradient clipping**.
+        with torch.autocast(device_type=device.type) if use_amp else nullcontext():
-        grad_scaler.unscale_(optimizer)
+            loss, output_dict = policy.forward(batch)
+        grad_scaler.scale(loss).backward()
 
+        grad_scaler.unscale_(optimizer)
         grad_norm = torch.nn.utils.clip_grad_norm_(
             policy.parameters(),
             grad_clip_norm,
             error_if_nonfinite=False,
         )
 
-        # Optimizer's gradients are already unscaled, so scaler.step does not unscale them,
-        # although it still skips optimizer.step() if the gradients contain infs or NaNs.
         with lock if lock is not None else nullcontext():
             grad_scaler.step(optimizer)
-        # Updates the scale for next iteration.
         grad_scaler.update()
-
         optimizer.zero_grad()
 
     # Step through pytorch scheduler at every batch instead of epoch
@@ -168,16 +161,11 @@ def train(cfg: TrainPipelineConfig):
         cfg: A `TrainPipelineConfig` object containing all training configurations.
     """
     cfg.validate()
-
-    # Only log config on main process when using accelerate
-    # For now we don't know if we're using accelerate yet, so we'll log this always
-    # and fix the duplicate later if needed
     logging.info(pformat(cfg.to_dict()))
 
-    # Initialize Accelerate if requested
+    # Initialize Accelerate if enabled
     accelerator = None
     if cfg.use_accelerate:
-        # Configure DDP to handle unused parameters
         ddp_kwargs = DistributedDataParallelKwargs(find_unused_parameters=True)
         accelerator = Accelerator(
             gradient_accumulation_steps=cfg.gradient_accumulation_steps,
@@ -191,19 +179,17 @@ def train(cfg: TrainPipelineConfig):
             )
             logging.info(f"Training on {accelerator.num_processes} processes")
     else:
-        # Check device is available (original behavior)
         device = get_safe_torch_device(cfg.policy.device, log=True)
 
-    # Only create wandb logger on main process
+    # Only create wandb logger on main process when using accelerate
-    if cfg.wandb.enable and cfg.wandb.project:
     if accelerator is None or accelerator.is_main_process:
+        if cfg.wandb.enable and cfg.wandb.project:
             wandb_logger = WandBLogger(cfg)
         else:
             wandb_logger = None
+            logging.info(colored("Logs will be saved locally.", "yellow", attrs=["bold"]))
     else:
         wandb_logger = None
-        if accelerator is None or accelerator.is_main_process:
-            logging.info(colored("Logs will be saved locally.", "yellow", attrs=["bold"]))
 
     if cfg.seed is not None:
         set_seed(cfg.seed)
@@ -211,7 +197,6 @@ def train(cfg: TrainPipelineConfig):
     torch.backends.cudnn.benchmark = True
     torch.backends.cuda.matmul.allow_tf32 = True
 
-    if accelerator is None or accelerator.is_main_process:
     logging.info("Creating dataset")
     dataset = make_dataset(cfg)
 
@@ -220,11 +205,9 @@ def train(cfg: TrainPipelineConfig):
     # using the eval.py instead, with gym_dora environment and dora-rs.
     eval_env = None
     if cfg.eval_freq > 0 and cfg.env is not None:
-        if accelerator is None or accelerator.is_main_process:
         logging.info("Creating env")
         eval_env = make_env(cfg.env, n_envs=cfg.eval.batch_size, use_async_envs=cfg.eval.use_async_envs)
 
-    if accelerator is None or accelerator.is_main_process:
     logging.info("Creating policy")
     policy = make_policy(
         cfg=cfg.policy,
@@ -246,28 +229,35 @@ def train(cfg: TrainPipelineConfig):
 
     if accelerator is None or accelerator.is_main_process:
         logging.info("Creating optimizer and scheduler")
-    optimizer, lr_scheduler = make_optimizer_and_scheduler(cfg, policy)
 
-    grad_scaler = GradScaler(device.type, enabled=cfg.policy.use_amp)
+    # Scale scheduler parameters for multi-GPU training
+    if accelerator is not None and accelerator.num_processes > 1:
+        # With more GPUs, we process data faster, so scheduler should adapt faster
+        original_warmup_steps = cfg.scheduler.num_warmup_steps if cfg.scheduler else 0
+        original_decay_steps = cfg.scheduler.num_decay_steps if cfg.scheduler else 0
+
+        if cfg.scheduler is not None:
+            cfg.scheduler.num_warmup_steps = max(
+                1, cfg.scheduler.num_warmup_steps // accelerator.num_processes
+            )
+            cfg.scheduler.num_decay_steps = max(1, cfg.scheduler.num_decay_steps // accelerator.num_processes)
+
+            if accelerator.is_main_process:
+                logging.info(f"Scaled scheduler for {accelerator.num_processes} GPUs:")
+                logging.info(f"  Warmup steps: {original_warmup_steps} → {cfg.scheduler.num_warmup_steps}")
+                logging.info(f"  Decay steps: {original_decay_steps} → {cfg.scheduler.num_decay_steps}")
+
+    optimizer, lr_scheduler = make_optimizer_and_scheduler(cfg, policy)
+    grad_scaler = GradScaler(device.type, enabled=cfg.policy.use_amp and accelerator is None)
 
     step = 0  # number of policy updates (forward + backward + optim)
 
     if cfg.resume:
-        if accelerator is not None:
-            # Load accelerate-specific state if available
-            accelerate_state_path = cfg.checkpoint_path / "accelerate_state"
-            if accelerate_state_path.exists():
-                accelerator.load_state(str(accelerate_state_path))
-                if accelerator.is_main_process:
-                    logging.info("Loaded Accelerate state from checkpoint")
-
         step, optimizer, lr_scheduler = load_training_state(cfg.checkpoint_path, optimizer, lr_scheduler)
 
     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())
 
-    # Only log setup info on main process
-    if accelerator is None or accelerator.is_main_process:
     logging.info(colored("Output dir:", "yellow", attrs=["bold"]) + f" {cfg.output_dir}")
     if cfg.env is not None:
         logging.info(f"{cfg.env.task=}")
@@ -277,14 +267,6 @@ def train(cfg: TrainPipelineConfig):
     logging.info(f"{num_learnable_params=} ({format_big_number(num_learnable_params)})")
     logging.info(f"{num_total_params=} ({format_big_number(num_total_params)})")
 
-        # Log batch size and learning rate info
-        if accelerator is not None:
-            logging.info(f"Per-GPU batch size: {cfg.batch_size}")
-            logging.info(f"Effective batch size (total): {cfg.batch_size * accelerator.num_processes}")
-        else:
-            logging.info(f"Batch size: {cfg.batch_size}")
-        logging.info(f"Learning rate: {optimizer.param_groups[0]['lr']:.2e}")
-
     # create dataloader for offline training
     if hasattr(cfg.policy, "drop_n_last_frames"):
         shuffle = False
@@ -318,7 +300,6 @@ def train(cfg: TrainPipelineConfig):
             logging.info("Policy, optimizer, dataloader, and scheduler prepared with Accelerate")
 
     dl_iter = cycle(dataloader)
-
     policy.train()
 
     train_metrics = {
@@ -336,29 +317,8 @@ def train(cfg: TrainPipelineConfig):
         effective_batch_size, dataset.num_frames, dataset.num_episodes, train_metrics, initial_step=step
     )
 
-    if accelerator is None or accelerator.is_main_process:
     logging.info("Start offline training on a fixed dataset")
     for _ in range(step, cfg.steps):
-        # Handle gradient accumulation
-        if accelerator is not None:
-            with accelerator.accumulate(policy):
-                start_time = time.perf_counter()
-                batch = next(dl_iter)
-                batch = preprocessor(batch)
-                train_tracker.dataloading_s = time.perf_counter() - start_time
-
-                train_tracker, output_dict = update_policy(
-                    train_tracker,
-                    policy,
-                    batch,
-                    optimizer,
-                    cfg.optimizer.grad_clip_norm,
-                    grad_scaler=grad_scaler,
-                    lr_scheduler=lr_scheduler,
-                    use_amp=cfg.policy.use_amp,
-                    accelerator=accelerator,
-                )
-        else:
         start_time = time.perf_counter()
         batch = next(dl_iter)
         batch = preprocessor(batch)
@@ -385,8 +345,6 @@ def train(cfg: TrainPipelineConfig):
         is_eval_step = cfg.eval_freq > 0 and step % cfg.eval_freq == 0
 
         if is_log_step:
-            # Only log training metrics on main process
-            if accelerator is None or accelerator.is_main_process:
             logging.info(train_tracker)
             if wandb_logger:
                 wandb_log_dict = train_tracker.to_dict()
@@ -396,33 +354,8 @@ def train(cfg: TrainPipelineConfig):
             train_tracker.reset_averages()
 
         if cfg.save_checkpoint and is_saving_step:
-            if accelerator is None or accelerator.is_main_process:
             logging.info(f"Checkpoint policy after step {step}")
             checkpoint_dir = get_step_checkpoint_dir(cfg.output_dir, cfg.steps, step)
-
-            if accelerator is not None:
-                # Use accelerate's checkpointing - only saves on main process
-                accelerator.wait_for_everyone()  # Synchronize all processes
-                if accelerator.is_main_process:
-                    # Use unwrapped model for saving
-                    unwrapped_policy = accelerator.unwrap_model(policy)
-                    save_checkpoint(
-                        checkpoint_dir,
-                        step,
-                        cfg,
-                        unwrapped_policy,
-                        optimizer,
-                        lr_scheduler,
-                        preprocessor,
-                        postprocessor,
-                    )
-                    update_last_checkpoint(checkpoint_dir)
-                    if wandb_logger:
-                        wandb_logger.log_policy(checkpoint_dir)
-                # Save accelerate-specific state
-                accelerator.save_state(checkpoint_dir / "accelerate_state")
-            else:
-                # Original behavior for non-accelerate
             save_checkpoint(
                 checkpoint_dir, step, cfg, policy, optimizer, lr_scheduler, preprocessor, postprocessor
             )
@@ -431,21 +364,15 @@ def train(cfg: TrainPipelineConfig):
                 wandb_logger.log_policy(checkpoint_dir)
 
         if cfg.env and is_eval_step:
-            # Only evaluate on main process when using accelerate
-            if accelerator is None or accelerator.is_main_process:
             step_id = get_step_identifier(step, cfg.steps)
             logging.info(f"Eval policy at step {step}")
-
-                # Use unwrapped model for evaluation if using accelerate
-                eval_policy = accelerator.unwrap_model(policy) if accelerator is not None else policy
-
             with (
                 torch.no_grad(),
                 torch.autocast(device_type=device.type) if cfg.policy.use_amp else nullcontext(),
             ):
                 eval_info = eval_policy_all(
                     envs=eval_env,  # dict[suite][task_id] -> vec_env
-                        policy=eval_policy,
+                    policy=policy,
                     preprocessor=preprocessor,
                     postprocessor=postprocessor,
                     n_episodes=cfg.eval.n_episodes,
@@ -478,26 +405,12 @@ def train(cfg: TrainPipelineConfig):
                 wandb_logger.log_dict(wandb_log_dict, step, mode="eval")
                 wandb_logger.log_video(eval_info["overall"]["video_paths"][0], step, mode="eval")
 
-            # Synchronize all processes after evaluation
-            if accelerator is not None:
-                accelerator.wait_for_everyone()
-
     if eval_env:
         close_envs(eval_env)
-
-    if accelerator is None or accelerator.is_main_process:
     logging.info("End of training")
 
-    # Synchronize all processes before finishing
-    if accelerator is not None:
-        accelerator.wait_for_everyone()
-
     if cfg.policy.push_to_hub:
-        # Only push to hub from main process when using accelerate
+        policy.push_model_to_hub(cfg)
-        if accelerator is None or accelerator.is_main_process:
-            # Use unwrapped model for hub pushing if using accelerate
-            hub_policy = accelerator.unwrap_model(policy) if accelerator is not None else policy
-            hub_policy.push_model_to_hub(cfg)
         preprocessor.push_to_hub(cfg.policy.repo_id)
         postprocessor.push_to_hub(cfg.policy.repo_id)