Substantially reduce memory usage in _update_causal_mask for large batches by using .expand instead of .repeat
Author: nqgl
Pull request: yes

# What does this PR do?
Reduce memory usage in _update_causal_mask when batch size is large
<!--
Congratulations! You've made it this far! You're not quite done yet though.

Once merged, your PR is going to appear in the release notes with the title you set, so make sure it's a great title that fully reflects the extent of your awesome contribution.

Then, please replace this with a description of the change and which issue is fixed (if applicable). Please also include relevant motivation and context. List any dependencies (if any) that are required for this change.

Once you're done, someone will review your PR shortly (see the section "Who can review?" below to tag some potential reviewers). They may suggest changes to make the code even better. If no one reviewed your PR after a week has passed, don't hesitate to post a new comment @-mentioning the same persons---sometimes notifications get lost.
-->

<!-- Remove if not applicable -->


# Explanation
Forward passes in Gemma (or Llama) with large batch sizes use a large amount of VRAM, much of it allocated for the causal mask (for example, with Gemma, I have seen 16gb allocation attempts for batch size 128 forward passes on the causal mask). 


- more trivial change: I'm pretty sure this ordering of operations is slightly faster and uses less memory with the .to()s in cases where dtype does not already match, 
- main change: switching the .repeat to a .expand since I do not expect the causal mask to usually recieve in-place edits, if true there is no need to make copies of it so .expand is preferred

