Update README.md

README.md

@@ -137,9 +137,7 @@ from transformers import AutoModelForCausalLM, AutoTokenizer, TorchAoConfig
 
 model_id = "google/gemma-3-12b-it"
 model_to_quantize = "google/gemma-3-12b-it"
-
-
-from torchao.quantization import Int4WeightOnlyConfig, quantize_
+from torchao.quantization import Int4WeightOnlyConfig, quantize_, ModuleFqnToConfig
 from torchao.prototype.awq import (
     AWQConfig,
 )
@@ -150,14 +148,24 @@ model = AutoModelForCausalLM.from_pretrained(
     torch_dtype=torch.bfloat16,
 )
 tokenizer = AutoTokenizer.from_pretrained(model_id)
-
+def get_quant_config(linear_config):
+    return ModuleFqnToConfig({
+        r"re:language_model\.model\.layers\..+\.mlp\..+_proj": linear_config,
+        r"re:language_model\.model\.layers\..+\.self_attn\..+_proj": linear_config,
+        r"re:model\.language_model\.layers\..+\.mlp\..+_proj": linear_config,
+        r"re:model\.language_model\.layers\..+\.self_attn\..+_proj": linear_config,
+    })
 # AWQ only works for H100 INT4 so far
 base_config = Int4WeightOnlyConfig(group_size=128)
-quant_config = AWQConfig(base_config, step="prepare")
+linear_config = AWQConfig(base_config, step="prepare")
+quant_config = get_quant_config(linear_config)
 quantize_(
     model,
     quant_config,
 )
+tasks = ["mmlu_philosophy"]
+calibration_limit=30
+max_seq_length=2048
 TransformerEvalWrapper(
     model=model,
     tokenizer=tokenizer,
@@ -166,22 +174,25 @@ TransformerEvalWrapper(
     tasks=tasks,
     limit=calibration_limit,
 )
-quant_config = AWQConfig(base_config, step="convert")
+linear_config = AWQConfig(base_config, step="convert")
+quant_config = get_quant_config(linear_config)
 quantize_(model, quant_config)
-
 quantized_model = model
-quant_config = AWQConfig(base_config, step="prepare_for_loading")
+linear_config = AWQConfig(base_config, step="prepare_for_loading")
+quant_config = get_quant_config(linear_config)
 quantized_model.config.quantization_config = TorchAoConfig(quant_config)
-
-
 # Push to hub
 USER_ID = "YOUR_USER_ID"
 MODEL_NAME = model_id.split("/")[-1]
 save_to = f"{USER_ID}/{MODEL_NAME}-AWQ-INT4"
 quantized_model.push_to_hub(save_to, safe_serialization=False)
 tokenizer.push_to_hub(save_to)
-
 # Manual Testing
+quantized_model = AutoModelForCausalLM.from_pretrained(
+    save_to,
+    device_map="cuda:0",
+    torch_dtype=torch.bfloat16,
+)
 prompt = "Hey, are you conscious? Can you talk to me?"
 messages = [
     {