Issue converting Flax model to Pytorch

[BirgerMoell]

When using the following script to convert a trained flax model to pytorch, the model seems to perform extremely poorly.

from transformers import RobertaForMaskedLM
model = RobertaForMaskedLM.from_pretrained("./", from_flax=True)
model.save_pretrained("./")

from transformers import RobertaForMaskedLM, FlaxRobertaForMaskedLM
import numpy as np
import torch 

model_fx = FlaxRobertaForMaskedLM.from_pretrained("birgermoell/roberta-swedish")
model_pt = RobertaForMaskedLM.from_pretrained("birgermoell/roberta-swedish", 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_fx(input_ids).logits
print(logits_fx)

Comparing gives the following input.

tensor([[[  1.7789, -13.5291, -11.2138,  ...,  -5.2875,  -9.3274,  -4.7912],
         [  2.3076, -13.4161, -11.1511,  ...,  -5.3181,  -9.0602,  -4.6083],
         [  2.6451, -13.4425, -11.0671,  ...,  -5.2838,  -8.8323,  -4.2280],
         ...,
         [  1.9009, -13.6516, -11.2348,  ...,  -4.9726,  -9.3278,  -4.6060],
         [  2.0522, -13.5394, -11.2804,  ...,  -4.9960,  -9.1956,  -4.5691],
         [  2.2570, -13.5093, -11.2640,  ...,  -4.9986,  -9.1292,  -4.3310]],

        [[  1.7789, -13.5291, -11.2138,  ...,  -5.2875,  -9.3274,  -4.7912],
         [  2.3076, -13.4161, -11.1511,  ...,  -5.3181,  -9.0602,  -4.6083],
         [  2.6451, -13.4425, -11.0671,  ...,  -5.2838,  -8.8323,  -4.2280],
         ...,
         [  1.9009, -13.6516, -11.2348,  ...,  -4.9726,  -9.3278,  -4.6060],
         [  2.0522, -13.5394, -11.2804,  ...,  -4.9960,  -9.1956,  -4.5691],
         [  2.2570, -13.5093, -11.2640,  ...,  -4.9986,  -9.1292,  -4.3310]]],
       grad_fn=<AddBackward0>)
[[[  0.1418128  -14.170926   -11.12649    ...  -7.542998   -10.79537
    -9.382975  ]
  [  1.7505689  -13.178099   -10.356588   ...  -6.794136   -10.567211
    -8.6670065 ]
  [  2.0270724  -13.522658   -10.372475   ...  -7.0110755  -10.396935
    -8.419178  ]
  ...
  [  0.19080782 -14.390833   -11.399942   ...  -7.469897   -10.715849
    -9.234054  ]
  [  1.3052869  -13.332332   -10.702984   ...  -6.9498534  -10.813769
    -8.608736  ]
  [  1.6442876  -13.226774   -10.59941    ...  -7.0290956  -10.693554
    -8.457008  ]]

 [[  0.1418128  -14.170926   -11.12649    ...  -7.542998   -10.79537
    -9.382975  ]
  [  1.7505689  -13.178099   -10.356588   ...  -6.794136   -10.567211
    -8.6670065 ]
  [  2.0270724  -13.522658   -10.372475   ...  -7.0110755  -10.396935
    -8.419178  ]
  ...
  [  0.19080782 -14.390833   -11.399942   ...  -7.469897   -10.715849
    -9.234054  ]
  [  1.3052869  -13.332332   -10.702984   ...  -6.9498534  -10.813769
    -8.608736  ]
  [  1.6442876  -13.226774   -10.59941    ...  -7.0290956  -10.693554
    -8.457008  ]]]

[patrickvonplaten]

Running the following command:

from transformers import RobertaForMaskedLM, FlaxRobertaForMaskedLM
import numpy as np
import torch 

model_fx = FlaxRobertaForMaskedLM.from_pretrained("birgermoell/roberta-swedish")
model_pt = RobertaForMaskedLM.from_pretrained("birgermoell/roberta-swedish", 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_fx(input_ids).logits
print(logits_fx)

should give more or less identical results

[patrickvonplaten]

Just corrected the pt weights. If you run:

from transformers import RobertaForMaskedLM, FlaxRobertaForMaskedLM
import numpy as np
import torch 

model_fx = FlaxRobertaForMaskedLM.from_pretrained("birgermoell/roberta-swedish")
model_pt = RobertaForMaskedLM.from_pretrained("birgermoell/roberta-swedish")
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_fx(input_ids).logits
print(logits_fx)

You should see equal results. The checkpoint was somehow incorrectly converted.

[patrickvonplaten]

Note that one should convert checkpoints with:

from transformers import RobertaForMaskedLM

model = RobertaForMaskedLM.from_pretrained("...", from_flax=True)
model.save_pretrained("./")

and not the AutoModel.... classes.

Also it's important to realize that the lm head layer is actually tied to the input word embedding layer which is why Flax just doesn't save those weights. Then when converting those weights to PyTorch, PyTorch says there are missing but since the weights are tied PyTorch would have overwritten those weights anyways with the input embeddings which is why it the warning:

Some weights of RobertaForMaskedLM were not initialized from the Flax model and are newly initialized: ['lm_head.decoder.bias', 'lm_head.decoder.weight']

doesn't matter.

[patrickvonplaten]

@BirgerMoell Also note that your local .git repository must be huge since you've essentially uploaded ~100 checkpoints of 500 MB each -> so your local .git stores 50 GB already I think.

[patrickvonplaten]

Widget seems to work: https://huggingface.co/birgermoell/roberta-swedish?text=Var+kan+jag+hitta+n%C3%A5gon+%3Cmask%3E+talar+engelska%3F

[BirgerMoell]

Awesome. Just to clarify. Once I'm done with training, this script should help me convert the model to pytorch.

from transformers import RobertaForMaskedLM

model = RobertaForMaskedLM.from_pretrained("...", from_flax=True)
model.save_pretrained("./")

[BirgerMoell]

@patrickvonplaten the uploaded model is still performing poorly so I'm not 100% the issue is fully resolved.

As you can see it outputs empty tokens.

[w11wo]

@patrickvonplaten the uploaded model is still performing poorly so I'm not 100% the issue is fully resolved.

As you can see it outputs empty tokens.

Hi @BirgerMoell, I'm training a RoBERTa model too using JAX during this community week -- model here. I got about 2.188 evaluation loss, yet the results are still somewhat jibberish despite the result. I think our models are, somehow, trained incorrectly? Or possibly require more data cleaning of some sort.

[BirgerMoell]

@w11wo Yeah. Something is definitely up. I think a good idea would be that people who work with similar models figure out a good way to clean the data and look at other things that might be wrong.

[jppaolim]

Facing same issue here, trained a model with Flax / Jax, then saved. When loading in Pytorch via "from Flax = True" , I have silly output despite training showing OK loss... Did you manage to find a solution or understand the issue ?

[w11wo]

Hi @jppaolim !

In my case, I loaded the earlier weights of the model (from the first few epochs), instead of the fully-trained model weights from the last training epoch. Loading the right model weights fixed it for me.

Another way to fix it might be training for longer.

Hope this helps! :)