Fix Gemma 7B checkpoint save

[ebsmothers]

Fixes #1122. See point (2) here.

Test:

tune run --nnodes 1 --nproc_per_node 4 lora_finetune_distributed --config \
gemma/7B_lora max_steps_per_epoch=5 epochs=1
...
INFO:torchtune.utils.logging:Model checkpoint of size 5.00 GB saved to /tmp/gemma-7b/hf_model_0001_0.pt
INFO:torchtune.utils.logging:Model checkpoint of size 4.98 GB saved to /tmp/gemma-7b/hf_model_0002_0.pt
INFO:torchtune.utils.logging:Model checkpoint of size 4.98 GB saved to /tmp/gemma-7b/hf_model_0003_0.pt
INFO:torchtune.utils.logging:Model checkpoint of size 2.11 GB saved to /tmp/gemma-7b/hf_model_0004_0.pt
INFO:torchtune.utils.logging:Adapter checkpoint of size 0.37 GB saved to /tmp/gemma-7b/adapter_0.pt
INFO:torchtune.utils.logging:Adapter checkpoint of size 0.37 GB saved to /tmp/gemma-7b/adapter_model.bin
INFO:torchtune.utils.logging:Adapter checkpoint of size 0.00 GB saved to /tmp/gemma-7b/adapter_config.json

[pytorch-bot]

🔗 Helpful Links

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

• 📄 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 7033610 with merge base 1af5135 ():
💚 Looks good so far! There are no failures yet. 💚

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

[joecummings]

show me the receipts

[joecummings]

show me the receipts

(Screenshot of successful model + adapter saving)

[joecummings]

show me the receipts

(Screenshot of successful model + adapter saving)

Wasn't sure if I was being clear

[joecummings]

show me the receipts

(Screenshot of successful model + adapter saving)

Wasn't sure if I was being clear

Does it make sense?

[ebsmothers]

show me the receipts

(Screenshot of successful model + adapter saving)

Wasn't sure if I was being clear

Does it make sense?

No

[ebsmothers]

show me the receipts

(Screenshot of successful model + adapter saving)

Wasn't sure if I was being clear

Does it make sense?

No

But I updated the test plan