NF4 quantization of linear layers without LoRA applied

[winglian]

Context

What is the purpose of this PR? Is it to

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

Please link to any issues this PR addresses. #1093

Changelog

Reverts #658 to bring back FrozenNF4Linear. When quantize_base is set to true, all base weights for linear layers are quantized, even if they do not have LoRA applied to them.

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

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

[pytorch-bot]

🔗 Helpful Links

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

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

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

[joecummings]

cc @msaroufim

[janeyx99]

Clarification question: does thid mean that quantize_base is the boolean with which one specifies the q in qLoRA, which is globally applied to all projs? (versus how LoRA could be applied to a subset of the parameters of q, k, v and not o).

[winglian]

running tune run lora_finetune_single_device --config llama3/8B_qlora_single_device
but only with lora_attn_modules: ['q_proj'] instead of all the attention linear layers, uses 15810MiB on main and 15000MiB on this branch.

EDIT: setting lora_attn_modules active across all 4 modules uses 15044MiB in this branch, which is expected to be more than the 15000MiB due to the optimizer, and also less than main because all the weights are quantized.

[janeyx99]

Rest of PR looks fine; we should implement FrozenNF4Linear as cleanly as possible though

[janeyx99]

Could we avoid accessing .data of self.weight?

I'm curious if there's a way to use swap_tensors here to accomplish this more cleanly. https://pytorch.org/docs/stable/generated/torch.utils.swap_tensors.html cc @mikaylagawarecki any ideas?

[mikaylagawarecki]

@winglian Could you explain the rationale for the .data access?

[ebsmothers]

Yeah can we do something similar to what's currently done in LoRALinear for this?

[winglian]

This is simply reverting back the code that was deleted in https://github.com/pytorch/torchtune/pull/658/files#diff-74b0d911936ebd1d0e216004577afed27b84b6bfdff9c6a9a1a28f6fac054850L45

Removing the .data does seem to work as well, so I've checked in that change.

[janeyx99]

Also, what test cases can we write to verify functionality?

[ebsmothers]

Overall the changes look reasonable to me. Re @janeyx99's testing comment, a couple things I think we should add here:

(1) a unit test for nf4_linear (can probably look at #465 for ideas)
(2) some kind of e2e test. I am thinking: run one of our QLoRA recipes before and after the change, confirm that (a) the peak memory reduces as expected, and (b) we see no regression in eval metrics on the resulting fine-tuned checkpoint.

It's also likely that this will break some of our existing QLoRA tests. So e.g. the values here would need to be updated

[ebsmothers]

Yeah can we do something similar to what's currently done in LoRALinear for this?