Only merge model weights in LoRA recipe when save_adapter_weights_only=False
#1476
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/torchtune/1476
Note: Links to docs will display an error until the docs builds have been completed. ✅ No FailuresAs of commit 4991014 with merge base be8f1e7 ( This comment was automatically generated by Dr. CI and updates every 15 minutes. |
facebook-github-bot
added
the
CLA Signed
Codecov ReportAttention: Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #1476 +/- ##
===========================================
- Coverage 73.12% 27.00% -46.13%
===========================================
Files 289 290 +1
Lines 14175 14252 +77
===========================================
- Hits 10366 3849 -6517
- Misses 3809 10403 +6594 ☔ View full report in Codecov by Sentry. |
| @@ -0,0 +1,182 @@ | |||
| # Copyright (c) Meta Platforms, Inc. and affiliates. | |||
There was a problem hiding this comment.
I was just afraid someone would ask me to run the recipe manually, so this felt like less effort.
| adapter_state_dict = { | ||
| k: v | ||
| for k, v in self._model.state_dict().items() | ||
| if adapter_key_filter(k) | ||
| } |
There was a problem hiding this comment.
Doesn't this still materialize the entire state dict though?
| alpha=self._lora_alpha, | ||
| ) | ||
| checkpoint_dict.update({training.MODEL_KEY: merged_state_dict}) | ||
| if not self._save_adapter_weights_only: |
There was a problem hiding this comment.
I think @pbontrager had looked at something with this previously, but we still create the full state dict on CPU in this case then, right? To actually reduce the memory footprint to just the adapter weights we'd need to change the call to get_full_model_state_dict above (btw it's ok to say "let's do this in a follow-up")
There was a problem hiding this comment.
A couple questions around whether we can get around materializing the full model state dict when save_adapter_weights_only=True. (I think it'll be trickier for the distributed recipe, but would be nice to sort out for single-device.) Other than that looks good!
Context
What is the purpose of this PR? Is it to
Please link to any issues this PR addresses.
Changelog
When saving LoRA checkpoints, we only need to copy and merge the full model state dict when
save_adapter_weights_only=False. For distributed recipes, this change will only skip the merging logic.I've also added tests for resuming recipe state with LoRA DPO.
Test plan
Please make sure to do each of the following if applicable to your PR. (If you're not sure about any one of these just ask and we will happily help. We also have a contributing page for some guidance on contributing.)
pre-commit install)pytest testspytest tests -m integration_testUX
If your function changed a public API, please add a dummy example of what the user experience will look like when calling it.
Example of docstring:
torchtune/torchtune/modules/vision_transformer.py
Line 285 in 6a7951f
Example in our docs: https://pytorch.org/torchtune/main/tutorials/qat_finetune.html#applying-qat-to-llama3-models