[DataLoader2] Saving and restoring initial seed generator

[NivekT]

Stack from ghstack:

• [DataLoader2] Adding support for naive checkpointing #1119
• -> [DataLoader2] Saving and restoring initial seed generator #1124

Reland of #998 with added guard while loading randomness state in DataLoader2 for backward compatibility

Changes to DataLoader2:

• Modifying state_dict to store randomness_state, which includes:
  • _seed: int
  • _reset_seed: bool - flag indicating whether _seed needs to be set
  • _seed_generator - the latest version at the time when state_dict is called
  • _initial_seed_generator - the versopm that is saved at the beginning of very epoch
• Modifying from_state and load_state_dict to restore randomness_state
• Adding a method _restore_checkpoint_beginning_of_epoch
  • This sets self._seed_generator = self._initial_seed_generator, allowing users to re-create an epoch from the beginning.

---

Considerations

Storing the randomness states provide more flexibility for users to restore as they see fit. The decision to do that should not be controversial.

I decided to make add a new method for checkpointing at the beginning of the epoch, ensure that users are not confused about what randomness is restored by default.

The basic idea is that we want to allow users to restore dl2._seed_generator to the previously saved version. From that point on, they can create a new __iter__ and continue from the beginning of the epoch.

• Note that since _seed and _reset_seed are also saved, if the users were planning to use a different seed or if there was a need to re-seed, those remain valid after restoring the checkpoint.
• Finally, if users change their mind at any point (after restoring) and want to manual set seed. That seed will override any other behavior and the seed will be used.

Differential Revision: D44748514

[NivekT]

@NivekT has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator.

[facebook-github-bot]

@NivekT merged this pull request in 851e26a.