About torch.nn.utils clipping functions

[carmocca]

❓ Questions and Help

Hello,

I noticed xla patches the torch.nn.utils.clip_grad_norm_ function. Allegedly due to performance issues. From https://github.com/pytorch/xla/blob/master/TROUBLESHOOTING.md:

Use torch.where to substitute control flow when applicable. E.g. The control flow with item() used in clip_grad_norm_ is problematic and impacts performance, so we have patched clip_grad_norm_ by calling torch.where instead, which gives us a dramatic performance improvement.

PyTorch's implementation of clip_grad_norm_ was updated in pytorch/pytorch#32020 (which made it into 1.5.0) so the computation no longer relies on .item()

Does that mean xla's patch is no longer necessary on torch >= 1.5.0?

Additionally, are there any known issues for torch.nn.utils.clip_grad_value_? I am assuming that's not the case since there are no .item() calls, but could not find any confirmation anywhere.

Thanks!

[JackCaoG]

Thanks! We meant to test this out. I agree that torch.nn.utils.clip_grad_norm_ seems fine now. We need to double check the speed on TPU before removing the patch. @ailzhang tried this once and run into a weird error last time.

[stale]

This issue has been automatically marked as stale because it has not had recent activity. It will be closed if no further activity occurs. Thank you for your contributions.

[carmocca]

It would be nice to revisit this.

In Lightning, we are seeing this error when torch_xla is imported

E       RuntimeError: The norm of order 2.0 for a gradient from `parameters` is non-finite, so it cannot be clipped. This error can be disabled with `error_if_nonfinite=False`

For tests that run and pass normally if torch_xla had not been imported.

As a workaround, I'm doing:

if hasattr(torch.nn.utils.clip_grad_norm_, "_orig"):
    # hacky workaround to https://github.com/pytorch/xla/issues/2884: undo xla patching on import
    torch.nn.utils.clip_grad_norm_ = torch.nn.utils.clip_grad_norm_._orig

(in Lightning-AI/pytorch-lightning#17519)

This global patching is particularly problematic because it is done regardless of whether you actually end up using XLA at all.