[torchao] Support quantization configs using module swap

[jerryzh168]

Summary:

Current torchao integration quantizes the weights by wrapping weights in a top level linear module and use quantize_ to quantize it, this works for quantization methods that do inplace changes to the weight itself, such as int4, float8, since these only do inplace changes to the linear module itself

but there are quantization configs that would need module swap, such as awq, that's not supported, in order to support these, we wrap the linear in nn.Sequential so it is no longer a top level module and can be swapped to another module.

Test Plan:
uplodated an awq checkpoint: https://huggingface.co/torchao-testing/Phi-4-mini-instruct-int4wo-awq-0.13-dev and we test by loading the checkpoint

python tests/quantization/test_torchao.py

Reviewers:

Subscribers:

Tasks:

Tags:

[gemini-code-assist]

Code Review

This pull request enables support for torchao quantization methods that require module swaps, like AWQ, by wrapping the linear layer in an nn.Sequential. The change is well-contained and accompanied by a relevant test case. I've suggested one improvement to make the code more robust against future changes in torchao or different quantization configurations.

[gemini-code-assist]

After quantization, the module may have been swapped. Instead of directly accessing dummy_linear[0].weight, retrieve the weight by inspecting the module's parameters. This avoids making fragile assumptions about the internal structure of the quantized module, which may change in future torchao versions.

    dummy_linear[0].weight = param
    quantize_(dummy_linear, torchao_config)
    # After quantization, the module may have been swapped.
    # We retrieve the single parameter, which is the quantized weight.
    params = list(dummy_linear.parameters())
    assert len(params) == 1, (
        "Expected the dummy module to have exactly one parameter after "
        f"quantization, but found {len(params)}."
    )
    return params[0].data

[github-actions]

👋 Hi! Thank you for contributing to the vLLM project.

💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in #pr-reviews, coordinate on features in #feat- channels, or join special interest groups in #sig- channels.

Just a reminder: PRs would not trigger full CI run by default. Instead, it would only run fastcheck CI which starts running only a small and essential subset of CI tests to quickly catch errors. You can run other CI tests on top of those by going to your fastcheck build on Buildkite UI (linked in the PR checks section) and unblock them. If you do not have permission to unblock, ping simon-mo or khluu to add you in our Buildkite org.

Once the PR is approved and ready to go, your PR reviewer(s) can run CI to test the changes comprehensively before merging.

To run CI, PR reviewers can either: Add ready label to the PR or enable auto-merge.

🚀

[mgoin]

Looks reasonable to me, thanks!

[houseroad]

Looks good.

[DarkLight1337]

Can you merge from main to fix the CI failures?

[jerryzh168]

the current error is because we just landed AWQ updates in torchao recently and it's not picked up by nightly yet
there is actually one more update to awq I plan to do next week, we can land this one after that