Add fp32 support for QLoRA

[rohan-varma]

Context

• QLoRA is currently coupled to bf16, but some older HW types don't support bf16, and we'd ideally like to enable at least one memory efficient finetuning solution for these HW arches. For example, T4s are 16GB and don't support bf16. This PR enables QLoRA to run compute + checkpoint in fp32 instead of bf16, to eliminate the hard coupling of bf16 to QLoRA training.

Changelog

• Remove bf16 assumptions from nf4
• Remove bf16 assumptions from LoRALinear and upstream
• Generalize a few functions to be less coupled to bf16
• Fix tests

NOTE

• We've currently forked over the LinearNF4 from torchao while changes to LinearNF4 land in ao are in progress. Will revert back to using ao's implementation asap, but we need the changes in this forked version to decouple support from fp32. This is a temporary (~days) mitigation.

Test plan

• Modified unittests - computation in fp32, checkpointing, parity is covered. fp32 coverage for qlora is at the same coverage as bf16.
• Verified manually that gradients are computed in fp32 (via inspecting .grad field)
• Checkpoints are saved in fp32:

• Run recipe: tune lora_finetune_single_device --config llama2/7B_qlora_single_device dtype=fp32 epochs=1
• Loss is comparable to QLoRA and LoRA-bf16 (see QLoRA #478 for those curves):

• Memory - increase over QLoRA bf16 is +20% peak memory allocated, +3% reserved memory

Memory Stats::
 GPU peak memory allocation: 6.98 GB
 GPU peak memory reserved: 9.57 GB
 GPU peak memory active: 6.98 GB

• Eval result:

[pytorch-bot]

🔗 Helpful Links

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

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

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

[ebsmothers]

Honestly I don't really see a better way to do this

[ebsmothers]

remove

[ebsmothers]

Looks great!