Implement incremental state transfer from worker to dataloader process #1262

pytorch / data

A PyTorch repo for data loading and utilities to be shared by the PyTorch domain libraries.

BSD 3-Clause "New" or "Revised" License

1.08k stars 142 forks source link

Implement incremental state transfer from worker to dataloader process #1262 #1264

Closed gokulavasan closed 1 month ago

gokulavasan commented 1 month ago

Summary: Transfer of full state from worker to main process can be expensive if the state is very large. In order to alleviate the cost of state transfer, maintain a flattened state map in the worker process and send only the key/value that have changed to the main dataloader process where the full state is reconstructed with the changes received from the workers.

Test Plan: Unit tests

Reviewers:

Subscribers:

Tasks:

Tags:Please read through our contribution guide prior to creating your pull request.

Note that there is a section on requirements related to adding a new DataPipe.

Fixes #{issue number}

Changes

-

pytorch-bot[bot] commented 1 month ago

:link: Helpful Links

:test_tube: See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/data/1264

:page_facing_up: Preview Python docs built from this PR

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

:heavy_exclamation_mark: 1 Active SEVs

There are 1 currently active SEVs. If your PR is affected, please view them below:

Upgrade MacOS runner to 14

:white_check_mark: No Failures

As of commit 441891329949461a52360f0b866648ec8d6101bb with merge base b474083ebbe0712fe075d8ac15698c99d070e616 (): :green_heart: Looks good so far! There are no failures yet. :green_heart:

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

gokulavasan commented 1 month ago

@andrewkho Thanks for the review! Addressed/responded to the comments, please take a look again, thanks!