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

fix(train): download dataset once on the global main process

Browse Source 
Gate the training dataset download on the global is_main_process (download once to the
shared dataset root, barrier, then every other rank reads the already-populated copy)
instead of per-node is_local_main_process. LeRobotDataset skips its snapshot_download
when try_load() succeeds, so no rank re-downloads. Assumes the dataset root / HF cache is
on storage shared across nodes.

Co-authored-by: Cursor <cursoragent@cursor.com>
This commit is contained in:
pepijn
2026-06-11 17:21:35 +00:00
parent 3f6909fb63
commit 926fb9c31e
1 changed files with 7 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/lerobot/scripts/lerobot_train.py
+7 -8
View File
@@ -234,18 +234,17 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
torch.backends.cudnn.benchmark = True torch.backends.cudnn.benchmark = True
torch.backends.cuda.matmul.allow_tf32 = True torch.backends.cuda.matmul.allow_tf32 = True
# Dataset loading synchronization: each node's local main process downloads first to avoid # Dataset loading synchronization: the global main process downloads once to the shared
# race conditions (the global main process only exists on node 0, so gating on it would let # dataset root, then a barrier lets every other rank read the already-populated copy.
# all ranks of the other nodes download and build the Arrow cache concurrently). # LeRobotDataset skips its snapshot_download when try_load() succeeds, so no rank re-downloads.
if accelerator.is_local_main_process: if is_main_process:
if is_main_process: logging.info("Creating dataset")
logging.info("Creating dataset")
dataset = make_dataset(cfg) dataset = make_dataset(cfg)
accelerator.wait_for_everyone() accelerator.wait_for_everyone()
# Now all other processes can safely load the dataset from the local cache # Other ranks read from the shared copy populated by the main process.
if not accelerator.is_local_main_process: if not is_main_process:
dataset = make_dataset(cfg) dataset = make_dataset(cfg)
# Create environment used for evaluating checkpoints during training on simulation data. # Create environment used for evaluating checkpoints during training on simulation data.