Fine tuning with FSDP v2 #44
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
Custom SPMD sharding is more complicated and does not offer signiificant performance advantages over FSDP v2 (that uses SPMD). Example is updated and two helpers have been added to optimum.tpu.
SPMD sharding is not used anymore in training, unless FSDP v2 is used, that is supported by transformers and pytorch 2.3. So this is now deleted.
|
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. |
No special need in modeling, since there is global mesh and shard marking is done in the Trainer class now.
Trying to make optimum tpu smart enough to deduce fine tune args for supported models.
For gemma 7b a bigger system might be necessary to avoid OOMs.
Also if model is not matched, raise a clear error.
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.
What does this PR do?
After some discussions, we figured out that tuning with SPMD API as it was done in #39 might not have been the preferred way, and we moved to FSPD v2. This implementation uses SPMD underneath, so performance is very similar, with the benefit of being much simpler to integrate.
Llama and Gemma models have been tested and there are now two examples, in form of script and notebook.