Update README.md (#4)

- Update README.md (628b89a76d3199661bc5cf3d4ae19ed436d37c01)


Co-authored-by: Scott Roy <[email protected]>

README.md

@@ -30,24 +30,33 @@ We used following code to get the quantized model:
 
 ```
 from transformers import (
-    AutoModelForCausalLM,
-    AutoProcessor,
-    AutoTokenizer,
-    TorchAoConfig,
+  AutoModelForCausalLM,
+  AutoProcessor,
+  AutoTokenizer,
+  TorchAoConfig,
 )
 from torchao.quantization.quant_api import (
+    IntxWeightOnlyConfig,
     Int8DynamicActivationIntxWeightConfig,
+    AOPerModuleConfig
 )
-from torchao.quantization.granularity import PerGroup
+from torchao.quantization.granularity import PerGroup, PerAxis
 import torch
 
 model_id = "microsoft/Phi-4-mini-instruct"
+
+embedding_config = IntxWeightOnlyConfig(
+    weight_dtype=torch.int8,
+    granularity=PerAxis(0),
+)
 linear_config = Int8DynamicActivationIntxWeightConfig(
     weight_dtype=torch.int4,
     weight_granularity=PerGroup(32),
+    weight_scale_dtype=torch.bfloat16,
 )
-quantization_config = TorchAoConfig(quant_type=linear_config)
-quantized_model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype="auto", device_map="auto", quantization_config=quantization_config)
+quant_config = AOPerModuleConfig({"_default": linear_config, "model.embed_tokens": embedding_config})
+quantization_config = TorchAoConfig(quant_type=quant_config, include_embedding=True)
+quantized_model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.float32, device_map="auto", quantization_config=quantization_config)
 tokenizer = AutoTokenizer.from_pretrained(model_id)
 
 # Push to hub
@@ -91,7 +100,7 @@ torch.save(state_dict, "phi4-mini-8dq4w.pt")
 The response from the manual testing is:
 
 ```
-Hello! As an AI, I don't have consciousness in the way humans do, but I'm here and ready to assist you. How can I help you today?
+Hello! As an AI, I don't have consciousness in the way humans do, but I am fully operational and here to assist you. How can I help you today?
 ```
 
 # Model Quality
@@ -100,7 +109,7 @@ We rely on [lm-evaluation-harness](https://github.com/EleutherAI/lm-evaluation-h
 
 ## baseline
 ```
-lm_eval --model hf --model_args pretrained=microsoft/Phi-4-mini-instruct --tasks hellaswag --device cuda:0 --batch_size 8
+lm_eval --model hf --model_args pretrained=microsoft/Phi-4-mini-instruct --tasks hellaswag --device cuda:0 --batch_size 64
 ```
 
 ## 8dq4w
@@ -111,7 +120,7 @@ from lm_eval.utils import (
     make_table,
 )
 
-lm_eval_model = lm_eval.models.huggingface.HFLM(pretrained=quantized_model, batch_size=8)
+lm_eval_model = lm_eval.models.huggingface.HFLM(pretrained=quantized_model, batch_size=64)
 results = evaluator.simple_evaluate(
     lm_eval_model, tasks=["hellaswag"], device="cuda:0", batch_size="auto"
 )
@@ -123,7 +132,7 @@ print(make_table(results))
 |                                  | Phi-4 mini-Ins | phi4-mini-8dq4w | 
 | **Popular aggregated benchmark** |             |                   |
 | **Reasoning**                    |             |                   |
-| HellaSwag                        | 54.57        | 53.19            |
+| HellaSwag                        | 54.57        | 53.24            |
 | **Multilingual**                 |             |                   |
 | **Math**                         |             |                   |
 | **Overall**                      | **TODO**    | **TODO**          |