From 674c990a3935f0c35b1b7c127f557110874f87eb Mon Sep 17 00:00:00 2001 From: pepijn Date: Fri, 12 Jun 2026 09:24:32 +0000 Subject: [PATCH] feat(streaming): default episode pool 1024 and wire streaming into lerobot-train Raise the default episode_pool_size to 1024 (DatasetConfig + StreamingLeRobotDataset) for better default shuffle quality at scale. Streaming is now a first-class option of the main train script: when cfg.dataset.streaming is set, the dataloader is not handed to accelerate (the dataset is already rank-disjoint via split_dataset_by_node, so IterableDatasetShard would drop (N-1)/N of each rank's stream), batches are moved to device manually, and the episode-aware sampler is skipped. Remove the standalone examples/scaling/train_streaming_multinode.py example in favor of this wiring. Co-authored-by: Cursor --- examples/scaling/train_streaming_multinode.py | 192 ------------------ src/lerobot/configs/default.py | 2 +- src/lerobot/datasets/streaming_dataset.py | 8 +- src/lerobot/scripts/lerobot_train.py | 18 +- 4 files changed, 19 insertions(+), 201 deletions(-) delete mode 100644 examples/scaling/train_streaming_multinode.py diff --git a/examples/scaling/train_streaming_multinode.py b/examples/scaling/train_streaming_multinode.py deleted file mode 100644 index 5beb29b55..000000000 --- a/examples/scaling/train_streaming_multinode.py +++ /dev/null @@ -1,192 +0,0 @@ -# Copyright 2025 The HuggingFace Inc. team. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""Distributed, resumable streaming training on a large HF-hosted dataset. - -This example shows how to train (or just stress the data pipeline) over a multi-TB dataset that never -touches local disk, scaling across GPUs and nodes with Accelerate. It demonstrates the large-scale -streaming features of :class:`StreamingLeRobotDataset`: - -- per-rank sharding via ``split_dataset_by_node`` (each GPU streams disjoint data; ``rank``/``world_size`` - are auto-resolved from the Accelerate state, so nothing needs to be passed explicitly); -- DataLoader-worker shard splitting (no duplicate frames within a rank); -- native `datasets` resume: the loader checkpoints stream state via ``state_dict()`` (``torchdata`` StatefulDataLoader when available, so ``num_workers > 0`` resumes too); -- an explicit video-decoder cache size so the working set of open decoders does not thrash. - -Launch with Accelerate (single node, N GPUs): - - accelerate launch --num_processes=8 examples/scaling/train_streaming_multinode.py \ - --repo_id=lerobot/droid_1.0.1 --batch_size=64 - -Multinode runs launch the same script with your cluster's accelerate/SLURM setup. - -Pass ``--dummy`` to skip the model entirely and measure pure dataloading throughput. -""" - -import argparse -import time -from pathlib import Path - -import torch -from accelerate import Accelerator -from torch.utils.data import DataLoader - -from lerobot.datasets import LeRobotDatasetMetadata, StreamingLeRobotDataset -from lerobot.utils.constants import ACTION - - -def parse_args() -> argparse.Namespace: - parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument("--repo_id", type=str, default="lerobot/droid_1.0.1") - parser.add_argument( - "--root", type=str, default=None, help="Local/prewarmed dataset root (else stream from Hub)." - ) - parser.add_argument("--output_dir", type=str, default="outputs/train/streaming_multinode") - parser.add_argument("--steps", type=int, default=1000) - parser.add_argument("--batch_size", type=int, default=64, help="Per-process batch size.") - parser.add_argument("--num_workers", type=int, default=8) - parser.add_argument( - "--episode_pool_size", - type=int, - default=64, - help="Whole episodes open per consumer (randomness knob).", - ) - parser.add_argument("--video_decoder_cache_size", type=int, default=None) - parser.add_argument("--n_action_steps", type=int, default=16, help="Action-chunk length (delta horizon).") - parser.add_argument("--save_freq", type=int, default=200) - parser.add_argument("--log_freq", type=int, default=20) - parser.add_argument("--resume_from", type=str, default=None, help="Checkpoint dir to resume from.") - parser.add_argument("--dummy", action="store_true", help="Skip the model; measure dataloading only.") - return parser.parse_args() - - -def make_dataloader( - args: argparse.Namespace, meta: LeRobotDatasetMetadata -) -> tuple[DataLoader, StreamingLeRobotDataset]: - # Supervise an action chunk; delta_timestamps drive the SARM-style temporal window. - delta_timestamps = {ACTION: [t / meta.fps for t in range(args.n_action_steps)]} - # rank / world_size are resolved automatically from the Accelerate state inside the dataset. - dataset = StreamingLeRobotDataset( - args.repo_id, - root=args.root, - delta_timestamps=delta_timestamps, - episode_pool_size=args.episode_pool_size, - video_decoder_cache_size=args.video_decoder_cache_size, - tolerance_s=1e-3, - ) - # torchdata's StatefulDataLoader checkpoints each worker's dataset state through the - # dataset's native state_dict protocol, making resume work with num_workers > 0. Fall back - # to the plain DataLoader (resume then requires num_workers=0). - try: - from torchdata.stateful_dataloader import StatefulDataLoader - - loader_cls = StatefulDataLoader - except ImportError: - loader_cls = DataLoader - loader = loader_cls( - dataset, - batch_size=args.batch_size, - num_workers=args.num_workers, - pin_memory=True, - drop_last=True, - prefetch_factor=2 if args.num_workers > 0 else None, - ) - return loader, dataset - - -def main() -> None: - args = parse_args() - accelerator = Accelerator() - output_dir = Path(args.output_dir) - if accelerator.is_main_process: - output_dir.mkdir(parents=True, exist_ok=True) - - meta = LeRobotDatasetMetadata(args.repo_id, root=args.root) - loader, dataset = make_dataloader(args, meta) - - if args.dummy: - model = optimizer = None - else: - from lerobot.policies.act import ACTConfig, ACTPolicy - from lerobot.utils.feature_utils import dataset_to_policy_features - - features = dataset_to_policy_features(meta.features) - output_features = {k: ft for k, ft in features.items() if k == ACTION} - input_features = {k: ft for k, ft in features.items() if k not in output_features} - cfg = ACTConfig(input_features=input_features, output_features=output_features) - model = ACTPolicy(cfg) - optimizer = torch.optim.Adam(model.parameters(), lr=1e-4) - # Do NOT prepare the dataloader: the dataset is already rank-disjoint via - # split_dataset_by_node, and accelerate's IterableDatasetShard would keep only every - # world_size-th batch of it (silently training on 1/N of the data while decoding all - # of it). Batches are moved to the device manually in the loop. - model, optimizer = accelerator.prepare(model, optimizer) - - # Resume: native datasets stream state, saved per rank. With torchdata's StatefulDataLoader - # the state covers every worker; with the plain DataLoader it is exact for num_workers=0. - can_checkpoint_loader = hasattr(loader, "state_dict") - if args.resume_from is not None: - state_path = Path(args.resume_from) / f"dataset_state_rank{accelerator.process_index}.pt" - state = torch.load(state_path, weights_only=False) # plain dict of stream offsets # nosec B614 - if can_checkpoint_loader: - loader.load_state_dict(state) - else: - dataset.load_state_dict(state) - accelerator.print(f"Resumed dataset stream from {state_path}") - - step = 0 - frames_seen = 0 - window_start = time.perf_counter() - done = False - while not done: - for batch in loader: - if model is not None: - batch = {k: (v.to(accelerator.device) if torch.is_tensor(v) else v) for k, v in batch.items()} - loss, _ = model.forward(batch) - accelerator.backward(loss) - optimizer.step() - optimizer.zero_grad() - - step += 1 - frames_seen += args.batch_size - if step % args.log_freq == 0: - elapsed = time.perf_counter() - window_start - fps_per_proc = (args.log_freq * args.batch_size) / max(elapsed, 1e-9) - total_fps = fps_per_proc * accelerator.num_processes - accelerator.print( - f"step {step} | {fps_per_proc:.1f} frames/s/proc | {total_fps:.1f} frames/s total" - + ("" if model is None else f" | loss {loss.item():.3f}") - ) - window_start = time.perf_counter() - - if step % args.save_freq == 0: - ckpt = output_dir / f"checkpoint-{step}" - if accelerator.is_main_process: - ckpt.mkdir(parents=True, exist_ok=True) - accelerator.wait_for_everyone() - # Every rank saves its own stream state: shard positions differ per rank. - state = loader.state_dict() if can_checkpoint_loader else dataset.state_dict() - torch.save(state, ckpt / f"dataset_state_rank{accelerator.process_index}.pt") - if model is not None and accelerator.is_main_process: - accelerator.unwrap_model(model).save_pretrained(ckpt) - - if step >= args.steps: - done = True - break - - accelerator.print(f"End of training: {step} steps, ~{frames_seen} frames/proc") - - -if __name__ == "__main__": - main() diff --git a/src/lerobot/configs/default.py b/src/lerobot/configs/default.py index 08fdda209..6cf095502 100644 --- a/src/lerobot/configs/default.py +++ b/src/lerobot/configs/default.py @@ -42,7 +42,7 @@ class DatasetConfig: # Whole episodes each streaming consumer keeps open to shuffle across (the randomness knob). # Larger mixes more episodes per batch at the cost of cold-start latency; RAM stays small because # the pool holds tabular rows only. Ignored when streaming is False. - streaming_episode_pool_size: int = 64 + streaming_episode_pool_size: int = 1024 def __post_init__(self) -> None: if self.episodes is not None: diff --git a/src/lerobot/datasets/streaming_dataset.py b/src/lerobot/datasets/streaming_dataset.py index 7d4f9fa8a..a3c337dfb 100644 --- a/src/lerobot/datasets/streaming_dataset.py +++ b/src/lerobot/datasets/streaming_dataset.py @@ -79,7 +79,7 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset): dataset = StreamingLeRobotDataset( repo_id="your-dataset-repo-id", delta_timestamps={"action": [0.0, 0.1, 0.2]}, - episode_pool_size=64, + episode_pool_size=1024, ) for sample in dataset: ... @@ -97,7 +97,7 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset): revision: str | None = None, force_cache_sync: bool = False, streaming: bool = True, - episode_pool_size: int | None = 64, + episode_pool_size: int | None = 1024, frame_shuffle_buffer_size: int | None = None, buffer_size: int | None = None, max_num_shards: int | None = None, @@ -128,7 +128,7 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset): episode_pool_size (int, optional): Whole episodes each consumer keeps open to shuffle across — the randomness knob. Larger mixes more episodes per batch (closer to map-style uniform) at the cost of cold-start latency and frame-buffer RAM. - Defaults to 64. + Defaults to 1024. frame_shuffle_buffer_size (int | None, optional): Frame-level shuffle buffer after the episode pool. Defaults to ``episode_pool_size x average episode length`` (capped), which matches the pool's mixing radius. @@ -178,7 +178,7 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset): self.shuffle = shuffle self.streaming = streaming - self.episode_pool_size = max(1, episode_pool_size) if episode_pool_size else 64 + self.episode_pool_size = max(1, episode_pool_size) if episode_pool_size else 1024 self._return_uint8 = return_uint8 self.rank, self.world_size = self._resolve_distributed(rank, world_size) diff --git a/src/lerobot/scripts/lerobot_train.py b/src/lerobot/scripts/lerobot_train.py index 3d210f00b..fc57e5602 100644 --- a/src/lerobot/scripts/lerobot_train.py +++ b/src/lerobot/scripts/lerobot_train.py @@ -387,7 +387,7 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None): logging.info(f"{num_total_params=} ({format_big_number(num_total_params)})") # create dataloader for offline training - if hasattr(active_cfg, "drop_n_last_frames"): + if hasattr(active_cfg, "drop_n_last_frames") and not cfg.dataset.streaming: shuffle = False # A dedicated generator (rather than the global torch RNG) lets accelerator.prepare # synchronize the shuffle permutation across ranks, keeping batch shards disjoint even @@ -426,9 +426,16 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None): # Prepare everything with accelerator accelerator.wait_for_everyone() - policy, optimizer, dataloader, lr_scheduler = accelerator.prepare( - policy, optimizer, dataloader, lr_scheduler - ) + if cfg.dataset.streaming: + # The streaming IterableDataset is already rank-disjoint via split_dataset_by_node, so we must + # NOT hand the dataloader to accelerate: its IterableDatasetShard would keep only every + # world_size-th batch of each rank's already-disjoint stream (silently training on 1/N of the + # data while decoding all of it). Batches are moved to the device manually in the loop below. + policy, optimizer, lr_scheduler = accelerator.prepare(policy, optimizer, lr_scheduler) + else: + policy, optimizer, dataloader, lr_scheduler = accelerator.prepare( + policy, optimizer, dataloader, lr_scheduler + ) dl_iter = cycle(dataloader) policy.train() @@ -468,6 +475,9 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None): for _ in range(step, cfg.steps): start_time = time.perf_counter() batch = next(dl_iter) + if cfg.dataset.streaming: + # The streaming dataloader is not accelerate-prepared (see above), so move to device here. + batch = {k: (v.to(device, non_blocking=True) if torch.is_tensor(v) else v) for k, v in batch.items()} for cam_key in dataset.meta.camera_keys: if cam_key in batch and batch[cam_key].dtype == torch.uint8: batch[cam_key] = batch[cam_key].to(dtype=torch.float32) / 255.0