wassname/vllm
1
0
Fork 0
mirror of https://github.com/wassname/vllm.git synced 2026-07-24 13:20:45 +08:00
Code Issues Packages Projects Releases Wiki Activity

[TPU] Add Load-time W8A16 quantization for TPU Backend (#7005)

Browse Source
This commit is contained in:
Siyuan Liu
2024-08-08 18:35:49 -07:00
committed by GitHub
parent 5923532e15
commit 0fa14907da
4 changed files with 135 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files
vllm/model_executor/model_loader/loader.py
+9 -8
View File
@@ -94,14 +94,15 @@ def _get_quantization_config(
"""Get the quantization config."""
if model_config.quantization is not None:
quant_config = get_quant_config(model_config, load_config)
capability = current_platform.get_device_capability()
capability = capability[0] * 10 + capability[1]
if capability < quant_config.get_min_capability():
raise ValueError(
f"The quantization method {model_config.quantization} is not "
"supported for the current GPU. "
f"Minimum capability: {quant_config.get_min_capability()}. "
f"Current capability: {capability}.")
if not is_tpu():
capability = current_platform.get_device_capability()
capability = capability[0] * 10 + capability[1]
if capability < quant_config.get_min_capability():
raise ValueError(
f"The quantization method {model_config.quantization} "
"is not supported for the current GPU. "
f"Minimum capability: {quant_config.get_min_capability()}. "
f"Current capability: {capability}.")
supported_dtypes = quant_config.get_supported_act_dtypes()
if model_config.dtype not in supported_dtypes:
raise ValueError(