Mistral kills the process by taking too many RAM

[deep-diver]

Mistral kills the process by taking too many RAM

[tirthasheshpatel]

Is this an expected behavior? or could there be bugs?

Mistral loads in bfloat16 by default. I noticed this caused issues in Colab with the JAX backbend (TensorFlow runs fine though with 16.5 GB of RAM).

The workaround to use the dtype set using keras.mixed_precision module is to pass dtype=None to the from_preset method:

mistral_lm = keras_nlp.models.MistralCausalLM.from_preset('mistral_instruct_7b_en', preprocessor=preprocessor, dtype=None)

Let me know if this lowers the RAM usage. This will be fixed in the next release.

[mattdangerw]

Thanks for the bug! Just synced up with @tirthasheshpatel. We want to change two things here

1. By default, mistral should follow global keras default settings. So keras.mixed_precision.set_global_policy("mixed...") -> variables load as float32. keras.config.set_floatx("bfloat16") -> variables load at bfloat16.
2. There is a bug with the jax backend only, where generation for mistral is consuming significantly too much CPU and GPU memory. It's a one liner fix on our side I think.

These are both simple but important fixes, we should have a patch fix for this in a couple days. Thanks @deep-diver!