Views
No views yet
1from vllm.assets.audio import AudioAsset
2from vllm import LLM, SamplingParams
3
4# prepare model
5llm = LLM(
6 model="neuralmagic/whisper-medium-quantized.w4a16",
7 max_model_len=448,
8 max_num_seqs=400,
9 limit_mm_per_prompt={"audio": 1},
10)
11
12# prepare inputs
13inputs = { # Test explicit encoder/decoder prompt
14 "encoder_prompt": {
15 "prompt": "",
16 "multi_modal_data": {
17 "audio": AudioAsset("winning_call").audio_and_sample_rate,
18 },
19 },
20 "decoder_prompt": "<|startoftranscript|>",
21}
22
23# generate response
24print("========== SAMPLE GENERATION ==============")
25outputs = llm.generate(inputs, SamplingParams(temperature=0.0, max_tokens=64))
26print(f"PROMPT : {outputs[0].prompt}")
27print(f"RESPONSE: {outputs[0].outputs[0].text}")
28print("==========================================")python quantize.py --model_path openai/whisper-medium --quant_path "output_dir/whisper-medium-quantized.w4a16" --calib_size 1024 --group_size 64 --dampening_frac 0.01 --actorder weight1import torch
2import argparse
3from datasets import load_dataset
4from transformers import WhisperProcessor
5from llmcompressor import oneshot
6from llmcompressor.modifiers.quantization import GPTQModifier
7from llmcompressor.transformers.tracing import TraceableWhisperForConditionalGeneration
8import os
9from compressed_tensors.quantization import QuantizationArgs, QuantizationType, QuantizationStrategy, ActivationOrdering, QuantizationScheme
10from llmcompressor.modifiers.smoothquant import SmoothQuantModifier
11
12parser = argparse.ArgumentParser()
13parser.add_argument('--model_path', type=str)
14parser.add_argument('--quant_path', type=str)
15parser.add_argument('--calib_size', type=int, default=256)
16parser.add_argument('--dampening_frac', type=float, default=0.1)
17parser.add_argument('--observer', type=str, default="minmax")
18parser.add_argument('--actorder', type=str, default="dynamic")
19parser.add_argument('--group_size', type=int, default=128)
20parser.add_argument('--save_dir', type=str, required=True)
21
22
23args = parser.parse_args()
24model_id = args.model_path
25
26model = TraceableWhisperForConditionalGeneration.from_pretrained(
27 model_id,
28 device_map="auto",
29 torch_dtype="auto",
30)
31model.config.forced_decoder_ids = None
32processor = WhisperProcessor.from_pretrained(model_id)
33
34# Configure processor the dataset task.
35processor.tokenizer.set_prefix_tokens(language="en", task="transcribe")
36
37# Select calibration dataset.
38DATASET_ID = "MLCommons/peoples_speech"
39DATASET_SUBSET = "test"
40DATASET_SPLIT = "test"
41
42# Select number of samples for calibration. 512 samples is a good place to start.
43# Increasing the number of samples can improve accuracy.
44
45NUM_CALIBRATION_SAMPLES = args.calib_size
46MAX_SEQUENCE_LENGTH = 2048
47dampening_frac=args.dampening_frac
48actorder_arg=args.actorder
49group_size=args.group_size
50
51# Load dataset and preprocess.
52ds = load_dataset(
53 DATASET_ID,
54 DATASET_SUBSET,
55 split=f"{DATASET_SPLIT}[:{NUM_CALIBRATION_SAMPLES}]",
56 trust_remote_code=True,
57)
58
59def preprocess(example):
60 return {
61 "array": example["audio"]["array"],
62 "sampling_rate": example["audio"]["sampling_rate"],
63 "text": " " + example["text"].capitalize(),
64 }
65
66ds = ds.map(preprocess, remove_columns=ds.column_names)
67
68# Process inputs.
69def process(sample):
70 inputs = processor(
71 audio=sample["array"],
72 sampling_rate=sample["sampling_rate"],
73 text=sample["text"],
74 add_special_tokens=True,
75 return_tensors="pt",
76 )
77
78 inputs["input_features"] = inputs["input_features"].to(dtype=model.dtype)
79 inputs["decoder_input_ids"] = inputs["labels"]
80 del inputs["labels"]
81
82 return inputs
83
84ds = ds.map(process, remove_columns=ds.column_names)
85
86# Define a oneshot data collator for multimodal inputs.
87def data_collator(batch):
88 assert len(batch) == 1
89 return {key: torch.tensor(value) for key, value in batch[0].items()}
90
91ignore=["lm_head"]
92
93# Recipe
94recipe = GPTQModifier(
95 targets="Linear",
96 config_groups={
97 "config_group": QuantizationScheme(
98 targets=["Linear"],
99 weights=QuantizationArgs(
100 num_bits=4,
101 type=QuantizationType.INT,
102 strategy=QuantizationStrategy.GROUP,
103 group_size=group_size,
104 symmetric=True,
105 dynamic=False,
106 actorder=getattr(ActivationOrdering, actorder_arg.upper()),
107 ),
108 ),
109 },
110 sequential_targets=["WhisperEncoderLayer", "WhisperDecoderLayer"],
111 ignore=["re:.*lm_head"],
112 update_size=NUM_CALIBRATION_SAMPLES,
113 dampening_frac=dampening_frac
114)
115
116# Apply algorithms.
117oneshot(
118 model=model,
119 dataset=ds,
120 recipe=recipe,
121 max_seq_length=MAX_SEQUENCE_LENGTH,
122 num_calibration_samples=NUM_CALIBRATION_SAMPLES,
123 data_collator=data_collator,
124)
125
126
127# Save to disk compressed.
128save_name = f"{model_id.split('/')[-1]}-quantized.w4a16"
129save_path = os.path.join(args.save_dir, save_name)
130print("Saving model:", save_path)
131model.save_pretrained(save_path, save_compressed=True)
132processor.save_pretrained(save_path)lmms-eval \
--model=whisper_vllm \
--model_args="pretrained=neuralmagic-ent/whisper-medium-quantized.w4a16" \
--batch_size 64 \
--output_path <output_file_path> \
--tasks librispeechlmms-eval \
--model=whisper_vllm \
--model_args="pretrained=neuralmagic-ent/whisper-medium-quantized.w4a16" \
--batch_size 64 \
--output_path <output_file_path> \
--tasks fleurs| Benchmark | Split | BF16 | W4A16 | Recovery (%) |
|---|---|---|---|---|
| LibriSpeech (WER) | test-clean | 2.8269 | 2.9419 | 96.10% |
| test-other | 6.4445 | 6.5561 | 98.30% | |
| Fleurs (X→en, WER) | cmn_hans_cn | 13.3371 | 18.4749 | 72.2% |
| en | 4.6004 | 4.7926 | 95.99% | |
| yue_hant_hk | 9.8107 | 13.4349 | 73.02% |