[Flax] Error converting model to PyTorch from Flax

[w11wo]

Hi, I followed the causal language modeling in Flax tutorial notebook provided here in Colab. And at the end of the training, I'd like to get a working PyTorch model from the JAX/Flax weights, hence I did this:

from transformers import GPT2LMHeadModel

mdl_path = "w11wo/sundanese-gpt2-base"

pt_model = GPT2LMHeadModel.from_pretrained(mdl_path, from_flax=True)
pt_model.save_pretrained(mdl_path)

But during the conversion, it raised this error

/usr/local/lib/python3.7/dist-packages/transformers/modeling_flax_pytorch_utils.py:201: UserWarning: The given NumPy array is not writeable, and PyTorch does not support non-writeable tensors. This means you can write to the underlying (supposedly non-writeable) NumPy array using the tensor. You may want to copy the array to protect its data or make it writeable before converting it to a tensor. This type of warning will be suppressed for the rest of this program. (Triggered internally at  /pytorch/torch/csrc/utils/tensor_numpy.cpp:180.)
  pt_model_dict[flax_key] = torch.from_numpy(flax_tensor)
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
<ipython-input-44-08a8ff6e575c> in <module>()
      7 tokenizer.save_pretrained(mdl_path)
      8 
----> 9 pt_model = GPT2LMHeadModel.from_pretrained(mdl_path, from_flax=True)
     10 pt_model.save_pretrained(mdl_path)
     11 

2 frames
/usr/local/lib/python3.7/dist-packages/transformers/modeling_flax_pytorch_utils.py in load_flax_weights_in_pytorch_model(pt_model, flax_state)
    199                 # add weight to pytorch dict
    200                 flax_tensor = np.asarray(flax_tensor) if not isinstance(flax_tensor, np.ndarray) else flax_tensor
--> 201                 pt_model_dict[flax_key] = torch.from_numpy(flax_tensor)
    202                 # remove from missing keys
    203                 missing_keys.remove(flax_key)

TypeError: can't convert np.ndarray of type bfloat16. The only supported types are: float64, float32, float16, complex64, complex128, int64, int32, int16, int8, uint8, and bool.

I think this issue occurred because the model I instantiated used bfloat16 -- just as the tutorial showed. Specifically this block

from transformers import FlaxAutoModelForCausalLM

model = FlaxAutoModelForCausalLM.from_config(config, seed=training_seed, dtype=jnp.dtype("bfloat16"))

I'd like to know if there's a workaround to this problem. Thanks!

[w11wo]

Seeing how we are able to convert float32 weights to bfloat16 as shown here, I tried to the reverse, i.e. converting bfloat16 to float32. I executed the following lines:

from transformers import GPT2LMHeadModel, GPT2Tokenizer, FlaxGPT2LMHeadModel

pretrained = "w11wo/sundanese-gpt2-base"
tmp_path = "sundanese-gpt-base"

model = FlaxGPT2LMHeadModel.from_pretrained(pretrained)
tokenizer = GPT2Tokenizer.from_pretrained(pretrained)

def to_f32(t):
    return jax.tree_map(lambda x: x.astype(jnp.float32) if x.dtype == jnp.bfloat16 else x, t)

model.params = to_f32(model.params)
model.save_pretrained(tmp_path)

pt_model = GPT2LMHeadModel.from_pretrained(tmp_path, from_flax=True)
pt_model.save_pretrained(tmp_path)

And the model converted returned the following messages,

All Flax model weights were used when initializing GPT2LMHeadModel.

Some weights of GPT2LMHeadModel were not initialized from the Flax model and are newly initialized: ['transformer.h.5.attn.masked_bias', 'transformer.h.4.attn.masked_bias', 'transformer.h.6.attn.bias', 'transformer.h.9.attn.bias', 'transformer.h.7.attn.bias', 'transformer.h.2.attn.masked_bias', 'transformer.h.6.attn.masked_bias', 'transformer.h.0.attn.bias', 'transformer.h.0.attn.masked_bias', 'transformer.h.1.attn.masked_bias', 'transformer.h.8.attn.masked_bias', 'transformer.h.10.attn.bias', 'transformer.h.8.attn.bias', 'transformer.h.4.attn.bias', 'transformer.h.9.attn.masked_bias', 'transformer.h.1.attn.bias', 'transformer.h.10.attn.masked_bias', 'transformer.h.3.attn.masked_bias', 'transformer.h.11.attn.bias', 'lm_head.weight', 'transformer.h.3.attn.bias', 'transformer.h.2.attn.bias', 'transformer.h.5.attn.bias', 'transformer.h.7.attn.masked_bias', 'transformer.h.11.attn.masked_bias']
You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.

