Repository: pytorch
Issue: [torch.compile][HF] torch.compile issue tracker for torch.compile on forward method of Meta Llama model 
Labels: triaged,oncall: pt2,module: dynamo

### 🐛 Describe the bug

Script

~~~
from transformers import AutoTokenizer, AutoModelForCausalLM
import torch
import os
os.environ["TOKENIZERS_PARALLELISM"] = "false"  # To prevent long warnings :)


torch._dynamo.config.inline_inbuilt_nn_modules = True

tokenizer = AutoTokenizer.from_pretrained("meta-llama/Meta-Llama-3-8B")
model = AutoModelForCausalLM.from_pretrained("meta-llama/Meta-Llama-3-8B", device_map="auto", torch_dtype=torch.float16)

model.generation_config.cache_implementation = "static"

model.forward = torch.compile(model.forward, mode="reduce-overhead")
# backend="eager")
input_text = "The theory of special relativity states "
input_ids = tokenizer(input_text, return_tensors="pt").to(model.device)

outputs = model.generate(**input_ids)
print(tokenizer.batch_decode(outputs, skip_special_tokens=True))
~~~

There are 2 issues I found
* Recompilation because of stride change - We can probably build on top of https://github.com/pytorch/pytorch/pull/130232 and update `mark_dynamic` to work with strides as well. I think HF team is ok with the source code change to avoid this single recompilation as well.

~~~
DEBUG:torch._dynamo.guards.__recompiles:Recompiling function forward in /home/anijain/local/transformers/src/transformers/models/llama/modeling_llama.py:1021
    triggered by the following guard failure(s):
    - 0/0: tensor 'L['input_ids']' stride mismatch at index 0. expected 9, actual 1
~~~

* Cudgraph skipping 

~~~
skipping cudagraphs due to mutated inputs (64 instances). Found from :
   File "/home/anijain/local/transformers/src/transformers/models/llama/modeling_llama.py", line 1069, in forward
    outputs = self.model(
  File "/home/anijain/local/transformers/src/transformers/models/llama/modeling_llama.py", line 873, in forward
    layer_outputs = decoder_layer(
  File "/home/anijain/local/transformers/src/transformers/models/llama/modeling_llama.py", line 612, in forward
    hidden_states, self_attn_weights, present_key_value = self.self_attn(
  File "/home/anijain/local/transformers/src/transformers/models/llama/modeling_llama.py", line 522, in forward
    key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)
  File "/home/anijain/local/transformers/src/transformers/cache_utils.py", line 865, in update
    k_out[:, :, cache_position] = key_states

~~~

There are 2 problems here
* HF marks some tensors as static address to avoid this. This should not have been required in the first place. But, because of this code, cudagraphs are happy on main (you can see mark static address in [PR](https://github.com/huggingface/transformers/pull/32102))
* When I turn on inlining of inbuilt nn modules (which will soon happen on main), I see this skipping. Dont know why though.


### Error logs

_No response_

### Minified repro

_No response_

### Versions

NA

cc @ezyang @chauhang @penguinwu @voznesenskym @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @chenyang78 @kadeng @amjames