Add torchao mps ops

[manuelcandales]

This PR adds the quantization scheme linear:afpwx. It quantizes only the weights in a groupwise manner with a specified bitwidth and groupsize. It takes arguments bitwidth (1, 2, 3, 4, 5, 6, 7) and groupsize (32, 64, 128, 256).

To use linear:afpwx, you must first set up the torchao mps experimental kernels. These will only work on a device with Apple Silicon.

From the torchchat root directory, run

bash torchchat/utils/scripts/build_torchao_ops.sh mps

Notice that this quantization scheme is currently implemented only for device mps.

python3 torchchat.py generate stories110M --device mps --dtype float32 --quantize '{"linear:afpwx": {"bitwidth": 4, "groupsize": 256}}' --prompt "Once upon a time,"

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/torchchat/1415

• 📄 Preview Python docs built from this PR

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

⏳ 1 Pending, 2 Unrelated Failures

As of commit c164f88 with merge base 4dc2f89 ():

BROKEN TRUNK - The following jobs failed but were present on the merge base:

👉 Rebase onto the `viable/strict` branch to avoid these failures

• pull / test-cpu-aoti (aarch64, stories15M) (gh) (trunk failure)
  Process completed with exit code 1.
• pull / test-cpu-aoti (x86_64, stories15M) (gh) (trunk failure)
  Process completed with exit code 1.

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

[metascroy]

Should we keep the naming convention that torchchat uses (of "a" followed by type and "w" followed by type). This started with a8w4dq before I added any kernels. In your case this would be something like afpwx?

[manuelcandales]

Ok, this makes sense.

[metascroy]

Do these only work with eager? If so, explicitly say that in the set-up section?

[manuelcandales]

The metal lowbit kernels run with ExecuTorch as well (llama runner can use them). However, my aim in this torchchat PR was only to enable eager. I plan to have a follow up PR to enable them via the torchchat ET path as well. But I prefer to keep it modular.

[manuelcandales]

I added a sentence in the setup section, clarifying that currently torchchat can only use them on Eager mode

[metascroy]

Can you add a CI test for the MPS kernels to make sure they install and run?

See https://github.com/pytorch/torchchat/blob/main/.github/workflows/pull.yml#L1060 as an example.

[manuelcandales]

done!