transformers
transformers copied to clipboard
huggingface
Extend save_pretrained to offloaded models
What does this PR do?
Fixes #20072 and addresses the second part of https://github.com/huggingface/peft/issues/868
Models with offloaded weights are currently incompatible with save_pretrained. This PR allows large models that are loaded onto the gpu and cpu to be saved, which is particularly useful for big models that have undergone merging and unloading via https://github.com/huggingface/peft/pull/1063.
The implementation is to iterate through modules and onload parameters to the execution device (typically gpu) before sending the appropriate elements of the state dict to the cpu in-place, where the final state dictionary is assembled and saved.
Before submitting
- [ ] This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
- [ x] Did you read the contributor guideline, Pull Request section?
- [ ] Was this discussed/approved via a Github issue or the forum? Please add a link to it if that's the case.
- [ ] Did you make sure to update the documentation with your changes? Here are the documentation guidelines, and here are tips on formatting docstrings.
- [ ] Did you write any new necessary tests?
Still working on the tests (some small models are not compatible with offloading due to architectural considerations) but am happy to submit a colab version with a large model in the meantime:)
Who can review?
Anyone! @pacman100
Overall I can see this being a pretty nice idea. Made some nits to improve. cc @SunMarc for your thoughts as well :)
Would it be possible to eventually offload this to accelerate?
Yes, we can definitely put it in a get_offloaded_model_state_dict() function in accelerate and use it here.
Thanks very much for reviewing everyone! @muellerzr @SunMarc your suggestions are much appreciated, I have gone ahead and implemented pretty much all of them.
A test for save_pretrained applied to a model with offloaded parameters has been added to test_modeling_utils.py. I previously used an equality test on the state dictionaries themselves and can add that too if that would be desirable.
@SunMarc sounds good to me! Thanks for the review and for giving pointers on where the method would live as well. I will open a PR in accelerate for adding the function get_state_dict_offloaded_model() with the proper integrations.
@ArthurZucker thanks much for the review!
Just to make sure I understand correctly, are you referring to adding compatibility for an offloaded model to be saved when it does not fit in cpu memory, with a possible method being onloading the state dictionary in parts and saving these parts as shards?
I certainly agree that it would be best to be able to save models that don't fit in cpu memory, rather than just gpu memory as is currently the case.
As a general question to guide this, should these changes be added here or in accelerate? Being that most of this PR's additions were sent to https://github.com/huggingface/accelerate/pull/2156 we could implement this compatibility feature either in transformers/modeling_utils or in accelerator.accelerate.
It's more like it does not make sense for me to support save_pretrained for offloaded models if we don't support the most common use case for offload which is that the RAM is not big enough to hold the model (you can't save from the GPU anyway) ! This could come in a followup PR, but it clearly makes more sense to support both instead of overloading the ram even if we have a big ram. Will also be faster
We should be good to go for saving models too large to fit on cpu memory (although I am still cleaning up the code and adding to the test). The approach is to map shard keys to modules and onload each shard's key just before saving that shard, before offloading once again.
Safe serialization provides a wrinkle in that meta tensors are by default recognized as shared tensors. I am not sure if there is a good way around this without onloading all tensors, which is a no-go if we are trying to preserve memory. Open to suggestions on this, however.
This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.
Please note that issues that do not follow the contributing guidelines are likely to be ignored.
![github-actions[bot] avatar](https://avatars.githubusercontent.com/in/15368?v=4 "Published by a pub.dev verified publisher"){kind=small}
Code has now been refactored and the test is updated, all are passing now.
I have run the test externally with small max_shard_size to confirm that the model only loads the specified amount into memory at any one time (checked via linux tools). A memory trace (ie via tracemalloc or memory_profiler) can be added to the test if that is considered desirable too.
Went ahead and added cpu memory checks via psutil to the test. The idea is that the model.save_pretrained function will onload no more memory than max_shard_size from the disk to the cpu at any given time during the saving process, accounting for a margin of error that is somewhat model-dependent.
This is perhaps easier to check with a large model, but I think this test (for a very small model) is still valid as-is: if we increase the shard size substantially (or use the default shard size which is larger than the model size) the memory check fails.
Thanks a lot, I'm just waiting for @muellerzr to run a first pass and will review as well 🤗
No problem, sounds good to me! I have also performed tests on disk storage used and there does not appear to be any change in storage usage compared to the head branch. I am not sure if we need a disk storage check added to the current test, however.
After some re-examination this approach seems fairly robust and straightforward with the exception of the sharded memory allocation. It appears to me that the current block allocation using id_tensor_storage() will not suffice for saving offloaded modules (hence the extra conditions added to avoid classifying 'meta' device tensors as shared tensors), but I may be overlooking something there @muellerzr
Added a warning for the user to check that the free memory exceeds the shard_size, which may not be the case for `device_map="auto" and would cause saving to fail otherwise without a clear reason. I have tested on some recent models with more exotic architectures (mixtral 8x7b etc.) and everything looks good there too.
cc @ArthurZucker
Thanks very much for your review @SunMarc ! I would be happy with moving some of this to accelerate, that seems like a natural place especially for the code you pointed to. Nice one thinking of using find_tied_parameters too, that should solve the problem of getting tied weights on the meta device. I should have time to go through your suggestions in the coming days and will add a PR in accelerate for get_state_dict_from_offload as well.
I have finished sending the accelerate -specific code to that library and the necessary integrations here, and after testing with some large models everything looks good.
Just for everyone's awareness, it looks like the psutils memory check in the test_save_offloaded_model test is not entirely reliable (fails around one in five times on my local machine). This appears to be due to the inherent noise in background process memory allocation compared to the small test model used, as adapting the test to much larger models removes this issue. I am not sure that we would want to use a large model, however, as this makes the test quite slow.
@SunMarc @ArthurZucker
This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.
Please note that issues that do not follow the contributing guidelines are likely to be ignored.
@SunMarc thanks very much for taking a look! No worries, I have been very busy too and would not have had much time to work on this before now anyways. I will plan make time to go through your suggestions tomorrow and will let you know if I can't make the finishing touches myself, in which case you would be more than welcome to do so
@SunMarc went ahead and implemented your suggestions, thanks! I was not sure which version of accelerate we should check for (the current version if we merge PR #2619 asap or the next version?) however.
The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.
Thank you again @blbadger for your patience and your work ! I really appreciate your contribution 🔥 Congrats on merging this amazing feature !
Happy to contribute! Thanks very much @SunMarc for shepherding this through and @amyeroberts @muellerzr @ArthurZucker for your reviews.