mirror of
https://github.com/huggingface/lerobot.git
synced 2026-07-26 11:16:00 +00:00
128 lines
6.1 KiB
Plaintext
128 lines
6.1 KiB
Plaintext
# Training Dataset Streaming
|
|
|
|
Training-time dataset streaming lets `lerobot-train` consume a LeRobotDataset v3 without first
|
|
downloading its complete Parquet and video payload. Enable it through the existing public switch:
|
|
|
|
```bash
|
|
lerobot-train \
|
|
--dataset.repo_id=OWNER/DATASET \
|
|
--dataset.streaming=true \
|
|
--policy.type=act \
|
|
--output_dir=outputs/train/act_streaming
|
|
```
|
|
|
|
This feature is independent of `--dataset.streaming_encoding=true`. `streaming_encoding` controls
|
|
how videos are written while recording; `dataset.streaming` controls how an existing dataset is read
|
|
during training. Recording and rollout encoders are not used by this training path.
|
|
|
|
## How an epoch is read
|
|
|
|
Each distributed rank and DataLoader worker owns a disjoint set of complete episodes. Within that
|
|
set, an exact-coverage sampler mixes a bounded pool of episodes while visiting every selected frame
|
|
once per epoch. Parquet columns and compressed MP4 byte ranges are prefetched from the same episode
|
|
admission frontier. Temporal history and future windows are resolved inside the complete episode,
|
|
including the same boundary padding masks as map-style loading.
|
|
|
|
The default map-style `LeRobotDataset` behavior is unchanged when `--dataset.streaming=false`.
|
|
|
|
## MP4 sidecars and the first run
|
|
|
|
Video streaming uses a small MP4 index sidecar. Dataset initialization first checks the
|
|
revision-keyed local cache, then looks for a valid published sidecar. If neither is available,
|
|
LeRobot builds the sidecar locally while holding a process lock and installs it atomically. A
|
|
failed or interrupted build does not replace the previous valid file.
|
|
|
|
Training is read-only: it never uploads a sidecar or modifies the dataset repository. On a cluster
|
|
with node-local caches, the first job may build once per node. A shared LeRobot cache avoids that
|
|
duplication.
|
|
|
|
Dataset maintainers can build a sidecar ahead of time:
|
|
|
|
```bash
|
|
uv run python scripts/build_mp4_sidecar.py \
|
|
--repo-id=OWNER/DATASET \
|
|
--revision=COMMIT_SHA \
|
|
--data-root=hf://datasets/OWNER/DATASET@COMMIT_SHA \
|
|
--output=/tmp/dataset-mp4-sidecar.npz
|
|
```
|
|
|
|
Publication is always explicit. Add `--push` only after validating the complete-dataset sidecar.
|
|
Subset sidecars cannot be published.
|
|
|
|
## Configuration
|
|
|
|
The production defaults are:
|
|
|
|
| Option | Default | Meaning |
|
|
| ----------------------------- | ------: | --------------------------------------------------- |
|
|
| `streaming_episode_pool_size` | 32 | Complete episodes mixed by each worker |
|
|
| `streaming_prefetch_episodes` | 8 | Episodes fetched ahead of the active pool |
|
|
| `streaming_byte_budget_gb` | 8 | Maximum synthesized MP4 bytes per worker |
|
|
| `streaming_data_root` | unset | Optional local, Hub, bucket, or fsspec payload root |
|
|
|
|
Memory limits are per DataLoader worker. For example, four workers with the default byte budget can
|
|
use up to 32 GiB for compressed episode video bytes, plus Parquet batches, decoders, and framework
|
|
overhead. Start with fewer workers or a smaller budget on memory-constrained hosts:
|
|
|
|
```bash
|
|
lerobot-train \
|
|
--dataset.repo_id=OWNER/DATASET \
|
|
--dataset.streaming=true \
|
|
--dataset.streaming_episode_pool_size=16 \
|
|
--dataset.streaming_prefetch_episodes=4 \
|
|
--dataset.streaming_byte_budget_gb=4 \
|
|
--num_workers=2 \
|
|
--policy.type=act \
|
|
--output_dir=outputs/train/act_streaming
|
|
```
|
|
|
|
If metadata remains in a dataset repository while payload files are mirrored elsewhere, set
|
|
`--dataset.streaming_data_root`. Supported values include local paths, revision-qualified
|
|
`hf://datasets/...` roots, `hf://buckets/...` roots, and fsspec URLs.
|
|
|
|
## Resume and shuffle migration
|
|
|
|
The earlier streaming reader used a bounded row shuffle buffer. The episode reader instead has
|
|
deterministic exact-coverage ordering derived from the seed and epoch. Checkpoint resume restores the
|
|
per-rank sample offset using the checkpoint batch size. Changing distributed world size or
|
|
DataLoader worker count changes episode ownership; changing batch size changes the batch boundaries.
|
|
For sample-exact comparisons, resume with the same world size, worker count, and batch size.
|
|
|
|
## Benchmarking
|
|
|
|
Measure the integrated path on the same hosts used for training:
|
|
|
|
```bash
|
|
uv run python scripts/bench_streaming_dataset.py \
|
|
--repo-id=OWNER/DATASET \
|
|
--revision=COMMIT_SHA \
|
|
--batch-size=16 \
|
|
--num-workers=4 \
|
|
--summary-json=streaming-benchmark.json
|
|
```
|
|
|
|
The output records source revisions, startup and first-batch latency, steady-state throughput,
|
|
batch-wait percentiles, duplicate indices, memory high-water marks, and the exact settings. Run a
|
|
separate end-to-end training A/B to include policy compute, device transfer, and optimizer time.
|
|
Do not compare results unless the code revision, dataset revision, hardware, and settings match.
|
|
|
|
For lower-level byte-fetch, retry, decoder-open, and refill diagnostics, build the sidecar explicitly
|
|
and run `scripts/bench_episode_byte_cache.py` with the same `--repo-id`, `--revision`,
|
|
`--data-root`, and `--sidecar-path`. Treat this as a stage profiler; throughput claims should come
|
|
from the integrated dataset benchmark and end-to-end training A/B.
|
|
|
|
## Troubleshooting
|
|
|
|
- **The first batch takes a long time:** check logs for a sidecar build. Reuse a shared cache or
|
|
publish a validated complete sidecar explicitly.
|
|
- **A sidecar lock times out:** another process may still be indexing the same revision. Confirm it
|
|
is healthy before removing a stale lock.
|
|
- **A rank owns no data:** reduce the number of ranks or workers so every rank owns at least one
|
|
selected episode.
|
|
- **The byte budget is exceeded:** lower the episode pool, increase the per-worker byte budget, or
|
|
use fewer/larger episodes per worker.
|
|
- **Remote reads cannot authenticate:** verify the normal Hugging Face token or fsspec credentials
|
|
are available in every worker environment. Credentials are never embedded in the sidecar.
|
|
- **Refill stalls are high:** compare p95/p99 batch wait, reduce network contention, raise prefetch
|
|
gradually, and verify that the dataset sidecar matches the exact revision.
|