admin/lerobot
1
0
Fork 0
mirror of https://github.com/huggingface/lerobot.git synced 2026-07-24 10:16:09 +00:00
Code Issues Packages Projects Releases Wiki Activity

nccl

Browse Source
This commit is contained in:
Pepijn
2026-02-21 18:44:35 +01:00
parent 87a91b4b08
commit 40f4386e4a
1 changed files with 5 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/lerobot/scripts/lerobot_train.py
+5 -4
View File
@@ -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 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 # We set find_unused_parameters=True to handle models with conditional computation
if accelerator is None: 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) ddp_kwargs = DistributedDataParallelKwargs(find_unused_parameters=True)
# Accelerate auto-detects the device based on the available hardware and ignores the policy.device setting. init_kwargs = InitProcessGroupKwargs(timeout=timedelta(seconds=int(os.environ.get("NCCL_TIMEOUT", 600))))
# Force the device to be CPU when policy.device is set to CPU.
force_cpu = cfg.policy.device == "cpu" force_cpu = cfg.policy.device == "cpu"
accelerator = Accelerator( accelerator = Accelerator(
step_scheduler_with_optimizer=False, step_scheduler_with_optimizer=False,
kwargs_handlers=[ddp_kwargs], kwargs_handlers=[ddp_kwargs, init_kwargs],
cpu=force_cpu, cpu=force_cpu,
) )