[Port] TensorFlow implementation of Mistral
Author: ariG23498
Pull request: yes
Labels: name: TensorFlow

name: New model

name: single-model-run-slow

name: run-slow

# What does this PR do?

This PR adds the TensorFlow implementation of Mistral.

## Before submitting
- [x] Did you read the [contributor guideline](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md#create-a-pull-request),
      Pull Request section?
- [x] Did you make sure to update the documentation with your changes? Here are the
      [documentation guidelines](https://github.com/huggingface/transformers/tree/main/docs), and
      [here are tips on formatting docstrings](https://github.com/huggingface/transformers/tree/main/docs#writing-source-documentation).


## Who can review?

@Rocketknight1


Comments:
@Rocketknight1 at this stage if I run the following
```python
from transformers.models.mistral import TFMistralModel

model = TFMistralModel.from_pretrained("mistralai/Mistral-7B-v0.1", from_pt=True)
```

I get the following stack trace.

```
Traceback (most recent call last):
  File "/home/arig23498/Documents/git-repos/transformers/demo.py", line 3, in <module>
    model = TFMistralModel.from_pretrained("mistralai/Mistral-7B-v0.1", from_pt=True)
  File "/home/arig23498/Documents/git-repos/transformers/src/transformers/modeling_tf_utils.py", line 2889, in from_pretrained
    return load_pytorch_checkpoint_in_tf2_model(
  File "/home/arig23498/Documents/git-repos/transformers/src/transformers/modeling_tf_pytorch_utils.py", line 198, in load_pytorch_checkpoint_in_tf2_model
    return load_pytorch_weights_in_tf2_model(
  File "/home/arig23498/Documents/git-repos/transformers/src/transformers/modeling_tf_pytorch_utils.py", line 238, in load_pytorch_weights_in_tf2_model
    pt_state_dict = {k: v.numpy() for k, v in pt_state_dict.items()}
  File "/home/arig23498/Documents/git-repos/transformers/src/transformers/modeling_tf_pytorch_utils.py", line 238, in <dictcomp>
    pt_state_dict = {k: v.numpy() for k, v in pt_state_dict.items()}
TypeError: Got unsupported ScalarType BFloat16
```

What am I missing here?

Hi @ariG23498, the cause there is most likely that Numpy doesn't support `bfloat16` dtypes, and so the code fails because there is no direct conversion from Torch -> TF, it has to go through a Numpy array.

We'll probably need a patch in `transformers` for this - I'll see if I can do it quickly!

PR is open at #29755 

PR merged! If you rebase, those loads should now work.

I was unable to test the port with the 7B model and reached out to @Rocketknight1 for a fix. He mentioned it would be wise to reduce the model size and host a random tiny model on hf hub and use the model to check with the initial port.

The code that I have used (as guided by @Rocketknight1)
```python
from transformers import AutoModel, AutoConfig
# Initializing a Mistral 7B style configuration
configuration = AutoConfig.from_pretrained("mistralai/Mistral-7B-v0.1")
# Reduce the size of the model
configuration.num_hidden_layers = 2
# Initializing a tiny model from the reduced configuration
model = AutoModel.from_config(configuration)
# Push the model to hub
model.push_to_hub("tiny-random-mistral")
```

The `tiny-random-mistral` is hosted [here](ariG23498/tiny-random-mistral).

Using the tiny model I was able to by-pass the OOM and the process killed issues.

Note: Also having rebased the PR I was able to use the *safetensors* and *bfloat16* weights inside the TF implementation. :partying_face:

Running the following code:
```python
import os
os.environ["CUDA_VISIBLE_DEVICES"] = "-1"

from transformers.models.mistral import TFMistralForCausalLM
model = TFMistralForCausalLM.from_pretrained("ariG23498/tiny-random-mistral-for-causal-lm")
```

Results in the follwoing naming error:

```shell
Some weights of the PyTorch model were not used when initializing the TF 2.0 model TFMistralForCausalLM: ['model.layers.1.mlp.down_proj.weight', 'model.layers.0.input_layernorm.weight', 'model.layers.1.self_attn.q_proj.weight', 'model.layers.1.self_attn.v_proj.weight', 'model.layers.1.self_attn.o_proj.weight', 'model.layers.1.mlp.gate_proj.weight', 'model.layers.0.mlp.up_proj.weight', 'model.layers.0.self_attn.k_proj.weight', 'model.layers.0.mlp.down_proj.weight', 'model.layers.0.self_attn.o_proj.weight', 'model.layers.1.input_layernorm.weight', 'lm_head.weight', 'model.layers.0.self_attn.q_proj.weight', 'model.layers.1.self_attn.k_proj.weight', 'model.layers.0.mlp.gate_proj.weight', 'model.norm.weight', 'model.layers.0.post_attention_layernorm.weight', 'model.layers.1.mlp.up_proj.weight', 'model.layers.0.self_attn.v_proj.weight', 'model.layers.1.post_attention_layernorm.weight']
- This IS expected if you are initializing TFMistralForCausalLM from a PyTorch model trained on another task or with another architecture (e.g. initializing a TFBertForSequenceClassification model from a BertForPreTraining model).
- This IS NOT expected if you are initializing TFMistralForCausalLM from a PyTorch model that you expect to be exactly identical (e.g. initializing a TFBertForSequenceClassification model from a BertForSequenceClassification model).
Some weights or buffers of the TF 2.0 model TFMistralForCausalLM were not initialized from the PyTorch model and are newly initialized: ['weight', 'weight', 'weight', 'weight', 'weight']
You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.
```

I have made sure that I add the names of each layer/module in the TF port. @Rocketknight1 do you think I am missing something here?

Note: I have commented out the `build()` methods containing *dense* layers. I need to inspect the shapes better to un-comment that section. If I use `[None, Non, units]` as the input shape to the dense layers, it results in a mismatched shape for the weights.

@ariG23498 yes, errors like these almost always indicate that weights haven't been built! My guess is that since a lot of weights in each layer are missing, the problem is one of the outer build() methods. Remember, the build() methods have to nest correctly - when you call the outermost `model.build()` or `model.build_in_name_scope()`, that should call all of the build methods for sublayers too! Take a look at any of the other TF models in the codebase if you want to see examples.

Also, when building `Dense` layers, remember that the argument you need to pass is the layer input dimension, not the output dimension. You don't need a full shape, so just `(input_dim,)` is good enough. If you're getting shape errors, try putting a `breakpoint()` before the layer is called in `call()` to see what shape is actually being passed to it, and use that knowledge to set its `build()` method.

@Rocketknight1 I was able to port all the weights from PyTorch to TensorFlow following your advice. I wanted your opinion on a small modification that I have made.

```diff
class TFMistralRMSNorm(tf.keras.layers.Layer):
    def __init__(self, hidden_size, eps=1e-6, **kwargs):
        """
        TFMistralRMSNorm is equivalent to T5LayerNorm
        """
        super().__init__(**kwargs)
        self.hidden_size = hidden_size
        self.variance_epsilon = eps

    def build(self, input_shape=None):
-        with tf.name_scope("weight"):
        self.weight = self.add_weight(
            name="weight",
            shape=self.hidden_size,
            initializer="ones",
        )
        if self.built:
            return
        self.built = True

    def call(self, hidden_states):
        input_dtype = hidden_states.dtype
        hidden_states = tf.cast(hidden_states, tf.float32)
        variance = tf.reduce_mean(tf.square(hidden_states), axis=-1, keepdims=True)
        hidden_states = tf.divide(hidden_states, tf.sqrt(variance + self.variance_epsilon))
        return self.weight * tf.cast(hidden_states, input_dtype)
```
Removing the named scope helped in building the varaibles with the same name as that of PyTorch. Adding the named scope made the variable names `<layer_name>.weight.weight`. WDYT?

@ariG23498 Yes, correct! When you `self.add_weight()` in `build()` then you don't need to define an extra name scope with the same name.

Hi @Rocketknight1 currently I think the `TFMistralForCausalLM` needs some review. I am unable to get this part aligned with the PyTorch implementation.

Could you highlight something that is very evident to you -- which might be failing the tests?

@Rocketknight1 I tried to investigate the failing `test_save_load` test.

Modifying the test like so
```py
    def test_save_load(self):
        config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()

        for model_class in self.all_model_classes:
            model1 = model_class(config, name="model1")
            outputs = model1(self._prepare_for_class(inputs_dict, model_class))

            with tempfile.TemporaryDirectory() as tmpdirname:
                model1.save_pretrained(tmpdirname, saved_model=False)

                # the config file (and the generation config file, if it can generate) should be saved
                self.assertTrue(os.path.exists(os.path.join(tmpdirname, CONFIG_NAME)))
                self.assertEqual(
                    model1.can_generate(), os.path.exists(os.path.join(tmpdirname, GENERATION_CONFIG_NAME))
                )

                model2 = model_class.from_pretrained(tmpdirname, name="model2")
                after_outputs = model2(self._prepare_for_class(inputs_dict, model_class))

                for model1_layers, model2_layers in zip(model1.layers, model2.layers):
                    for layers1, layers2 in zip(model1_layers.layers, model2_layers.layers):
                        for var1, var2 in zip(layers1.trainable_variables, layers2.trainable_variables):
                            if not tf.reduce_all(tf.equal(var1, var2)):
                                print(f"not equal {var1.name} {var2.name}")
                            else:
                                print(f"equal {var1.name} {var2.name}")

                self.assert_outputs_same(after_outputs, outputs)
```

Outputs
```shell
not equal model1/model/layers.0/self_attn/q_proj/kernel:0 model2/model/layers.0/self_attn/q_proj/kernel:0
not equal model1/model/layers.0/self_attn/k_proj/kernel:0 model2/model/layers.0/self_attn/k_proj/kernel:0
not equal model1/model/layers.0/self_attn/v_proj/kernel:0 model2/model/layers.0/self_attn/v_proj/kernel:0
not equal model1/model/layers.0/self_attn/o_proj/kernel:0 model2/model/layers.0/self_attn/o_proj/kernel:0
not equal model1/model/layers.0/mlp/gate_proj/kernel:0 model2/model/layers.0/mlp/gate_proj/kernel:0
not equal model1/model/layers.0/mlp/up_proj/kernel:0 model2/model/layers.0/mlp/up_proj/kernel:0
not equal model1/model/layers.0/mlp/down_proj/kernel:0 model2/model/layers.0/mlp/down_proj/kernel:0
equal model1/model/layers.0/input_layernorm/weight:0 model2/model/layers.0/input_layernorm/weight:0
equal model1/model/layers.0/post_attention_layernorm/weight:0 model2/model/layers.0/post_attention_layernorm/weight:0
not equal model1/model/layers.1/self_attn/q_proj/kernel:0 model2/model/layers.1/self_attn/q_proj/kernel:0
not equal model1/model/layers.1/self_attn/k_proj/kernel:0 model2/model/layers.1/self_attn/k_proj/kernel:0
not equal model1/model/layers.1/self_attn/v_proj/kernel:0 model2/model/layers.1/self_attn/v_proj/kernel:0
not equal model1/model/layers.1/self_attn/o_proj/kernel:0 model2/model/layers.1/self_attn/o_proj/kernel:0
not equal model1/model/layers.1/mlp/gate_proj/kernel:0 model2/model/layers.1/mlp/gate_proj/kernel:0
not equal model1/model/layers.1/mlp/up_proj/kernel:0 model2/model/layers.1/mlp/up_proj/kernel:0
not equal model1/model/layers.1/mlp/down_proj/kernel:0 model2/model/layers.1/mlp/down_proj/kernel:0
equal model1/model/layers.1/input_layernorm/weight:0 model2/model/layers.1/input_layernorm/weight:0
equal model1/model/layers.1/post_attention_layernorm/weight:0 model2/model/layers.1/post_attention_layernorm/weight:0
```

Do you think we should move into investigative the `save_pretrained` and `from_pretrained` methods?

Documenting the fix:
The problem was with missign `self.built=Ture` in two layer build methods. This resulted in building the weights from scratch without loading it from the saved model.

@Rocketknight1 all tests pass. I am ready for review!

Hi @ariG23498, we're almost there! I just ran the doctests and there's a couple of small failures, but it should just be small stuff like missing return values in the docstrings - you can see the logs in the `Build PR Documentation` CI.

@Rocketknight1 I have made the updates. The tests pass. :heavy_check_mark: 

The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/transformers/pr_29708). All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

Looks clean to me now! Is there anything else you want to finish up, or are you happy for me to merge it?

> Looks clean to me now! Is there anything else you want to finish up, or are you happy for me to merge it?

@Rocketknight1  You don't want to use the feature PR slow CI ...? 😭 

I am okay with merging right after @ydshieh's question is being answered.

@ydshieh I'm sorry!! Trying it now!

@Rocketknight1 

No worry!

github label `single-model-run-slow` is no longer used. It should be `run-slow`. And a new commit with `[run-slow] mistral` has to be pushed from @ariG23498 side after the label is  added to this PR.

I will push the commit as soon as the label is added!

I added !

Hey @Rocketknight1 @ydshieh could you guide me to fix the currently failing tests. I am not able to figure out why the tests fail by looking at the logs.

Hi @ariG23498 @Rocketknight1 You can focus on the following 1 test failure. The others are irrelevant to this PR.

```
FAILED tests/models/mistral/test_modeling_tf_mistral.py::TFMistralModelTest::test_save_load_after_resize_token_embeddings - TypeError: Exception encountered when calling layer 'tf_mistral_for_causal_lm_35' (type TFMistralForCausalLM).

'NoneType' object is not callable

Call arguments received by layer 'tf_mistral_for_causal_lm_35' (type TFMistralForCausalLM):
  • input_ids=tf.Tensor(shape=(13, 7), dtype=int32)
  • attention_mask=tf.Tensor(shape=(13, 7), dtype=int32)
  • position_ids=None
  • past_key_values=None
  • inputs_embeds=None
  • labels=None
  • use_cache=None
  • output_attentions=None
  • output_hidden_states=None
  • return_dict=None

```

@ydshieh Do you know why some of the torch tests fail in the slow CI for Mistral even though they pass in the nightly tests?

@ydshieh do let me know if you need something from my side.

@ariG23498 as long as the failures are all in torch tests I think we can merge anyway - can you do a final rebase to rerun the tests and make sure nothing will break, and then I'll merge?

@Rocketknight1 Are you able to merge in this case (with some failures)?

Good from my side.

@Rocketknight1 I don't think the failing tests are realted to mistral -- could you verify?

@aritra yes, this is a dependency issue where the wrong version of `tensorflow-probability` is being installed in the CI. Investigating!

@ariG23498 this is a CI issue and the fix is being pushed soon. After that we can re-run the tests and hopefully merge this PR!

@Rocketknight1 I have rebased for the tests to pass.

Review done - merging this now, since the slow CI failures are all in the Torch model (that's already in the library)