Update README.md

README.md

@@ -48,9 +48,17 @@ intxwo = IntxWeightOnlyConfig(weight_dtype=torch.int8, granularity=PerAxis(0))
 qconfig_dict = {
     # highest priority
     "model.decoder.layers.3.self_attn.q_proj": int4wo,
+    "model.decoder.layers.3.self_attn.k_proj": int4wo,
+    "model.decoder.layers.3.self_attn.v_proj": int4wo,
+    # vllm
+    "model.decoder.layers.3.self_attn.qkv_proj": int4wo,
+
     "model.decoder.layers.*.self_attn.q_proj": float8dyn,
     "model.decoder.layers.*.self_attn.k_proj": float8dyn,
-    "model.decoder.layers.*.self_attn.v_proj": None,
+    "model.decoder.layers.*.self_attn.v_proj": float8dyn,
+    # vllm
+    "model.decoder.layers.*.self_attn.qkv_proj": float8dyn,
+
     "_default": intxwo,
 }
 quant_config = ModuleFqnToConfig(qconfig_dict)
@@ -65,19 +73,23 @@ print("quantized model:", quantized_model)
 tokenizer = AutoTokenizer.from_pretrained(model_id)
 for i in range(12):
     if i == 3:
-        print("type:", quantized_model.model.decoder.layers[i].self_attn.q_proj.weight)
         assert isinstance(quantized_model.model.decoder.layers[i].self_attn.q_proj.weight, Int4TilePackedTo4dTensor)
+        assert isinstance(quantized_model.model.decoder.layers[i].self_attn.k_proj.weight, Int4TilePackedTo4dTensor)
+        assert isinstance(quantized_model.model.decoder.layers[i].self_attn.v_proj.weight, Int4TilePackedTo4dTensor)
     else:
         assert isinstance(quantized_model.model.decoder.layers[i].self_attn.q_proj.weight, Float8Tensor)
-    assert isinstance(quantized_model.model.decoder.layers[i].self_attn.k_proj.weight, Float8Tensor)
-    assert not isinstance(quantized_model.model.decoder.layers[i].self_attn.v_proj.weight, Float8Tensor)
+        assert isinstance(quantized_model.model.decoder.layers[i].self_attn.k_proj.weight, Float8Tensor)
+        assert isinstance(quantized_model.model.decoder.layers[i].self_attn.v_proj.weight, Float8Tensor)
     assert isinstance(quantized_model.model.decoder.layers[i].self_attn.out_proj.weight, IntxUnpackedToInt8Tensor)
 
-# Push to hub
+# # # Push to hub
 MODEL_NAME = model_id.split("/")[-1]
 save_to = f"torchao-testing/{MODEL_NAME}-ModuleFqnToConfig-v1-regex-0.14.0.dev"
 quantized_model.push_to_hub(save_to, safe_serialization=False)
 tokenizer.push_to_hub(save_to)
+# quantized_model.save_pretrained(save_to, safe_serialization=False)
+# tokenizer.save_pretrained(save_to)
+
 
 # Manual Testing
 prompt = "What are we having for dinner?"
@@ -143,10 +155,12 @@ print("quantized model:", quantized_model)
 for i in range(12):
     if i == 3:
         assert isinstance(quantized_model.model.decoder.layers[i].self_attn.q_proj.weight, Int4TilePackedTo4dTensor)
+        assert isinstance(quantized_model.model.decoder.layers[i].self_attn.k_proj.weight, Int4TilePackedTo4dTensor)
+        assert isinstance(quantized_model.model.decoder.layers[i].self_attn.v_proj.weight, Int4TilePackedTo4dTensor)
     else:
         assert isinstance(quantized_model.model.decoder.layers[i].self_attn.q_proj.weight, Float8Tensor)
-    assert isinstance(quantized_model.model.decoder.layers[i].self_attn.k_proj.weight, Float8Tensor)
-    assert not isinstance(quantized_model.model.decoder.layers[i].self_attn.v_proj.weight, Float8Tensor)
+        assert isinstance(quantized_model.model.decoder.layers[i].self_attn.k_proj.weight, Float8Tensor)
+        assert isinstance(quantized_model.model.decoder.layers[i].self_attn.v_proj.weight, Float8Tensor)
     assert isinstance(quantized_model.model.decoder.layers[i].self_attn.out_proj.weight, IntxUnpackedToInt8Tensor)
 
 tokenizer = AutoTokenizer.from_pretrained(model_name)
@@ -154,10 +168,6 @@ tokenizer = AutoTokenizer.from_pretrained(model_name)
 input_ids = tokenizer(input_text, return_tensors="pt").to(device)
 
 output = quantized_model.generate(**input_ids, max_new_tokens=max_new_tokens)
-EXPECTED_OUTPUT = [
-    "What are we having for dinner?\n\nJessica: (smiling)",
-    "What are we having for dinner?\n\nJess: (smiling) I",
-]
-# self.assertTrue(tokenizer.decode(output[0], skip_special_tokens=True) in EXPECTED_OUTPUT)
+print(tokenizer.decode(output[0], skip_special_tokens=True))
 
 ```