d

src/lerobot/scripts/train.py

@@ -167,6 +167,10 @@ 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
@@ -177,15 +181,24 @@ def train(cfg: TrainPipelineConfig):
             mixed_precision=cfg.mixed_precision,
         )
         device = accelerator.device
-        logging.info(f"Accelerate initialized with device: {device}, mixed_precision: {cfg.mixed_precision}")
+        if accelerator.is_main_process:
+            logging.info(
+                f"Accelerate initialized with device: {device}, mixed_precision: {cfg.mixed_precision}"
+            )
+            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:
             wandb_logger = WandBLogger(cfg)
         else:
             wandb_logger = None
+    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:
@@ -194,6 +207,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")
     dataset = make_dataset(cfg)
 
@@ -202,9 +216,11 @@ 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,
@@ -224,6 +240,7 @@ def train(cfg: TrainPipelineConfig):
         policy_cfg=cfg.policy, pretrained_path=cfg.policy.pretrained_path, **processor_kwargs
     )
 
+    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)
@@ -236,6 +253,7 @@ def train(cfg: TrainPipelineConfig):
             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)
@@ -243,6 +261,8 @@ def train(cfg: TrainPipelineConfig):
     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=}")
@@ -281,6 +301,7 @@ def train(cfg: TrainPipelineConfig):
         policy, optimizer, dataloader, lr_scheduler = accelerator.prepare(
             policy, optimizer, dataloader, lr_scheduler
         )
+        if accelerator.is_main_process:
             logging.info("Policy, optimizer, dataloader, and scheduler prepared with Accelerate")
 
     dl_iter = cycle(dataloader)
@@ -299,6 +320,7 @@ def train(cfg: TrainPipelineConfig):
         cfg.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
@@ -347,6 +369,8 @@ 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()
@@ -356,6 +380,7 @@ 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)
 
@@ -443,8 +468,14 @@ def train(cfg: TrainPipelineConfig):
 
     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
         if accelerator is None or accelerator.is_main_process: