Repository: keras-team/keras-nlp

Issue:
Mistral kills the process by taking too many RAM
```python
preprocessor = keras_nlp.models.MistralCausalLMPreprocessor.from_preset(
    "mistral_instruct_7b_en",
    sequence_length=128,
)
mistral_lm = keras_nlp.models.MistralCausalLM.from_preset(
    "mistral_instruct_7b_en", preprocessor=preprocessor
)

output = mistral_lm.generate("My trip to Yosemite was", max_length=64)
print("\nMistral output:")
print(output)
```

<img width="1323" alt="Screenshot 2024-02-23 at 3 01 32 AM" src="https://github.com/keras-team/keras-nlp/assets/26025527/f6972b4c-98dd-4c9d-8d8f-bcb12fd842b8">

I was running Mistral model on Colab environment w/ A100(40GB) and 80GB RAM. I loaded up the model successfully. However, when generate text, the RAM usage hit the peak, and the runtime got restarted.

Is this an expected behavior? or could there be bugs?



Discussion:
> 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.
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!

Resolving patch:
diff --git a/keras_nlp/models/mistral/mistral_causal_lm.py b/keras_nlp/models/mistral/mistral_causal_lm.py
--- a/keras_nlp/models/mistral/mistral_causal_lm.py
+++ b/keras_nlp/models/mistral/mistral_causal_lm.py
@@ -190,6 +190,7 @@ def next(prompt, cache, index):
             mask=padding_mask,
             end_token_id=end_token_id,
             hidden_states=hidden_states,
+            model=self,
         )
 
         # Compute an output padding mask with the token ids we updated.
diff --git a/keras_nlp/models/mistral/mistral_presets.py b/keras_nlp/models/mistral/mistral_presets.py
--- a/keras_nlp/models/mistral/mistral_presets.py
+++ b/keras_nlp/models/mistral/mistral_presets.py
@@ -23,7 +23,7 @@
             "path": "mistral",
             "model_card": "https://github.com/mistralai/mistral-src/blob/main/README.md",
         },
-        "kaggle_handle": "kaggle://keras/mistral/keras/mistral_7b_en/3",
+        "kaggle_handle": "kaggle://keras/mistral/keras/mistral_7b_en/6",
     },
     "mistral_instruct_7b_en": {
         "metadata": {
@@ -33,6 +33,6 @@
             "path": "mistral",
             "model_card": "https://github.com/mistralai/mistral-src/blob/main/README.md",
         },
-        "kaggle_handle": "kaggle://keras/mistral/keras/mistral_instruct_7b_en/3",
+        "kaggle_handle": "kaggle://keras/mistral/keras/mistral_instruct_7b_en/6",
     },
 }
diff --git a/tools/checkpoint_conversion/convert_mistral_checkpoints.py b/tools/checkpoint_conversion/convert_mistral_checkpoints.py
--- a/tools/checkpoint_conversion/convert_mistral_checkpoints.py
+++ b/tools/checkpoint_conversion/convert_mistral_checkpoints.py
@@ -11,14 +11,12 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-import datetime
 import gc
-import json
 import os
-import pathlib
+import shutil
+import tempfile
 import traceback
 
-import keras
 import numpy as np
 import requests
 from absl import app
@@ -27,10 +25,10 @@
 from transformers import AutoTokenizer
 from transformers import MistralForCausalLM
 
-import keras_nlp
 from keras_nlp.models import MistralBackbone
 from keras_nlp.models import MistralCausalLMPreprocessor
 from keras_nlp.models import MistralTokenizer
