Unnecessary KV Cache Updates During Training Mode
Author: Hannibal046
Pull request: no
Labels: name: bug

### System Info

- `transformers` version: 4.48.0
- Platform: Linux-5.4.0-58-generic-x86_64-with-glibc2.31
- Python version: 3.11.10
- Huggingface_hub version: 0.26.3
- Safetensors version: 0.4.5
- Accelerate version: 1.1.1
- Accelerate config:    not found
- PyTorch version (GPU?): 2.5.1+cu118 (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-PCIE-40GB

### Who can help?

@ArthurZucker @muellerz @SunMarc 

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

I've identified an issue regarding unnecessary KV cache updates during training, which affects all current LLM models in the library and impacts both memory efficiency and torch.compile compatibility.

Taking `src/transformers/models/llama/modeling_llama.py` as an example:

1. The `past_key_values` is set as a `DynamicCache` object even during training:
https://github.com/huggingface/transformers/blob/15bd3e61f8d3680ca472c9314ad07584d20f7b81/src/transformers/models/llama/modeling_llama.py#L547-L548

2. This leads to unnecessary memory allocation for storing KV cache during the forward pass:

https://github.com/huggingface/transformers/blob/15bd3e61f8d3680ca472c9314ad07584d20f7b81/src/transformers/models/llama/modeling_llama.py#L273-L276

3. Additionally, this causes issues with `torch.compile`, resulting in multiple recompilations due to layer_idx guards (my internal test) :
```
[rank6]:V0113 13:51:42.251000 31893 site-packages/torch/_dynamo/guards.py:2813] [10/6] [__recompiles] Recompiling function forward in /weka-jd/prod/deepseek/permanent/shared/chengxin/workspace/research_project/lingua/apps/main/hf_transformers/modeling_transformer.py:315
[rank6]:V0113 13:51:42.251000 31893 site-packages/torch/_dynamo/guards.py:2813] [10/6] [__recompiles]     triggered by the following guard failure(s):
[rank6]:V0113 13:51:42.251000 31893 site-packages/torch/_dynamo/guards.py:2813] [10/6] [__recompiles]     - 10/5: L['self']._modules['self_attn'].layer_idx == 5              
[rank6]:V0113 13:51:42.251000 31893 site-packages/torch/_dynamo/guards.py:2813] [10/6] [__recompiles]     - 10/4: L['self']._modules['self_attn'].layer_idx == 4              
[rank6]:V0113 13:51:42.251000 31893 site-packages/torch/_dynamo/guards.py:2813] [10/6] [__recompiles]     - 10/3: L['self']._modules['self_attn'].layer_idx == 3              
[rank6]:V0113 13:51:42.251000 31893 site-packages/torch/_dynamo/guards.py:2813] [10/6] [__recompiles]     - 10/2: L['self']._modules['self_attn'].layer_idx == 2              
[rank6]:V0113 13:51:42.251000 31893 site-packages/torch/_dynamo/guards.py:2813] [10/6] [__recompiles]     - 10/1: L['self']._modules['self_attn'].layer_idx == 1              
[rank6]:V0113 13:51:42.251000 31893 site-packages/torch/_dynamo/guards.py:2813] [10/6] [__recompiles]     - 10/0: L['self']._modules['self_attn'].layer_idx == 0              
```

### Workaround
A temporary solution is to explicitly set `use_cache=False` during training:
```
loss = model.forward(input_ids, labels, use_cache=False)
```

### Expected behavior

no kv cache during training

Comments:
Or we could change from:
https://github.com/huggingface/transformers/blob/15bd3e61f8d3680ca472c9314ad07584d20f7b81/src/transformers/models/llama/modeling_llama.py#L532

to:
```python
if use_cache is None:
   use_cache = False if self.training else self.config.use_cache
```

cc @gante for cache as well

@Hannibal046 👋 

@Rocketknight1 and I chatted offline and we can't think of a reason why the cache should stay active while training. ~I'm going to update to the pattern as you wrote -- if the user doesn't manually specify the cache, only use the config file value (which is `True` by default) if not training~ (see next comment)

@Hannibal046 @Rocketknight1 

Correction: we may indeed want to pass `past_key_values` with `self.training=True`, in fine-tuning routines like prefix tuning (thank you @BenjaminBossan for reminding me). As such, the change will NOT be implemented.

In fact, I'm going to add a tiny integration test that will fail if we decide to disallow caches at training time.

Hi @gante,

Thanks for the quick response! While I'm not entirely certain about why prefix tuning requires caching during training, I believe there's an issue with the current past_key_values implementation. 

For reference, in previous Hugging Face versions (such as transformers v4.30.2), past_key_values was consistently set to None during training:
https://github.com/huggingface/transformers/blob/66fd3a8d626a32989f4569260db32785c6cbf42a/src/transformers/models/llama/modeling_llama.py#L517-L519

https://github.com/huggingface/transformers/blob/66fd3a8d626a32989f4569260db32785c6cbf42a/src/transformers/models/llama/modeling_llama.py#L199-L201

I don't see why we should compromise the general use case for a special case. Most applications don't require caching during training, and the current approach has two significant drawbacks:
1. Unnecessary GPU memory consumption
2. Compatibility issues with `torch.compile`

If specific cases like prefix tuning require caching, they can simply pass `use_cache=True` in the forward method. This approach would better serve both general and special use cases while maintaining code clarity and efficiency.

@Hannibal046 

The old cache design had many drawbacks, including now allowing to specify different types of cache or preventing compilation with text generation. The new design came at the cost of minor changes like the root cause for this issue, but the benefits far outweigh the costs.

In 🤗 `transformers`, our goal is to enable a wide set of possibilities with respect to model usage. We also strive to enable backward compatibility. It's impossible to optimize in every direction, so sometimes we have to sacrifice some speed or UX. 

I agree it would be ideal to work optimally without any flag, but let's consider the alternatives:
1. We disable cache by default at train time -> backward compatibility issues with PEFT and almost certainly other use cases we don't have visibility. Code will crash unless the library owners update their code, and set new version restrictions;
2. [current] We enable cache by default at train time -> If the user enabled gradient checkpointing, which is very common nowadays, there will be no drawback at all as the cache will be disabled. For the other users: minor speed issues and incompatibility with advanced features like compilation, that can be solved with `use_cache=False`

I hope the benefits of not changing the default are clear 😉 

Thanks for your detailed explanation!