Add quant recipe

README.md

@@ -40,6 +40,141 @@ optimum-cli export executorch \
   --output_dir ./smollm3_3b
 ```
 
+# Quantization Recipe
+
+First need to install the required packages:
+```Shell
+pip install git+https://github.com/huggingface/transformers@main
+pip install torchao
+```
+
+## Untie Embedding Weights
+We want to quantize the embedding and lm_head differently.  Since those layers are tied, we first need to untie the model:
+
+```Py
+from transformers import (
+  AutoModelForCausalLM,
+  AutoProcessor,
+  AutoTokenizer,
+)
+import torch
+
+model_id = "HuggingFaceTB/SmolLM3-3B"
+untied_model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype="auto", device_map="auto")
+tokenizer = AutoTokenizer.from_pretrained(model_id)
+
+print(untied_model)
+from transformers.modeling_utils import find_tied_parameters
+print("tied weights:", find_tied_parameters(untied_model))
+if getattr(untied_model.config.get_text_config(decoder=True), "tie_word_embeddings"):
+    setattr(untied_model.config.get_text_config(decoder=True), "tie_word_embeddings", False)
+
+untied_model._tied_weights_keys = []
+untied_model.lm_head.weight = torch.nn.Parameter(untied_model.lm_head.weight.clone())
+
+print("tied weights:", find_tied_parameters(untied_model))
+
+USER_ID = "YOUR_USER_ID"
+MODEL_NAME = model_id.split("/")[-1]
+save_to = f"{USER_ID}/{MODEL_NAME}-untied-weights"
+
+untied_model.push_to_hub(save_to)
+tokenizer.push_to_hub(save_to)
+
+# or save locally
+save_to_local_path = f"{MODEL_NAME}-untied-weights"
+untied_model.save_pretrained(save_to_local_path)
+tokenizer.save_pretrained(save_to)
+```
+
+Note: to `push_to_hub` you need to run
+```Shell
+pip install -U "huggingface_hub[cli]"
+huggingface-cli login
+```
+and use a token with write access, from https://huggingface.co/settings/tokens
+
+## Quantization
+
+We used following code to get the quantized model:
+
+```Py
+from transformers import (
+  AutoModelForCausalLM,
+  AutoProcessor,
+  AutoTokenizer,
+  TorchAoConfig,
+)
+from torchao.quantization.quant_api import (
+    IntxWeightOnlyConfig,
+    Int8DynamicActivationIntxWeightConfig,
+    ModuleFqnToConfig,
+    quantize_,
+)
+from torchao.quantization.granularity import PerGroup, PerAxis
+import torch
+
+# we start from the model with untied weights
+model_id = "HuggingFaceTB/SmolLM3-3B"
+USER_ID = "YOUR_USER_ID"
+MODEL_NAME = model_id.split("/")[-1]
+untied_model_id = f"{USER_ID}/{MODEL_NAME}-untied-weights"
+untied_model_local_path = f"{MODEL_NAME}-untied-weights"
+
+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,
+)
+quant_config = ModuleFqnToConfig({"_default": linear_config, "model.embed_tokens": embedding_config})
+quantization_config = TorchAoConfig(quant_type=quant_config, include_embedding=True, untie_embedding_weights=True, modules_to_not_convert=[])
+
+# either use `untied_model_id` or `untied_model_local_path`
+quantized_model = AutoModelForCausalLM.from_pretrained(untied_model_id, torch_dtype=torch.float32, device_map="auto", quantization_config=quantization_config)
+tokenizer = AutoTokenizer.from_pretrained(model_id)
+
+# Push to hub
+MODEL_NAME = model_id.split("/")[-1]
+save_to = f"{USER_ID}/{MODEL_NAME}-8da4w"
+quantized_model.push_to_hub(save_to, safe_serialization=False)
+tokenizer.push_to_hub(save_to)
+
+# Manual testing
+prompt = "Hey, are you conscious? Can you talk to me?"
+messages = [
+    {
+        "role": "system",
+        "content": "",
+    },
+    {"role": "user", "content": prompt},
+]
+templated_prompt = tokenizer.apply_chat_template(
+    messages,
+    tokenize=False,
+    add_generation_prompt=True,
+)
+print("Prompt:", prompt)
+print("Templated prompt:", templated_prompt)
+inputs = tokenizer(
+    templated_prompt,
+    return_tensors="pt",
+).to("cuda")
+generated_ids = quantized_model.generate(**inputs, max_new_tokens=128)
+output_text = tokenizer.batch_decode(
+    generated_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False
+)
+print("Response:", output_text[0][len(prompt):])
+```
+
+The response from the manual testing is:
+
+```
+Okay, the user is asking if I can talk to them. First, I need to clarify that I can't communicate like a human because I don't have consciousness or emotions. I'm an AI model created by Hugging Face. 
+```
 
 # Disclaimer
 PyTorch has not performed safety evaluations or red teamed the quantized models. Performance characteristics, outputs, and behaviors may differ from the original models. Users are solely responsible for selecting appropriate use cases, evaluating and mitigating for accuracy, safety, and fairness, ensuring security, and complying with all applicable laws and regulations.