Skip to content

feat(data): shard-backed dataset, H5-to-shard converter, planner_shards package - #416

Closed
Dionysus326 wants to merge 1 commit into
new-architecture/mainfrom
feat/new-arch-shard-dataset
Closed

Dionysus326 wants to merge 1 commit into
new-architecture/mainfrom
feat/new-arch-shard-dataset

Conversation

@Dionysus326

Copy link
Copy Markdown

Summary

Adds a tar-shard data path to the new-architecture training tree, next to the existing HDF5 loader:

  • packages/planner_shards — the versioned shard dataset library (packer, manifests, key-sets, DDP-safe loader) as a uv workspace package. It is a copy of the data_pipeline / shard_* modules from the tier4-main tree (feat(data_pipeline): versioned shard dataset (WebDataset-style tars + parquet manifest) with DDP loader #390) with imports renamed, because both trees use the package name diffusion_planner. Deliberate temporary duplication: neither PR can depend on the other while both are drafts; once one lands, the other switches to a dependency.
  • ShardPlannerDataset / build_shard_dataloader (diffusion_planner.data) — iterates one rank's share of a packed dataset and applies the same frame transforms as PlannerDataset. Rank and world size come from the torchrun environment.
  • configs/train/dataloader/shards.yaml — same transform stack as default.yaml; select with train/dataloader@dataloader=shards.
  • scripts/train/train.py — when the loader wraps ShardPlannerDataset, it is not passed through accelerator.prepare (it is already sharded per rank) and batches are moved to the device in the loop. The HDF5 path is unchanged.
  • scripts/dataset/convert_h5_to_shards.py (diffusion_planner.data.h5_to_shards) — converts the frames of a Parquet index into shards: one member per frame, one partition per rosbag, staged in chunks through a scratch directory and packed by the unchanged packer; writes a key-set, scrubs, and spot-checks members bit-exact against the H5 source.

Testing

  • packages/diffusion_planner/tests/data/test_shard_planner_dataset.py — two ranks cover every frame exactly once, bit-exact; transforms apply in order on writable arrays; loader length matches the plan; loader settings must match the plan.
  • packages/diffusion_planner/tests/data/test_h5_to_shards.py — chunked conversion of a synthetic H5 layout, key-set size, intermediate and final versions, bit-exact verification, staging cleanup.
  • packages/planner_shards/tests — the library's own suite (108 tests) with imports renamed.
  • Hydra smoke: train/train composed with train/dataloader@dataloader=shards instantiates and yields batches.
  • Existing tests/data suite still passes.

Measured on an internal 8-GPU box: whole training epochs over identical frames were 1.5x faster through shards than through the HDF5 loader at its default worker count; details are internal. Not a convergence comparison.

Notes for review

  • planner_shards is not yet in the pyright include list; the copied library predates type-checking and will be cleaned up when it becomes a shared dependency.
  • The converter stages frames as compressed .npz before packing (the packer's existing input contract). A direct H5-to-member writer is a follow-up.
  • Manifest metadata columns for the H5 frame metadata (frame time, ego speed, turn indicator, map ids) are not carried into the shards yet; that schema needs agreement before it is added.

…ds package

Add a tar-shard data path beside the HDF5 loader:

- packages/planner_shards: the versioned shard dataset library (packer,
  manifests, key-sets, DDP-safe loader) as a workspace package; a copy of
  the tier4-main data_pipeline/shard_* modules with imports renamed, since
  both trees use the diffusion_planner package name.
- ShardPlannerDataset / build_shard_dataloader apply the PlannerDataset
  transforms to frames read from shards; rank and world size come from the
  torchrun environment.
- configs/train/dataloader/shards.yaml mirrors default.yaml's transforms.
- train.py keeps a shard-backed loader out of accelerator.prepare (it is
  already sharded per rank) and moves batches to the device itself.
- convert_h5_to_shards.py stages the frames of a Parquet index in chunks,
  packs them with the unchanged packer, writes a key-set, scrubs, and
  spot-checks members bit-exact against the H5 source.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant