| |
| """Quantize a Whisper model with llmcompressor. |
| |
| This script loads a Whisper checkpoint, prepares English transcription |
| calibration samples from MLCommons People's Speech, applies either a W8A8 or |
| W4A8 GPTQ quantization recipe to linear layers while leaving the language model |
| head unquantized, and saves the compressed model plus processor files under the |
| requested output directory. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import shutil |
| from pathlib import Path |
| from typing import Any, Sequence |
|
|
| import torch |
| from compressed_tensors.quantization import ( |
| QuantizationArgs, |
| QuantizationScheme, |
| QuantizationStrategy, |
| QuantizationType, |
| ) |
| from datasets import load_dataset |
| from llmcompressor import oneshot |
| from llmcompressor.modifiers.quantization import GPTQModifier |
| from transformers import WhisperForConditionalGeneration, WhisperProcessor |
| from transformers.utils import cached_file |
|
|
|
|
| DATASET_ID = "MLCommons/peoples_speech" |
| DATASET_SUBSET = "test" |
| DATASET_SPLIT = "test" |
| MAX_SEQUENCE_LENGTH = 2048 |
| IGNORE_TARGETS = ["proj_out"] |
| MODEL_INPUT_KEYS = {"input_features", "decoder_input_ids"} |
| PROCESSOR_ARTIFACT_FILES = [ |
| "added_tokens.json", |
| "merges.txt", |
| "normalizer.json", |
| "preprocessor_config.json", |
| "special_tokens_map.json", |
| "vocab.json", |
| ] |
|
|
|
|
| def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace: |
| parser = argparse.ArgumentParser( |
| description="Quantize a Whisper model using W4A8 GPTQ.", |
| formatter_class=argparse.ArgumentDefaultsHelpFormatter, |
| ) |
| parser.add_argument( |
| "--model_path", |
| required=True, |
| help="Hugging Face model ID or local model directory to quantize.", |
| ) |
| parser.add_argument( |
| "--num_calibration_samples", |
| type=int, |
| default=1024, |
| help="Number of samples to load from the calibration dataset.", |
| ) |
| parser.add_argument( |
| "--save_dir", |
| type=Path, |
| required=True, |
| help="Parent directory where the generated model is saved.", |
| ) |
| parser.add_argument( |
| "--dampening_frac", |
| type=float, |
| default=0.01, |
| help="GPTQ dampening fraction.", |
| ) |
| return parser.parse_args(argv) |
|
|
|
|
| def load_model_and_processor(model_id: str): |
| model = WhisperForConditionalGeneration.from_pretrained( |
| model_id, |
| device_map="auto", |
| dtype="auto", |
| ) |
| model = model.to(torch.float32) |
| model.config.forced_decoder_ids = None |
|
|
| processor = WhisperProcessor.from_pretrained(model_id) |
| processor.tokenizer.set_prefix_tokens(language="en", task="transcribe") |
| return model, processor |
|
|
| def resolve_source_processor_artifacts(model_id: str) -> dict[str, Path]: |
| sources = {} |
| for filename in PROCESSOR_ARTIFACT_FILES: |
| source = cached_file(model_id, filename) |
| if source is None: |
| raise FileNotFoundError( |
| f"Could not find required processor artifact {filename!r} for {model_id!r}" |
| ) |
| sources[filename] = Path(source) |
| return sources |
|
|
| def copy_source_processor_artifacts( |
| source_artifacts: dict[str, Path], |
| save_path: Path, |
| ) -> None: |
| for filename, source in source_artifacts.items(): |
| destination = save_path / filename |
| if source.resolve() != destination.resolve(): |
| shutil.copy2(source, destination) |
|
|
|
|
| def save_processor_artifacts( |
| processor: Any, |
| source_artifacts: dict[str, Path], |
| save_path: Path, |
| ) -> None: |
| processor.save_pretrained(save_path) |
| processor.feature_extractor.save_pretrained(save_path) |
| processor.tokenizer.save_pretrained(save_path) |
| copy_source_processor_artifacts(source_artifacts, save_path) |
|
|
|
|
|
|
| def preprocess_example(example: dict[str, Any]) -> dict[str, Any]: |
| return { |
| "array": example["audio"]["array"], |
| "sampling_rate": example["audio"]["sampling_rate"], |
| "text": " " + example["text"].capitalize(), |
| } |
|
|
|
|
| def process_sample( |
| sample: dict[str, Any], |
| processor: Any, |
| model_dtype: Any, |
| ) -> dict[str, Any]: |
| inputs = processor( |
| audio=sample["array"], |
| sampling_rate=sample["sampling_rate"], |
| text=sample["text"], |
| add_special_tokens=True, |
| return_tensors="pt", |
| ) |
|
|
| inputs["input_features"] = inputs["input_features"].to(dtype=model_dtype) |
| inputs["decoder_input_ids"] = inputs.pop("labels") |
| return inputs |
|
|
|
|
| def load_calibration_dataset(num_samples: int, processor: Any, model_dtype: Any): |
| dataset = load_dataset( |
| DATASET_ID, |
| DATASET_SUBSET, |
| split=f"{DATASET_SPLIT}[:{num_samples}]", |
| trust_remote_code=True, |
| ) |
| dataset = dataset.map(preprocess_example, remove_columns=dataset.column_names) |
| dataset = dataset.map( |
| lambda sample: process_sample(sample, processor, model_dtype), |
| remove_columns=dataset.column_names, |
| ) |
| return dataset.add_column("input_ids", [[0]] * len(dataset)) |
|
|
|
|
| def data_collator(batch: list[dict[str, Any]]) -> dict[str, Any]: |
| assert len(batch) == 1 |
| return { |
| key: torch.tensor(value) |
| for key, value in batch[0].items() |
| if key in MODEL_INPUT_KEYS |
| } |
|
|
|
|
| def build_recipe(args: argparse.Namespace): |
|
|
| weights_args = QuantizationArgs( |
| num_bits=4, |
| type=QuantizationType.INT, |
| strategy=QuantizationStrategy.CHANNEL, |
| symmetric=True, |
| dynamic=False, |
| observer="memoryless_minmax", |
| ) |
| input_acts = QuantizationArgs( |
| num_bits=8, |
| type=QuantizationType.INT, |
| strategy=QuantizationStrategy.TOKEN, |
| symmetric=True, |
| dynamic=True, |
| observer=None, |
| ) |
| scheme = QuantizationScheme( |
| targets=["Linear"], |
| weights=weights_args, |
| input_activations=input_acts, |
| ) |
| return [ |
| GPTQModifier( |
| config_groups={"group_0": scheme}, |
| targets=["Linear"], |
| actorder="static", |
| dampening_frac=args.dampening_frac, |
| ignore=IGNORE_TARGETS, |
| ) |
| ] |
|
|
|
|
| def save_path_for(model_id: str, save_dir: Path) -> Path: |
| model_name = Path(model_id.rstrip("/")).name |
| return save_dir / f"{model_name}-quantized.w4a8" |
|
|
| def main(argv: Sequence[str] | None = None) -> None: |
| args = parse_args(argv) |
| model, processor = load_model_and_processor(args.model_path) |
| processor_artifacts = resolve_source_processor_artifacts(args.model_path) |
| dataset = load_calibration_dataset( |
| args.num_calibration_samples, |
| processor, |
| model.dtype, |
| ) |
| recipe = build_recipe(args) |
|
|
| save_path = save_path_for(args.model_path, args.save_dir) |
| print(f"Saving model: {save_path}") |
| oneshot( |
| model=model, |
| processor=processor, |
| dataset=dataset, |
| recipe=recipe, |
| max_seq_length=MAX_SEQUENCE_LENGTH, |
| num_calibration_samples=args.num_calibration_samples, |
| data_collator=data_collator, |
| pipeline="basic", |
| output_dir=save_path, |
| ) |
|
|
| save_processor_artifacts(processor, processor_artifacts, save_path) |
|
|
| if __name__ == "__main__": |
| main() |
|
|
