Use cyclic distributed padding for oversubscribed datasets - #2903
Closed
hjinnkim wants to merge 6 commits into
Closed
Use cyclic distributed padding for oversubscribed datasets#2903hjinnkim wants to merge 6 commits into
hjinnkim wants to merge 6 commits into
Conversation
Explicit dataset oversubscription can increase a bucket's repeat count, but those repeats are not materialized before the bucket is sliced across ranks. The existing fixed-step padding gate could therefore leave DP ranks empty. - In data_backend/factory.py, enable padding for epoch-driven runs or when allow_dataset_oversubscription is explicit. Ordinary fixed-step runs keep their existing no-padding behavior. - In metadata/backends/base.py, fill short and completely empty Context Parallel shards from the final item in the non-empty global bucket. The old local-only repetition had no source item for an empty shard; this also matches Accelerate's non-CP last-item policy. - In training/trainer.py, forward the same policy when epoch rollover rebuilds and re-splits buckets. Otherwise the omitted argument defaults to false and silently drops the initial split policy.
Rank-local RNG state could permute the same bucket differently before each rank sliced it, causing duplicate and missing samples. Resume could then apply saved progress to a newly rebuilt partition that did not match the checkpoint. - In metadata/backends/base.py, derive every rank's bucket order from one configured or main-rank-generated seed, broadcast it, and shuffle a copied bucket with a dataset-and-bucket-specific RNG. Avoid a second full-list copy when the trim limit keeps the complete bucket. - In multiaspect/sampler.py, save the exact local partition together with world size, rank, effective DP size, DP rank, and CP size. Validate that topology and restore the partition before applying seen or exhausted progress. - Reject topology mismatches, and distributed legacy partition state without topology metadata, before mutating the live sampler because those partitions cannot be mapped safely to the current ranks.
Explicit oversubscription still inherited the fixed-step no-padding gate, and existing padding concentrated all extra entries on one final sample. Context Parallel and non-CP paths also applied padding at different stages. - In data_backend/factory.py, enable padding for epoch-driven or explicitly oversubscribed runs while preserving ordinary fixed-step behavior. - In metadata/backends/base.py, establish one shared bucket seed and use a deterministic dataset-and-bucket RNG so every rank pads and slices the same order. - Cyclically extend the global bucket to a multiple of effective DP size before splitting, disable Accelerate's final-item padding, and feed the same padded list to CP and non-CP paths. Repeats are distributed across the prefix rather than concentrated on one sample. - Skip padding arithmetic for empty buckets because they have no valid source item to repeat. - In training/trainer.py, forward the same policy when epoch rollover rebuilds buckets so the post-rebuild distribution matches the initial split.
…e-oversubscribed-step-padding
Skip empty buckets during repeat validation so an empty bucket cannot trigger a division-by-zero calculation.\n\nPad the shuffled global bucket to an effective-DP multiple with a cyclic prefix before rank partitioning. This spreads duplicated samples across the bucket instead of repeating only the final item.\n\nDisable the older CP tail-padding block and Accelerate per-rank padding after global padding, keeping DDP, FSDP2, and context-parallel partitions on the same policy.\n\nThis commit is stacked on the oversubscription policy and deterministic rank-shuffle fixes so the three changes merge without overlapping reimplementations.
Preserve the published PR07 history while merging the exact oversubscription-policy and deterministic-rank-shuffle prerequisite commits. Resolve the overlapping metadata backend implementation to the prerequisite-aware cyclic-padding version: empty buckets bypass padding arithmetic, one shared shuffled global bucket is extended by a cyclic prefix, and both Context Parallel and ordinary distributed paths consume the same pre-padded list without Accelerate tail padding. This makes the branch merge cleanly after PR01 through PR06 and leaves only the cyclic-padding enhancement once the prerequisite commits are present upstream.
Contributor
Author
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
1 participant
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.
Use cyclic distributed padding for oversubscribed datasets
Dependencies
series can be reviewed together.
draft until both prerequisite PRs are merged.
mainand verify that
Files changedcontains only the cyclic-padding delta beforemarking Use cyclic distributed padding for oversubscribed datasets #2903 ready for review.
Summary
entries on one sample; Context Parallel and non-CP paths also padded at
different stages.
DDP/FSDP2 or Context Parallel splitting.
Prerequisite behavior
simpletuner/helpers/data_backend/factory.py: enable padding for epoch-drivenruns or explicit oversubscription while preserving ordinary step-bounded
behavior.
simpletuner/helpers/metadata/backends/base.py: establish one deterministicdataset-and-bucket order shared by every rank.
simpletuner/helpers/training/trainer.py: forward the same padding policy afterepoch-rollover bucket rebuilds.
Changes in this PR
simpletuner/helpers/metadata/backends/base.py: preserve empty buckets withoutentering repeat or padding arithmetic.
size before splitting.
list is already padded.
Result
padding is enabled.
sample; CP and non-CP use the same global policy.
[A, B, C, D], final-itempadding produces
[A, B, C, D, D, D, D, D]: four added copies ofD, fiveDoccurrences total. A shared shuffled order such as[C, A, D, B]nowbecomes
[C, A, D, B, C, A, D, B], spreading the four added entries acrossthe bucket prefix.
DP ranks, CP size 2, epoch rollover, empty buckets, and seed reproducibility;
the related 110-test regression suite passed. PR01 through PR06 followed by
this stacked PR merged without conflicts, and the two-rank distributed
integration smoke passed.
Environment
3.13.142.13.0+cu1301.14.00.39.05.13.1590.48.01