Add ability to restart on new epoch #383
Merged
Conversation
This file contains 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
You can set the epoch via the option `--epoch=[INTEGER]`. This automatically handles changing the data order each epoch by setting the data seed to `seed + epoch`. So `--epoch` is the only flag you need to set when restarting on a new epoch. Everything else in the config can stay the same. Note that we count epochs starting from 0. So to start the 2nd epoch you would add the flag `--epoch=1`.
2015aroras
reviewed
Nov 27, 2023
olmo/train.py
Outdated
| @@ -147,40 +152,47 @@ def load_trainer_state_dict(self, state_dict: Dict[str, Any]) -> None: | |||
| ] | |||
|
|
|||
| # Dataset / dataloader position. | |||
| checkpoint_epoch = state_dict.get("epoch", 0) | |||
| self.global_step = state_dict["global_step"] | |||
| self.global_data_step = state_dict["global_data_step"] | |||
| self.global_train_examples_seen = state_dict.get( # newer addition | |||
There was a problem hiding this comment.
Is there any reason to keep self.global_train_examples_seen? You can perform this state_dict.get() backwards compatibility check without keeping the global_train_examples_seen variable.
There was a problem hiding this comment.
Yea good point. Turns out we don't need global_data_step either. 4d6e61c
`global_train_examples_seen` and `global_data_step` no longer needed
2015aroras
reviewed
Nov 27, 2023
olmo/train.py
Outdated
| "global_train_tokens_seen", | ||
| self.global_data_step * self.cfg.global_train_batch_size * self.cfg.model.max_sequence_length, | ||
| state_dict.get("global_data_step", 0) # for backwards compatibility |
There was a problem hiding this comment.
This will result in throughput/total_tokens being reset to 0 if global_train_tokens_seen and global_data_step are both not present. Maybe state_dict.get("global_data_step", self.global_step) is safer?
There was a problem hiding this comment.
Absolutely, good catch: b8ca94d
2015aroras
approved these changes
Nov 27, 2023
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
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.
You can set the epoch via the option
--epoch=[INTEGER]. Thisautomatically handles changing the data order each epoch by setting the
data seed to
seed + epoch. So--epochis the only flag you need toset when restarting on a new epoch. Everything else in the config can
stay the same.
Note that we count epochs starting from 0. So to start the 2nd epoch you would add
the flag
--epoch=1.I cherry-picked this commit from #350, which has now started its 2nd epoch.