I suspect that there are conversion errors, but am clueless as to the reason behind the error.

Further, I tried using the converted PyTorch model for text-generation via the following pipeline:

from transformers import pipeline

nlp = pipeline(
    "text-generation",
    model=tmp_path,
    tokenizer=tokenizer
)

nlp("Nami abdi Budi")
>> Setting `pad_token_id` to `eos_token_id`:50256 for open-end generation.
[{'generated_text': "Testipit ayeuna ayeuna409 pillowsKrKr bunderan Platform sipil poé Summary Sakitu Enak maké bunderan sipilthough papatong hubungiashmina Stock protésnyieun248duta terhadap Mm �issinginjeumdoi ' 'Dealer Studio gunarior floridabodas ' békénAndroid Holland majalah dot mangaruhanumpahinting"}]

Unsurprisingly, the output is very much jibberish, despite the model trained down to about 3.66 validation loss. I think the conversion is somehow incorrect, hence the incorrect porting of weights, and thus the jibberish output.

[patrickvonplaten]

I will take a look later today!

[patrickvonplaten]

Also cc @patil-suraj

[w11wo]

Thanks @patrickvonplaten! Really appreciate what you and your 🤗 team are doing!

[w11wo]

One more thing, I trained RoBERTa using a different script for the Flax community week; hub repo here. It seems that converting the model to PyTorch raises a possibly related error:

from transformers import RobertaForMaskedLM
model = RobertaForMaskedLM.from_pretrained("flax-community/indonesian-roberta-base", from_flax=True)

/usr/local/lib/python3.7/dist-packages/transformers/modeling_flax_pytorch_utils.py:201: UserWarning: The given NumPy array is not writeable, and PyTorch does not support non-writeable tensors. This means you can write to the underlying (supposedly non-writeable) NumPy array using the tensor. You may want to copy the array to protect its data or make it writeable before converting it to a tensor. This type of warning will be suppressed for the rest of this program. (Triggered internally at  /pytorch/torch/csrc/utils/tensor_numpy.cpp:180.)
  pt_model_dict[flax_key] = torch.from_numpy(flax_tensor)
All Flax model weights were used when initializing RobertaForMaskedLM.

Some weights of RobertaForMaskedLM were not initialized from the Flax model and are newly initialized: ['lm_head.decoder.bias', 'roberta.embeddings.position_ids', 'lm_head.decoder.weight']
You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.

[patrickvonplaten]

Seeing how we are able to convert float32 weights to bfloat16 as shown here, I tried to the reverse, i.e. converting bfloat16 to float32. I executed the following lines:

from transformers import GPT2LMHeadModel, GPT2Tokenizer, FlaxGPT2LMHeadModel

pretrained = "w11wo/sundanese-gpt2-base"
tmp_path = "sundanese-gpt-base"

model = FlaxGPT2LMHeadModel.from_pretrained(pretrained)
tokenizer = GPT2Tokenizer.from_pretrained(pretrained)

def to_f32(t):
    return jax.tree_map(lambda x: x.astype(jnp.float32) if x.dtype == jnp.bfloat16 else x, t)

model.params = to_f32(model.params)
model.save_pretrained(tmp_path)

pt_model = GPT2LMHeadModel.from_pretrained(tmp_path, from_flax=True)
pt_model.save_pretrained(tmp_path)

And the model converted returned the following messages,

All Flax model weights were used when initializing GPT2LMHeadModel.

