#!/usr/bin/env python3 """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()