Weight tying is broken on TPUs leading to silent errors

[ibeltagy]

🐛 Bug

PyTorch/XLA documentation mentions here that weight tying should happen after moving tensors to XLA, otherwise the tensors are copied. This is a silent error that can easily go undetected (thanks to @matt-peters for pointing it out), and it would be good if PL guards the user against it. Notice that weight tying is pretty common in today's models not a corner case.

Code sample

The following code snippet shows how to detect that this issue is happening and how to guard against it.

import pl
class MyPLModel(pl.LightningModule):
    def to(self, *args, **kwargs):
        param_count_before_moving_to_device = len(list(self.parameters()))  # 
        super().to(*args, **kwargs)
        if self.trainer.use_tpu:
            # need to re-tie the weights after moving to XLA!
            self.tie_weights()  # a new function that the user needs to implement
        param_count_after_moving_to_device = len(list(self.parameters()))
        assert param_count_before_moving_to_device == param_count_after_moving_to_device

[stale]

This issue has been automatically marked as stale because it hasn't had any recent activity. This issue will be closed in 7 days if no further activity occurs. Thank you for your contributions, Pytorch Lightning Team!

[ibeltagy]

@zcain117, might be worth checking. It is a subtle issue and users are likely going to miss it.

[Borda]

@ibeltagy is the problem still on master? 🐰
cc: @zcain117 @lezwon

[ibeltagy]

Most probably yes. It needs a small API change to support the LightningModule.tie_weights function to call from LightningModule.to. I don't know if there's a way to fix it without such API change.

[Borda]

@zcain117 @lezwon ^^