Move non-NF4 tensor to device prior to quantization on copy

[ebsmothers]

Addressing #642 based on @gau-nernst's suggestion here. Tested in torchtune that (a) init time and memory are similar to torchtune's inplace copy here and (b) FSDP2 recipes still succeed (since they don't with torchtune copy enabled).

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/ao/737

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

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

[msaroufim]

Stamping so we can merge fast to unblock but one thing that might help here is a test ideally here since we'd rather catch issues early or at the very least in tune since right now if another load time regression happens we won't catch it either

[ebsmothers]

Stamping so we can merge fast to unblock but one thing that might help here is a test ideally here since we'd rather catch issues early or at the very least in tune since right now if another load time regression happens we won't catch it either

Yeah good point. We don't currently have any perf-related testing but it actually shouldn't be too hard to add something like that on our end. Honestly just general sanity checks for QLoRA around load time, peak memory, etc will do us a lot of good on this front