[ghstack-poisoned]

[ghstack-poisoned]

Changes to `DataLoader2`:
- Modifying `state_dict` to store the `initial_seed_generator` that is saved at the beginning of an epoch.
- Modifying `from_state` and `load_state_dict` to restore `initial_seed_generator` if the user sets the parameter to `True`
- Within `__iter__, skips over the re-seeding process if no manual seed has been specified AND the `seed_generator` was explicitly restored.


---
### Consideration

I decided to make modification to the existing APIs. Alternatively, we can create a new method.

The basic idea is that we want to allow users to restore `dl2._seed_generator` to the previously saved version, at the same time, we need to skip over the logic that re-do seeding in `__iter__` (hence the new variable `_skip_iteration_seeding` is needed.

I see 2 main scenarios:
1. Users want to restore DataPipe and ReadingService but not the initial state of RNG
    - I think lots of current users (including some internals) are in this category.
    - This should work by default because `restore_initial_seed_generator=False` unless user explicitly change it
2. Users actively want to restore DP, RS, and initial state of RNG
    - Users will need to set an extra variable to `True` and we will make sure `_skip_iteration_seeding=True` so no re-seeding will happen in the first subsequent call of `__iter__`

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.


[ghstack-poisoned]

Changes to `DataLoader2`:
- Modifying `state_dict` to store the `initial_seed_generator` that is saved at the beginning of an epoch.
- Modifying `from_state` and `load_state_dict` to restore `initial_seed_generator` if the user sets the parameter to `True`
- Within `__iter__, skips over the re-seeding process if no manual seed has been specified AND the `seed_generator` was explicitly restored.


---
### Consideration

I decided to make modification to the existing APIs. Alternatively, we can create a new method.

