[NF4] Add quantize_() API support for NF4

[gau-nernst]

Background

When I was working on INT8 training stuff a while back, I realized there was no issue to make NF4 training work with quantize_() API (i.e. swap plain tensor with tensor subclass). This is because autograd+compile still works correctly for __torch_function__ (might not be the case in the past when NF4 was implemented) -> eliminate the need for custom nn.Module for downstream users as well as need for explicitly call to linear_nf4().

Usage

from torchao import quantize_
from torchao.dtypes.nf4 import nf4_weight_only

model = ...
quantize_(model, nf4_weight_only())  # works for both training and inference. NF4 weight is non-trainable

This would also compose nicely with LoRA, making QLoRA implementation more seamless. e.g.

from torch import nn
import torch.nn.functional as F

class LoRALinear(nn.Module):
    ...
    def forward(self, x):
        out = F.linear(x, self.weight)  # self.weight can be NF4Tensor
        out = out + self.lora_b(self.lora_a(x)) * (self.alpha / self.rank)
        return out

lora_layer = LoRALinear(...)
quantize_(
    lora_layer,
    nf4_weight_only(),
    filter_fn=lambda module, name: isinstance(module, LoRALinear),
)  # now this becomes QLoRA

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

❌ 1 New Failure

As of commit f9a5f94 with merge base f99b667 ():

NEW FAILURE - The following job has failed:

• Run Regression Tests / test (CPU Nightly, linux.4xlarge, --pre torch --index-url https://download.pytorch.org/whl/nightl... / linux-job (gh)
  test/prototype/test_parametrization.py::TestFakeSparsity::test_jit_trace

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

[suryanshgupta9933]

This would be really neat.