Restore from master weights (& allow restoring from a checkpoint of different precision) #702
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
…rministically by also saving RNG state of last update)
ademeure
force-pushed
the
restore_from_master_weights
branch
from
6387e66
to
4d77ece
Compare
karpathy
reviewed
Jul 28, 2024
| @@ -61,6 +61,16 @@ __global__ void adamw_kernel3(Tp* params_memory, float* master_params_memory, Tg | |||
| ); | |||
| } | |||
|
|
|||
| template <typename Tp> | |||
| __global__ void init_from_master_kernel(Tp* params_memory, float* master_params_memory, size_t num_parameters, | |||
| ptrdiff_t w_stride, ptrdiff_t s_stride, unsigned int seed, bool check_identical) { | |||
There was a problem hiding this comment.
oops, that used to be for debug lots that I removed, thanks
karpathy
reviewed
Jul 28, 2024
| @@ -1532,7 +1554,7 @@ int main(int argc, char *argv[]) { | |||
| gpt2_init_common(&model); | |||
| if (resuming == 1) { | |||
| // if `-y 1` was set, then we are resuming from the latest checkpoint | |||
| gpt2_build_from_checkpoint(&model, filename_buffer); | |||
| gpt2_build_from_checkpoint(&model, filename_buffer, true); | |||
There was a problem hiding this comment.
at this point model.use_master_weights is not yet initialized with use_master_weights, that happens below, and is just the default (true). This variable is used inside gpt2_build_from_checkpoint and this is probably a bug?
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.
This is fully deterministic for new checkpoints where the new rng_state_last_update is saved, so that stochastic rounding from master weights is done with the exact same seeds (while restoring the actual final rng_state again afterwards, in case anything else changed it between that update and saving the checkpoint).
In the case where we are resuming from a checkpoint (not just a regular model file) and we have master weights, this simply skips loading the weights from the checkpoint completely, so it doesn't matter if they are not even the right number of bytes.
It should be useful to check if FP32 helps with runs exploding, and going forward it will allow FP8 runs to not care too much about what format the non-master weights are saved, so we don't need to worry about changes breaking compatibility etc...