Mitchish65 #388
Merged
Mitchish65 #388
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
|
With the latest fix, I can load an unsharded checkpoint this way. It takes a long time, and it's very inefficient (because every rank will load everything but then discard most of what it just loaded), but it does work. Unfortunately I seem to have broken how the optimizer state works. So this is not complete. |
|
In fact, with the way this works, we can now load a model that's so big it wouldn't fit into CPU memory. Not that we need to do this. But we could. |
For example, suppose you want to keep your checkpoints at every 1000 steps, but you also want to save a temporary checkpoint every 100 steps in case your job fails. In that case you would set `save_interval=1000` and `save_interval_ephemeral=100`.
…ints. But also still contains the old code.
epwalsh
approved these changes
Jan 4, 2024
| fsdp_model.load_state_dict(state_dict_to_load) | ||
| del state_dict_to_load | ||
| with torch.no_grad(): | ||
| # fill everything with NaN, so we can check afterwards that every parameter has been restored |
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.
More checkpointing formats
The problem that's being solved here is, how do we restore the 65B model from an unsharded checkpoint? The existing way works, I think, only by accident. If it works at all. It's possible that this never worked, or maybe only with sharding strategies that don't actually shard the model (like
NO_SHARDorSHARD_GRAD_OPorwrapping_strategy=null).So this new way uses
FSDP.apply()to copy tensors from a state dict into the model. This part is pretty straightforward. That's whatapply()is for.The part that isn't straightforward is the detour through the safetensors format. Safetensors is brilliant. It lets you create a
Dict[str, Tensor], where the tensors are memory mapped files. It loads up a 500GB files in seconds (because, of course, it doesn't actually read the tensor bytes until later). So this PR contains a script that can read an existing unsharded checkpoint (in .pt format), and write it to disk in safetensors format (.safetensors). This can be done on CPU, though you need a lot of memory to do it. When reading a .pt file, we check whether there happens to be a .safetensors file with the same name, and if so, we load that instead.One more problem is that state dicts are not
Dict[str, Tensor]. State dicts can contain inner dicts, and optimizer state dicts contain even more crazy stuff. So there is a mapper in this PR that maps crazy state dicts to well-formedDict[str, Tensor]s and back. This sacrifices human interpretability of the files, but retains the lazy-loading memmap goodness from safetensors.TODO