No speed-up of model.generate() with StaticCache + torch.compile in 4.39.3
Author: learning-chip
Pull request: no
Labels: name: Generation

name: Cache

### System Info

torch==2.2.2
transformers==4.39.3

Platform: RTX 4090 or RTX A6000 rent on vast.ai

### Who can help?

 @ArthurZucker @gan

### Information

- [ ] The official example scripts
- [X] 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

Below is the complete benchmark script. The compile setting follows https://github.com/huggingface/transformers/pull/29791. The generate setting and benchmark metric follows [gpt-fast/generate.py](https://github.com/pytorch-labs/gpt-fast/blob/main/generate.py)

```python
import time
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, StaticCache

import torch
import torch._dynamo.config
import torch._inductor.config

torch._inductor.config.coordinate_descent_tuning = True
torch._inductor.config.triton.unique_kernel_names = True
torch._inductor.config.fx_graph_cache = True

# set same as gpt-fast default config
compile = True
device = "cuda"
precision = torch.bfloat16
num_samples = 5
max_new_tokens = 200
start = -1 if compile else 0
torch.manual_seed(1234)

def device_sync(device):
    if "cuda" in device:
        torch.cuda.synchronize(device)
    elif ("cpu" in device) or ("mps" in device):
        pass
    else:
        print(f"device={device} is not yet suppported")

model_name = "openlm-research/open_llama_7b"
tokenizer = AutoTokenizer.from_pretrained(model_name)

t0 = time.time()
model = AutoModelForCausalLM.from_pretrained(
    model_name, device_map={"": device}, torch_dtype=precision, attn_implementation='sdpa'
)
device_sync(device=device)
t = time.time() - t0
print(f"Time to load model: {t:.02f} seconds")

prompt = "Hello, my name is"
inputs = tokenizer(prompt, return_tensors='pt').to(device)
prompt_length = inputs["input_ids"].size(1)
max_seq_length = max_new_tokens + prompt_length

model._setup_cache(StaticCache, max_batch_size=1, max_cache_len=max_seq_length)

if compile:
    model = torch.compile(model, mode="reduce-overhead", fullgraph=True)

aggregate_metrics = {
    'tokens_per_sec': []
}

for i in range(start, num_samples):
    device_sync(device)
    t0 = time.perf_counter()
    output = model.generate(**inputs, max_new_tokens=max_new_tokens, do_sample=True, top_k=200, temperature=0.8)
    device_sync(device)
    t = time.perf_counter() - t0

    print(tokenizer.decode(output[0]))

    tokens_generated = output.size(1) - prompt_length
    tokens_sec = tokens_generated / t
    if i == -1:
        print(f"Compilation time: {time.perf_counter() - t0:.2f} seconds")
        continue
    
    aggregate_metrics['tokens_per_sec'].append(tokens_sec)
    print(f"tokens generated: {tokens_generated}")
    print(f"Time for inference {i + 1}: {t:.02f} sec total, {tokens_sec:.02f} tokens/sec")

print(f"Average tokens/sec: {torch.mean(torch.tensor(aggregate_metrics['tokens_per_sec'])).item():.2f}")
print(f"Memory used: {torch.cuda.max_memory_reserved() / 1e9:.02f} GB")
```

