[Torchax] Support Running AWQ models

[kyuyeunk]

Description

• Following PR adds support for running loading & running AWQ models in torchax
  • Please note that this requires a fix in vLLM repo to be submitted as well: [Bugfix] Fix packed_factor missing attribute error vllm#23902
• Using model Qwen/Qwen2.5-32B-Instruct-AWQ, I have verified that it'ss running correctly.
  • Numeric validation
    • Command TPU_BACKEND_TYPE=jax MODEL_IMPL_TYPE=vllm python3 examples/offline_inference.py --model=Qwen/Qwen2.5-32B-Instruct-AWQ --tensor_parallel_size=8 --task=generate --max_model_len=1024 --download_dir=/mnt/disks/persist
    • Result: http://gpaste/6205721003425792
  • Performance validation (degradation over baseline is expected)
    • Baseline (Qwen/Qwen2.5-32B-Instruct): 12.61 reqs/s
    • AWQ (Qwen/Qwen2.5-32B-Instruct-AWQ): 11.34 reqs/s
• Limitation / Future work
  • AWQ uses w4a16 subchannel quantization, and for initial implementation, I have utilized XLA instead of kernel. However, there is known performance issue with subchannel quantization in XLA. Therefore, we need to implement w4a16 subchannel support into the kernel for optimal performance.

Tests

pytest -v tests/models/vllm/layers/test_awq.py

Checklist

Before submitting this PR, please make sure:

• I have performed a self-review of my code.
• I have necessary comments in my code, particularly in hard-to-understand areas.
• I have made or will make corresponding changes to any relevant documentation.

[github-actions]

Description

Start with a short description of what the PR does and how this is a change from
the past.

The rest of the description includes relevant details and context, examples:

• why is this change being made,
• the problem being solved and any relevant context,
• why this is a good solution,
• some information about the specific implementation,
• shortcomings of the solution and possible future improvements.

If the change fixes a bug or a Github issue, please include a link, e.g.,:
FIXES: b/123456
FIXES: #123456

Tests

Please describe how you tested this change, and include any instructions and/or
commands to reproduce.

Checklist

Before submitting this PR, please make sure:

• I have performed a self-review of my code.
• I have necessary comments in my code, particularly in hard-to-understand areas.
• I have made or will make corresponding changes to any relevant documentation.

[vanbasten23]

add this test "pytest -v tests/models/vllm/layers/test_awq.py" to the CI?

[kyuyeunk]

This gets run in Jax unit test: https://github.com/vllm-project/tpu_commons/blob/4470371755e316bf7c3a5e5940133923d7ebc325/.buildkite/pipeline_jax.yml#L76-L88

In a (failed...) CI, I have confirmed that is the case: https://buildkite.com/tpu-commons/tpu-commons-ci/builds/2546#01990bbf-4525-4e7c-ad33-667076f9482d

Is it a better idea to add this in torch ci?: https://github.com/vllm-project/tpu_commons/blob/main/.buildkite/pipeline_torch.yml

[vanbasten23]

I see. Thanks. If it's framework agnostic, then it's fine to leave it as is.

[kyuyeunk]

vllm-project/vllm#23902 has been merged. This PR is ready to be submitted.

[kyuyeunk]

This gets run in Jax unit test: https://github.com/vllm-project/tpu_commons/blob/4470371755e316bf7c3a5e5940133923d7ebc325/.buildkite/pipeline_jax.yml#L76-L88

In a (failed...) CI, I have confirmed that is the case: https://buildkite.com/tpu-commons/tpu-commons-ci/builds/2546#01990bbf-4525-4e7c-ad33-667076f9482d

Is it a better idea to add this in torch ci?: https://github.com/vllm-project/tpu_commons/blob/main/.buildkite/pipeline_torch.yml