Tpu save

[lezwon]

What does this PR do?

Fixes #2700
fixes #2303
fixes #3660

• Move the XLA tensors within a checkpoint to CPU before saving.
• Lazy check if TPU device exists
• Removed calling accelerator.barrier() on TPU when calling .test() as the multiprocessing begins only in .fit()
• Moved TPU device check to run when Trainer initialized.

Related to #3044

Before submitting

• Was this discussed/approved via a Github issue? (no need for typos and docs improvements)
• Did you read the contributor guideline, Pull Request section?
• Did you make sure your PR does only one thing, instead of bundling different changes together? Otherwise, we ask you to create a separate PR for every change.
• Did you make sure to update the documentation with your changes?
• Did you write any new necessary tests?
• Did you verify new and existing tests pass locally with your changes?
• If you made a notable change (that affects users), did you update the CHANGELOG?

PR review

• Is this pull request ready for review? (if not, please submit in draft mode)

Anyone in the community is free to review the PR once the tests have passed.
If we didn't discuss your PR in Github issues there's a high chance it will not be merged.

Did you have fun?

Make sure you had fun coding 🙃

[codecov]

Codecov Report

Merging #4309 (eef0849) into master (0c763b2) will decrease coverage by 0%.
The diff coverage is 100%.

@@          Coverage Diff           @@
##           master   #4309   +/-   ##
======================================
- Coverage      93%     93%   -0%     
======================================
  Files         124     124           
  Lines        9320    9303   -17     
======================================
- Hits         8640    8616   -24     
- Misses        680     687    +7     

[williamFalcon]

@lezwon does that move to cpu really belong in this file?
Why is this needed for TPUs?
You realize this will affect more than just TPUs no? this is why i don't like making these changes outside the accelerator.

If it turns out this is the correct place for this, please add a on_save() function to each accelerator. Then make all of them no op:

def on_save(self, ...):
    pass

And add the change ONLY to the TPU accelerator.

def on_save(self, ...):
   your_changes()

Going forward, any changes that are accelerator specific should not be done inside methods like this. Instead, each accelerator needs to implement this method and then called like:

self.accelerator.on_save()

The reason is that we are trying to break up all the underlying accelerator code so they are easier to debug and changes to an accelerator won't break all the others.

[lezwon]

@lezwon does that move to cpu really belong in this file?
Why is this needed for TPUs?
You realize this will affect more than just TPUs no? this is why i don't like making these changes outside the accelerator.

If it turns out this is the correct place for this, please add a on_save() function to each accelerator. Then make all of them no op:

def on_save(self, ...):
    pass

And add the change ONLY to the TPU accelerator.

def on_save(self, ...):
   your_changes()

Going forward, any changes that are accelerator specific should not be done inside methods like this. Instead, each accelerator needs to implement this method and then called like:

self.accelerator.on_save()

The reason is that we are trying to break up all the underlying accelerator code so they are easier to debug and changes to an accelerator won't break all the others.

The XLA guide recommends users to move the tensors to CPU before saving so that they can be loaded on non-TPU devices. Lightning users have faced this issue wherein they have trained a model on a TPU and arent able to use it on a GPU or CPU. Hence we need to move these tensors to CPU before saving. I get your point about separating such code and making it accelerator specific. Will refactor this into on_save. 👍

[mergify]

This pull request is now in conflict... :(

[SeanNaren]

hey @lezwon any activity here? were you able to reproduce this using the boring model btw? Don't mind picking this up

[lezwon]

hey @SeanNaren, I've been making some refactors to make this change TPU specific. WIll push an updated branch soon :)

[Borda]

lgtm

[williamFalcon]

where did this go? cc @tchaton

[tchaton]

I guess useless imports.

[Borda]

shall not be there at all in the first place, @SeanNaren

[tchaton]

I guess useless imports.