Repository: huggingface/transformers
Pull request: Break weight tying when quantizing input embedding

Summary:
Currently when we try to quantize input_embedding for some models, the output embedding (lm_head) will also be quantized the same way, since they are tied, and this may not be what we want. To break the tie, we added the option to allow people to
1. load unquantized weight
2. tie weights
3. quantize

so that the tie will be broken

Test Plan:
```
from transformers import (
  AutoModelForCausalLM,
  AutoProcessor,
  AutoTokenizer,
  TorchAoConfig,
)
from torchao.quantization.quant_api import (
    IntxWeightOnlyConfig,
    Int8DynamicActivationIntxWeightConfig,
    AOPerModuleConfig
)
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,
)
quant_config = AOPerModuleConfig({"_default": linear_config, "model.embed_tokens": embedding_config})
quantization_config = TorchAoConfig(quant_type=quant_config, include_embedding=True, untie_embedding_weights=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)

print(quantized_model)
print("embed_tokens.weight:", quantized_model.model.embed_tokens.weight)
print("lm head weight:", quantized_model.lm_head.weight)
from transformers.modeling_utils import find_tied_parameters
print(find_tied_parameters(quantized_model))

# 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):])

```

output

```
Phi3ForCausalLM(
  (model): Phi3Model(
    (embed_tokens): Embedding(200064, 3072, padding_idx=199999)
    (layers): ModuleList(
      (0-31): 32 x Phi3DecoderLayer(
        (self_attn): Phi3Attention(
          (o_proj): Linear(in_features=3072, out_features=3072, weight=LinearActivationQuantizedTensor(activation=<function _int8_asymm_per_token_quant at 0x7f3973d8f250>, weight=AffineQuantizedTensor(shape=torch.Size([3072, 3072]), block_size=(1, 32), device=cuda:0, _layout=QDQLayout(), tensor_impl_dtype=torch.int8, quant_min=-8, quant_max=7)))
          (qkv_proj): Linear(in_features=3072, out_features=5120, weight=LinearActivationQuantizedTensor(activation=<function _int8_asymm_per_token_quant at 0x7f3973d8f250>, weight=AffineQuantizedTensor(shape=torch.Size([5120, 3072]), block_size=(1, 32), device=cuda:0, _layout=QDQLayout(), tensor_impl_dtype=torch.int8, quant_min=-8, quant_max=7)))
        )
        (mlp): Phi3MLP(
          (gate_up_proj): Linear(in_features=3072, out_features=16384, weight=LinearActivationQuantizedTensor(activation=<function _int8_asymm_per_token_quant at 0x7f3973d8f250>, weight=AffineQuantizedTensor(shape=torch.Size([16384, 3072]), block_size=(1, 32), device=cuda:0, _layout=QDQLayout(), tensor_impl_dtype=torch.int8, quant_min=-8, quant_max=7)))
          (down_proj): Linear(in_features=8192, out_features=3072, weight=LinearActivationQuantizedTensor(activation=<function _int8_asymm_per_token_quant at 0x7f3973d8f250>, weight=AffineQuantizedTensor(shape=torch.Size([3072, 8192]), block_size=(1, 32), device=cuda:0, _layout=QDQLayout(), tensor_impl_dtype=torch.int8, quant_min=-8, quant_max=7)))
          (activation_fn): SiLU()
        )
        (input_layernorm): Phi3RMSNorm((3072,), eps=1e-05)
        (post_attention_layernorm): Phi3RMSNorm((3072,), eps=1e-05)
        (resid_attn_dropout): Dropout(p=0.0, inplace=False)
        (resid_mlp_dropout): Dropout(p=0.0, inplace=False)
      )
    )
    (norm): Phi3RMSNorm((3072,), eps=1e-05)
    (rotary_emb): Phi3RotaryEmbedding()
  )
  (lm_head): Linear(in_features=3072, out_features=200064, bias=False)
)
embed_tokens.weight: AffineQuantizedTensor(tensor_impl=QDQTensorImpl(data=tensor([[-20,   4,  13,  ...,   8,  -5,  -3],
        [ -2,   1,  13,  ...,   0, -18,  15],
        [  1,   2,  11,  ...,  15,   0,  18],
        ...,
        [  0,  -2,   7,  ...,   4,  10,  12],
        [  0,  -2,   7,  ...,   4,  10,  12],
        [  0,  -2,   7,  ...,   4,  10,  12]], device='cuda:0',
       dtype=torch.int8)... , scale=tensor([0.0083, 0.0099, 0.0115,  ..., 0.0009, 0.0009, 0.0009], device='cuda:0')... , zero_point=tensor([0, 0, 0,  ..., 0, 0, 0], device='cuda:0', dtype=torch.int8)... , _layout=QDQLayout()), block_size=(1, 3072), shape=torch.Size([200064, 3072]), device=cuda:0, dtype=torch.float32, requires_grad=False)
lm head weight: Parameter containing:
tensor([[-0.1689,  0.0317,  0.1060,  ...,  0.0635, -0.0378, -0.0260],
        [-0.0233,  0.0072,  0.1299,  ...,  0.0013, -0.1748,  0.1465],
        [ 0.0159,  0.0206,  0.1260,  ...,  0.1748, -0.0027,  0.2041],
        ...,
        [ 0.0002, -0.0020,  0.0062,  ...,  0.0038,  0.0095,  0.0113],
        [ 0.0002, -0.0020,  0.0062,  ...,  0.0038,  0.0095,  0.0113],
        [ 0.0002, -0.0020,  0.0062,  ...,  0.0038,  0.0095,  0.0113]],
       device='cuda:0')
[]
Prompt: Hey, are you conscious? Can you talk to me?
Templated prompt: 


Hey, are you conscious? Can you talk to me?


Response: Hello! As an AI, I don't have consciousness in the way humans do, but I'm fully operational and here to assist you. How can I help you today?
```
Reviewers:

