diff --git a/src/lerobot/scripts/lerobot_train.py b/src/lerobot/scripts/lerobot_train.py index 890225b37..64d0884ea 100644 --- a/src/lerobot/scripts/lerobot_train.py +++ b/src/lerobot/scripts/lerobot_train.py @@ -172,15 +172,16 @@ def train(cfg: TrainPipelineConfig, accelerator: Accelerator | None = None): # We set step_scheduler_with_optimizer=False to prevent accelerate from adjusting the lr_scheduler steps based on the num_processes # We set find_unused_parameters=True to handle models with conditional computation if accelerator is None: - from accelerate.utils import DistributedDataParallelKwargs + from datetime import timedelta + + from accelerate.utils import DistributedDataParallelKwargs, InitProcessGroupKwargs ddp_kwargs = DistributedDataParallelKwargs(find_unused_parameters=True) - # Accelerate auto-detects the device based on the available hardware and ignores the policy.device setting. - # Force the device to be CPU when policy.device is set to CPU. + init_kwargs = InitProcessGroupKwargs(timeout=timedelta(seconds=int(os.environ.get("NCCL_TIMEOUT", 600)))) force_cpu = cfg.policy.device == "cpu" accelerator = Accelerator( step_scheduler_with_optimizer=False, - kwargs_handlers=[ddp_kwargs], + kwargs_handlers=[ddp_kwargs, init_kwargs], cpu=force_cpu, )