+from keras_nlp.utils.preset_utils import save_to_preset
 
 PRESET_MAP = {
     "mistral_7b_en": "mistralai/Mistral-7B-v0.1",
@@ -227,124 +225,74 @@ def main(_):
     preset = FLAGS.preset
     hf_preset = PRESET_MAP[preset]
 
-    # === Create the save directories ===
-    model_dir = pathlib.Path(__file__).parent / f"{preset}"
-    tokenizer_dir = model_dir / "assets" / "tokenizer"
-    if not model_dir.exists():
-        os.makedirs(model_dir)
-    if not tokenizer_dir.exists():
-        os.makedirs(tokenizer_dir)
+    # === Create the temporary save directories ===
+    temp_dir = tempfile.mkdtemp()
 
-    # === Load the Huggingface model ===
-    hf_model = MistralForCausalLM.from_pretrained(hf_preset)
-    hf_tokenizer = AutoTokenizer.from_pretrained(hf_preset)
-    hf_model.eval()
-    print("\n-> Huggingface model and tokenizer loaded")
-
-    # === Load the KerasNLP model ===
-    keras_nlp_config = dict(
-        vocabulary_size=hf_model.config.vocab_size,
-        hidden_dim=hf_model.config.hidden_size,
-        num_layers=hf_model.config.num_hidden_layers,
-        num_query_heads=hf_model.config.num_attention_heads,
-        num_key_value_heads=hf_model.config.num_key_value_heads,
-        intermediate_dim=hf_model.config.intermediate_size,
-        sliding_window=hf_model.config.sliding_window,
-        layer_norm_epsilon=hf_model.config.rms_norm_eps,
-        rope_max_wavelength=hf_model.config.rope_theta,
-        dtype="float32",
-    )
-    keras_nlp_model = MistralBackbone(**keras_nlp_config)
-
-    # === Download the tokenizer from Huggingface model card ===
-    spm_path = (
-        f"https://huggingface.co/{hf_preset}/resolve/main/tokenizer.model"
-    )
-    response = requests.get(spm_path)
-    if not response.ok:
-        raise ValueError(f"Couldn't fetch {preset}'s tokenizer.")
-    tokenizer_path = tokenizer_dir / "vocabulary.spm"
-    with open(tokenizer_path, "wb") as tokenizer_file:
-        tokenizer_file.write(response.content)
-    keras_nlp_tokenizer = MistralTokenizer(str(tokenizer_path.absolute()))
-    print("\n-> Keras 3 model and tokenizer loaded.")
-
-    # === Port the weights ===
-    convert_checkpoints(keras_nlp_model, hf_model)
-    print("\n-> Weight transfer done.")
-
-    # === Check that the models and tokenizers outputs match ===
-    test_tokenizer(keras_nlp_tokenizer, hf_tokenizer)
-    test_model(keras_nlp_model, keras_nlp_tokenizer, hf_model, hf_tokenizer)
-    print("\n-> Tests passed!")
-
-    # === Save the model weights in float32 format ===
-    keras_nlp_model.save_weights(
-        str((model_dir / "model.weights.h5").absolute())
-    )
-    print("\n-> Saved the model weights in float16")
-
-    del keras_nlp_model, hf_model
-    gc.collect()
-
-    keras_nlp_config["dtype"] = "float16"
-
-    # === Save the weights again in float16 ===
-    keras_nlp_model = MistralBackbone(**keras_nlp_config)
-    keras_nlp_model.load_weights(
-        str((model_dir / "model.weights.h5").absolute())
-    )
-    keras_nlp_model.save_weights(
-        str((model_dir / "model.weights.h5").absolute())
-    )
-    print("-> Saved the model weights in float16")
-
-    # === Save the model config ===
-    keras_nlp_config["dtype"] = "bfloat16"
-    model_config = {
-        "module": "keras_nlp.src.models.mistral.mistral_backbone",
-        "class_name": "MistralBackbone",
-        "config": {**keras_nlp_config},
-        "registered_name": "keras_nlp>MistralBackbone",
-        "assets": [],
-        "weights": "model.weights.h5",
-    }
-    model_config_json = json.dumps(model_config)
-    with open(model_dir / "config.json", "w") as model_config_file:
-        model_config_file.write(model_config_json)
-    print("\n-> Saved model config")
-
-    # === Save the tokenizer config ===
-    tokenizer_config = {
-        "module": "keras_nlp.src.models.mistral.Mistral_tokenizer",
-        "class_name": "MistralTokenizer",
-        "config": {
-            "name": "mistral_tokenizer",
-            "trainable": True,
-            "dtype": "int32",
-            "proto": None,
-            "sequence_length": None,
-        },
-        "registered_name": "keras_nlp>MistralTokenizer",
-        "assets": ["assets/tokenizer/vocabulary.spm"],
-        "weights": None,
-    }
-    tokenizer_config_json = json.dumps(tokenizer_config)
-    with open(model_dir / "tokenizer.json", "w") as tokenizer_config_file:
-        tokenizer_config_file.write(tokenizer_config_json)
-    print("\n-> Saved tokenizer config")
+    try:
+        # === Load the Huggingface model ===
+        hf_model = MistralForCausalLM.from_pretrained(hf_preset)
+        hf_tokenizer = AutoTokenizer.from_pretrained(hf_preset)
+        hf_model.eval()
+        print("\n-> Huggingface model and tokenizer loaded")
+
+        # === Load the KerasNLP model ===
+        backbone_kwargs = dict(
+            vocabulary_size=hf_model.config.vocab_size,
+            hidden_dim=hf_model.config.hidden_size,
+            num_layers=hf_model.config.num_hidden_layers,
+            num_query_heads=hf_model.config.num_attention_heads,
+            num_key_value_heads=hf_model.config.num_key_value_heads,
+            intermediate_dim=hf_model.config.intermediate_size,
+            sliding_window=hf_model.config.sliding_window,
+            layer_norm_epsilon=hf_model.config.rms_norm_eps,
+            rope_max_wavelength=hf_model.config.rope_theta,
+            dtype="float32",
+        )
+        keras_nlp_model = MistralBackbone(**backbone_kwargs)
 
-    # === Save metadata ===
-    metadata_config = {
-        "keras_version": keras.__version__,
-        "keras_nlp_version": keras_nlp.__version__,
-        "parameter_count": keras_nlp_model.count_params(),
-        "date_saved": datetime.datetime.utcnow().strftime("%Y-%m-%d@%H:%M:%S"),
-    }
-    metadata_config_json = json.dumps(metadata_config)
-    with open(model_dir / "metadata.json", "w") as metadata_config_file:
-        metadata_config_file.write(metadata_config_json)
-    print("\n-> Saved metadata")
+        # === Download the tokenizer from Huggingface model card ===
+        spm_path = (
+            f"https://huggingface.co/{hf_preset}/resolve/main/tokenizer.model"
+        )
+        response = requests.get(spm_path)
+        if not response.ok:
+            raise ValueError(f"Couldn't fetch {preset}'s tokenizer.")
+        tokenizer_path = os.path.join(temp_dir, "vocabulary.spm")
+        with open(tokenizer_path, "wb") as tokenizer_file:
+            tokenizer_file.write(response.content)
+        keras_nlp_tokenizer = MistralTokenizer(tokenizer_path)
+        print("\n-> Keras 3 model and tokenizer loaded.")
+
+        # === Port the weights ===
+        convert_checkpoints(keras_nlp_model, hf_model)
+        print("\n-> Weight transfer done.")
+
+        # === Check that the models and tokenizers outputs match ===
+        test_tokenizer(keras_nlp_tokenizer, hf_tokenizer)
+        test_model(keras_nlp_model, keras_nlp_tokenizer, hf_model, hf_tokenizer)
+        print("\n-> Tests passed!")
+
+        # === Save the model weights in float32 format ===
+        keras_nlp_model.save_weights(os.path.join(temp_dir, "model.weights.h5"))
+        print("\n-> Saved the model weights in float32")
+
+        del keras_nlp_model, hf_model
+        gc.collect()
+
+        # === Save the weights again in float16 ===
+        backbone_kwargs["dtype"] = "float16"
+        keras_nlp_model = MistralBackbone(**backbone_kwargs)
+        keras_nlp_model.load_weights(os.path.join(temp_dir, "model.weights.h5"))
+        save_to_preset(keras_nlp_model, preset)
+        print("\n-> Saved the model preset in float16")
+
+        # === Save the tokenizer ===
+        save_to_preset(
+            keras_nlp_tokenizer, preset, config_filename="tokenizer.json"
+        )
+        print("\n-> Saved the tokenizer")
+    finally:
+        shutil.rmtree(temp_dir)
 
 
 if __name__ == "__main__":


Test patch: