AttributeError: 'BoringModel' object has no attribute 'require_backward_grad_sync' when using manual optimization with TPU

[evanatyourservice]

🐛 Bug

Hello!

When using manual optimization with TPU, I am getting an AttributeError: 'BoringModel' object has no attribute 'require_backward_grad_sync'. When I replace self.manual_backward(loss) with loss.backward() things seem to work, but I am not sure if this is a safe or sustainable workaround. It seems the error happens at the self.manual_backward(loss) step in training_step. Any help would be much appreciated.

Please reproduce using the BoringModel

Here is the notebook reproducing the error:

https://colab.research.google.com/drive/1LPYgtUAiHd1OXuTK6I1WkRaCUQScxEPg?usp=sharing

Environment

WARNING:root:TPU has started up successfully with version pytorch-1.8

• CUDA:
  • GPU:
  • available: False
  • version: 10.1
• Packages:
  • numpy: 1.19.5
  • pyTorch_debug: False
  • pyTorch_version: 1.8.0+cu101
  • pytorch-lightning: 1.3.0dev
  • tqdm: 4.41.1
• System:
  • OS: Linux
  • architecture:
    • 64bit
  • processor: x86_64
  • python: 3.7.10
  • version: Proposal for help #1 SMP Thu Jul 23 08:00:38 PDT 2020

Installed torch-xla using:
!pip install cloud-tpu-client==0.10 https://storage.googleapis.com/tpu-pytorch/wheels/torch_xla-1.8-cp37-cp37m-linux_x86_64.whl
to match colab defaults

[tchaton]

Dear @evanatyourservice,

Manual Optimization doesn't work with TPU yet. Sorry for the inconvenience.
loss.backward() might work, but I would have to investigate deeper.

Best,
T.C

[evanatyourservice]

Ok, no worries. I don't have the time to work on this myself, so I will set up my own TPU notebook.

[SeanNaren]

Hey @evanatyourservice @7970 may have fixed this issue, could you retry using master?

[evanatyourservice]

The notebook I linked to above doesn't run, it still gets an error:

Exception in device=TPU:4: 'LightningDistributedModule' object has no attribute 'require_backward_grad_sync'

but it's been a while since I made it so I'm not sure if the notebook code is wrong or lightning's.

I ended up making my own TPU notebook so I stopped using lightning but I was trying to implement sharpness aware minimization if anyone wanted to try this with lightning