Subscribers:

Tasks:

Tags:

Changed files:
filename: src/transformers/quantizers/quantizer_torchao.py
patch: @@ -247,6 +247,16 @@ def create_quantized_param(
             module._parameters[tensor_name] = torch.nn.Parameter(
                 param_value, requires_grad=param_value.requires_grad
             ).to(device=target_device)
+            # if we are quantizing tied parameters, to avoid tying the quantized weights
+            # the correct order to do it is
+            # 1. load the weight to model
+            # 2. run tie_weights to populate the weights
+            # 3. quantize
+            input_embed = model.get_input_embeddings()
+            if self.quantization_config.untie_embedding_weights and id(module) == id(input_embed):
+                model.tie_weights()
+                setattr(model.config.get_text_config(decoder=True), "tie_word_embeddings", False)
+
             # handle AOPerModuleConfig, introduced in torchao 0.11.0+
             if self.quantization_config._get_ao_version() > version.Version("0.10.0"):
                 from torchao.quantization import AOPerModuleConfig

filename: src/transformers/utils/quantization_config.py
patch: @@ -1555,6 +1555,7 @@ class TorchAoConfig(QuantizationConfigMixin):
     modules_to_not_convert: Optional[List]
     quant_type_kwargs: Dict[str, Any]
     include_embedding: bool
+    untie_embedding_weights: bool
 
     """This is a config class for torchao quantization/sparsity techniques.
 
@@ -1569,6 +1570,9 @@ class TorchAoConfig(QuantizationConfigMixin):
         inlcude_embedding (`bool`, default to `False`):
             Whether to include embedding in quantization or not, input embedding will be removed from
             the module_not_to_convert list as well if this flag is set.
+        untie_embedding_weights (`bool`, default to `False`):
+            Whether to untie the weights when we are quantizing input embedding weights that is tied
+            to other weights.
         kwargs (`Dict[str, Any]`, *optional*):
             The keyword arguments for the chosen type of quantization, for example, int4_weight_only quantization supports two keyword arguments
             `group_size` and `inner_k_tiles` currently. More API examples and documentation of arguments can be found in
@@ -1614,13 +1618,15 @@ def __init__(
         quant_type: Union[str, "AOBaseConfig"],  # noqa: F821
         modules_to_not_convert: Optional[List] = None,
         include_embedding: bool = False,
+        untie_embedding_weights: bool = False,
         **kwargs,
     ):
         self.quant_method = QuantizationMethod.TORCHAO
         self.quant_type = quant_type
         self.modules_to_not_convert = modules_to_not_convert
         self.quant_type_kwargs = kwargs.get("quant_type_kwargs", kwargs)
         self.include_embedding = include_embedding
+        self.untie_embedding_weights = untie_embedding_weights
         self.post_init()
 
     @staticmethod