Some weights of GPT2LMHeadModel were not initialized from the Flax model and are newly initialized: ['transformer.h.5.attn.masked_bias', 'transformer.h.4.attn.masked_bias', 'transformer.h.6.attn.bias', 'transformer.h.9.attn.bias', 'transformer.h.7.attn.bias', 'transformer.h.2.attn.masked_bias', 'transformer.h.6.attn.masked_bias', 'transformer.h.0.attn.bias', 'transformer.h.0.attn.masked_bias', 'transformer.h.1.attn.masked_bias', 'transformer.h.8.attn.masked_bias', 'transformer.h.10.attn.bias', 'transformer.h.8.attn.bias', 'transformer.h.4.attn.bias', 'transformer.h.9.attn.masked_bias', 'transformer.h.1.attn.bias', 'transformer.h.10.attn.masked_bias', 'transformer.h.3.attn.masked_bias', 'transformer.h.11.attn.bias', 'lm_head.weight', 'transformer.h.3.attn.bias', 'transformer.h.2.attn.bias', 'transformer.h.5.attn.bias', 'transformer.h.7.attn.masked_bias', 'transformer.h.11.attn.masked_bias']
You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.

I suspect that there are conversion errors, but am clueless as to the reason behind the error.

Further, I tried using the converted PyTorch model for text-generation via the following pipeline:

from transformers import pipeline

nlp = pipeline(
    "text-generation",
    model=tmp_path,
    tokenizer=tokenizer
)

nlp("Nami abdi Budi")
>> Setting `pad_token_id` to `eos_token_id`:50256 for open-end generation.
[{'generated_text': "Testipit ayeuna ayeuna409 pillowsKrKr bunderan Platform sipil poé Summary Sakitu Enak maké bunderan sipilthough papatong hubungiashmina Stock protésnyieun248duta terhadap Mm �issinginjeumdoi ' 'Dealer Studio gunarior floridabodas ' békénAndroid Holland majalah dot mangaruhanumpahinting"}]

Unsurprisingly, the output is very much jibberish, despite the model trained down to about 3.66 validation loss. I think the conversion is somehow incorrect, hence the incorrect porting of weights, and thus the jibberish output.

I think this is actually the correct way how you should convert bfloat16 to float32. Having followed this way of conversion can you verify that the PT and Flax model give similar outputs (ideally on CPU since TPU uses approximations)

[patrickvonplaten]

One more thing, I trained RoBERTa using a different script for the Flax community week; hub repo here. It seems that converting the model to PyTorch raises a possibly related error:

from transformers import RobertaForMaskedLM
model = RobertaForMaskedLM.from_pretrained("flax-community/indonesian-roberta-base", from_flax=True)

/usr/local/lib/python3.7/dist-packages/transformers/modeling_flax_pytorch_utils.py:201: UserWarning: The given NumPy array is not writeable, and PyTorch does not support non-writeable tensors. This means you can write to the underlying (supposedly non-writeable) NumPy array using the tensor. You may want to copy the array to protect its data or make it writeable before converting it to a tensor. This type of warning will be suppressed for the rest of this program. (Triggered internally at  /pytorch/torch/csrc/utils/tensor_numpy.cpp:180.)
  pt_model_dict[flax_key] = torch.from_numpy(flax_tensor)
All Flax model weights were used when initializing RobertaForMaskedLM.

Some weights of RobertaForMaskedLM were not initialized from the Flax model and are newly initialized: ['lm_head.decoder.bias', 'roberta.embeddings.position_ids', 'lm_head.decoder.weight']
You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.

This is expected can you check with #12554

[patrickvonplaten]

from transformers import GPT2LMHeadModel, GPT2Tokenizer, FlaxGPT2LMHeadModel
import torch
import numpy as np
import jax
import jax.numpy as jnp

pretrained = "w11wo/sundanese-gpt2-base"
tmp_path = "./"

model = FlaxGPT2LMHeadModel.from_pretrained(pretrained)

def to_f32(t):
    return jax.tree_map(lambda x: x.astype(jnp.float32) if x.dtype == jnp.bfloat16 else x, t)

model.params = to_f32(model.params)
model.save_pretrained(tmp_path)

model_pt = GPT2LMHeadModel.from_pretrained(tmp_path, from_flax=True)

input_ids = np.asarray(2 * [128 * [0]], dtype=np.int32)
input_ids_pt = torch.tensor(input_ids)

logits_pt = model_pt(input_ids_pt).logits
print(logits_pt)
logits_fx = model(input_ids).logits
print(logits_fx)

=> here you can see that you are correctly converting the weighst. The two models give the same results when doing a forward pass. This probably means that the training didn't work very well. Also not that vanilla generation often doesn't work very well. One should use do_sample=True

