model.generate cannot handle past_key_values correctly
Author: zhuango
Pull request: no
Labels: name: bug

### System Info

- `transformers` version: 4.42.4
- Platform: Linux-5.15.0-94-generic-x86_64-with-glibc2.35
- Python version: 3.10.12
- Huggingface_hub version: 0.23.5
- Safetensors version: 0.4.3
- Accelerate version: 0.32.1
- Accelerate config:    not found
- PyTorch version (GPU?): 2.3.1+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?: No
- Using GPU in script?: Yes
- GPU type: NVIDIA A100-SXM4-80GB

### Who can help?

@ArthurZucker @zucchini-nlp  @gante

### 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

Here is my scripts.

```python
import torch
from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
    TextStreamer,
)

model_name_or_path='./glm-4-9b-chat'
device = "cuda"

tokenizer = AutoTokenizer.from_pretrained(model_name_or_path,trust_remote_code=True)

model = AutoModelForCausalLM.from_pretrained(
    model_name_or_path,
    torch_dtype=torch.bfloat16,
    low_cpu_mem_usage=True,
    trust_remote_code=True
).to(device).eval()

streamer = TextStreamer(tokenizer, skip_prompt=True, decode_kwargs=dict(skip_special_tokens=True))
gen_kwargs = {"max_length": 2500, "do_sample": True, "top_k": 1, "streamer":streamer, "return_dict_in_generate": True}
past_key_values = None
inputs = None

system_message = "你是一个名为 GLM-4 的人工智能助手。你是基于智谱AI训练的语言模型 GLM-4 模型开发的，你的任务是针对用户的问题和要求提供适当的答复和支持。"
history = [{"role": "system", "content": system_message}]

past_key_values = None

while True:
    query = input("Human: ")
    if len(query.strip()) == 0:
        history = [{"role": "system", "content": system_message}]

        continue
        
    history.append({
        'role': "user",
        "content": query
    })
    inputs = tokenizer.apply_chat_template(history,
                add_generation_prompt=True,
                tokenize=True,
                return_tensors="pt",
                return_dict=True
                )
    inputs = inputs.to(device)
    print(inputs['input_ids'].shape)

    with torch.no_grad():
        print("Assistant:")
        outputs = model.generate(**inputs, **gen_kwargs, past_key_values=past_key_values)
        past_key_values = outputs['past_key_values']
        
        outputs = outputs['sequences'][:, inputs['input_ids'].shape[1]:]
        outputs = tokenizer.decode(outputs[0], skip_special_tokens=True)
        history.append({
            'role': "assistant",
            "content": outputs
        })
```

The first round chat works fine, but the second round (which has past_key_values passed into model.generate) runs into error:
```
Traceback (most recent call last):
  File "/home/ubuntu/llm/generation.txt", line 52, in <module>
    outputs = model.generate(**inputs, **gen_kwargs, past_key_values=past_key_values)
  File "/home/ubuntu/.local/lib/python3.10/site-packages/torch/utils/_contextlib.py", line 115, in decorate_context
    return func(*args, **kwargs)
  File "/home/ubuntu/.local/lib/python3.10/site-packages/transformers/generation/utils.py", line 1914, in generate
    result = self._sample(
  File "/home/ubuntu/.local/lib/python3.10/site-packages/transformers/generation/utils.py", line 2651, in _sample
    outputs = self(
  File "/home/ubuntu/.local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1532, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/home/ubuntu/.local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1541, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/ubuntu/.cache/huggingface/modules/transformers_modules/glm-4-9b-chat/modeling_chatglm.py", line 997, in forward
    transformer_outputs = self.transformer(
  File "/home/ubuntu/.local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1532, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/home/ubuntu/.local/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1541, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/ubuntu/.cache/huggingface/modules/transformers_modules/glm-4-9b-chat/modeling_chatglm.py", line 882, in forward
    full_attention_mask = self.get_masks(input_ids, past_key_values, padding_mask=attention_mask)
  File "/home/ubuntu/.cache/huggingface/modules/transformers_modules/glm-4-9b-chat/modeling_chatglm.py", line 783, in get_masks
    full_attention_mask = full_attention_mask * padding_mask.unsqueeze(1)
RuntimeError: The size of tensor a (130) must match the size of tensor b (69) at non-singleton dimension 2
```

### Expected behavior

It shoud generate new tokens with kv cache.

Comments:
hey! Does the error occur in models that are integrated into the library, e.g. Llama? From the error msg seems like it's related to the way ChatGLM is implemented

@zucchini-nlp Thanks for the replay. I just tested on LLama, it runs correctly. So, like you mentioned, it's a bug of ChatGLM  implementation. I will get some feedback from them.