this made some >50% reduction in VRAM used by Gemma in my experiences with gemma-2b-it with batch sizes 128-512 (don't have an exact measure on this)

This needs tests run on it, and I need someone more familiar with hf transformers to tell me whether there are components that expect unique memory addresses from causal mask along the batch dim (in which case this pull would cause issues). Concerned if the .repeat may have been there for a good reason :grin: 

## Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR. 
@ArthurZucker 
@younesbelkada 



Comments:
### Note on 2x calls to .to

In the original, 
```python
        causal_mask = (
            self.causal_mask[None, None, :, :]
            .repeat(batch_size, 1, 1, 1)
            .to(dtype) # <--- 1st .to
            * min_dtype
        )
        causal_mask = causal_mask.to(dtype=dtype, device=device) # <--- 2nd .to
```
I was unclear about why there were two separate .to()s, and my instinct was to change that to something like this:
```python
        causal_mask = (
            self.causal_mask[None, None, :, :]
            .repeat(batch_size, 1, 1, 1)
            .to(dtype=dtype, device=device) # <--- combined .to
            * min_dtype
        )
```
but I figure it is likely to have been written that way for a reason, so I took the cautious route and kept the pattern around in the edits, in case I'm overlooking some important consideration. So, the current PR is like this:
```python
        causal_mask = (
            (
                self.causal_mask[None, None, :, :]
                .to(dtype=dtype) # <--- 1st .to
                * min_dtype
            )
            .to(dtype=dtype, device=device) # <--- 2nd .to
            .expand(batch_size, 1, -1, -1)
        )
```

However, if the .to()s are fine to combine, I will edit it to:
```python
        causal_mask = (
            (
                self.causal_mask[None, None, :, :]
                .to(dtype=dtype, device=device) # <--- combined .to
                * min_dtype
            )
            .expand(batch_size, 1, -1, -1)
        )
```

(I manually formatted these snippets and maximally blew up the expressions, in order to have the parts that stayed the same continue to look the same. Hopefully this increased rather than decreased readability.)

Whether two `to` are needed I am not sure, let's use only one and make sure the tests pass

Sounds good, will do! Thanks for the feedback!

No failed tests on Gemma, regardless of one vs two calls to .to

The slow tests ran faster than I expected though, which mildly concerned me. So, I also checked my ~/.cache/hub to make sure that Gemma got downloaded, and I don't see it there.  Maybe that's expected though. In case not, here's outputs:

<details><summary>output on Gemma</summary>

```
root@C.9931582:~/transformers$ RUN_SLOW=1 python3 -m pytest tests/models/gemma/test_modeling_gemma.py                                               
===================================== test session starts =====================================                                                                                               
platform linux -- Python 3.10.13, pytest-7.4.4, pluggy-1.4.0                                   
rootdir: /root/transformers                                                                    
configfile: pyproject.toml                                                                                                                                                                    
plugins: hypothesis-6.98.10, xdist-3.5.0, timeout-2.2.0, dash-2.16.0, anyio-4.3.0    
configfile: pyproject.toml                                                                                                                                                           [50/1564]plugins: hypothesis-6.98.10, xdist-3.5.0, timeout-2.2.0, dash-2.16.0, anyio-4.3.0                                                                                                             
collected 141 items                                                                                                                                                                           
                                               
tests/models/gemma/test_modeling_gemma.py .................................sssssssss... [ 31%]                                                                                                
.............................ssss..ss...ssssss.sssssssssss..sssss.sss..........s......s [ 93%]                                                                                                s........                                                                               [100%]
                                                                                                                                                                                              
====================================== warnings summary =======================================                                                                                               
../../opt/conda/lib/python3.10/site-packages/_pytest/config/__init__.py:1373                   
  /opt/conda/lib/python3.10/site-packages/_pytest/config/__init__.py:1373: PytestConfigWarning: Unknown config option: doctest_glob                                                           
                                                                                                                                                                                              
    self._warn_or_fail_if_strict(f"Unknown config option: {key}\n")                                                                                                                                                                                                                          
src/transformers/deepspeed.py:23                                                                                                                                                              
  /root/transformers/src/transformers/deepspeed.py:23: FutureWarning: transformers.deepspeed module is deprecated and will be removed in a future version. Please import deepspeed modules dir
ectly from transformers.integrations                                                           
    warnings.warn(                                                                                                                                                                            
                                                                                                                                                                                              tests/models/gemma/test_modeling_gemma.py::GemmaModelTest::test_beam_sample_generate           
tests/models/gemma/test_modeling_gemma.py::GemmaModelTest::test_beam_sample_generate_dict_output                                                                                              
  /root/transformers/src/transformers/generation/utils.py:3359: UserWarning: `max_length` is deprecated in this function, use `stopping_criteria=StoppingCriteriaList([MaxLengthCriteria(max_l
ength=max_length)])` instead.                                                                  
    warnings.warn(                                                                             
                                                                                                                                                                                              tests/models/gemma/test_modeling_gemma.py::GemmaModelTest::test_beam_search_generate           
tests/models/gemma/test_modeling_gemma.py::GemmaModelTest::test_beam_search_generate_dict_output                                                                                              
tests/models/gemma/test_modeling_gemma.py::GemmaModelTest::test_beam_search_generate_dict_outputs_use_cache                                                                                   
  /root/transformers/src/transformers/generation/utils.py:2970: UserWarning: `max_length` is deprecated in this function, use `stopping_criteria=StoppingCriteriaList([MaxLengthCriteria(max_l
ength=max_length)])` instead.                                                                                                                                                                 
    warnings.warn(                                                                             
                                                                                                                                                                                              tests/models/gemma/test_modeling_gemma.py::GemmaModelTest::test_constrained_beam_search_generate                                                                                              
tests/models/gemma/test_modeling_gemma.py::GemmaModelTest::test_constrained_beam_search_generate_dict_output                                                                                  
  /root/transformers/src/transformers/generation/utils.py:4101: UserWarning: `max_length` is deprecated in this function, use `stopping_criteria=StoppingCriteriaList([MaxLengthCriteria(max_l
ength=max_length)])` instead.                                                                                                                                                                 
    warnings.warn(                                                                                                                                                                            
                                                                                               
tests/models/gemma/test_modeling_gemma.py::GemmaModelTest::test_equivalence_flax_to_pt                                                                                                          /root/transformers/src/transformers/modeling_flax_pytorch_utils.py:460: UserWarning: The given NumPy array is not writable, and PyTorch does not support non-writable tensors. This means wr
iting to this tensor will result in undefined behavior. You may want to copy the array to protect its data or make it writable before converting it to a tensor. This type of warning will be 
suppressed for the rest of this program. (Triggered internally at /opt/conda/conda-bld/pytorch_1708025847130/work/torch/csrc/utils/tensor_numpy.cpp:206.)                                     
    pt_model_dict[flax_key] = torch.from_numpy(flax_tensor)                                    
                                                                                                                                                                                              
tests/models/gemma/test_modeling_gemma.py::GemmaModelTest::test_fast_init_context_manager      
  /root/transformers/tests/test_modeling_common.py:464: FutureWarning: `torch.testing.assert_allclose()` is deprecated since 1.12 and will be removed in a future release. Please use `torch.testing.assert_close()` instead. You can find detailed upgrade instructions in https://github.com/pytorch/pytorch/issues/61844.                      
    torch.testing.assert_allclose(init_instance.linear.bias, expected_bias, rtol=1e-3, atol=1e-4)                                                                                             
                                                                                               
tests/models/gemma/test_modeling_gemma.py::GemmaModelTest::test_fast_init_context_manager      
  /root/transformers/tests/test_modeling_common.py:467: FutureWarning: `torch.testing.assert_allclose()` is deprecated since 1.12 and will be removed in a future release. Please use `torch.t
esting.assert_close()` instead. You can find detailed upgrade instructions in https://github.com/pytorch/pytorch/issues/61844.                                                                
  /root/transformers/tests/test_modeling_common.py:467: FutureWarning: `torch.testing.assert_allclose()` is deprecated since 1.12 and will be removed in a future release. Please use [0/1564]esting.assert_close()` instead. You can find detailed upgrade instructions in https://github.com/pytorch/pytorch/issues/61844.                                                                
    torch.testing.assert_allclose(                                                                                                                                                            
                                               
tests/models/gemma/test_modeling_gemma.py::GemmaModelTest::test_generate_from_inputs_embeds_decoder_only                                                                                      
  /root/transformers/src/transformers/generation/utils.py:1181: UserWarning: Using the model-agnostic default `max_length` (=20) to control the generation length. We recommend setting `max_new_tokens` to control the maximum length of the generation.                                   
    warnings.warn(                                                                                                                                                                            
                                                                                                                                                                                              
tests/models/gemma/test_modeling_gemma.py::GemmaModelTest::test_greedy_generate                
tests/models/gemma/test_modeling_gemma.py::GemmaModelTest::test_greedy_generate_dict_outputs                                                                                                  
tests/models/gemma/test_modeling_gemma.py::GemmaModelTest::test_greedy_generate_dict_outputs_use_cache                                                                                        
  /root/transformers/src/transformers/generation/utils.py:2361: UserWarning: `max_length` is deprecated in this function, use `stopping_criteria=StoppingCriteriaList([MaxLengthCriteria(max_length=max_length)])` instead.                                                                  
    warnings.warn(                                                                                                                                                                            
                                                                                                                                                                                              
tests/models/gemma/test_modeling_gemma.py::GemmaModelTest::test_group_beam_search_generate     
tests/models/gemma/test_modeling_gemma.py::GemmaModelTest::test_group_beam_search_generate_dict_output                                                                                        
  /root/transformers/src/transformers/generation/utils.py:3699: UserWarning: `max_length` is deprecated in this function, use `stopping_criteria=StoppingCriteriaList([MaxLengthCriteria(max_length=max_length)])` instead.                                                                  
    warnings.warn(                                                                                                                                                                            
                                                                                                                                                                                              
tests/models/gemma/test_modeling_gemma.py::GemmaModelTest::test_sample_generate                
tests/models/gemma/test_modeling_gemma.py::GemmaModelTest::test_sample_generate_dict_output    
  /root/transformers/src/transformers/generation/utils.py:2651: UserWarning: `max_length` is deprecated in this function, use `stopping_criteria=StoppingCriteriaList([MaxLengthCriteria(max_length=max_length)])` instead.                                                                  
    warnings.warn(                                                                                                                                                                            
                                                                                                                                                                                              
tests/models/gemma/test_modeling_gemma.py::GemmaModelTest::test_torchscript_output_attentions                                                                                                 
tests/models/gemma/test_modeling_gemma.py::GemmaModelTest::test_torchscript_output_hidden_state                                                                                               
tests/models/gemma/test_modeling_gemma.py::GemmaModelTest::test_torchscript_simple             
  /root/transformers/src/transformers/modeling_utils.py:4199: FutureWarning: `_is_quantized_training_enabled` is going to be deprecated in transformers 4.39.0. Please use `model.hf_quantizer.is_trainable` instead                                                                                                                                                                        
    warnings.warn(                                                                                                                                                                            
                                                                                                                                                                                              
tests/models/gemma/test_modeling_gemma.py::GemmaModelTest::test_torchscript_output_attentions                                                                                                 
tests/models/gemma/test_modeling_gemma.py::GemmaModelTest::test_torchscript_output_hidden_state                                                                                               
tests/models/gemma/test_modeling_gemma.py::GemmaModelTest::test_torchscript_simple             
  /root/transformers/src/transformers/models/gemma/modeling_gemma.py:969: TracerWarning: Converting a tensor to a Python boolean might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs!                                          
    if seq_length > self.causal_mask.shape[-1]:                                                                                                                                               
                                                                                                                                                                                              
tests/models/gemma/test_modeling_gemma.py::GemmaModelTest::test_torchscript_output_attentions  
tests/models/gemma/test_modeling_gemma.py::GemmaModelTest::test_torchscript_output_hidden_state                                                                                               
tests/models/gemma/test_modeling_gemma.py::GemmaModelTest::test_torchscript_simple             
  /root/transformers/src/transformers/models/gemma/modeling_gemma.py:280: TracerWarning: Converting a tensor to a Python boolean might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs!
    if attn_output.size() != (bsz, self.num_heads, q_len, self.head_dim):                                                                                                                     
                                                                                               
-- Docs: https://docs.pytest.org/en/stable/how-to/capture-warnings.html                        
======================== 98 passed, 43 skipped, 29 warnings in 25.44s =========================                                                                                               
```
</details> 



I got two errors on Llama, but it looks to be the same whether I run my branch or run the main repo, so these maybe are an environment setup issue or something?

<details><summary>Llama test summary running on main branch (type 1) </summary> 

```
================================================================================== short test summary info ===================================================================================FAILED tests/models/llama/test_modeling_llama.py::LlamaModelTest::test_eager_matches_sdpa_inference_1_bfloat16 - AssertionError: False is not true : padding_side=left, use_mask=False, batch_size=5, enable_kernels=True: mean relative difference: 2.197e-02, torch atol = 0.01, torch rtol = 0.03
FAILED tests/models/llama/test_modeling_llama.py::LlamaIntegrationTest::test_compile_static_cache - OSError: Can't load tokenizer for 'meta-llama/Llama-2-7b-hf'. If you were trying to load it from 'https://huggingface.co/models', make sure you don't have a local directory with the sam...
FAILED tests/models/llama/test_modeling_llama.py::CodeLlamaIntegrationTest::test_model_7b_logits - AssertionError: Lists differ: ['
▁<PRE> def remove_non_ascii(s: str) -> st[893 chars]ID>'] != ['
 <PRE> def remove_non_ascii(s: str) -> st[893 chars]ID>']
================================================================== 3 failed, 102 passed, 44 skipped, 29 warnings in 36.66s ===================================================================
```
</details>

<details><summary>Llama test summary running on main branch (type 2) </summary> 

```
================================================================================== short test summary info ===================================================================================
FAILED tests/models/llama/test_modeling_llama.py::LlamaIntegrationTest::test_compile_static_cache - OSError: Can't load tokenizer for 'meta-llama/Llama-2-7b-hf'. If you were trying to load it from 'https://huggingface.co/models', make sure you don't have a local directory with the sam...
FAILED tests/models/llama/test_modeling_llama.py::CodeLlamaIntegrationTest::test_model_7b_logits - AssertionError: Lists differ: ['
▁<PRE> def remove_non_ascii(s: str) -> st[893 chars]ID>'] != ['
 <PRE> def remove_non_ascii(s: str) -> st[893 chars]ID>']
================================================================== 2 failed, 103 passed, 44 skipped, 29 warnings in 35.92s ===================================================================
```
</details>

(Looks like there's a bit of nondeterminism in whether the bfloat case gives an error, doesn't seem super concerning but maybe a/r tol wants to be adjusted)



<details><summary>Llama test running on my branch with errors (matching main branch case Type 2 all 5+ times I ran the tests) </summary> 

```
================================================================================== short test summary info ===================================================================================
FAILED tests/models/llama/test_modeling_llama.py::LlamaIntegrationTest::test_compile_static_cache - OSError: Can't load tokenizer for 'meta-llama/Llama-2-7b-hf'. If you were trying to load it from 'https://huggingface.co/models', make sure you don't have a local directory with the sam...
FAILED tests/models/llama/test_modeling_llama.py::CodeLlamaIntegrationTest::test_model_7b_logits - AssertionError: Lists differ: ['
▁<PRE> def remove_non_ascii(s: str) -> st[893 chars]ID>'] != ['
 <PRE> def remove_non_ascii(s: str) -> st[893 chars]ID>']
================================================================== 2 failed, 103 passed, 44 skipped, 29 warnings in 37.14s ===================================================================
```
</details>

So, it looks like I'm having some trouble with the tokenizers. If that's not an issue that makes the rest of the tests unsound, and the fast Gemma tests are normal, then this should be good to go.

### Requesting input on comments for non-obvious operations

Following the "use comments only for non-obvious optimizations" principle, I wonder about if I should add a comment on the expand and slightly clarify the second comment, since both operations are potentially non-obvious

[link to commit with proposed change](https://github.com/nqgl/transformers/commit/4e0779e9c169fea9968a28abe1c947d2a196df19)

<details><summary>
As image:
</summary>

![image](https://github.com/huggingface/transformers/assets/35577239/9d14700e-e8d1-42ec-9522-8e4ccd1ea97a)
</details>


<details><summary>
as plain text:
</summary>


```python
        causal_mask = (
            (self.causal_mask[None, None, :, :].to(dtype=dtype, device=device) * min_dtype)
            .expand(batch_size, 1, -1, -1)  # expand instead of repeat to conserve memory
        )
        if attention_mask is not None and attention_mask.dim() == 2:
            causal_mask = causal_mask.clone()  # copy to create contiguous memory for in-place edit
```

</details>

Let's just fix the CIs! 🤗 

Great :)
Changes accepted.

Thanks for the fix!

FYI @gante, @fxmarty