Add model+loss compile for full finetune single device

[gau-nernst]

Context

What is the purpose of this PR? Is it to

• add a new feature
• fix a bug
• update tests and/or documentation
• other (please add here)

Please link to any issues this PR addresses. #1228

Changelog

What are the changes made in this PR?

Add model+loss compile for full finetune single device. Basically the same as #1296

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.)

• run pre-commit hooks and linters (make sure you've first installed via pre-commit install)
• add unit tests for any new functionality
• update docstrings for any new or updated methods or classes
• run unit tests via pytest tests
• run recipe tests via pytest tests -m integration_test
• manually run any new or modified recipes with sufficient proof of correctness
• include relevant commands and any other artifacts in this summary (pastes of loss curves, eval results, etc.)

Qwen2-1.5B with 8-bit Adam bs=4

tune run full_finetune_single_device --config qwen2/1.5B_full_single_device optimizer._component_=bitsandbytes.optim.AdamW8bit optimizer_in_bwd=False compile=True metric_logger._component_=torchtune.utils.metric_logging.WandBLogger batch_size=4 log_peak_memory_stats=True

tok/s almost doubles, similar to @ebsmothers results for Llama2-7B with QLoRA on A100. Also much better memory usage (I believe the logits are the culprit here). When bs=8, "compile model only" OOM, while "compile model+loss" is ok.

All runs are done with 4070Ti SUPER 16GB.

[pytorch-bot]

🔗 Helpful Links

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

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

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

[SalmanMohammadi]

Nice to see you taking advantage of the small model configs : )
LGTM. Thank you for adding this.