[w11wo]

One more thing, I trained RoBERTa using a different script for the Flax community week; hub repo here. It seems that converting the model to PyTorch raises a possibly related error:

from transformers import RobertaForMaskedLM
model = RobertaForMaskedLM.from_pretrained("flax-community/indonesian-roberta-base", from_flax=True)

/usr/local/lib/python3.7/dist-packages/transformers/modeling_flax_pytorch_utils.py:201: UserWarning: The given NumPy array is not writeable, and PyTorch does not support non-writeable tensors. This means you can write to the underlying (supposedly non-writeable) NumPy array using the tensor. You may want to copy the array to protect its data or make it writeable before converting it to a tensor. This type of warning will be suppressed for the rest of this program. (Triggered internally at  /pytorch/torch/csrc/utils/tensor_numpy.cpp:180.)
  pt_model_dict[flax_key] = torch.from_numpy(flax_tensor)
All Flax model weights were used when initializing RobertaForMaskedLM.

Some weights of RobertaForMaskedLM were not initialized from the Flax model and are newly initialized: ['lm_head.decoder.bias', 'roberta.embeddings.position_ids', 'lm_head.decoder.weight']
You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.

This is expected can you check with #12554

Checked, and it does return the same results. Thanks!

[w11wo]

from transformers import GPT2LMHeadModel, GPT2Tokenizer, FlaxGPT2LMHeadModel
import torch
import numpy as np
import jax
import jax.numpy as jnp

pretrained = "w11wo/sundanese-gpt2-base"
tmp_path = "./"

model = FlaxGPT2LMHeadModel.from_pretrained(pretrained)

def to_f32(t):
    return jax.tree_map(lambda x: x.astype(jnp.float32) if x.dtype == jnp.bfloat16 else x, t)

model.params = to_f32(model.params)
model.save_pretrained(tmp_path)

pt_model = GPT2LMHeadModel.from_pretrained(tmp_path, from_flax=True)

input_ids = np.asarray(2 * [128 * [0]], dtype=np.int32)
input_ids_pt = torch.tensor(input_ids)

logits_pt = model_pt(input_ids_pt).logits
print(logits_pt)
logits_fx = model(input_ids).logits
print(logits_fx)

This doesn't seem to return the same results on my end.

All Flax model weights were used when initializing GPT2LMHeadModel.

Some weights of GPT2LMHeadModel were not initialized from the Flax model and are newly initialized: ['transformer.h.9.attn.masked_bias', 'transformer.h.8.attn.bias', 'transformer.h.11.attn.masked_bias', 'transformer.h.2.attn.bias', 'transformer.h.7.attn.masked_bias', 'transformer.h.6.attn.masked_bias', 'transformer.h.6.attn.bias', 'transformer.h.3.attn.masked_bias', 'transformer.h.10.attn.bias', 'transformer.h.4.attn.masked_bias', 'transformer.h.1.attn.bias', 'transformer.h.9.attn.bias', 'transformer.h.5.attn.masked_bias', 'transformer.h.11.attn.bias', 'transformer.h.7.attn.bias', 'transformer.h.4.attn.bias', 'transformer.h.0.attn.masked_bias', 'transformer.h.1.attn.masked_bias', 'transformer.h.0.attn.bias', 'lm_head.weight', 'transformer.h.3.attn.bias', 'transformer.h.5.attn.bias', 'transformer.h.2.attn.masked_bias', 'transformer.h.8.attn.masked_bias', 'transformer.h.10.attn.masked_bias']
You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.
tensor([[[ 3.3714, -8.7829, -7.0013,  ..., -0.9126, -5.1224,  1.1825],
         [ 3.3988, -8.7616, -7.0675,  ..., -0.8201, -5.0986,  1.1000],
         [ 3.3717, -8.7876, -7.0300,  ..., -0.8022, -5.0759,  1.1368],
         ...,
         [ 3.0262, -8.7904, -7.3067,  ..., -1.0611, -5.3096,  1.0842],
         [ 3.0953, -8.7540, -7.2462,  ..., -1.0426, -5.2823,  1.1086],
         [ 3.0298, -8.7541, -7.3130,  ..., -1.0767, -5.3349,  1.0631]],

        [[ 3.3714, -8.7829, -7.0013,  ..., -0.9126, -5.1224,  1.1825],
         [ 3.3988, -8.7616, -7.0675,  ..., -0.8201, -5.0986,  1.1000],
         [ 3.3717, -8.7876, -7.0300,  ..., -0.8022, -5.0759,  1.1368],
         ...,
         [ 3.0262, -8.7904, -7.3067,  ..., -1.0611, -5.3096,  1.0842],
         [ 3.0953, -8.7540, -7.2462,  ..., -1.0426, -5.2823,  1.1086],
         [ 3.0298, -8.7541, -7.3130,  ..., -1.0767, -5.3349,  1.0631]]],
       grad_fn=<AddBackward0>)
