Refactor NF4Tensor to use dispatch utils

[jerryzh168]

Summary:
As a follow up of #542, we can simplify the code of nf4tensor by using the dispatch utils as well.

Test Plan:
python test/dtypes/test_nf4.py

Reviewers:

Subscribers:

Tasks:

Tags:

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

❌ 6 New Failures

As of commit 452901a with merge base afde175 ():

NEW FAILURES - The following jobs have failed:

• Run Regression Tests / test (CPU 2.2.2, linux.4xlarge, torch==2.2.2 --index-url https://download.pytorch.org/whl/cpu "nu... / linux-job (gh)
  test/quantization/test_quant_api.py::TestQuantFlow::test_dynamic_quant_gpu_singleline
• Run Regression Tests / test (CPU 2.3, linux.4xlarge, torch==2.3.0 --index-url https://download.pytorch.org/whl/cpu, cpu) / linux-job (gh)
  test/quantization/test_quant_api.py::TestQuantFlow::test_dynamic_quant_gpu_singleline
• Run Regression Tests / test (CPU Nightly, linux.4xlarge, --pre torch --index-url https://download.pytorch.org/whl/nightl... / linux-job (gh)
  test/quantization/test_quant_api.py::TestQuantFlow::test_quantized_tensor_subclass_8da4w
• Run Regression Tests / test (CUDA 2.2.2, linux.g5.12xlarge.nvidia.gpu, torch==2.2.2 "numpy<2" , cuda, 12.1) / linux-job (gh)
  test/quantization/test_quant_api.py::TestQuantFlow::test_dynamic_quant_gpu_singleline
• Run Regression Tests / test (CUDA 2.3, linux.g5.12xlarge.nvidia.gpu, torch==2.3.0, cuda, 12.1) / linux-job (gh)
  test/quantization/test_quant_api.py::TestQuantFlow::test_dynamic_quant_gpu_singleline
• Run Regression Tests / test (CUDA Nightly, linux.g5.12xlarge.nvidia.gpu, --pre torch --index-url https://download.pytorc... / linux-job (gh)
  test/sparsity/test_sparse_api.py::TestQuantSemiSparse::test_quant_semi_sparse

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

[jerryzh168]

Note: it seems that some ops are implemented this way to work with QLoRA. from @weifengpy:

I remember they hack around state dict loading and has to call args[1]["dtype"] instead of kwargs["dtype"] so changing the behavior might break single-device QLoRA recipe

So I'll leave this refactor to nf4 people. cc @msaroufim