Add more options to the unshard_checkpoint function to help scale
#145
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
2015aroras
approved these changes
Jan 24, 2025
| chunks = [] | ||
| for chunk_num, key in enumerate(metadata.state_dict_metadata.keys()): | ||
| if key.startswith(f"{prefix}."): | ||
| chunks.append((path / f"chunk-{chunk_num:05d}.{ext}", [key])) |
There was a problem hiding this comment.
thought: it could be nice to name it as chunk-{key}, maybe with some cleanup of the key. Then one could easily examine the weights of specific parameters. Would possibly be slightly cleaner for hypothetical future converter logic too.
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.
I was getting worried about unsharding really big checkpoints like for the 32B, which we'll need to do soon. The main issue at the moment is that in order to unshard we need to load the entire model (or optimizer state) in memory, which clearly isn't scalable. So I've added an option to unshard the checkpoint into chunks of a given size which helps scale because only a single chunk (which could be as small as a single tensor) needs to load into memory at a time. Each chunk gets written to a unique file. I think HuggingFace does something similar.
Note: this is not supported for optimizer state yet. But, speaking of optimizer state, this PR also adds a function called
load_keys()for loading (and unsharding) specific keys from a checkpoint. So if you want to inspect part of the optimizer state, you could use that function without having to unshard the whole optimizer state.