[Port] TensorFlow implementation of Mistral

[ariG23498]

What does this PR do?

This PR adds the TensorFlow implementation of Mistral.

Before submitting

• Did you read the contributor guideline,
  Pull Request section?
• Did you make sure to update the documentation with your changes? Here are the
  documentation guidelines, and
  here are tips on formatting docstrings.

Who can review?

@Rocketknight1

[ariG23498]

@Rocketknight1 at this stage if I run the following

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?

[Rocketknight1]

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!

[Rocketknight1]

PR is open at #29755

[Rocketknight1]

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

[ariG23498]

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)

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.

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. 🥳

[ariG23498]

Running the following code:

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:

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.

[Rocketknight1]

@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.

[ariG23498]

@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.

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?

[Rocketknight1]

@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.

[Rocketknight1]

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

[ydshieh]

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 ...? 😭

[ariG23498]

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

[Rocketknight1]

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

[ydshieh]

@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.

[ariG23498]

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

[ydshieh]

I added !

[ariG23498]

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.

[ydshieh]

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

[Rocketknight1]

@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?

[ariG23498]

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

[Rocketknight1]

@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?

[ydshieh]

@Rocketknight1 Are you able to merge in this case (with some failures)?

Good from my side.

[ariG23498]

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

[Rocketknight1]

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

[Rocketknight1]

@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!

[ariG23498]

@Rocketknight1 I have rebased for the tests to pass.

[Rocketknight1]

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