Problem with restoring GPT-2 weights
Author: mgrankin
Pull request: no
Labels: name: wontfix

Hello, I've been debugging an issue for a while and it seem it's a model-specific issue. 

I'm training GPT-2 on a TPU and I can't save and restore it. It looks like there is a code that silently changes parameter values right in `load_state_dict()`.

```
    print(state_dict['transformer.wte.weight'])
    print(state_dict['transformer.wte.weight'].shape)
    
    cpu_model = model_class(config=config)
    cpu_model.load_state_dict(state_dict)
    print(cpu_model.state_dict()['transformer.wte.weight'])
    print(cpu_model.state_dict()['transformer.wte.weight'].shape)
```

```
tensor([[-0.1101, -0.0393,  0.0331,  ..., -0.1364,  0.0151,  0.0453],
        [ 0.0417, -0.0488,  0.0485,  ...,  0.0827,  0.0097,  0.0454],
        [-0.1275,  0.0479,  0.1841,  ...,  0.0899, -0.1297, -0.0879],
        ...,
        [-0.0439, -0.0579,  0.0103,  ...,  0.1113,  0.0919, -0.0724],
        [ 0.1846,  0.0156,  0.0444,  ..., -0.0974,  0.0785, -0.0211],
        [ 0.0471, -0.0284,  0.0492,  ...,  0.0048,  0.1511,  0.1202]])
torch.Size([50257, 768])
tensor([[-0.1317, -0.0305,  0.0339,  ..., -0.1310,  0.0113,  0.0262],
        [ 0.0413, -0.0491,  0.0451,  ...,  0.0930, -0.0019,  0.0457],
        [-0.1465,  0.0565,  0.1839,  ...,  0.0962, -0.1339, -0.1074],
        ...,
        [-0.0432, -0.0628,  0.0088,  ...,  0.1002,  0.1045, -0.0654],
        [ 0.1725,  0.0160,  0.0444,  ..., -0.0944,  0.0760, -0.0289],
        [ 0.0330, -0.0182,  0.0455,  ...,  0.0136,  0.1487,  0.0975]])
torch.Size([50257, 768])
```

For the context
https://github.com/pytorch/xla/issues/1245
https://discuss.pytorch.org/t/problem-with-model-accuracy-after-restore-on-tpu/59304/3

Full code is here 
https://github.com/mgrankin/ru_transformers/blob/9d52a4caef16df5b921c386f4841c879877d03a4/debug_lm_finetuning.py


Comments:
I found a bug, it's TPU related. For some reason, after I move the mode to TPU, using `model = model.to(device)`, the weights become decoupled. Then I save decoupled weights and during restore it ties them again. It loads correctly, it just doesn't expect tied weights to be different.

The workaround is to tie weights again after moving model to the TPU.

```
    model = model.to(args.device)
    model.tie_weights()
```



I'm sorry to reopen this issue, Davide Libenzi suggesting this is the model issue, not the PyTorch XLA issue. I'm a bit of a tired from debugging and I'm happy with the workaround. 

You can find details here https://github.com/pytorch/xla/issues/1245


Ok, do you think we should fix this upstream in our library?
I'm not super excited about overwriting PyTorch built-in `nn.Module.apply()` method.

It feels to me, the Pytorch/XLA is more appropriate place for fix, since Pytorch/Cuda have that behavior and the fix will make two libraries consistent. 
But I don't feel competent neither in Pytorch/XLA nor in Transformers to insist. 
It would be great to have somebody from Transformers to talk to Pytorch/XLA over this issue. 

Ok, we'll try to push this up-stream

The PyTorch community decided it's more appropriate to tie weights after moving the model to the device (TPU/GPU/CPU). I believe it's worth to fix the model accordingly. 

https://github.com/pytorch/xla/issues/1245#issuecomment-552559970
https://github.com/pytorch/xla/pull/1335

This issue has been automatically marked as stale because it has not had recent activity. It will be closed if no further activity occurs. Thank you for your contributions.


Hello, stale bot, it would be great to keep it open.

This issue has been automatically marked as stale because it has not had recent activity. It will be closed if no further activity occurs. Thank you for your contributions.