-
Notifications
You must be signed in to change notification settings - Fork 241
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Speed up model loading for generate #709
base: main
Are you sure you want to change the base?
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/torchtune/709
Note: Links to docs will display an error until the docs builds have been completed. ❌ 1 New FailureAs of commit add24c6 with merge base ada5224 (): NEW FAILURE - The following job has failed:
This comment was automatically generated by Dr. CI and updates every 15 minutes. |
Any way I can know from the CI logs what is my lint mistake so I can fix it? |
@albanD thanks so much for putting this up! I'll take a more detailed look tomorrow, but to answer your lint question - you can do the following:
This will fix all of the issues for you. |
Thanks for the PR @albanD! Tbh we have already had a fraught relationship with meta device initialization 😅 (see e.g. #317, #418, #514). Our latest status is that we deliberately sacrifice a bit on time-to-first-batch for the sake of keeping code in the model components agnostic to meta device. But generation is an interesting case since the total runtime is much lower than on a finetune with FSDP (which is what we were focusing on previously). Out of curiosity, what is the speedup of meta device vs just initializing directly on GPU in this case? |
I would need to check once I go back on the machine in question. @kartikayk I saw that but I don't have pre-commit in my environment :p |
model = config.instantiate(model_cfg) | ||
|
||
model.load_state_dict(model_state_dict, assign=True) |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
for quantized models we'd need to load after we do quantization I think
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
We should consider doing this https://docs-preview.pytorch.org/pytorch/tutorials/2824/recipes/recipes/swap_tensors.html?highlight=swap_tensor once we can use 2.3+ in AO/here.
This has not been extensively tested (only mistral 7b) and more of a proposal!
This change does the follow:
This makes the model loading almost instant on my machine.