HIGGS Quantization not working properly
Author: Godofnothing
Pull request: no
Labels: name: bug

### System Info

**Environment**
```
- `transformers` version: 4.48.2
- Platform: Linux-5.4.210-39.1.pagevecsize-x86_64-with-glibc2.27
- Python version: 3.11.10
- Huggingface_hub version: 0.26.2
- Safetensors version: 0.4.5
- Accelerate version: 1.1.1
- Accelerate config:    not found
- PyTorch version (GPU?): 2.4.0+cu121 (True)
- Tensorflow version (GPU?): not installed (NA)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using distributed or parallel set-up in script?: <fill in>
- Using GPU in script?: <fill in>
- GPU type: NVIDIA A100-SXM4-80GB
- fast_hadamard_transform           1.0.4.post1
```

### Who can help?

@BlackSamorez 
@SunMarc 
@ArthurZucker 

### Information

- [ ] The official example scripts
- [x] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [ ] My own task or dataset (give details below)

### Reproduction

Recently, in the [PR](https://github.com/huggingface/transformers/pull/34997)  HIGGS quantization from the paper [Pushing the Limits of Large Language Model Quantization via the Linearity Theorem](https://arxiv.org/abs/2411.17525) was introduced.

But when attempting to load the quantized `Llama-3.1-8B-Instruct `model in this format as follows:
```python
model_name = "meta-llama/Llama-3.1-8B-Instruct"
quantization_config = HiggsConfig(bits=4, p=2)
model = AutoModelForCausalLM.from_pretrained(
    model_name, 
    device_map="auto", 
    torch_dtype=torch.bfloat16, 
    low_cpu_mem_usage=True,
    quantization_config=quantization_config
)
model.config.use_cache = False
```
And doing forward pass with dummy inputs
```python
inputs = torch.randint(0, model.config.vocab_size, device="cuda", size=(8,))
with torch.no_grad():
    outputs = model(inputs)
```
I get the following error in the RoPE:
```bash
File ~/miniconda3/envs/llm/lib/python3.11/site-packages/transformers/models/llama/modeling_llama.py:271, in LlamaAttention.forward(self, hidden_states, position_embeddings, attention_mask, past_key_value, cache_position, **kwargs)
    [268](https://vscode-remote+ssh-002dremote-002bultramar.vscode-resource.vscode-cdn.net/home/dkuznedelev/FLUTE_Playground/~/miniconda3/envs/llm/lib/python3.11/site-packages/transformers/models/llama/modeling_llama.py:268) value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)
    [270](https://vscode-remote+ssh-002dremote-002bultramar.vscode-resource.vscode-cdn.net/home/dkuznedelev/FLUTE_Playground/~/miniconda3/envs/llm/lib/python3.11/site-packages/transformers/models/llama/modeling_llama.py:270) cos, sin = position_embeddings
--> [271](https://vscode-remote+ssh-002dremote-002bultramar.vscode-resource.vscode-cdn.net/home/dkuznedelev/FLUTE_Playground/~/miniconda3/envs/llm/lib/python3.11/site-packages/transformers/models/llama/modeling_llama.py:271) query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
    [273](https://vscode-remote+ssh-002dremote-002bultramar.vscode-resource.vscode-cdn.net/home/dkuznedelev/FLUTE_Playground/~/miniconda3/envs/llm/lib/python3.11/site-packages/transformers/models/llama/modeling_llama.py:273) if past_key_value is not None:
    [274](https://vscode-remote+ssh-002dremote-002bultramar.vscode-resource.vscode-cdn.net/home/dkuznedelev/FLUTE_Playground/~/miniconda3/envs/llm/lib/python3.11/site-packages/transformers/models/llama/modeling_llama.py:274)     # sin and cos are specific to RoPE models; cache_position needed for the static cache
    [275](https://vscode-remote+ssh-002dremote-002bultramar.vscode-resource.vscode-cdn.net/home/dkuznedelev/FLUTE_Playground/~/miniconda3/envs/llm/lib/python3.11/site-packages/transformers/models/llama/modeling_llama.py:275)     cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position}

File ~/miniconda3/envs/llm/lib/python3.11/site-packages/transformers/models/llama/modeling_llama.py:169, in apply_rotary_pos_emb(q, k, cos, sin, position_ids, unsqueeze_dim)
    [167](https://vscode-remote+ssh-002dremote-002bultramar.vscode-resource.vscode-cdn.net/home/dkuznedelev/FLUTE_Playground/~/miniconda3/envs/llm/lib/python3.11/site-packages/transformers/models/llama/modeling_llama.py:167) cos = cos.unsqueeze(unsqueeze_dim)
    [168](https://vscode-remote+ssh-002dremote-002bultramar.vscode-resource.vscode-cdn.net/home/dkuznedelev/FLUTE_Playground/~/miniconda3/envs/llm/lib/python3.11/site-packages/transformers/models/llama/modeling_llama.py:168) sin = sin.unsqueeze(unsqueeze_dim)
--> [169](https://vscode-remote+ssh-002dremote-002bultramar.vscode-resource.vscode-cdn.net/home/dkuznedelev/FLUTE_Playground/~/miniconda3/envs/llm/lib/python3.11/site-packages/transformers/models/llama/modeling_llama.py:169) q_embed = (q * cos) + (rotate_half(q) * sin)
    [170](https://vscode-remote+ssh-002dremote-002bultramar.vscode-resource.vscode-cdn.net/home/dkuznedelev/FLUTE_Playground/~/miniconda3/envs/llm/lib/python3.11/site-packages/transformers/models/llama/modeling_llama.py:170) k_embed = (k * cos) + (rotate_half(k) * sin)
    [171](https://vscode-remote+ssh-002dremote-002bultramar.vscode-resource.vscode-cdn.net/home/dkuznedelev/FLUTE_Playground/~/miniconda3/envs/llm/lib/python3.11/site-packages/transformers/models/llama/modeling_llama.py:171) return q_embed, k_embed

RuntimeError: The size of tensor a (32) must match the size of tensor b (128) at non-singleton dimension 3
```

### Expected behavior

I would expect successful forward pass through the quantized model.

Comments:
Did you manage to fix your issue @Godofnothing ?

@SunMarc yes, the input format was incorrect. 
The standard way to produce inputs via tokenization of text prompt works fine as well as `generate`. 
Therefore, I closed the issue.

> ### System Info
> 
> **Environment**
> 
> ```
> - `transformers` version: 4.48.2
> - Platform: Linux-5.4.210-39.1.pagevecsize-x86_64-with-glibc2.27
> - Python version: 3.11.10
> - Huggingface_hub version: 0.26.2
> - Safetensors version: 0.4.5
> - Accelerate version: 1.1.1
> - Accelerate config:    not found
> - PyTorch version (GPU?): 2.4.0+cu121 (True)
> - Tensorflow version (GPU?): not installed (NA)
> - Flax version (CPU?/GPU?/TPU?): not installed (NA)
> - Jax version: not installed
> - JaxLib version: not installed
> - Using distributed or parallel set-up in script?: <fill in>
> - Using GPU in script?: <fill in>
> - GPU type: NVIDIA A100-SXM4-80GB
> - fast_hadamard_transform           1.0.4.post1
> ```
> 
> ### Who can help?
> 
> [@BlackSamorez](https://github.com/BlackSamorez) [@SunMarc](https://github.com/SunMarc) [@ArthurZucker](https://github.com/ArthurZucker)
> ### Information
> 
>     * [ ]  The official example scripts[x]  My own modified scripts
> 
> 
> ### Tasks
> 
>     * [ ]  An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)[ ]  My own task or dataset (give details below)
> 
> 
> ### Reproduction
> 
> Recently, in the [PR](https://github.com/huggingface/transformers/pull/34997) HIGGS quantization from the paper [Pushing the Limits of Large Language Model Quantization via the Linearity Theorem](https://arxiv.org/abs/2411.17525) was introduced.
> 
> But when attempting to load the quantized `Llama-3.1-8B-Instruct `model in this format as follows:
> 
> model_name = "meta-llama/Llama-3.1-8B-Instruct"
> quantization_config = HiggsConfig(bits=4, p=2)
> model = AutoModelForCausalLM.from_pretrained(
>     model_name, 
>     device_map="auto", 
>     torch_dtype=torch.bfloat16, 
>     low_cpu_mem_usage=True,
>     quantization_config=quantization_config
> )
> model.config.use_cache = False
> 
> And doing forward pass with dummy inputs
> 
> inputs = torch.randint(0, model.config.vocab_size, device="cuda", size=(8,))
> with torch.no_grad():
>     outputs = model(inputs)
> 
> I get the following error in the RoPE:
> 
> File ~/miniconda3/envs/llm/lib/python3.11/site-packages/transformers/models/llama/modeling_llama.py:271, in LlamaAttention.forward(self, hidden_states, position_embeddings, attention_mask, past_key_value, cache_position, **kwargs)
>     [268](https://vscode-remote+ssh-002dremote-002bultramar.vscode-resource.vscode-cdn.net/home/dkuznedelev/FLUTE_Playground/~/miniconda3/envs/llm/lib/python3.11/site-packages/transformers/models/llama/modeling_llama.py:268) value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)
>     [270](https://vscode-remote+ssh-002dremote-002bultramar.vscode-resource.vscode-cdn.net/home/dkuznedelev/FLUTE_Playground/~/miniconda3/envs/llm/lib/python3.11/site-packages/transformers/models/llama/modeling_llama.py:270) cos, sin = position_embeddings
> --> [271](https://vscode-remote+ssh-002dremote-002bultramar.vscode-resource.vscode-cdn.net/home/dkuznedelev/FLUTE_Playground/~/miniconda3/envs/llm/lib/python3.11/site-packages/transformers/models/llama/modeling_llama.py:271) query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
>     [273](https://vscode-remote+ssh-002dremote-002bultramar.vscode-resource.vscode-cdn.net/home/dkuznedelev/FLUTE_Playground/~/miniconda3/envs/llm/lib/python3.11/site-packages/transformers/models/llama/modeling_llama.py:273) if past_key_value is not None:
>     [274](https://vscode-remote+ssh-002dremote-002bultramar.vscode-resource.vscode-cdn.net/home/dkuznedelev/FLUTE_Playground/~/miniconda3/envs/llm/lib/python3.11/site-packages/transformers/models/llama/modeling_llama.py:274)     # sin and cos are specific to RoPE models; cache_position needed for the static cache
>     [275](https://vscode-remote+ssh-002dremote-002bultramar.vscode-resource.vscode-cdn.net/home/dkuznedelev/FLUTE_Playground/~/miniconda3/envs/llm/lib/python3.11/site-packages/transformers/models/llama/modeling_llama.py:275)     cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position}
> 
> File ~/miniconda3/envs/llm/lib/python3.11/site-packages/transformers/models/llama/modeling_llama.py:169, in apply_rotary_pos_emb(q, k, cos, sin, position_ids, unsqueeze_dim)
>     [167](https://vscode-remote+ssh-002dremote-002bultramar.vscode-resource.vscode-cdn.net/home/dkuznedelev/FLUTE_Playground/~/miniconda3/envs/llm/lib/python3.11/site-packages/transformers/models/llama/modeling_llama.py:167) cos = cos.unsqueeze(unsqueeze_dim)
>     [168](https://vscode-remote+ssh-002dremote-002bultramar.vscode-resource.vscode-cdn.net/home/dkuznedelev/FLUTE_Playground/~/miniconda3/envs/llm/lib/python3.11/site-packages/transformers/models/llama/modeling_llama.py:168) sin = sin.unsqueeze(unsqueeze_dim)
> --> [169](https://vscode-remote+ssh-002dremote-002bultramar.vscode-resource.vscode-cdn.net/home/dkuznedelev/FLUTE_Playground/~/miniconda3/envs/llm/lib/python3.11/site-packages/transformers/models/llama/modeling_llama.py:169) q_embed = (q * cos) + (rotate_half(q) * sin)
>     [170](https://vscode-remote+ssh-002dremote-002bultramar.vscode-resource.vscode-cdn.net/home/dkuznedelev/FLUTE_Playground/~/miniconda3/envs/llm/lib/python3.11/site-packages/transformers/models/llama/modeling_llama.py:170) k_embed = (k * cos) + (rotate_half(k) * sin)
>     [171](https://vscode-remote+ssh-002dremote-002bultramar.vscode-resource.vscode-cdn.net/home/dkuznedelev/FLUTE_Playground/~/miniconda3/envs/llm/lib/python3.11/site-packages/transformers/models/llama/modeling_llama.py:171) return q_embed, k_embed
> 
> RuntimeError: The size of tensor a (32) must match the size of tensor b (128) at non-singleton dimension 3
> 
> ### Expected behavior
> 
> I would expect successful forward pass through the quantized model.

why?      ?????????????????????????????????????????????????????????????