mirror of
https://github.com/huggingface/lerobot.git
synced 2026-07-18 23:41:55 +00:00
Pepijn
1050c2fb6c
Replace the shard/Backtrackable/decoded-shuffle-buffer internals with an
episode pool: each (rank x worker) consumer keeps episode_pool_size whole
episodes' tabular rows in RAM and emits uniformly random frames across
them. delta_timestamps windows become exact in-RAM slices with correct
boundary padding (the Backtrackable machinery and its lookback/lookahead
ceilings are gone), and video is decoded only when a sample is emitted,
so pool memory stays tabular-sized instead of buffer_size decoded
samples.
- Prefetch-on-admit: when streaming from a remote source, each pooled
episode's video files download to a local cache in the background
(refcounted, since v3 packs several episodes per file; deleted on
eviction), so decode-on-exit reads local bytes instead of paying
network seek latency.
- Per-consumer RNG derived from (seed, epoch, rank, worker): consumers
decorrelated, runs reproducible, epochs reshuffle automatically.
- Deterministic fast-forward resume: load_state_dict takes the trainer's
{batches_consumed, batch_size}; each worker re-derives its own skip
from the DataLoader's round-robin batch assignment and replays
tabular-only (no decode). Exact within an epoch, works with
num_workers > 0, and the same state file serves every rank. Replaces
the per-shard HF state_dict approach, which lived in worker processes
and could not be captured from the trainer.
- Shard-cap default removed (max_num_shards=None uses every parquet
shard); runtime warnings for non-divisible world sizes (datasets
degrades to read-everything splitting) and workers left without
shards.
- episode_pool_size replaces buffer_size (deprecated, ignored with a
warning); decoder cache sized to the pool working set, capped at 128.
Legacy order-replication tests asserted the old buffer algorithm
step-by-step and are rewritten as behavior contracts (exactly-once
coverage, per-seed determinism, epoch reshuffle). Value-level parity
tests against the map-style dataset pass unchanged.
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
41 lines
1.4 KiB
Bash
41 lines
1.4 KiB
Bash
#!/bin/bash
|
|
#SBATCH --job-name=bench_stream
|
|
#SBATCH --nodes=2
|
|
#SBATCH --ntasks-per-node=1
|
|
#SBATCH --gpus-per-node=8
|
|
#SBATCH --cpus-per-task=96
|
|
#SBATCH --exclusive
|
|
#SBATCH --time=02:00:00
|
|
#SBATCH --output=logs/%x-%j.out
|
|
|
|
# Per-node dataloading benchmark for StreamingLeRobotDataset across 1-2 nodes. Each node runs an
|
|
# independent dummy-consumer benchmark; per-node throughput should be independent (separate network).
|
|
# Results are written per (node, source, mode) under --out_dir.
|
|
#
|
|
# Submit with: sbatch slurm/benchmark_streaming_robocasa.sh
|
|
# Override the source label for cold/warm bucket runs: SOURCE=warmed_bucket sbatch slurm/benchmark_streaming_robocasa.sh
|
|
|
|
set -euo pipefail
|
|
|
|
REPO_ID=${REPO_ID:-pepijn223/robocasa_pretrain_human300_v4}
|
|
SOURCE=${SOURCE:-hub}
|
|
OUT_DIR=${OUT_DIR:-benchmarks/streaming/results}
|
|
|
|
export HF_HOME=${HF_HOME:-$SCRATCH/hf_home}
|
|
export TOKENIZERS_PARALLELISM=false
|
|
|
|
# One benchmark process per node (each saturates the node's DataLoader workers + network independently).
|
|
srun --kill-on-bad-exit=1 bash -c '
|
|
for MODE in single sarm; do
|
|
python benchmarks/streaming/benchmark_streaming.py \
|
|
--repo_id '"$REPO_ID"' \
|
|
--source '"$SOURCE"' \
|
|
--mode $MODE \
|
|
--batch_size 64 \
|
|
--num_workers 12 \
|
|
--episode_pool_size 64 \
|
|
--num_batches 300 \
|
|
--out_dir '"$OUT_DIR"'/node${SLURM_NODEID}
|
|
done
|
|
'
|