With either `compile=True` or `compile=False`, the throughput is always ~33 token/s.
Complete logs on A6000:
- [hf_generate_eager.log](https://github.com/huggingface/transformers/files/14876744/hf_generate_eager.log)
- [hf_generate_compile.log](https://github.com/huggingface/transformers/files/14876745/hf_generate_compile.log)

For comparison, same benchmark setting using https://github.com/pytorch-labs/gpt-fast:

```bash
export MODEL_REPO=openlm-research/open_llama_7b
python ./generate.py --checkpoint_path checkpoints/$MODEL_REPO/model.pth --prompt "Hello, my name is" # eager
python ./generate.py --compile  checkpoints/$MODEL_REPO/model.pth ./checkpoints/open_llama_7b/model.pth  --prompt "Hello, my name is" # compile
```

Here the eager version gives ~30 token/s while compiled version gives 50 token/s.

- [gptfast_eager.log](https://github.com/huggingface/transformers/files/14876867/gptfast_eager.log)
- [gptfast_compile.log](https://github.com/huggingface/transformers/files/14876866/gptfast_compile.log)

### Expected behavior

StaticCache + torch.compile should also achieve ~2x speed-up as gpt-fast

Comments:
Looking at the complete logs, gpt-fast compile step takes 30s (this is second time with `fx_graph_cache`; the initial compile takes ~100s). But HF's compile step only takes 6.5 seconds. Any TorchInductor settings that are missing in my above benchmark script? 

Thanks for providing a detailed reproducer! 
I re-run my custom benchmark and I am getting: `compiled static cache   (1,  512) - Took : 0.8398 sec total, 8.398 ms per forward, 119.08 tokens/s` 
Something like this:

![image](https://github.com/huggingface/transformers/assets/48595927/e20a4f09-5960-4ab8-b302-87cf7c14ce42)

Should work.
cc @gante and @fxmarty 


> I re-run my custom benchmark

I am missing the `cache_implementation="static"`, maybe that's the difference🤔


I found the problem -- using `torch.compile(model.forward)` instead of `torch.compile(model)` will properly trigger the compilation.




With this updated script, I can confirm `torch.compile` speed-up:

Expand here:
<details>

```python
import os
import time
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, StaticCache

import torch
import torch._dynamo.config
import torch._inductor.config

torch._inductor.config.coordinate_descent_tuning = True
torch._inductor.config.triton.unique_kernel_names = True
torch._inductor.config.fx_graph_cache = True

os.environ["TOKENIZERS_PARALLELISM"] = "true"

print(torch.version.cuda)

# set same as gpt-fast default config
compile = False
device = "cuda"
precision = torch.float16 # torch.bfloat16
num_samples = 5
max_new_tokens = 100
start = -1 if compile else 0
torch.manual_seed(1234)

def device_sync(device):
    if "cuda" in device:
        torch.cuda.synchronize(device)
    elif ("cpu" in device) or ("mps" in device):
        pass
    else:
        print(f"device={device} is not yet suppported")

model_name = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
# model_name = "openlm-research/open_llama_7b"
tokenizer = AutoTokenizer.from_pretrained(model_name)

t0 = time.time()
model = AutoModelForCausalLM.from_pretrained(
    model_name, device_map={"": device}, torch_dtype=precision, attn_implementation='sdpa'
)
device_sync(device=device)
t = time.time() - t0
print(f"Time to load model: {t:.02f} seconds")

# prompt = ["Hello, my name is"]
prompt = ["The quick brown"]
inputs = tokenizer(prompt, return_tensors='pt').to(device)
prompt_length = inputs["input_ids"].size(1)
max_seq_length = max_new_tokens + prompt_length

# model._setup_cache(StaticCache, max_batch_size=1, max_cache_len=max_seq_length)

if compile:
    # model = torch.compile(model, mode="reduce-overhead", fullgraph=True)
    model.forward = torch.compile(model.forward, mode="reduce-overhead", fullgraph=True)

aggregate_metrics = {
    'tokens_per_sec': []
}

for i in range(start, num_samples):
    device_sync(device)
    t0 = time.perf_counter()
    output = model.generate(
        **inputs, max_new_tokens=max_new_tokens, cache_implementation="static",
        do_sample=False #, top_k=200, temperature=0.8
    )
    device_sync(device)
    t = time.perf_counter() - t0

    print(tokenizer.decode(output[0]))

    tokens_generated = output.size(1) - prompt_length
    tokens_sec = tokens_generated / t
    if i == -1:
        print(f"Compilation time: {time.perf_counter() - t0:.2f} seconds")
        continue
    
    aggregate_metrics['tokens_per_sec'].append(tokens_sec)
    print(f"tokens generated: {tokens_generated}")
    print(f"Time for inference {i + 1}: {t:.02f} sec total, {tokens_sec:.02f} tokens/sec")

print(f"Average tokens/sec: {torch.mean(torch.tensor(aggregate_metrics['tokens_per_sec'])).item():.2f}")
print(f"Memory used: {torch.cuda.max_memory_reserved() / 1e9:.02f} GB")
```

</details>

Complete logs:
- [run_hf_eager.log](https://github.com/huggingface/transformers/files/14885576/run_hf_eager.log)
- [run_hf_compile.log](https://github.com/huggingface/transformers/files/14885578/run_hf_compile.log)

Eager gives 44 token/s while compile gives 200+ token/s !

However, the compile pass seems not very robust regarding input/output sequence length. When I used a different prompt `prompt = ["Hello, my name is"]` for the above benchmark script, I got: `RuntimeError: attn_bias is not correctly aligned (strideH). attn_bias.stride(1) = 636, and should be a multiple of 8.`

Complete error log:
<details>

```python
12.1
Time to load model: 2.01 seconds
Traceback (most recent call last):
  File "/home/./hf_bench.py", line 66, in <module>
    output = model.generate(
  File "/opt/conda/lib/python3.10/site-packages/torch/utils/_contextlib.py", line 115, in decorate_context
    return func(*args, **kwargs)
  File "/opt/conda/lib/python3.10/site-packages/transformers/generation/utils.py", line 1527, in generate
    result = self._greedy_search(
  File "/opt/conda/lib/python3.10/site-packages/transformers/generation/utils.py", line 2411, in _greedy_search
    outputs = self(
  File "/opt/conda/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1511, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/opt/conda/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1520, in _call_impl
    return forward_call(*args, **kwargs)
  File "/opt/conda/lib/python3.10/site-packages/torch/_dynamo/eval_frame.py", line 489, in _fn
    return fn(*args, **kwargs)
  File "/opt/conda/lib/python3.10/site-packages/transformers/models/llama/modeling_llama.py", line 1148, in forward
    @add_start_docstrings_to_model_forward(LLAMA_INPUTS_DOCSTRING)
  File "/opt/conda/lib/python3.10/site-packages/torch/_dynamo/eval_frame.py", line 489, in _fn
    return fn(*args, **kwargs)
  File "/opt/conda/lib/python3.10/site-packages/torch/_dynamo/external_utils.py", line 17, in inner
    return fn(*args, **kwargs)
  File "/opt/conda/lib/python3.10/site-packages/torch/_functorch/aot_autograd.py", line 901, in forward
    return compiled_fn(full_args)
  File "/opt/conda/lib/python3.10/site-packages/torch/_functorch/_aot_autograd/utils.py", line 81, in g
    return f(*args)
  File "/opt/conda/lib/python3.10/site-packages/torch/_functorch/_aot_autograd/runtime_wrappers.py", line 94, in runtime_wrapper
    all_outs = call_func_at_runtime_with_args(
  File "/opt/conda/lib/python3.10/site-packages/torch/_functorch/_aot_autograd/utils.py", line 105, in call_func_at_runtime_with_args
    out = normalize_as_list(f(args))
  File "/opt/conda/lib/python3.10/site-packages/torch/_functorch/_aot_autograd/jit_compile_runtime_wrappers.py", line 118, in rng_functionalization_wrapper
    return compiled_fw(args)
  File "/opt/conda/lib/python3.10/site-packages/torch/_inductor/codecache.py", line 864, in __call__
    return self.get_current_callable()(inputs)
  File "/opt/conda/lib/python3.10/site-packages/torch/_inductor/compile_fx.py", line 665, in run
    return compiled_fn(new_inputs)
  File "/opt/conda/lib/python3.10/site-packages/torch/_inductor/cudagraph_trees.py", line 380, in deferred_cudagraphify
    fn, out = cudagraphify(model, inputs, new_static_input_idxs, *args, **kwargs)
  File "/opt/conda/lib/python3.10/site-packages/torch/_inductor/cudagraph_trees.py", line 408, in cudagraphify
    return manager.add_function(
  File "/opt/conda/lib/python3.10/site-packages/torch/_inductor/cudagraph_trees.py", line 1941, in add_function
    return fn, fn(inputs)
  File "/opt/conda/lib/python3.10/site-packages/torch/_inductor/cudagraph_trees.py", line 1755, in run
    out = self._run(new_inputs, function_id)
  File "/opt/conda/lib/python3.10/site-packages/torch/_inductor/cudagraph_trees.py", line 1796, in _run
    return self.run_eager(new_inputs, function_id)
  File "/opt/conda/lib/python3.10/site-packages/torch/_inductor/cudagraph_trees.py", line 1911, in run_eager
    return node.run(new_inputs)
  File "/opt/conda/lib/python3.10/site-packages/torch/_inductor/cudagraph_trees.py", line 611, in run
    out = self.wrapped_function.model(new_inputs)
  File "/opt/conda/lib/python3.10/site-packages/torch/_inductor/codecache.py", line 892, in _run_from_cache
    return compiled_graph.compiled_artifact(inputs)
  File "/tmp/torchinductor_root/pt/cptfwf3oza4rtyjjjspyuyn27gezumkor3wxuzntzqfnogv7jb33.py", line 1522, in call
    buf16 = aten._scaled_dot_product_efficient_attention(buf12, buf13, buf14, buf15, False)
  File "/opt/conda/lib/python3.10/site-packages/torch/_ops.py", line 755, in __call__
    return self._op(*args, **(kwargs or {}))
  File "/opt/conda/lib/python3.10/site-packages/torch/utils/_device.py", line 77, in __torch_function__
    return func(*args, **kwargs)
  File "/opt/conda/lib/python3.10/site-packages/torch/_ops.py", line 755, in __call__
    return self._op(*args, **(kwargs or {}))
RuntimeError: attn_bias is not correctly aligned (strideH). attn_bias.stride(1) = 636, and should be a multiple of 8.
```

</details>

The above RuntimeError goes away if I explicitly set `max_cache_len` to be power of 2 for `_setup_cache`, rather than relying on `cache_implementation="static"` to figure-out the cache size.

- [run_hf_compile(1).log](https://github.com/huggingface/transformers/files/14885720/run_hf_compile.1.log)
- [run_hf_compile(2).log](https://github.com/huggingface/transformers/files/14885908/run_hf_compile.2.log)

Complete updated script (sampling also works fine): 

<details>

```python
import os
import time
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, StaticCache

import torch
import torch._dynamo.config
import torch._inductor.config

# torch._inductor.config.coordinate_descent_tuning = True
# torch._inductor.config.triton.unique_kernel_names = True
# torch._inductor.config.fx_graph_cache = True

os.environ["TOKENIZERS_PARALLELISM"] = "true"

print(torch.version.cuda)

# set same as gpt-fast default config
compile = True
device = "cuda"
precision = torch.float16 # torch.bfloat16
num_samples = 5
max_new_tokens = 100
start = -1 if compile else 0
torch.manual_seed(1234)

def device_sync(device):
    if "cuda" in device:
        torch.cuda.synchronize(device)
    elif ("cpu" in device) or ("mps" in device):
        pass
    else:
        print(f"device={device} is not yet suppported")

model_name = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
# model_name = "openlm-research/open_llama_7b"
tokenizer = AutoTokenizer.from_pretrained(model_name)

t0 = time.time()
model = AutoModelForCausalLM.from_pretrained(
    model_name, device_map={"": device}, torch_dtype=precision, attn_implementation='sdpa'
)
device_sync(device=device)
t = time.time() - t0
print(f"Time to load model: {t:.02f} seconds")

prompt = ["Hello, my name is"]
# prompt = ["The quick brown"]
inputs = tokenizer(prompt, return_tensors='pt').to(device)
prompt_length = inputs["input_ids"].size(1)
max_seq_length = 512 # max_new_tokens + prompt_length

model._setup_cache(StaticCache, max_batch_size=1, max_cache_len=max_seq_length)

if compile:
    # model = torch.compile(model, mode="reduce-overhead", fullgraph=True)
    model.forward = torch.compile(model.forward, mode="reduce-overhead", fullgraph=True)

aggregate_metrics = {
    'tokens_per_sec': []
}

for i in range(start, num_samples):
    device_sync(device)
    t0 = time.perf_counter()
    output = model.generate(
        **inputs, max_new_tokens=max_new_tokens, # cache_implementation="static",
        do_sample=True, top_k=200, temperature=0.8
    )
    device_sync(device)
    t = time.perf_counter() - t0

    print(tokenizer.decode(output[0]))

    tokens_generated = output.size(1) - prompt_length
    tokens_sec = tokens_generated / t
    if i == -1:
        print(f"Compilation time: {time.perf_counter() - t0:.2f} seconds")
        continue
    
    aggregate_metrics['tokens_per_sec'].append(tokens_sec)
    print(f"tokens generated: {tokens_generated}")
    print(f"Time for inference {i + 1}: {t:.02f} sec total, {tokens_sec:.02f} tokens/sec")

print(f"Average tokens/sec: {torch.mean(torch.tensor(aggregate_metrics['tokens_per_sec'])).item():.2f}")
print(f"Memory used: {torch.cuda.max_memory_reserved() / 1e9:.02f} GB")
```

</details>

Would be good to have more comprehensive unit tests regarding sequence length, batch size, compile settings...


@learning-chip I am indeed getting the same error while I used not to (`attn_bias is not correctly aligned (strideH). attn_bias.stride(1) = 119, and should be a multiple of 8.`).

This looks a lot like a bug in PyTorch + compile + SDPA mem-efficient backend cc @drisspg


> SDPA mem-efficient backend

Indeed like a xformers problem: https://github.com/vllm-project/vllm/pull/834


I can take a look on Monday

> I found the problem -- using `torch.compile(model.forward)` instead of `torch.compile(model)` will properly trigger the compilation.

@learning-chip This is very surprising to me, considering what is written here:
https://huggingface.co/docs/transformers/main/perf_torch_compile#benefits-of-torchcompile

I'm not sure anymore what's the recommended way to use `torch.compile` with transformers.

> I'm not sure anymore what's the recommended way to use `torch.compile` with transformers.

GPT-fast also only compiles the model forward call:
https://github.com/pytorch-labs/gpt-fast/blob/c21a88962b02ee54b74999078e81be0fd24ac2af/generate.py#L310-L319


The doc is being updated #26058 and this should work as well:
```python 
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

model_id  = "meta-llama/Llama-2-7b-chat-hf"
model     = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.float16, attn_implementation="sdpa").cuda().eval();
tokenizer = AutoTokenizer.from_pretrained(model_id) 
tokenizer.add_bos_token = False

model.forward = torch.compile(model.forward, mode="reduce-overhead", fullgraph=True)

inputs = tokenizer(["
 
 Write an essay about large language models 
"], return_tensors="pt").to(model.device)
for _ in range(3):
	gen_out = model.generate(**inputs, do_sample=False, cache_implementation="static", max_new_tokens=100, pad_token_id=tokenizer.pad_token_id, temperature=None, top_p=None, use_cache=False)
print(tokenizer.decode(gen_out[0]))

inputs = tokenizer(["
 
 How to make a chocolate cake? 
"], return_tensors="pt").to(model.device)
gen_out = model.generate(**inputs, do_sample=False, cache_implementation="static", max_new_tokens=100, pad_token_id=tokenizer.pad_token_id, temperature=None, top_p=None, use_cache=False)
print(tokenizer.decode(gen_out[0]))
```

taken from #30351

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.

Closing as this works.

Hi! Is there any way to save the results of compilation? Or we have to go through the long compilation process every time we run the script?

Normally compilation should be save directly by torch!