[[[-0.18300235 -0.3880248   0.51604265 ... -0.03581356  0.33767372
    0.63595504]
  [-0.22266355 -0.43806082  0.54173964 ... -0.09416972  0.1144447
    0.46392882]
  [-0.21165054 -0.41117024  0.52854717 ... -0.10048242 -0.01106828
    0.21045035]
  ...
  [-0.13923354 -0.11146995  0.2263387  ... -0.06064492  0.6304022
   -0.27594602]
  [-0.20900789 -0.1306307   0.22801363 ... -0.07289732  0.6694305
   -0.2215142 ]
  [-0.17881839 -0.11881532  0.2094207  ... -0.06347632  0.6529123
   -0.2000624 ]]

 [[-0.18300235 -0.3880248   0.51604265 ... -0.03581356  0.33767372
    0.63595504]
  [-0.22266355 -0.43806082  0.54173964 ... -0.09416972  0.1144447
    0.46392882]
  [-0.21165054 -0.41117024  0.52854717 ... -0.10048242 -0.01106828
    0.21045035]
  ...
  [-0.13923354 -0.11146995  0.2263387  ... -0.06064492  0.6304022
   -0.27594602]
  [-0.20900789 -0.1306307   0.22801363 ... -0.07289732  0.6694305
   -0.2215142 ]
  [-0.17881839 -0.11881532  0.2094207  ... -0.06347632  0.6529123
   -0.2000624 ]]]

[patrickvonplaten]

from transformers import GPT2LMHeadModel, GPT2Tokenizer, FlaxGPT2LMHeadModel
import torch
import numpy as np
import jax
import jax.numpy as jnp

pretrained = "w11wo/sundanese-gpt2-base"
tmp_path = "./"

model = FlaxGPT2LMHeadModel.from_pretrained(pretrained)

def to_f32(t):
    return jax.tree_map(lambda x: x.astype(jnp.float32) if x.dtype == jnp.bfloat16 else x, t)

model.params = to_f32(model.params)
model.save_pretrained(tmp_path)

pt_model = GPT2LMHeadModel.from_pretrained(tmp_path, from_flax=True)

input_ids = np.asarray(2 * [128 * [0]], dtype=np.int32)
input_ids_pt = torch.tensor(input_ids)

logits_pt = model_pt(input_ids_pt).logits
print(logits_pt)
logits_fx = model(input_ids).logits
print(logits_fx)

This doesn't seem to return the same results on my end.

All Flax model weights were used when initializing GPT2LMHeadModel.