The basic idea is that we want to allow users to restore `dl2._seed_generator` to the previously saved version, at the same time, we need to skip over the logic that re-do seeding in `__iter__` (hence the new variable `_skip_iteration_seeding` is needed.

I see 2 main scenarios:
1. Users want to restore DataPipe and ReadingService but not the initial state of RNG
    - I think lots of current users (including some internals) are in this category.
    - This should work by default because `restore_initial_seed_generator=False` unless user explicitly change it
2. Users actively want to restore DP, RS, and initial state of RNG
    - Users will need to set an extra variable to `True` and we will make sure `_skip_iteration_seeding=True` so no re-seeding will happen in the first subsequent call of `__iter__`

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.


[ghstack-poisoned]

Changes to `DataLoader2`:
- Modifying `state_dict` to store the `initial_seed_generator` that is saved at the beginning of an epoch.
- Modifying `from_state` and `load_state_dict` to restore `initial_seed_generator` if the user sets the parameter to `True`
- Within `__iter__, skips over the re-seeding process if no manual seed has been specified AND the `seed_generator` was explicitly restored.


---
### Consideration

I decided to make modification to the existing APIs. Alternatively, we can create a new method.

The basic idea is that we want to allow users to restore `dl2._seed_generator` to the previously saved version, at the same time, we need to skip over the logic that re-do seeding in `__iter__` (hence the new variable `_skip_iteration_seeding` is needed.

I see 2 main scenarios:
1. Users want to restore DataPipe and ReadingService but not the initial state of RNG
    - I think lots of current users (including some internals) are in this category.
    - This should work by default because `restore_initial_seed_generator=False` unless user explicitly change it
2. Users actively want to restore DP, RS, and initial state of RNG
    - Users will need to set an extra variable to `True` and we will make sure `_skip_iteration_seeding=True` so no re-seeding will happen in the first subsequent call of `__iter__`

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.


[ghstack-poisoned]

Changes to `DataLoader2`:
- Modifying `state_dict` to store the `initial_seed_generator` that is saved at the beginning of an epoch.
- Modifying `from_state` and `load_state_dict` to restore `initial_seed_generator` if the user sets the parameter to `True`
- Within `__iter__, skips over the re-seeding process if no manual seed has been specified AND the `seed_generator` was explicitly restored.


---
### Consideration

I decided to make modification to the existing APIs. Alternatively, we can create a new method.

The basic idea is that we want to allow users to restore `dl2._seed_generator` to the previously saved version, at the same time, we need to skip over the logic that re-do seeding in `__iter__` (hence the new variable `_skip_iteration_seeding` is needed.

I see 2 main scenarios:
1. Users want to restore DataPipe and ReadingService but not the initial state of RNG
    - I think lots of current users (including some internals) are in this category.
    - This should work by default because `restore_initial_seed_generator=False` unless user explicitly change it
2. Users actively want to restore DP, RS, and initial state of RNG
    - Users will need to set an extra variable to `True` and we will make sure `_skip_iteration_seeding=True` so no re-seeding will happen in the first subsequent call of `__iter__`

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.


[ghstack-poisoned]

Changes to `DataLoader2`:
- Modifying `state_dict` to store the `initial_seed_generator` that is saved at the beginning of an epoch.
- Modifying `from_state` and `load_state_dict` to restore `initial_seed_generator` if the user sets the parameter to `True`
- Within `__iter__, skips over the re-seeding process if no manual seed has been specified AND the `seed_generator` was explicitly restored.


---
### Consideration

I decided to make modification to the existing APIs. Alternatively, we can create a new method.

The basic idea is that we want to allow users to restore `dl2._seed_generator` to the previously saved version, at the same time, we need to skip over the logic that re-do seeding in `__iter__` (hence the new variable `_skip_iteration_seeding` is needed.

I see 2 main scenarios:
1. Users want to restore DataPipe and ReadingService but not the initial state of RNG
    - I think lots of current users (including some internals) are in this category.
    - This should work by default because `restore_initial_seed_generator=False` unless user explicitly change it
2. Users actively want to restore DP, RS, and initial state of RNG
    - Users will need to set an extra variable to `True` and we will make sure `_skip_iteration_seeding=True` so no re-seeding will happen in the first subsequent call of `__iter__`

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.


[ghstack-poisoned]

Changes to `DataLoader2`:
- Modifying `state_dict` to store the `initial_seed_generator` that is saved at the beginning of an epoch.
- Modifying `from_state` and `load_state_dict` to restore `initial_seed_generator` if the user sets the parameter to `True`
- Within `__iter__, skips over the re-seeding process if no manual seed has been specified AND the `seed_generator` was explicitly restored.


---
### Consideration

I decided to make modification to the existing APIs. Alternatively, we can create a new method.

The basic idea is that we want to allow users to restore `dl2._seed_generator` to the previously saved version, at the same time, we need to skip over the logic that re-do seeding in `__iter__` (hence the new variable `_skip_iteration_seeding` is needed.

I see 2 main scenarios:
1. Users want to restore DataPipe and ReadingService but not the initial state of RNG
    - I think lots of current users (including some internals) are in this category.
    - This should work by default because `restore_initial_seed_generator=False` unless user explicitly change it
2. Users actively want to restore DP, RS, and initial state of RNG
    - Users will need to set an extra variable to `True` and we will make sure `_skip_iteration_seeding=True` so no re-seeding will happen in the first subsequent call of `__iter__`

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.


[ghstack-poisoned]

Changes to `DataLoader2`:
- Modifying `state_dict` to store the `initial_seed_generator` that is saved at the beginning of an epoch.
- Modifying `from_state` and `load_state_dict` to restore `initial_seed_generator` if the user sets the parameter to `True`
- Within `__iter__, skips over the re-seeding process if no manual seed has been specified AND the `seed_generator` was explicitly restored.


---
### Consideration

I decided to make modification to the existing APIs. Alternatively, we can create a new method.

The basic idea is that we want to allow users to restore `dl2._seed_generator` to the previously saved version, at the same time, we need to skip over the logic that re-do seeding in `__iter__` (hence the new variable `_skip_iteration_seeding` is needed.

I see 2 main scenarios:
1. Users want to restore DataPipe and ReadingService but not the initial state of RNG
    - I think lots of current users (including some internals) are in this category.
    - This should work by default because `restore_initial_seed_generator=False` unless user explicitly change it
2. Users actively want to restore DP, RS, and initial state of RNG
    - Users will need to set an extra variable to `True` and we will make sure `_skip_iteration_seeding=True` so no re-seeding will happen in the first subsequent call of `__iter__`

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.


[ghstack-poisoned]

Changes to `DataLoader2`:
- Modifying `state_dict` to store the `initial_seed_generator` that is saved at the beginning of an epoch.
- Modifying `from_state` and `load_state_dict` to restore `initial_seed_generator` if the user sets the parameter to `True`
- Within `__iter__, skips over the re-seeding process if no manual seed has been specified AND the `seed_generator` was explicitly restored.


---
### Consideration

I decided to make modification to the existing APIs. Alternatively, we can create a new method.

The basic idea is that we want to allow users to restore `dl2._seed_generator` to the previously saved version, at the same time, we need to skip over the logic that re-do seeding in `__iter__` (hence the new variable `_skip_iteration_seeding` is needed.

I see 2 main scenarios:
1. Users want to restore DataPipe and ReadingService but not the initial state of RNG
    - I think lots of current users (including some internals) are in this category.
    - This should work by default because `restore_initial_seed_generator=False` unless user explicitly change it
2. Users actively want to restore DP, RS, and initial state of RNG
    - Users will need to set an extra variable to `True` and we will make sure `_skip_iteration_seeding=True` so no re-seeding will happen in the first subsequent call of `__iter__`

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.


[ghstack-poisoned]

Changes to `DataLoader2`:
- Modifying `state_dict` to store the `initial_seed_generator` that is saved at the beginning of an epoch.
- Modifying `from_state` and `load_state_dict` to restore `initial_seed_generator` if the user sets the parameter to `True`
- Within `__iter__, skips over the re-seeding process if no manual seed has been specified AND the `seed_generator` was explicitly restored.


---
### Consideration

I decided to make modification to the existing APIs. Alternatively, we can create a new method.

The basic idea is that we want to allow users to restore `dl2._seed_generator` to the previously saved version, at the same time, we need to skip over the logic that re-do seeding in `__iter__` (hence the new variable `_skip_iteration_seeding` is needed.

I see 2 main scenarios:
1. Users want to restore DataPipe and ReadingService but not the initial state of RNG
    - I think lots of current users (including some internals) are in this category.
    - This should work by default because `restore_initial_seed_generator=False` unless user explicitly change it
2. Users actively want to restore DP, RS, and initial state of RNG
    - Users will need to set an extra variable to `True` and we will make sure `_skip_iteration_seeding=True` so no re-seeding will happen in the first subsequent call of `__iter__`

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.


[ghstack-poisoned]

Changes to `DataLoader2`:
- Modifying `state_dict` to store the `initial_seed_generator` that is saved at the beginning of an epoch.
- Modifying `from_state` and `load_state_dict` to restore `initial_seed_generator` if the user sets the parameter to `True`
- Within `__iter__, skips over the re-seeding process if no manual seed has been specified AND the `seed_generator` was explicitly restored.


---
### Consideration

I decided to make modification to the existing APIs. Alternatively, we can create a new method.

The basic idea is that we want to allow users to restore `dl2._seed_generator` to the previously saved version, at the same time, we need to skip over the logic that re-do seeding in `__iter__` (hence the new variable `_skip_iteration_seeding` is needed.

I see 2 main scenarios:
1. Users want to restore DataPipe and ReadingService but not the initial state of RNG
    - I think lots of current users (including some internals) are in this category.
    - This should work by default because `restore_initial_seed_generator=False` unless user explicitly change it
2. Users actively want to restore DP, RS, and initial state of RNG
    - Users will need to set an extra variable to `True` and we will make sure `_skip_iteration_seeding=True` so no re-seeding will happen in the first subsequent call of `__iter__`

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.


[ghstack-poisoned]

Changes to `DataLoader2`:
- Modifying `state_dict` to store the `initial_seed_generator` that is saved at the beginning of an epoch.
- Modifying `from_state` and `load_state_dict` to restore `initial_seed_generator` if the user sets the parameter to `True`
- Within `__iter__, skips over the re-seeding process if no manual seed has been specified AND the `seed_generator` was explicitly restored.


---
### Consideration

I decided to make modification to the existing APIs. Alternatively, we can create a new method.

The basic idea is that we want to allow users to restore `dl2._seed_generator` to the previously saved version, at the same time, we need to skip over the logic that re-do seeding in `__iter__` (hence the new variable `_skip_iteration_seeding` is needed.

I see 2 main scenarios:
1. Users want to restore DataPipe and ReadingService but not the initial state of RNG
    - I think lots of current users (including some internals) are in this category.
    - This should work by default because `restore_initial_seed_generator=False` unless user explicitly change it
2. Users actively want to restore DP, RS, and initial state of RNG
    - Users will need to set an extra variable to `True` and we will make sure `_skip_iteration_seeding=True` so no re-seeding will happen in the first subsequent call of `__iter__`

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.


[ghstack-poisoned]

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.


[ghstack-poisoned]

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.


[ghstack-poisoned]

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.


[ghstack-poisoned]

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.


[ghstack-poisoned]

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: [D44390519](https://our.internmc.facebook.com/intern/diff/D44390519)

[ghstack-poisoned]

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: [D44390519](https://our.internmc.facebook.com/intern/diff/D44390519)

[ghstack-poisoned]