# quark_gemma4_int8.pyimport os, torchfrom transformers import AutoTokenizer, Gemma4ForConditionalGenerationfrom quark.torch import ModelQuantizerfrom quark.torch.quantization.config.config import ( QTensorConfig, QuantizationConfig, Config, Dtype,)from quark.torch.quantization.config.type import ( RoundType, ScaleType, QSchemeType,)from quark.torch.quantization.observer import PerChannelMinMaxObserver MODEL_IN = "google/gemma-4-31B-it"MODEL_OUT = "./Gemma-4-31B-it-Quark-W8A8-INT8" tokenizer = AutoTokenizer.from_pretrained(MODEL_IN, trust_remote_code=True)model = Gemma4ForConditionalGeneration.from_pretrained( MODEL_IN, torch_dtype=torch.bfloat16, device_map="auto", trust_remote_code=True,) weight_spec = QTensorConfig( dtype=Dtype.int8, observer_cls=PerChannelMinMaxObserver, symmetric=True, is_dynamic=False, qscheme=QSchemeType.per_channel, ch_axis=0, round_method=RoundType.round, scale_type=ScaleType.float,)input_spec = QTensorConfig( dtype=Dtype.int8, observer_cls=PerChannelMinMaxObserver, symmetric=True, is_dynamic=True, qscheme=QSchemeType.per_channel, ch_axis=1, round_method=RoundType.round, scale_type=ScaleType.float,) q_cfg = Config( global_quant_config=QuantizationConfig( input_tensors=input_spec, weight=weight_spec, ), exclude=[ "lm_head", "*embed_tokens*", "*vision_tower*", "*embed_vision*", ],) quantizer = ModelQuantizer(q_cfg, multi_device=True)model = quantizer.quantize_model(model, dataloader=None) # PTQ, no calibration data needed for dynamic actquantizer.freeze(model) quantizer.export_model( model, MODEL_OUT, pack_method="order", weight_format="real_quantized", custom_mode="quark",)tokenizer.save_pretrained(MODEL_OUT)