Slow speed when inference LLaMA model with torchao
Author: ThisisBillhe
Pull request: no
Labels: name: Quantization

name: bug

### System Info

RTX3090 machine:
```
- `transformers` version: 4.45.2
- Platform: Linux-5.15.0-79-generic-x86_64-with-glibc2.31
- Python version: 3.10.15
- Huggingface_hub version: 0.26.1
- Safetensors version: 0.4.5
- Accelerate version: 1.0.1
- Accelerate config:    not found
- PyTorch version (GPU?): 2.5.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 GeForce RTX 3090
```

A100 machine:
```
- `transformers` version: 4.45.2
- Platform: Linux-4.19.91-014.15-kangaroo.alios7.x86_64-x86_64-with-glibc2.35
- Python version: 3.10.15
- Huggingface_hub version: 0.25.2
- Safetensors version: 0.4.5
- Accelerate version: 1.0.1
- Accelerate config:    not found
- PyTorch version (GPU?): 2.5.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
 ```

### Who can help?

@SunMarc 

### Information

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

### Tasks

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

### Reproduction

1. Quantize the Llama-3-8B [demo](https://huggingface.co/docs/transformers/main/quantization/torchao) in official document.
I got slower speed when using torchao on both A100 and 3090 GPUs. On 3090:
```
Setting `pad_token_id` to `eos_token_id`:None for open-end generation.
Setting `pad_token_id` to `eos_token_id`:None for open-end generation.
Setting `pad_token_id` to `eos_token_id`:None for open-end generation.
Setting `pad_token_id` to `eos_token_id`:None for open-end generation.
Setting `pad_token_id` to `eos_token_id`:None for open-end generation.
Setting `pad_token_id` to `eos_token_id`:None for open-end generation.
Setting `pad_token_id` to `eos_token_id`:None for open-end generation.
int4wo-128 model: 47.273
Loading checkpoint shards: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████| 4/4 [00:16<00:00,  4.02s/it]
Setting `pad_token_id` to `eos_token_id`:None for open-end generation.
Setting `pad_token_id` to `eos_token_id`:None for open-end generation.
Setting `pad_token_id` to `eos_token_id`:None for open-end generation.
Setting `pad_token_id` to `eos_token_id`:None for open-end generation.
Setting `pad_token_id` to `eos_token_id`:None for open-end generation.
Setting `pad_token_id` to `eos_token_id`:None for open-end generation.
Setting `pad_token_id` to `eos_token_id`:None for open-end generation.
bf16 model: 20.715
```
On A100:
```
Setting `pad_token_id` to `eos_token_id`:None for open-end generation.
Setting `pad_token_id` to `eos_token_id`:None for open-end generation.
Setting `pad_token_id` to `eos_token_id`:None for open-end generation.
Setting `pad_token_id` to `eos_token_id`:None for open-end generation.
Setting `pad_token_id` to `eos_token_id`:None for open-end generation.
Setting `pad_token_id` to `eos_token_id`:None for open-end generation.
Setting `pad_token_id` to `eos_token_id`:None for open-end generation.
int4wo-128 model: 36.820
Loading checkpoint shards: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 4/4 [00:08<00:00,  2.07s/it]
Setting `pad_token_id` to `eos_token_id`:None for open-end generation.
Setting `pad_token_id` to `eos_token_id`:None for open-end generation.
Setting `pad_token_id` to `eos_token_id`:None for open-end generation.
Setting `pad_token_id` to `eos_token_id`:None for open-end generation.
Setting `pad_token_id` to `eos_token_id`:None for open-end generation.
Setting `pad_token_id` to `eos_token_id`:None for open-end generation.
Setting `pad_token_id` to `eos_token_id`:None for open-end generation.
bf16 model: 10.142
```
How can I get the decoding speedup from weight-only quantization?

### Expected behavior

Speedup with torchao quantized LLMs. 

Comments:
Hey @ThisisBillhe, did you compile the model ? Can you test this [script](https://gist.github.com/SunMarc/2d863109027c083e5ea4d1c9b4991388) instead ? We are compiling with fullgraph=True: `model.forward = torch.compile(model.forward, mode="reduce-overhead", fullgraph=True)`

> Hey @ThisisBillhe, did you compile the model ? Can you test this [script](https://gist.github.com/SunMarc/2d863109027c083e5ea4d1c9b4991388) instead ? We are compiling with fullgraph=True: `model.forward = torch.compile(model.forward, mode="reduce-overhead", fullgraph=True)`

Thanks for your reply! Your script is working properly on my 3090 machine. However, it got stuck on my A100 machine which is strange. I find that as long as I use "model.generation_config.cache_implementation = "static", the program will get stuck on my A100 machine no matter whether torch.compile is used. You can see my env in the front.

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.

@SunMarc  On a RTX 3090 machine, I followed the 🤗 official doc of [torchao](https://huggingface.co/docs/transformers/v4.47.1/en/quantization/torchao#torchao), also got slower inference speed for `Llama-3.1-8B-Instruct`, comparing int4wo vs bf16:

```
int4wo-128 model: 28.364
bf16 model: 24.682
```

```python
import torch
from transformers import TorchAoConfig, AutoModelForCausalLM, AutoTokenizer

model_name = "meta-llama/Llama-3.1-8B-Instruct"
# model_name = "meta-llama/Meta-Llama-3-8B"
# We support int4_weight_only, int8_weight_only and int8_dynamic_activation_int8_weight
# More examples and documentations for arguments can be found in https://github.com/pytorch/ao/tree/main/torchao/quantization#other-available-quantization-techniques
quantization_config = TorchAoConfig("int4_weight_only", group_size=128)
quantized_model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype="auto", device_map="auto", quantization_config=quantization_config)

tokenizer = AutoTokenizer.from_pretrained(model_name)
input_text = "What are we having for dinner?"
input_ids = tokenizer(input_text, return_tensors="pt").to("cuda")

# compile the quantized model to get speedup
import torchao
torchao.quantization.utils.recommended_inductor_config_setter()
quantized_model = torch.compile(quantized_model, mode="max-autotune")

output = quantized_model.generate(**input_ids, max_new_tokens=10)
print(tokenizer.decode(output[0], skip_special_tokens=True))

# benchmark the performance
import torch.utils.benchmark as benchmark

def benchmark_fn(f, *args, **kwargs):
    # Manual warmup
    for _ in range(5):
        f(*args, **kwargs)
        
    t0 = benchmark.Timer(
        stmt="f(*args, **kwargs)",
        globals={"args": args, "kwargs": kwargs, "f": f},
        num_threads=torch.get_num_threads(),
    )
    return f"{(t0.blocked_autorange().mean):.3f}"

MAX_NEW_TOKENS = 1000
print("int4wo-128 model:", benchmark_fn(quantized_model.generate, **input_ids, max_new_tokens=MAX_NEW_TOKENS))

bf16_model = AutoModelForCausalLM.from_pretrained(model_name, device_map="cuda", torch_dtype=torch.bfloat16)
bf16_model = torch.compile(bf16_model, mode="max-autotune")
print("bf16 model:", benchmark_fn(bf16_model.generate, **input_ids, max_new_tokens=MAX_NEW_TOKENS))
```

Below is my env info:
```
- `transformers` version: 4.47.1
- Platform: Linux-5.15.0-127-generic-x86_64-with-glibc2.31
- Python version: 3.11.11
- Huggingface_hub version: 0.27.0
- Safetensors version: 0.4.5
- Accelerate version: 1.2.1
- Accelerate config:    not found
- PyTorch version (GPU?): 2.5.1+cu124 (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?: No
- Using GPU in script?: Yes
- GPU type: NVIDIA GeForce RTX 3090
```

If using [your script](https://gist.github.com/SunMarc/2d863109027c083e5ea4d1c9b4991388), I got

| | Max memory (MB) | Throughput (tokens/sec) |
| -- | -- | -- |
| BF16 |  16156.354048 | 44.429162663703245
| int4wo-128 | 5914.308096 | 91.65667215102958

If initiating the model with your script and run the benchmark stated in 🤗 official doc, I got

```diff
- quantized_model = torch.compile(quantized_model, mode="max-autotune")
+ quantized_model.generation_config.cache_implementation = "static"
+ quantized_model.forward = torch.compile(quantized_model.forward, mode="reduce-overhead", fullgraph=True)
```

```
int4wo-128 model: 12.329
bf16 model: 24.674
```

Shall we update 🤗 official doc to guide others how to correctly get the speedup?

Another question: I'm curious why I got faster inference speed for bs=1, but not the case for bs=32?

For bs=1:
```
int4wo-128 model: 12.329
bf16 model: 24.674
```

But for bs=32:
```
int4wo-128 model: 90.628
bf16 model: 68.694
```

Below is the code I used for bs=32 
```diff
input_text = "What are we having for dinner?"

+ input_batch = [input_text] * 32
+ input_ids = tokenizer(input_batch, return_tensors="pt").to("cuda")
- input_ids = tokenizer(input_text, return_tensors="pt").to("cuda")
```


Thanks for the nice write-up @martin0258 ! 

> Shall we update 🤗 official doc to guide others how to correctly get the speedup? 

Yeah, definitely ! Would you like to update the example script ? That would be awesome ! Also, we introduced automatic compiling when the users call generate with `cache_implementation="static"` with this [PR](https://github.com/huggingface/transformers/pull/34923) ! 

cc @jerryzh168 for the batch_size question 

@SunMarc Sure! I just created this [PR](https://github.com/huggingface/transformers/pull/35490) to update the example script with the new automatic compilation.

Thanks @martin0258 and @SunMarc for updating the script. I thought the slowdown was due to some other issue related to SDPA at the time, but looks like it's not.

regarding batch size 32, it's expected that the current int4wo does not have speedup, and we have added a new type of int4 weight only quant: https://github.com/pytorch/ao/blob/main/torchao/quantization/README.md#gemlite-triton to address the issue, we can upstream this to transformers