Enable torchao.experimental EmbeddingQuantization

[Jack-Khuu]

🚀 The feature, motivation and pitch

Quantization is a technique used to reduce the speed, size, or memory requirements of a model and torchao is PyTorch's native quantization library for inference and training

There are new experimental quantizations in torchao that we would like to enable in torchchat. Specifically this task is for enabling EmbeddingQuantizer and SharedEmbeddingQuantizer.

Entrypoint:

torchchat/torchchat/utils/quantize.py

Line 101 in 1384f7d

def quantize_model(

Task: Using ExecuTorch as a reference (pytorch/executorch#9548) add support for EmbeddingQuantizer and SharedEmbeddingQuantizer.

cc: @metascroy, @manuelcandales

Alternatives

No response

Additional context

No response

RFC (Optional)

No response

[dillondesilva]

@Jack-Khuu I'd love to have a crack at this if possible! Would you mind assigning it to me?

[Jack-Khuu]

Totally, give it a shot

[Jack-Khuu]

Hi @dillondesilva, how's the task going? Any questions?

[dillondesilva]

Hey @Jack-Khuu - I've just been busy with mid-semester exams in the past week. Should have time to start sometime this week and will send questions soon 👍 Thanks for checking in!

[dillondesilva]

@Jack-Khuu Good news! Here's the PR -> #1525

I don't know if I've oversimplified it so feel free to correct me if I'm wrong but to enable the above experimental quantizers, I'm assuming all that was needed was:

1. A mapping in quantizer_class_dict to enable devs to specify when they are using one of the new experimental options (i.e. experimental:embedding and experimental:shared)
2. Logic to run the respective quantizers if they have been specified

[Jack-Khuu]

Sweet! I'll take a look