Skip to content

[DCP][HF] Add option to parallelize reads in HF Storage Reader#160205

Closed
ankitageorge wants to merge 5 commits intogh/ankitageorge/21/basefrom
gh/ankitageorge/21/head
Closed

[DCP][HF] Add option to parallelize reads in HF Storage Reader#160205
ankitageorge wants to merge 5 commits intogh/ankitageorge/21/basefrom
gh/ankitageorge/21/head

Conversation

@ankitageorge
Copy link
Contributor

@ankitageorge ankitageorge commented Aug 8, 2025

Parallelize reading of data behind thread_count argument to HFStorageReader
Test plan: ensure existing tests pass and run a job successfully with these changes

Stack from ghstack (oldest at bottom):

Differential Revision: D79478188

cc @H-Huang @awgu @wanchaol @fegin @fduwjj @wz337 @wconstab @d4l3k @pragupta

Differential Revision: [D79478188](https://our.internmc.facebook.com/intern/diff/D79478188/)

[ghstack-poisoned]
@pytorch-bot
Copy link

pytorch-bot bot commented Aug 8, 2025

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/160205

Note: Links to docs will display an error until the docs builds have been completed.

✅ No Failures

As of commit ecd2d16 with merge base e1a64b7 (image):
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.

ankitageorge added a commit that referenced this pull request Aug 8, 2025
Differential Revision: [D79478188](https://our.internmc.facebook.com/intern/diff/D79478188/)

ghstack-source-id: 300233804
Pull Request resolved: #160205
@pytorch-bot pytorch-bot bot added oncall: distributed Add this issue/PR to distributed oncall triage queue release notes: distributed (checkpoint) labels Aug 8, 2025
@facebook-github-bot
Copy link
Contributor

This pull request was exported from Phabricator. Differential Revision: D79478188

Differential Revision: [D79478188](https://our.internmc.facebook.com/intern/diff/D79478188/)

cc H-Huang awgu wanchaol fegin fduwjj wz337 wconstab d4l3k pragupta

[ghstack-poisoned]
ankitageorge added a commit that referenced this pull request Aug 11, 2025
Parallelize reading of data behind thread_count argument to HFStorageReader
Pull Request resolved: #160205


ghstack-source-id: 302209653
@exported-using-ghexport

Differential Revision: [D79478188](https://our.internmc.facebook.com/intern/diff/D79478188/)
@facebook-github-bot
Copy link
Contributor

This pull request was exported from Phabricator. Differential Revision: D79478188

@ankitageorge ankitageorge changed the title parallelize reads [DCP][HF] Add option to parallelize reads in HF Storage Reader Aug 11, 2025
@pytorch-bot pytorch-bot bot added the ciflow/trunk Trigger trunk jobs on your pull request label Aug 11, 2025
…ader"

Parallelize reading of data behind thread_count argument to HFStorageReader
Test plan: ensure existing tests pass and run a job successfully with these changes


Differential Revision: [D79478188](https://our.internmc.facebook.com/intern/diff/D79478188/)

cc H-Huang awgu wanchaol fegin fduwjj wz337 wconstab d4l3k pragupta

[ghstack-poisoned]
ankitageorge added a commit that referenced this pull request Aug 12, 2025
Parallelize reading of data behind thread_count argument to HFStorageReader
Pull Request resolved: #160205


ghstack-source-id: 302487332
@exported-using-ghexport

Differential Revision: [D79478188](https://our.internmc.facebook.com/intern/diff/D79478188/)
@facebook-github-bot
Copy link
Contributor

This pull request was exported from Phabricator. Differential Revision: D79478188

Copy link
Contributor

@meetv18 meetv18 left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM :)

Any way to modularize this and also implement it for FileSystemReader?

except queue.Empty:
pass

def read_data(self, plan: LoadPlan, planner: LoadPlanner) -> Future[None]:
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Is there a plan to introduce the multithreading for save path as well?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

it's already there


target_tensor.copy_(tensor)
planner.commit_tensor(req, target_tensor)
if self.thread_count <= 1 or len(per_file) <= 1:
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We need a test for this method.

…ader"

Parallelize reading of data behind thread_count argument to HFStorageReader
Test plan: ensure existing tests pass and run a job successfully with these changes


Differential Revision: [D79478188](https://our.internmc.facebook.com/intern/diff/D79478188/)

cc H-Huang awgu wanchaol fegin fduwjj wz337 wconstab d4l3k pragupta

[ghstack-poisoned]
ankitageorge added a commit that referenced this pull request Aug 21, 2025
Parallelize reading of data behind thread_count argument to HFStorageReader
Pull Request resolved: #160205


ghstack-source-id: 304719983
@exported-using-ghexport

Differential Revision: [D79478188](https://our.internmc.facebook.com/intern/diff/D79478188/)
@facebook-github-bot
Copy link
Contributor

This pull request was exported from Phabricator. Differential Revision: D79478188

…ader"

Parallelize reading of data behind thread_count argument to HFStorageReader
Test plan: ensure existing tests pass and run a job successfully with these changes


Differential Revision: [D79478188](https://our.internmc.facebook.com/intern/diff/D79478188/)

cc H-Huang awgu wanchaol fegin fduwjj wz337 wconstab d4l3k pragupta

[ghstack-poisoned]
ankitageorge added a commit that referenced this pull request Aug 21, 2025
Parallelize reading of data behind thread_count argument to HFStorageReader
Pull Request resolved: #160205


ghstack-source-id: 304745088
@exported-using-ghexport

Differential Revision: [D79478188](https://our.internmc.facebook.com/intern/diff/D79478188/)
@facebook-github-bot
Copy link
Contributor

This pull request was exported from Phabricator. Differential Revision: D79478188

@ankitageorge
Copy link
Contributor Author

@pytorchmergebot merge

@pytorchmergebot
Copy link
Collaborator

Merge started

Your change will be merged once all checks pass (ETA 0-4 Hours).

Learn more about merging in the wiki.

Questions? Feedback? Please reach out to the PyTorch DevX Team

Advanced Debugging
Check the merge workflow status
here

wincent8 pushed a commit to wincent8/pytorch that referenced this pull request Aug 22, 2025
…ch#160205)

Parallelize reading of data behind thread_count argument to HFStorageReader
Test plan: ensure existing tests pass and run a job successfully with these changes

Differential Revision: [D79478188](https://our.internmc.facebook.com/intern/diff/D79478188/)

Pull Request resolved: pytorch#160205
Approved by: https://github.com/meetv18
markc-614 pushed a commit to markc-614/pytorch that referenced this pull request Sep 17, 2025
…ch#160205)

Parallelize reading of data behind thread_count argument to HFStorageReader
Test plan: ensure existing tests pass and run a job successfully with these changes

Differential Revision: [D79478188](https://our.internmc.facebook.com/intern/diff/D79478188/)

Pull Request resolved: pytorch#160205
Approved by: https://github.com/meetv18
@github-actions github-actions bot deleted the gh/ankitageorge/21/head branch September 21, 2025 02:15
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

ciflow/trunk Trigger trunk jobs on your pull request fb-exported Merged oncall: distributed Add this issue/PR to distributed oncall triage queue release notes: distributed (checkpoint)

Projects

None yet

Development

Successfully merging this pull request may close these issues.

5 participants