Skip to content

Fix automatic oversubscription scheduling with gradient accumulation - #2918

Draft
hjinnkim wants to merge 2 commits into
bghira:mainfrom
hjinnkim:bugfix/auto-repeat-gradient-accumulation
Draft

Fix automatic oversubscription scheduling with gradient accumulation#2918
hjinnkim wants to merge 2 commits into
bghira:mainfrom
hjinnkim:bugfix/auto-repeat-gradient-accumulation

Conversation

@hjinnkim

Copy link
Copy Markdown
Contributor

Dependency

Depends on #2897 for occurrence-aware sampler and resume semantics. This branch is stacked on #2897 and will be rebased after it merges.

Additional issues found

  • Automatic oversubscription increased the repeat count without constructing the repeated schedule, so gradient accumulation could still leave insufficient rank-local samples.
  • repeats: 0 was treated as a manual override.
  • Empty buckets could enter repeat division.
  • Mutating the backend repeat value could affect later splits.
  • Conditioning metadata copying could discard the current rank-local source schedule and overwrite canonical target bucket state with rank-local data.

Resolution

  • Keep the automatic repeat calculation local to the split.
  • Build a batch-aligned cyclic schedule directly at rank-local size.
  • Treat only positive configured repeats as a manual override.
  • Skip empty buckets.
  • Preserve the source schedule and avoid persisting rank-local target bucket indices over canonical metadata.

Regression coverage for gradient accumulation, standard DP, Context Parallel, empty buckets, manual repeats, and conditioning metadata is included.

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