[Misc] Update GPTQ to use vLLMParameters (#7976)

This commit is contained in:
Dipika Sikka
2024-09-03 17:21:44 -04:00
committed by GitHub
parent dc0b6066ab
commit 2188a60c7e
6 changed files with 93 additions and 62 deletions

View File

@@ -1,5 +1,7 @@
import os
import torch
MAX_MODEL_LEN = 1024
MODEL_NAME = os.environ.get("MODEL_NAME",
"robertgshaw2/zephyr-7b-beta-channelwise-gptq")
@@ -8,9 +10,12 @@ QUANTIZATION = os.environ.get("QUANTIZATION", "gptq_marlin")
def test_weight_loading(vllm_runner):
"""
Test parameter weight loading with tp>1.
"""
with vllm_runner(model_name=MODEL_NAME,
revision=REVISION,
dtype="auto",
dtype=torch.half if QUANTIZATION == "gptq" else "auto",
quantization=QUANTIZATION,
max_model_len=MAX_MODEL_LEN,
tensor_parallel_size=2) as model: