diff --git a/src/lerobot/configs/train.py b/src/lerobot/configs/train.py index f2d07cd7f..60a4d81d5 100644 --- a/src/lerobot/configs/train.py +++ b/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 diff --git a/src/lerobot/optim/schedulers.py b/src/lerobot/optim/schedulers.py index 1ae5c0f24..106b2ecb6 100644 --- a/src/lerobot/optim/schedulers.py +++ b/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) diff --git a/src/lerobot/scripts/train.py b/src/lerobot/scripts/train.py index 90e7807ba..3ad4dd993 100644 --- a/src/lerobot/scripts/train.py +++ b/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 - with accelerator.autocast() if use_amp else nullcontext(): + # Use accelerate's autocast and backward + 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: - # Accelerate handles gradient scaling internally - grad_norm = accelerator.clip_grad_norm_(policy.parameters(), grad_clip_norm) - if grad_norm is None: - grad_norm = 0.0 + # Use accelerate's gradient clipping + 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 - # Unscale the gradient of the optimizer's assigned params in-place **prior to gradient clipping**. - grad_scaler.unscale_(optimizer) + # Original single-GPU path + with torch.autocast(device_type=device.type) if use_amp else nullcontext(): + 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 - if cfg.wandb.enable and cfg.wandb.project: - if accelerator is None or accelerator.is_main_process: + # Only create wandb logger on main process when using accelerate + 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,8 +197,7 @@ 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") + logging.info("Creating dataset") dataset = make_dataset(cfg) # Create environment used for evaluating checkpoints during training on simulation data. @@ -220,12 +205,10 @@ 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") + 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") + logging.info("Creating policy") policy = make_policy( cfg=cfg.policy, ds_meta=dataset.meta, @@ -246,44 +229,43 @@ 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=}") - logging.info(f"{cfg.steps=} ({format_big_number(cfg.steps)})") - logging.info(f"{dataset.num_frames=} ({format_big_number(dataset.num_frames)})") - logging.info(f"{dataset.num_episodes=}") - 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}") + logging.info(colored("Output dir:", "yellow", attrs=["bold"]) + f" {cfg.output_dir}") + if cfg.env is not None: + logging.info(f"{cfg.env.task=}") + logging.info(f"{cfg.steps=} ({format_big_number(cfg.steps)})") + logging.info(f"{dataset.num_frames=} ({format_big_number(dataset.num_frames)})") + logging.info(f"{dataset.num_episodes=}") + logging.info(f"{num_learnable_params=} ({format_big_number(num_learnable_params)})") + logging.info(f"{num_total_params=} ({format_big_number(num_total_params)})") # create dataloader for offline training if hasattr(cfg.policy, "drop_n_last_frames"): @@ -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,45 +317,24 @@ 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") + 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 + 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) - 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, - ) + 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, + ) # Note: eval and checkpoint happens *after* the `step`th training update has completed, so we # increment `step` here. @@ -385,121 +345,74 @@ 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() - if output_dict: - wandb_log_dict.update(output_dict) - wandb_logger.log_dict(wandb_log_dict, step) + logging.info(train_tracker) + if wandb_logger: + wandb_log_dict = train_tracker.to_dict() + if output_dict: + wandb_log_dict.update(output_dict) + wandb_logger.log_dict(wandb_log_dict, step) 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}") + 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 - ) - update_last_checkpoint(checkpoint_dir) - if wandb_logger: - wandb_logger.log_policy(checkpoint_dir) + save_checkpoint( + checkpoint_dir, step, cfg, policy, optimizer, lr_scheduler, preprocessor, postprocessor + ) + update_last_checkpoint(checkpoint_dir) + if wandb_logger: + 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, - preprocessor=preprocessor, - postprocessor=postprocessor, - n_episodes=cfg.eval.n_episodes, - videos_dir=cfg.output_dir / "eval" / f"videos_step_{step_id}", - max_episodes_rendered=4, - start_seed=cfg.seed, - max_parallel_tasks=cfg.env.max_parallel_tasks, - ) - # overall metrics (suite-agnostic) - aggregated = eval_info["overall"] - - # optional: per-suite logging - for suite, suite_info in eval_info.items(): - logging.info("Suite %s aggregated: %s", suite, suite_info) - - # meters/tracker - eval_metrics = { - "avg_sum_reward": AverageMeter("∑rwrd", ":.3f"), - "pc_success": AverageMeter("success", ":.1f"), - "eval_s": AverageMeter("eval_s", ":.3f"), - } - eval_tracker = MetricsTracker( - cfg.batch_size, dataset.num_frames, dataset.num_episodes, eval_metrics, initial_step=step + step_id = get_step_identifier(step, cfg.steps) + logging.info(f"Eval policy at step {step}") + 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=policy, + preprocessor=preprocessor, + postprocessor=postprocessor, + n_episodes=cfg.eval.n_episodes, + videos_dir=cfg.output_dir / "eval" / f"videos_step_{step_id}", + max_episodes_rendered=4, + start_seed=cfg.seed, + max_parallel_tasks=cfg.env.max_parallel_tasks, ) - eval_tracker.eval_s = aggregated.pop("eval_s") - eval_tracker.avg_sum_reward = aggregated.pop("avg_sum_reward") - eval_tracker.pc_success = aggregated.pop("pc_success") - if wandb_logger: - wandb_log_dict = {**eval_tracker.to_dict(), **eval_info} - wandb_logger.log_dict(wandb_log_dict, step, mode="eval") - wandb_logger.log_video(eval_info["overall"]["video_paths"][0], step, mode="eval") + # overall metrics (suite-agnostic) + aggregated = eval_info["overall"] - # Synchronize all processes after evaluation - if accelerator is not None: - accelerator.wait_for_everyone() + # optional: per-suite logging + for suite, suite_info in eval_info.items(): + logging.info("Suite %s aggregated: %s", suite, suite_info) + + # meters/tracker + eval_metrics = { + "avg_sum_reward": AverageMeter("∑rwrd", ":.3f"), + "pc_success": AverageMeter("success", ":.1f"), + "eval_s": AverageMeter("eval_s", ":.3f"), + } + eval_tracker = MetricsTracker( + cfg.batch_size, dataset.num_frames, dataset.num_episodes, eval_metrics, initial_step=step + ) + eval_tracker.eval_s = aggregated.pop("eval_s") + eval_tracker.avg_sum_reward = aggregated.pop("avg_sum_reward") + eval_tracker.pc_success = aggregated.pop("pc_success") + if wandb_logger: + wandb_log_dict = {**eval_tracker.to_dict(), **eval_info} + wandb_logger.log_dict(wandb_log_dict, step, mode="eval") + wandb_logger.log_video(eval_info["overall"]["video_paths"][0], step, mode="eval") 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() + logging.info("End of training") if cfg.policy.push_to_hub: - # Only push to hub from main process when using accelerate - 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) + policy.push_model_to_hub(cfg) + preprocessor.push_to_hub(cfg.policy.repo_id) + postprocessor.push_to_hub(cfg.policy.repo_id) def main():