Improve FP6-LLM 2+4bit weight splitting + user API

[gau-nernst]

Address #208

2+4bit weight splitting

Port https://github.com/pytorch/ao/blob/4ca3985be603e6496da7ec57adf1942c8b32a78e/torchao/csrc/fp6_llm/weight_prepacking.cpp to pure PyTorch.

FP16 weight, (8192, 8192). Ryzen 5600, 4070Ti SUPER

device	op	time (m/s)
CPU (num_threads=1)	original FP16->FP6 + original 2+4bit splitting	1380.96
CPU (num_threads=1)	new FP16->FP6 + original 2+4bit splitting	616.285
CPU (num_threads=1)	new FP16->FP6 + new 2+4bit splitting	587.965
CPU (num_threads=4)	original FP16->FP6 + original 2+4bit splitting	1213.08
CPU (num_threads=4)	new FP16->FP6 + original 2+4bit splitting	334.502
CPU (num_threads=4)	new FP16->FP6 + new 2+4bit splitting	246.911
CUDA	new FP16->FP6 + original 2+4bit splitting	257.539
CUDA	new FP16->FP6 + new 2+4bit splitting	1.05908

Note:

• original 2+4bit splitting only works on CPU. Thus, for the 2nd last row, FP16->FP6 is done on GPU, but 2+4bit splitting is done on CPU.

User API

from torchao.quantization.fp6_llm import convert_fp6_llm

model = ...
convert_fp6_llm(model)  # nn.Linear modules will be replaced with Fp6LlmLinear in-place

I opt for custom linear module instead of tensor subclass mainly because it's easier to implement.

Note:

• Fp6LlmLinear will cast input to FP16 and cast output to original dtype.

[pytorch-bot]

🔗 Helpful Links

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

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

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