import argparseimport torchimport osfrom datasets import load_datasetfrom transformers import WhisperProcessorfrom llmcompressor import oneshotfrom llmcompressor.modifiers.quantization import QuantizationModifierfrom llmcompressor.transformers.tracing import TraceableWhisperForConditionalGenerationfrom compressed_tensors.quantization import QuantizationType
# --- Args ---parser = argparse.ArgumentParser()parser.add_argument('--model_path', type=str, required=True)parser.add_argument('--quant_path', type=str, required=True)parser.add_argument('--observer', type=str, default="minmax")args = parser.parse_args()
# --- Load Model ---model = TraceableWhisperForConditionalGeneration.from_pretrained( args.model_path, device_map="auto", torch_dtype="auto",)model.config.forced_decoder_ids = Noneprocessor = WhisperProcessor.from_pretrained(args.model_path)
# --- Recipe (FP8 Dynamic) ---recipe = [ QuantizationModifier( targets="Linear", scheme="FP8_DYNAMIC", sequential_targets=["WhisperEncoderLayer", "WhisperDecoderLayer"], ignore=["re:.*lm_head"], )]
# --- Run oneshot ---oneshot( model=model, recipe=recipe, trust_remote_code_model=True,)
# --- Save ---os.makedirs(args.quant_path, exist_ok=True)model.save_pretrained(args.quant_path, save_compressed=True)processor.save_pretrained(args.quant_path)