[WIP] LanceDB dataloader#100
Open
AyushExel wants to merge 43 commits into
Open
Conversation
LanceDB drop-in replacements for the action (LeRobot/DROID), VLM (WebDataset/
LLaVA-OneVision), and local vision-SFT (Bridge) dataloaders, with offline
converters, equivalence tests, and fair (same-device, shuffled) benchmarks.
Measured wins (single node, 4x L40S, CPU decode, shuffled, local):
action 2.0-2.5x, vision-SFT 6.5x e2e, VLM 3.7x raw access, combined 2.75x.
Equivalence: action/token-ids bit-exact; pre-composed video PSNR ~32-37dB
(one-time lossy re-encode); validated to preserve the training signal.
Method: store a training-optimized representation Lance can serve but the
canonical LeRobot/WebDataset formats can't (pre-composed/pre-resized all-intra
per-episode clips as blob-v2; columnar random access + global shuffle via the
Permutation API). No per-frame JPEG (disk stays 0.35x the original video).
cosmos_framework/data/lance/ dataloaders + docs (README/RESULTS/WHY_BASE_CANT/
OPTIMIZATION_ROADMAP/VALIDATION)
tools/lance_datagen/ offline converters
benchmarks/lance/ throughput benchmarks
tests/data/lance/ equivalence tests
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Pre-composed clips are 0.35x the original DROID video (gop=1), not a blowup (per-frame JPEG would be 1.8x). Full table + GOP tradeoffs. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
bench_filtered.py: Lance predicate pushdown vs WebDataset stream-and-discard. At 10% selectivity Lance is ~122x faster and reads 0.1x the bytes (2213->221 MB); webdataset must read 100% always (sequential tar, no skip). Structural Lance win. Adds EXPERIMENTS.md (results) and CONVERSION_EXPLAINED.md (layman conversion guide). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
LoRA-SFT Qwen2.5-VL-3B with base vs Lance VLM loader (same init/seed/order/lr). base-vs-lance mean |Δloss|=3.5e-3 < base-vs-base2 control 5.3e-3; step-0 identical; eval losses match within nondeterminism. Lance loader is a training-equivalent drop-in. Adds benchmarks/lance/train_compare_vlm.py + TRAIN_EQUIVALENCE.md. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
4xL40S DDP LoRA-SFT Qwen2.5-VL-3B: base 29.7s/epoch vs lance 30.8s/epoch, both at 1.1% data-wait (compute-bound) and identical loss. Loader speedup => faster training only in the data-bound regime; honest finding + train_multigpu_time.py harness. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…x faster epoch 4xL40S, action loader from S3, tiny compute head (fast-GPU proxy) => data-bound: base 5.4s/epoch vs lance 3.1s/epoch (1.74x, 149->258 samp/s global). Confirms the loader speedup converts to faster training wall-clock once data-bound (H100/8x/S3), toward the 2.5x decode ceiling. Compute-bound heavy-model run was 1.0x. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Qwen2.5-VL-3B LoRA, 4 epochs/2400 steps, base/base2/lance. base-vs-lance <= base-vs-base2 nondeterminism floor on loss, eval, LoRA weights; held-out greedy generations identical 12/12. Dataloader confirmed correct over a full multi-epoch run. Adds train_equiv_real.py. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
A) Borrow from base: LanceDROIDComposedIterable (episode-shuffle stream) — per-episode decoder built once + reused vs RandomSampler rebuild. Measured S3 2.55x (35.3->90.0 samp/s), neutral local. B) Lance-optimal: batch_readahead in LanceVLMShuffleScan; AUDIT.md documents compact_files()/create_scalar_index/file_system-IPC recommendations + what we already do right. Adds bench_episode_shuffle.py. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…asure S3 data-bound demo unchanged at tiny scale (subset ~3 episodes fits decoder cache so RandomSampler never misses; episode-shuffle slightly slower from overhead). Episode-shuffle's 2.55x S3 win needs the many-episode regime (cache << episodes), per bench_episode_shuffle. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Per-episode take_blobs(indices=[row]) in a loop does sequential S3 round-trips; one batched take_blobs([all rows]) issues them concurrently (measured ~2.3x on the fetch). Random-access S3 dataloader throughput 44 -> 60.8 samp/s (1.0x -> 1.45x vs base). _ensure_decoders never evicts an episode needed by the current batch (fixes KeyError). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
The README 2.5x compared base-RandomSampler (~2x artificially slow); the production base uses episode-shuffle. Faithful: base-episode vs lance-episode = 1.89x local / 1.69x S3 (327 eps, cache 16, 8 workers, LANCE_IO_THREADS=256). lance-random 2.15x local but 1.09x S3-at-scale (re-fetches clips) -> episode-shuffle is the right pattern. Adds bench_action_faithful.py; AUDIT.md carries the lance-S3 concurrency checklist. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…imes) bench_combined_faithful.py now runs the 3-loader 1:1:1 mixer in three regimes: LOCAL (apples-to-apples), S3 (lance native vs stock base access), and DEFAULT-MIXED (each loader on its real default: action local, vision-SFT S3 via boto3, VLM HF-Hub streaming). Self-contained (no longer imports the removed bench_combined/bench_action). Adds the stock-faithful boto3 vision-SFT base and the HF-Hub streaming VLM base, plus base-random mode to bench_action_faithful for the 2x2. _env.sh sets the LD_LIBRARY_PATH torchcodec needs. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Remove benches superseded by the faithful variants (bench_action, bench_combined) and transient experiment scripts (bench_episode_shuffle, bench_vanilla_vs_lance, train_compare_vlm). Consolidate 9 docs -> 3 (README, RESULTS, VALIDATION); the useful content from AUDIT/CONVERSION_EXPLAINED/EXPERIMENTS/OPTIMIZATION_ROADMAP/TRAIN_EQUIVALENCE/ WHY_BASE_CANT is folded into the keepers. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
…DUCE.md Correct the earlier inflated/unfair claims. Combined dataloader speedup is ~2.6-3.1x across all three storage regimes, gated by the slowest (video) loader; per-loader: action ~1.9x (worker-count-dependent, not 2.5x; shuffle-mode-neutral locally), vision-SFT ~7.6x local, VLM raw batch-dependent (e2e ~1x). Retract the bogus 8.5x S3 number (a FUSE artifact in the vision-SFT base; the stock boto3 base collapses it to 2.64x). Record the methodology lessons. Add REPRODUCE.md: full standalone recipe (env, datasets, conversions, S3 setup, all 3 regimes, expected numbers) linked from the README. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
… training bench, consolidated docs - action/vision-SFT loaders auto-detect plain large_binary vs blob-v2; columnar take() is ~6.3x faster than take_blobs on S3 for <2MB clips (converters default --storage plain). - bench_combined_faithful: force spawn for all sub-loaders (fixes fork/spawn SIGABRT) and add --action-workers/--vlm-workers/--vsft-workers for per-loader rebalancing (the dominant lever). - new benches: bench_blob_levers, bench_take_vs_blobs, bench_cold_cache, train_combined_e2e (e2e training). - VLM equivalence test (records byte-identical vs the HF stream). - docs: consolidate all numbers into BENCHMARKS.md, add HOW_IT_WORKS.md + RUN_BENCHMARKS_H100.md, slim README to headlines, remove RESULTS.md. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
… store Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…paths) The docs referenced these but they lived only in scratchpad. Parameterized via env (REPO/DATA/FUSE/S/BUCKET/REGION/ALLOCS/LAYERS/WORKERS) with dev-box defaults so they run on another machine. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
LanceDB-backed replacements for the three Cosmos training dataloaders (DROID action,
vision-SFT, VLM), consolidated to the minimal needed surface:
- Loaders: cosmos_framework/data/lance/{action_dataset,vision_sft_dataset,vlm_dataset}.py
— composed/raw DROID, pre-resized vision-SFT, columnar VLM (random + chunked-shuffle).
Free the base's dead per-frame index for lower per-worker memory at scale.
- Benchmarks: benchmarks/lance/ — throughput (per-loader + combined, LOCAL/S3 via genuine
base loaders + explicit S3 standins) and memory (bench_memory + build_scaled_droid).
- Converters: tools/lance_datagen/.
- Tests: tests/data/lance/ — equivalence vs the genuine base loaders.
- Docs consolidated into a single cosmos_framework/data/lance/README.md.
Verified: equivalence tests pass (action bit-exact, vision-SFT vs genuine SFTDataset
token-exact, VLM byte-identical; composed action video within ~1.5% re-encode); throughput
~3.3-4.9x combined, vision-SFT ~7-8x; per-worker memory ~3x lower at scale.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
- Move all function-level imports to module scope across the loaders, converters, benchmarks, and tests (proper style; verified no circular imports and no fork/teardown regressions — equivalence + bench_vlm/bench_memory(fork) re-run clean). - test_action: fold the redundant single-item loop into the batch check. - base_standins: drop the unused `import time`. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…factor main moved add_special_tokens into sequence_packing.modalities (sequence_packing is now a package). Update the vision-SFT loader to import it from there, matching the base SFTDataset. Equivalence tests pass against merged main. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…l surface) - Remove LanceDROIDDataset (raw mp4-bytes action loader): it was never benchmarked and no shipped converter builds its table (only the composed table is produced). We keep the composed loader, whose labels are bit-exact and video is within ~1.5% re-encode. Drops its now-unused imports (torch.nn.functional, torchvision, lancedb, Permutation, _IMAGE_FEATURES) and its equivalence test; README equivalence wording corrected (action = labels bit-exact + video within re-encode tolerance, not bit-exact video). - Remove benchmarks/lance/train_combined_e2e.py + run_e2e.sh (e2e training example) to reduce surface; throughput + memory benches remain. Equivalence: composed action + vision-SFT + VLM all pass. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…used blob path
Per the integration brief, the loaders now do their columnar reads through the LanceDB
Permutation API instead of dropping to pylance lance.dataset.take:
- action composed + vision-SFT: open via lancedb, build Permutation.identity(table)
.select_columns(...).with_format("arrow"), read video_bytes via __getitems__. The
episode/metadata index is read the same way. Matches the VLM loader + the object-detection
reference pattern.
- Removed the blob auto-detect + pylance take_blobs fallback and the converters' --storage
blob option: media is stored plain large_binary only. An isolated S3 test confirmed plain
is fastest for our small (<~2MB) clips, while blob-v2 wins for larger payloads (>=~8-16MB)
*when read in parallel* (a serial take_blobs loop is latency-bound — that was the source of
the earlier inflated "blob is slow" number). Loaders carry a TODO to move to blob-v2 if
per-row clip sizes grow; README documents the crossover.
Equivalence: composed action + vision-SFT + VLM all pass.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…al-model validation
- _read_clip_bytes keys results by row (Permutation take returns sorted order)
- LanceVLMShuffleScan on the Permutation API (no pylance dependency)
- get_lance_{action,vlm} drop-in factories; forward_equivalence.py (real Cosmos3-Nano)
- refresh benchmarks (throughput up to 4.8x, ~2.7x lighter at scale); ruff format/cleanup
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
…e in README _rows is a freeable redundancy DROIDLeRobotDataset never reads (verified bit-identical batches with/without it) — not a fundamental Lance advantage; a _rows-freed base reaches memory parity. Real wins are throughput + S3. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
…dec, GOP, current-vs-Lance) Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
…, blob-v2 note Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
…hf-stream vs Lance) Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
…) + resolution-derivation note Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
…set); clarify streams vs channels Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
…parquet labels) Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Converter writes {table}_frames/_tasks/_episodes (labels dumped verbatim from the
base arrays); loader takes lance_uri only and rebuilds the base's compact arrays
from Lance — labels bit-exact both action spaces, no LeRobot tree at train time.
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Loader now subclasses the new DROIDLeRobotDataset: split/span index via the base's own helpers, _fetch_sample assembles the LeRobot-shaped sample dict from Lance, _compose_multi_view decodes the stored composed clip. Labels bit-exact (joint_pos + midtrain, incl. the new train/val split); converter dumps from the base's LeRobot table; benchmarks/standins/tests updated for the new base + vfm->generator renames. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Review-driven fixes: - VLM map take: key by row, map back to requested order (dup-safe) - VLM shuffle-scan / vision iterable: (rank, worker) sharding with torch.distributed fallback; reshuffle each pass - VLM factory honors n (base .select(range(n)) semantics) - vision: base caption post-processing (caption_suffix, CFG dropout, duration/resolution suffixes) -> dense captions token-exact too - vision: skip-not-crash on short windows / missing captions (base process_one_sample contract); clamp window frame indices like the base's sequential decode; raise on build/serve resolution mismatch - vision converter: use the base's own metadata load + flatten (duration/min-frames filters); reject unstored caption keys - action: support use_state for midtrain/ee_pose_delta (new base allows it); fail loudly on empty filter-dict match; converter guards single-shard roots Cleanup: drop unused LanceDROIDComposedIterable, _EpisodeShuffle (benches use the genuine ActionIterableShuffleDataset), obsolete build_scaled_droid + bench_memory _rows plumbing, dead attrs; export all loaders from __init__; docs updated to match; tests cover midtrain+use_state, dense captions, unsorted+dup VLM batches. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…on size row Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Walkthrough of the design and the equivalence invariants:
cosmos_framework/data/lance/design.mdcosmos_framework/data/lance/README.mdEquivalence
Wherever possible the loaders reuse the base code rather than reimplement it (the action loader subclasses
DROIDLeRobotDatasetand builds its split/span index with the base's ownsplit_episode_ids/build_episode_spans; vision reuses_select_caption, the base caption post-processing order, andtokenize_caption), so equivalence is structural:joint_posandmidtrain, with and withoutuse_state, including the train/val split.text_token_idstoken-exact for structured and dense captions (incl. CFG dropout and duration/resolution conditioning suffixes).lr=0, fixed seed, per-step loss: VLM exact 0.00%, action/vision within the re-encode tolerance while different samples differ ~10× more).Tests:
pytest tests/data/lance/compares every loader per-sample against the genuine base loader.Benchmarks
Setup: 327-episode subset of the public
lerobot/droid_1.0.1(materialized viatools/lance_datagen/prepare_droid_subset.py), Bridge vision-SFT jsonl, LLaVA-OneVisionfigureqa. The base side is always the genuine shipped loader with its production access pattern (episode-shuffle for action on both sides); for S3 the base gets a materialization standin (download-then-run) since it has no native S3 path. Full methodology in the README.Combined 3-loader throughput (samples/s, batch 16):
Per-loader, standalone (tuned workers):
Where the wins come from: action decodes one pre-composed clip instead of three runtime views + resize/concat; vision-SFT decodes a pre-resized short-GOP clip in-process instead of a per-sample ffmpeg subprocess; VLM gets row-level access and S3-friendly chunked-shuffle scans instead of Hub streaming. The composed action table is also 0.37× the size of the original three views (1.47 GB → 0.55 GB) despite all-intra encoding.
Memory
A bug in upstream caused memory/scalability issues where droid dataset class had unused _rows member that cached frames. It has also been fixed.
TODO:
large_binaryfor now; it will move to blob encoding (blob-v2) once the lancedb-level blob API is available. All reads go through the lancedbPermutationAPI — no pylance dependency.max_num_history_actions,val_temp_seg, and multi-shard roots are not supported (documented in design.md; the converter fails loudly on multi-shard roots).