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

feat(streaming): default episode pool 1024 and wire streaming into lerobot-train

Browse Source 
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 <cursoragent@cursor.com>
This commit is contained in:
pepijn
2026-06-12 09:24:32 +00:00
parent 38106ea6b4
commit 674c990a39
4 changed files with 19 additions and 201 deletions
Download Patch File Download Diff File Expand all files Collapse all files
examples/scaling/train_streaming_multinode.py
-192
View File
@@ -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()
src/lerobot/configs/default.py
+1 -1
View File
@@ -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:
src/lerobot/datasets/streaming_dataset.py
+4 -4
View File
@@ -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)
src/lerobot/scripts/lerobot_train.py
+14 -4
View File
@@ -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