transformers - 💡(How to fix) Fix Model quantized via sinq broken after save_pretrained and from_pretrained [3 comments, 3 participants]

from transformers import AutoProcessor, AutoModelForCausalLM, SinqConfig

def quantize_model(model_id: str, save_dst: str):
    # Quantize and save resulting model
    quant_cfg = SinqConfig(
        nbits=8,
        group_size=64,
        tiling_mode='2D',
        method='sinq',
        modules_to_not_convert=["lm_head", "model.audio_tower"],
    )

    model = AutoModelForCausalLM.from_pretrained(
        model_id,
        device_map = 'cpu',
        quantization_config=quant_cfg,
    )

    model.save_pretrained(save_dst)



def test_quantized(model_id: str, save_dst: str):
    # Loading quantized model and appropriate processor 
    processor = AutoProcessor.from_pretrained(model_id)
    model = AutoModelForCausalLM.from_pretrained(save_dst, device_map='cpu')

    # Testing quantized model
    chat = []
    chat.append({ 'role': 'user', 'content': 'this is a test.' })

    model_input = processor.apply_chat_template(chat, add_generation_prompt=True, tokenize=True, return_dict=True, return_tensors="pt").to(model.device)
    skip_length = model_input['input_ids'].shape[-1]

    new_tokens = model.generate(**model_input, max_new_tokens=2048)[0, skip_length:]
    
    message = processor.parse_response(processor.decode(new_tokens))

    print(message)


model_id = 'google/gemma-4-E4B-it'
save_dst = './gemma-4-E4B-it-sinq/'

quantize_model(model_id, save_dst)
test_quantized(model_id, save_dst)