Memory blows up when training large models on all TPU cores

🐛 Bug

I am training on 8 TPU cores but the memory blows up when the epoch ends.

To Reproduce

Try training a bert large on 8 TPU cores

Expected behavior

Second epoch should get started

Environment

Kaggle TPU

• PyTorch Lightning Version (e.g., 1.5.0):
• PyTorch Version (e.g., 1.10):
• Python version (e.g., 3.9):
• OS (e.g., Linux): Linux
• CUDA/cuDNN version:
• GPU models and configuration:
• How you installed PyTorch (conda, pip, source):
• If compiling from source, the output of torch.__config__.show():
• Any other relevant information:

Additional context

I will try to use GPU instead of TPUs

cc @kaushikb11 @rohitgr7 @awaelchli @ananthsub @ninginthecloud

[KrishPro]

Actually I know what causes this issue. Earlier i was directly using torch/xla. The memory blows when it tires to save the model at the end on epoch. Like if you would remove the check pointing code this works ok.

The issue is in saving large models trained on TPU (multi-core)

[KrishPro]

Any updates ? Maybe I should close this issue and open a duplicate on torch/xla ?