Skip to content

[torchao] Support quantization configs using module swap #21982

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 1 commit into
base: main
Choose a base branch
from
+24 −7
Open

[torchao] Support quantization configs using module swap #21982

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
15 changes: 15 additions & 0 deletions tests/quantization/test_torchao.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -75,5 +75,20 @@ def test_qwenvl_int8wo_model_loading_with_params(vllm_runner):
print(output)


@pytest.mark.skipif(not TORCHAO_AVAILABLE, reason="torchao is not available")
def test_phi4mini_int4wo_awq_model_loading_with_params(vllm_runner):
torch._dynamo.reset()
model_name = "torchao-testing/Qwen3-4B-int4wo-awq-0.13-dev"
with vllm_runner(model_name=model_name,
quantization="torchao",
dtype="bfloat16",
pt_load_map_location="cuda:0") as llm:
output = llm.generate_greedy(["The capital of France is"],
max_tokens=32)

assert output
print(output)


if __name__ == "__main__":
pytest.main([__file__])
16 changes: 9 additions & 7 deletions vllm/model_executor/layers/quantization/torchao.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -152,18 +152,20 @@ def torchao_quantize_param_data(param: torch.Tensor,
from torchao.quantization import quantize_

assert isinstance(torchao_config, AOBaseConfig), f"{torchao_config}"
"""
Avoid real weight allocation for faster load, since we will
"""
Avoid real weight allocation for faster load, since we will
end up setting it to param.
"""
with torch.device("meta"):
dummy_linear = torch.nn.Linear(param.shape[1],
param.shape[0],
bias=False)
# linear can't be top level module since quantize_ is inplace
# while some of our configs need to do module swap, and only non-top
# level modules support module swap
dummy_linear = torch.nn.Sequential(
torch.nn.Linear(param.shape[1], param.shape[0], bias=False))

dummy_linear.weight = param
dummy_linear[0].weight = param
quantize_(dummy_linear, torchao_config)
return dummy_linear.weight
return dummy_linear[0].weight
Comment on lines +166 to +168
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

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



class TorchAOLinearMethod(LinearMethodBase):
Expand Down