Multi GPU Training with Custom Training Steps

[mogwai]

What's the protocol for using custom training steps with lightning + multi gpu + dp that invoke grad scalers etc.? I'm currently seeing this error:

Original Traceback (most recent call last):
  File "/home/harry/sonantic/src/sonantic/tts/.venv/lib/python3.9/site-packages/torch/nn/parallel/parallel_apply.py", line 61, in _worker
    output = module(*input, **kwargs)
  File "/home/harry/sonantic/src/sonantic/tts/.venv/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1130, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/harry/sonantic/src/sonantic/tts/.venv/lib/python3.9/site-packages/pytorch_lightning/overrides/data_parallel.py", line 65, in forward
    output = super().forward(*inputs, **kwargs)
  File "/home/harry/sonantic/src/sonantic/tts/.venv/lib/python3.9/site-packages/pytorch_lightning/overrides/base.py", line 79, in forward
    output = self.module.training_step(*inputs, **kwargs)
  File "/home/harry/sonantic/src/sonantic/tts/nn/vits/vits.py", line 728, in training_step
    scaler.unscale_(optim_d)
  File "/home/harry/sonantic/src/sonantic/tts/.venv/lib/python3.9/site-packages/torch/cuda/amp/grad_scaler.py", line 272, in unscale_
    raise RuntimeError("unscale_() is being called after step().")
RuntimeError: unscale_() is being called after step().

This some psuedo code from the vits model.

And the lightning implementation is here:

def training_step(self, batch, batch_idx):
    ...
    # Discriminator Backwards
    optim_d.zero_grad()
    self.manual_backward(scaler.scale(loss_disc_all))
    scaler.unscale_(optim_d)
    grad_norm_d = commons.clip_grad_value_(self.discriminator.parameters(), None)
    scaler.step(optim_d)
    self.log("grad_norm_d", grad_norm_d, prog_bar=True)

[akihironitta]

If you have Trainer(precision=16), it's handled by Lightning already, so you don't need scale/unscale in your training_step. FYI, the documentation on manual optimization is here: https://pytorch-lightning.readthedocs.io/en/stable/common/optimization.html?highlight=optimization#id2