Expose flag for should_save_peft_only

[irenedea]

This will allow users to control whether they want their composer checkpoints from lora to have the full weights or just the adapter weights.

Note: This only applies to composer checkpoints, not HF checkpoints. HF checkpoints saved through hf_checkpointer.py will still only contain the adapters only.

Tests

should_save_peft_only: False

2024-07-15 21:27:50     844110 .metadata
2024-07-15 21:27:55 4028659956 __0_0.distcp
2024-07-15 21:27:55 4019819348 __1_0.distcp
2024-07-15 21:27:55 4019819348 __2_0.distcp
2024-07-15 21:27:55 4019819348 __3_0.distcp
2024-07-15 21:27:55 4019819348 __4_0.distcp
2024-07-15 21:27:55 4019819348 __5_0.distcp
2024-07-15 21:27:55 4019819348 __6_0.distcp
2024-07-15 21:27:55 4019819348 __7_0.distcp

should_save_peft_only: True

2024-07-15 21:35:06     291838 .metadata
2024-07-15 21:35:06   13185952 __0_0.distcp
2024-07-15 21:35:06    4345344 __1_0.distcp
2024-07-15 21:35:06    4345344 __2_0.distcp
2024-07-15 21:35:06    4345344 __3_0.distcp
2024-07-15 21:35:06    4345344 __4_0.distcp
2024-07-15 21:35:06    4345344 __5_0.distcp
2024-07-15 21:35:06    4345344 __6_0.distcp
2024-07-15 21:35:06    4345344 __7_0.distcp

[snarayan21]

code changes lgtm.
before merging, mind doing a save/resumption test with and without the flag? want to make sure we can resume from checkpoints with this enabled

lol offline discussion

[snarayan21]

lgtm