Mistral 7B doesn't load in Flax (sliding_window + RESOURCE_EXHAUSTED)
Author: dfdx
Pull request: no

### System Info

- `transformers` version: 4.38.1
- Platform: Linux-6.2.0-1019-azure-x86_64-with-glibc2.35
- Python version: 3.10.12
- Huggingface_hub version: 0.21.1
- Safetensors version: 0.4.2
- Accelerate version: 0.27.2
- Accelerate config:    not found
- PyTorch version (GPU?): 2.2.1+cu121 (True)
- Tensorflow version (GPU?): not installed (NA)
- Flax version (CPU?/GPU?/TPU?): 0.8.1 (gpu)
- Jax version: 0.4.25
- JaxLib version: 0.4.25
- Using GPU in script?: Yes (JAX's default behavior)
- Using distributed or parallel set-up in script?: No

GPU specs: NVIDIA-SMI 525.147.05   Driver Version: 525.147.05   CUDA Version: 12.0

### Who can help?

@sanchit-gandhi

### 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

Try to load Mistral 7B in Flax on A100, e.g.:

```python
import jax.numpy as jnp
from transformers import FlaxAutoModelForCausalLM

MODEL_ID = "mistralai/Mistral-7B-Instruct-v0.2"
model = FlaxAutoModelForCausalLM.from_pretrained(MODEL_ID, from_pt=True) 
```
Which gives `TypeError: bad operand type for unary -: 'NoneType'`:

```python
File ~/.pyenv/versions/3.10.6/lib/python3.10/site-packages/transformers/models/mistral/modeling_flax_mistral.py:242, in FlaxMistralAttention.setup(self)
    240 self.o_proj = nn.Dense(self.hidden_size, use_bias=False, dtype=self.dtype)
    241 casual_mask = make_causal_mask(jnp.ones((1, config.max_position_embeddings), dtype="bool"), dtype="bool")
--> 242 self.causal_mask = jnp.triu(casual_mask, k=-config.sliding_window)
    243 self.rotary_emb = FlaxMistralRotaryEmbedding(config, dtype=self.dtype)

TypeError: bad operand type for unary -: 'NoneType'
```

Setting `sliding_window` fixes this error, but the model still doesn't load due to OOM:

```python
# setting sliding_window to just 256 for test
model = FlaxAutoModelForCausalLM.from_pretrained(MODEL_ID, from_pt=True, sliding_window=256) 
```
which gives:

```
ValueError                                Traceback (most recent call last)
Cell In[3], line 1
----> 1 model = FlaxAutoModelForCausalLM.from_pretrained(MODEL_ID, from_pt=True, sliding_window=256)

File ~/.pyenv/versions/3.10.6/lib/python3.10/site-packages/transformers/models/auto/auto_factory.py:561, in _BaseAutoModelClass.from_pretrained(cls, pretrained_model_name_or_path, *model_args, **kwargs)
    559 elif type(config) in cls._model_mapping.keys():
    560     model_class = _get_model_class(config, cls._model_mapping)
--> 561     return model_class.from_pretrained(
    562         pretrained_model_name_or_path, *model_args, config=config, **hub_kwargs, **kwargs
    563     )
    564 raise ValueError(
    565     f"Unrecognized configuration class {config.__class__} for this kind of AutoModel: {cls.__name__}.\n"
    566     f"Model type should be one of {', '.join(c.__name__ for c in cls._model_mapping.keys())}."
    567 )

File ~/.pyenv/versions/3.10.6/lib/python3.10/site-packages/transformers/modeling_flax_utils.py:900, in FlaxPreTrainedModel.from_pretrained(cls, pretrained_model_name_or_path, dtype, config, cache_dir, ignore_mismatched_sizes, force_download, local_files_only, token, revision, *model_args, **kwargs)
    897     safetensors_from_pt = safetensors_metadata.get("format") == "pt"
    899 # init random models
--> 900 model = cls(config, *model_args, _do_init=_do_init, **model_kwargs)
    902 if from_pt or safetensors_from_pt:
    903     state = load_pytorch_checkpoint_in_flax_state_dict(model, resolved_archive_file, is_sharded)

File ~/.pyenv/versions/3.10.6/lib/python3.10/site-packages/transformers/models/mistral/modeling_flax_mistral.py:416, in FlaxMistralPreTrainedModel.__init__(self, config, input_shape, seed, dtype, _do_init, **kwargs)
    406 def __init__(
    407     self,
    408     config: MistralConfig,
   (...)
    413     **kwargs,
    414 ):
    415     module = self.module_class(config=config, dtype=dtype, **kwargs)
--> 416     super().__init__(config, module, input_shape=input_shape, seed=seed, dtype=dtype, _do_init=_do_init)

File ~/.pyenv/versions/3.10.6/lib/python3.10/site-packages/transformers/modeling_flax_utils.py:219, in FlaxPreTrainedModel.__init__(self, config, module, input_shape, seed, dtype, _do_init)
    215 self._is_initialized = _do_init
    217 if _do_init:
    218     # randomly initialized parameters
--> 219     random_params = self.init_weights(self.key, input_shape)
    220     params_shape_tree = jax.eval_shape(lambda params: params, random_params)
    221 else:

File ~/.pyenv/versions/3.10.6/lib/python3.10/site-packages/transformers/models/mistral/modeling_flax_mistral.py:426, in FlaxMistralPreTrainedModel.init_weights(self, rng, input_shape, params)
    423 params_rng, dropout_rng = jax.random.split(rng)
    424 rngs = {"params": params_rng, "dropout": dropout_rng}
--> 426 random_params = self.module.init(rngs, input_ids, attention_mask, position_ids, return_dict=False)["params"]
    428 if params is not None:
    429     random_params = flatten_dict(unfreeze(random_params))

    [... skipping hidden 9 frame]

File ~/.pyenv/versions/3.10.6/lib/python3.10/site-packages/transformers/models/mistral/modeling_flax_mistral.py:677, in FlaxMistralForCausalLMModule.__call__(self, input_ids, attention_mask, position_ids, deterministic, init_cache, output_attentions, output_hidden_states, return_dict)
    666 def __call__(
    667     self,
    668     input_ids,
   (...)
    675     return_dict: bool = True,
    676 ):
--> 677     outputs = self.model(
    678         input_ids,
    679         position_ids=position_ids,
    680         attention_mask=attention_mask,
    681         deterministic=deterministic,
    682         init_cache=init_cache,
    683         output_attentions=output_attentions,
    684         output_hidden_states=output_hidden_states,
    685         return_dict=return_dict,
    686     )
    688     hidden_states = outputs[0]
    689     lm_logits = self.lm_head(hidden_states)

    [... skipping hidden 2 frame]

File ~/.pyenv/versions/3.10.6/lib/python3.10/site-packages/transformers/models/mistral/modeling_flax_mistral.py:603, in FlaxMistralModule.__call__(self, input_ids, attention_mask, position_ids, deterministic, init_cache, output_attentions, output_hidden_states, return_dict)
    592 def __call__(
    593     self,
    594     input_ids,
   (...)
    601     return_dict: bool = True,
    602 ):
--> 603     input_embeds = self.embed_tokens(input_ids.astype("i4"))
    605     outputs = self.layers(
    606         input_embeds,
    607         position_ids=position_ids,
   (...)
    613         return_dict=return_dict,
    614     )
    616     hidden_states = outputs[0]

    [... skipping hidden 5 frame]

File ~/.pyenv/versions/3.10.6/lib/python3.10/site-packages/flax/linen/linear.py:976, in Embed.setup(self)
    975 def setup(self):
--> 976   self.embedding = self.param(
    977     'embedding',
    978     self.embedding_init,
    979     (self.num_embeddings, self.features),
    980     self.param_dtype,
    981   )

    [... skipping hidden 2 frame]

File ~/.pyenv/versions/3.10.6/lib/python3.10/site-packages/jax/_src/nn/initializers.py:159, in normal.<locals>.init(key, shape, dtype)
    155 def init(key: KeyArray,
    156          shape: core.Shape,
    157          dtype: DTypeLikeInexact = dtype) -> Array:
    158   dtype = dtypes.canonicalize_dtype(dtype)
--> 159   return random.normal(key, shape, dtype) * stddev

File ~/.pyenv/versions/3.10.6/lib/python3.10/site-packages/jax/_src/random.py:720, in normal(key, shape, dtype)
    718 dtype = dtypes.canonicalize_dtype(dtype)
    719 shape = core.as_named_shape(shape)
--> 720 return _normal(key, shape, dtype)

ValueError: RESOURCE_EXHAUSTED: Out of memory while trying to allocate 1310720384 bytes.
BufferAssignment OOM Debugging.
BufferAssignment stats:
             parameter allocation:         8B
              constant allocation:         0B
        maybe_live_out allocation:  500.00MiB
     preallocated temp allocation:    1.22GiB
  preallocated temp fragmentation:         0B (0.00%)
                 total allocation:    1.71GiB
              total fragmentation:  250.00MiB (14.29%)
Peak buffers:
        Buffer 1:
                Size: 500.00MiB
                Operator: op_name="jit(_normal)/jit(main)/jit(_normal_real)/mul" source_file="<ipython-input-1-6f42fc5de7b5>" source_line=5
                XLA Label: fusion
                Shape: f32[32000,4096]
                ==========================

        Buffer 2:
                Size: 250.00MiB
                Operator: op_name="jit(_normal)/jit(main)/jit(_normal_real)/jit(_uniform)/threefry2x32" source_file="<ipython-input-1-6f42fc5de7b5>" source_line=5
                XLA Label: custom-call
                Shape: u32[65536000]
                ==========================

        Buffer 3:
                Size: 250.00MiB
                Operator: op_name="jit(_normal)/jit(main)/jit(_normal_real)/jit(_uniform)/threefry2x32" source_file="<ipython-input-1-6f42fc5de7b5>" source_line=5
                XLA Label: custom-call
                Shape: u32[65536000]
                ==========================

        Buffer 4:
                Size: 250.00MiB
                XLA Label: fusion
                Shape: u32[65536000]
                ==========================

        Buffer 5:
                Size: 250.00MiB
                XLA Label: fusion
                Shape: u32[65536000]
                ==========================

        Buffer 6:
                Size: 250.00MiB
                Operator: op_name="jit(_normal)/jit(main)/jit(_normal_real)/jit(_uniform)/threefry2x32" source_file="<ipython-input-1-6f42fc5de7b5>" source_line=5
                XLA Label: fusion
                Shape: u32[65536000]
                ==========================

        Buffer 7:
                Size: 16B
                Operator: op_name="jit(_normal)/jit(main)/jit(_normal_real)/jit(_uniform)/threefry2x32" source_file="<ipython-input-1-6f42fc5de7b5>" source_line=5
                XLA Label: fusion
                Shape: (u32[65536000], u32[65536000])
                ==========================

        Buffer 8:
                Size: 16B
                Operator: op_name="jit(_normal)/jit(main)/jit(_normal_real)/jit(_uniform)/threefry2x32" source_file="<ipython-input-1-6f42fc5de7b5>" source_line=5
                XLA Label: custom-call
                Shape: (u32[65536000], u32[65536000])
                ==========================

        Buffer 9:
                Size: 16B
                XLA Label: fusion
                Shape: (u32[65536000], u32[65536000])
                ==========================

        Buffer 10:
                Size: 8B
                Entry Parameter Subshape: u32[2]
                ==========================
```

### Expected behavior

The model is loaded onto NVIDIA A100 GPU with 80Gb of VRAM. 

Comments:
Ok, it's all about default arguments. The following works:

```python
model = FlaxAutoModelForCausalLM.from_pretrained(
        "mistralai/Mistral-7B-Instruct-v0.2",
        from_pt=True,
        dtype=jnp.bfloat16,
        max_position_embeddings=4096,
        sliding_window=4096
    )
```