Some weights of GPT2LMHeadModel were not initialized from the Flax model and are newly initialized: ['transformer.h.9.attn.masked_bias', 'transformer.h.8.attn.bias', 'transformer.h.11.attn.masked_bias', 'transformer.h.2.attn.bias', 'transformer.h.7.attn.masked_bias', 'transformer.h.6.attn.masked_bias', 'transformer.h.6.attn.bias', 'transformer.h.3.attn.masked_bias', 'transformer.h.10.attn.bias', 'transformer.h.4.attn.masked_bias', 'transformer.h.1.attn.bias', 'transformer.h.9.attn.bias', 'transformer.h.5.attn.masked_bias', 'transformer.h.11.attn.bias', 'transformer.h.7.attn.bias', 'transformer.h.4.attn.bias', 'transformer.h.0.attn.masked_bias', 'transformer.h.1.attn.masked_bias', 'transformer.h.0.attn.bias', 'lm_head.weight', 'transformer.h.3.attn.bias', 'transformer.h.5.attn.bias', 'transformer.h.2.attn.masked_bias', 'transformer.h.8.attn.masked_bias', 'transformer.h.10.attn.masked_bias']
You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.
tensor([[[ 3.3714, -8.7829, -7.0013,  ..., -0.9126, -5.1224,  1.1825],
         [ 3.3988, -8.7616, -7.0675,  ..., -0.8201, -5.0986,  1.1000],
         [ 3.3717, -8.7876, -7.0300,  ..., -0.8022, -5.0759,  1.1368],
         ...,
         [ 3.0262, -8.7904, -7.3067,  ..., -1.0611, -5.3096,  1.0842],
         [ 3.0953, -8.7540, -7.2462,  ..., -1.0426, -5.2823,  1.1086],
         [ 3.0298, -8.7541, -7.3130,  ..., -1.0767, -5.3349,  1.0631]],

        [[ 3.3714, -8.7829, -7.0013,  ..., -0.9126, -5.1224,  1.1825],
         [ 3.3988, -8.7616, -7.0675,  ..., -0.8201, -5.0986,  1.1000],
         [ 3.3717, -8.7876, -7.0300,  ..., -0.8022, -5.0759,  1.1368],
         ...,
         [ 3.0262, -8.7904, -7.3067,  ..., -1.0611, -5.3096,  1.0842],
         [ 3.0953, -8.7540, -7.2462,  ..., -1.0426, -5.2823,  1.1086],
         [ 3.0298, -8.7541, -7.3130,  ..., -1.0767, -5.3349,  1.0631]]],
       grad_fn=<AddBackward0>)
[[[-0.18300235 -0.3880248   0.51604265 ... -0.03581356  0.33767372
    0.63595504]
  [-0.22266355 -0.43806082  0.54173964 ... -0.09416972  0.1144447
    0.46392882]
  [-0.21165054 -0.41117024  0.52854717 ... -0.10048242 -0.01106828
    0.21045035]
  ...
  [-0.13923354 -0.11146995  0.2263387  ... -0.06064492  0.6304022
   -0.27594602]
  [-0.20900789 -0.1306307   0.22801363 ... -0.07289732  0.6694305
   -0.2215142 ]
  [-0.17881839 -0.11881532  0.2094207  ... -0.06347632  0.6529123
   -0.2000624 ]]

 [[-0.18300235 -0.3880248   0.51604265 ... -0.03581356  0.33767372
    0.63595504]
  [-0.22266355 -0.43806082  0.54173964 ... -0.09416972  0.1144447
    0.46392882]
  [-0.21165054 -0.41117024  0.52854717 ... -0.10048242 -0.01106828
    0.21045035]
  ...
  [-0.13923354 -0.11146995  0.2263387  ... -0.06064492  0.6304022
   -0.27594602]
  [-0.20900789 -0.1306307   0.22801363 ... -0.07289732  0.6694305
   -0.2215142 ]
  [-0.17881839 -0.11881532  0.2094207  ... -0.06347632  0.6529123
   -0.2000624 ]]]

weird it gives the same results for me

[w11wo]

Here's my environment @patrickvonplaten, in case it'll help

• transformers version: 4.9.0.dev0
• Platform: Linux-5.4.104+-x86_64-with-Ubuntu-18.04-bionic
• Python version: 3.7.10
• PyTorch version (GPU?): 1.9.0+cu102 (False)
• Tensorflow version (GPU?): 2.5.0 (False)
• Flax version (CPU?/GPU?/TPU?): 0.3.4 (cpu)
• Jax version: 0.2.13
• JaxLib version: 0.1.66
• Using GPU in script?:
• Using distributed or parallel set-up in script?:

[patrickvonplaten]

That's my output:

