Fix save adapter weights only

[ebsmothers]

Fixes #1713

Summary

We should not be using get_adapter_params on models with activation checkpointing enabled. This is because get_adapter_params iterates over named_modules, not state_dict. This is a good thing from a memory perspective, but it also means that the state dict hooks that are usually triggered to unwrap activation checkpointing wrapper modules (or previously FSDP wrapper modules) do not get triggered. We can consider adding an option to get_adapter_params to do this in that utility cause this is a bit tricky. At the same time, that also feels hacky.

Actually this is really only an issue in our single-device recipe when save_adapter_weights_only=True -- otherwise we filter the state dict on the keys of self.adapter_params (this is correct because (a) self.adapter_params is defined before any AC wrapping and (b) we are using the state dict so the relevant hooks have fired).

But actually, we don't need to use get_adapter_params here at all: we already have references to all the adapter weights and the correct (non-AC-wrapped) keys in self.adapter_params as defined here. So we can just use that, and we can do it regardless of the value of save_adapter_weights_only (which makes the code a bit cleaner).

Test plan

Why didn't we catch this before? Well actually all of our existing recipes tests set enable_activation_checkpointing=False so we missed it. For that reason I've updated the test_recipe_state_on_resume tests to have enable_activation_checkpointing=True. Running those tests on main with the new values for enable_activation_checkpointing:

pytest tests/recipes/test_lora_finetune_single_device.py -m integration_test -k 'test_training_state_on_resume'
...
===== 1 failed, 1 passed, 9 deselected, 1 warning in 7.29s ======


pytest tests/recipes/test_lora_dpo_single_device.py -m integration_test -k 'test_training_st
ate_on_resume'
...
==== 1 failed, 1 passed, 1 deselected, 1 warning in 28.37s =======

If we instead run with the changes on this PR:

pytest tests/recipes/test_lora_finetune_distributed.py -m integration_test
...
========== 7 passed in 127.43s (0:02:07) ==========

pytest tests/recipes/test_lora_dpo_single_device.py -m integration_test
...
========== 3 passed, 2 warnings in 52.46s =======

pytest tests/recipes/test_lora_finetune_single_device.py -m integration_test
...
========== 11 passed, 2 warnings in 271.11s (0:04:31) ===========

Also, the command from #1713 that was failing before now passes:

tune run lora_finetune_single_device --config llama3_2/3B_lora_single_device max_steps_per_epoch=10 \
gradient_accumulation_steps=1 epochs=1 save_adapter_weights_only=True

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/torchtune/1764

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

✅ No Failures

As of commit 8357be3 with merge base a8a64ec ():
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[RdoubleA]

solid fix