[SDL] StatefulDistributedSampler add epoch to state_dict #1515
+13
−4
Conversation
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
meta-cla
bot
added
the
CLA Signed
aelavender
changed the title
aelavender
changed the title
|
@aelavender has imported this pull request. If you are a Meta employee, you can view this in D87398440. |
divyanshk
reviewed
Nov 19, 2025
| self.next_yielded = None | ||
|
|
||
| def __iter__(self): | ||
| print(f"Calling __iter__... {self.yielded=} {self.next_yielded=}") |
Contributor
Author
There was a problem hiding this comment.
Sorry, will remove
divyanshk
reviewed
Nov 19, 2025
| def load_state_dict(self, state_dict: Dict[str, Any]) -> None: | ||
| if self._EPOCH not in state_dict: | ||
| raise ValueError("Invalid state_dict") | ||
| self.set_epoch(state_dict[self._EPOCH]) |
Contributor
There was a problem hiding this comment.
If users are used to calling set_epoch in a certain in the main loop, then doing self.set_epoch(..) on load_state_dict might interfere with that ?
Contributor
Author
There was a problem hiding this comment.
If the user is trying to set_epoch to something other than what was previously saved, then they are likely going to get repeated indices in the current epoch, which seems like broken behaviour. Perhaps we should warn in that case.
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
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.
Partially fixes #1513
StatefulDistributedSamplerhas epoch-dependent behaviour (the sampling order depends on epoch), so we need to save and restore the correct epoch.When used with
StatefulDataLoaderthis is mostly an edge case, as users are expected to explicitly callset_epochbefore each epoch. But the order matters here: if the user forgets to callset_epoch()or calls__iter__()beforeset_epoch(), things will be incorrect. By putting the epoch in state we can avoid these mistakes.See also #1516.