# 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()