Some weights of GPT2LMHeadModel were not initialized from the Flax model and are newly initialized: ['transformer.h.9.attn.masked_bias', 'transformer.h.3.attn.bias', 'transformer.h.8.attn.masked_bias', 'transformer.h.2.attn.bias', 'transformer.h.4.attn.masked_bias', 'transformer.h.6.attn.masked_bias', 'transformer.h.2.attn.masked_bias', 'transformer.h.3.attn.masked_bias', 'transformer.h.11.attn.bias', 'transformer.h.0.attn.masked_bias', 'transformer.h.10.attn.masked_bias', 'lm_head.weight', 'transformer.h.11.attn.masked_bias', 'transformer.h.0.attn.bias', 'transformer.h.10.attn.bias', 'transformer.h.7.attn.masked_bias', 'transformer.h.4.attn.bias', 'transformer.h.1.attn.masked_bias', 'transformer.h.7.attn.bias', 'transformer.h.5.attn.masked_bias', 'transformer.h.9.attn.bias', 'transformer.h.5.attn.bias', 'transformer.h.6.attn.bias', 'transformer.h.8.attn.bias', 'transformer.h.1.attn.bias']
You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.
tensor([[[-0.1830, -0.3880,  0.5160,  ..., -0.0358,  0.3377,  0.6360],
         [-0.2227, -0.4381,  0.5417,  ..., -0.0942,  0.1144,  0.4639],
         [-0.2117, -0.4112,  0.5285,  ..., -0.1005, -0.0111,  0.2105],
         ...,
         [-0.1392, -0.1115,  0.2263,  ..., -0.0606,  0.6304, -0.2759],
         [-0.2090, -0.1306,  0.2280,  ..., -0.0729,  0.6694, -0.2215],
         [-0.1788, -0.1188,  0.2094,  ..., -0.0635,  0.6529, -0.2001]],

        [[-0.1830, -0.3880,  0.5160,  ..., -0.0358,  0.3377,  0.6360],
         [-0.2227, -0.4381,  0.5417,  ..., -0.0942,  0.1144,  0.4639],
         [-0.2117, -0.4112,  0.5285,  ..., -0.1005, -0.0111,  0.2105],
         ...,
         [-0.1392, -0.1115,  0.2263,  ..., -0.0606,  0.6304, -0.2759],
         [-0.2090, -0.1306,  0.2280,  ..., -0.0729,  0.6694, -0.2215],
         [-0.1788, -0.1188,  0.2094,  ..., -0.0635,  0.6529, -0.2001]]],
       grad_fn=<UnsafeViewBackward>)
2021-07-07 14:27:19.167439: W tensorflow/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libcudart.so.11.0'; dlerror: libcudart.so.11.0: cannot open shared object file: No such file or directory
[[[-0.18300146 -0.38216865  0.51648384 ... -0.03733653  0.33876732
    0.63263166]
  [-0.217145   -0.4353617   0.5423262  ... -0.09400246  0.11277512
    0.46030286]
  [-0.20957293 -0.4070356   0.52822715 ... -0.10254323 -0.01395297
    0.20638129]
  ...
  [-0.13773967 -0.10961649  0.22636533 ... -0.065424    0.6344244
   -0.27795315]
  [-0.20452675 -0.1309557   0.23095486 ... -0.07456343  0.66748977
   -0.21973252]
  [-0.17408411 -0.11795001  0.21041611 ... -0.06386398  0.6547011
   -0.1990372 ]]

 [[-0.18300146 -0.38216865  0.51648384 ... -0.03733653  0.33876732
    0.63263166]
  [-0.217145   -0.4353617   0.5423262  ... -0.09400246  0.11277512
    0.46030286]
  [-0.20957293 -0.4070356   0.52822715 ... -0.10254323 -0.01395297
    0.20638129]
  ...
  [-0.13773967 -0.10961649  0.22636533 ... -0.065424    0.6344244
   -0.27795315]
  [-0.20452675 -0.1309557   0.23095486 ... -0.07456343  0.66748977
   -0.21973252]
  [-0.17408411 -0.11795001  0.21041611 ... -0.06386398  0.6547011
   -0.1990372 ]]]

[patrickvonplaten]

My env is:

- `transformers` version: 4.9.0.dev0
- Platform: Linux-5.4.0-1043-gcp-x86_64-with-glibc2.29
- Python version: 3.8.10
- PyTorch version (GPU?): 1.9.0+cu102 (False)
- Tensorflow version (GPU?): 2.5.0 (False)
- Flax version (CPU?/GPU?/TPU?): 0.3.4 (tpu)
- Jax version: 0.2.16
- JaxLib version: 0.1.68
- Using GPU in script?: <fill in>
- Using distributed or parallel set-up in script?: <fill in>

[patrickvonplaten]

Can you update jax, flax and jaxlib and also check on TPU?