Mistral in Flax: generation is slow, JIT fails

### 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 (cpu) - Jax version: 0.4.25 - JaxLib version: 0.4.25 - Using GPU in script?: Yes (JAX default behavior) - Using distributed or parallel set-up in script?: No ### 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 On a VM/Docker with NVIDIA A100 run: ```python import jax import jax.numpy as jnp from transformers import FlaxAutoModelForCausalLM, AutoTokenizer MODEL_ID = "mistralai/Mistral-7B-Instruct-v0.2" model = FlaxAutoModelForCausalLM.from_pretrained( MODEL_ID, from_pt=True, dtype=jnp.bfloat16, max_position_embeddings=4096, # much smaller than the default value sliding_window=4096 ) tokenizer = AutoTokenizer.from_pretrained(MODEL_ID) texts = ["<s>[INST]Write a poem[/INST]"] input_ids = tokenizer(texts, return_tensors="np")["input_ids"] ``` With this setup, I'm able to generate some output using: ```python model.generate(input_ids, max_new_tokens=32) ``` But it takes 8.91s (after warmup) - longer than what I'd expect for total of 45 tokens. Obvious next step is to JIT-compile it: ```python jax.jit(model.generate, static_argnames=("max_new_tokens",))(input_ids, max_new_tokens=32) ``` But it fails with: ``` 0302 22:21:58.956701 19411 pjrt_stream_executor_client.cc:2804] Execution of replica 0 failed: INTERNAL: Failed to allocate 117440512 bytes for new constant -------------------------------------------------------------…