Update Gradient Clipping for TPU Accelerator

[kaushikb11]

What does this PR do?

Add updates to TPU Accelerator:

• Graident Clipping

Before submitting

• Was this discussed/approved via a GitHub issue? (not for typos and docs)
• 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?
• Did you make sure to update the documentation with your changes? (if necessary)
• Did you write any new necessary tests? (not for typos and docs)
• Did you verify new and existing tests pass locally with your changes?
• Did you update the CHANGELOG? (not for typos, docs, test updates, or internal minor changes/refactorings)

PR review

Anyone in the community is free to review the PR once the tests have passed.
Before you start reviewing make sure you have read Review guidelines. In short, see the following bullet-list:

• Is this pull request ready for review? (if not, please submit in draft mode)
• Check that all items from Before submitting are resolved
• Make sure the title is self-explanatory and the description concisely explains the PR
• Add labels and milestones (and optionally projects) to the PR so it can be classified

Did you have fun?

Make sure you had fun coding 🙃

[SeanNaren]

The patched XLA function does not do a check for grad_clip_val>0 so we might need to include this as done here:

https://github.com/PyTorchLightning/pytorch-lightning/blob/master/pytorch_lightning/plugins/precision/precision_plugin.py#L109-L110

In addition, will this mean we can get rid of the TODO here? https://github.com/PyTorchLightning/pytorch-lightning/blob/master/pytorch_lightning/plugins/precision/precision_plugin.py#L103

The implementation of clip gradients is TPU compliant I think, but I think that's not an issue.

[tchaton]

Good point !

[kaushikb11]

@SeanNaren @tchaton The issue for aitextgen?

The code hangs when it calls xm.save in Transformer’s save_pretrained method, when passed xm_save to the save_function param from aitextgen for TPUs.
Code: https://github.com/huggingface/transformers/blob/master/src/transformers/modeling_utils.py#L835

My hunch, xm.save has rendezvous in it. All the cores were not syncing up for rendezvous. Hence, the code was hanging.

Need to do the required changes on the aitextgen end.

[SeanNaren]

I know tests are annoying for TPUs, but could we add a test to ensure that clipping works? a simple test just to make sure the function is called like the sharded test:

@pytest.mark.parametrize("clip_val", [0, 10])
@RunIf(min_gpus=1, skip_windows=True, amp_native=True, fairscale=True)
@mock.patch('fairscale.optim.oss.OSS.clip_grad_norm')
def test_ddp_sharded_precision_16_clip_gradients(mock_oss_clip_grad_norm, clip_val, tmpdir):
    """
    Ensure that clip gradients is only called if the value is greater than 0.
    """
    model = BoringModel()
    trainer = Trainer(accelerator='ddp_sharded', gpus=1, precision=16, fast_dev_run=True, gradient_clip_val=clip_val)
    trainer.fit(model)
    if clip_val > 0:
        mock_oss_clip_grad_norm.assert_called()
    else:
        mock_oss_clip_grad_norm.assert_not_called()

[tchaton]

LGTM !

[awaelchli]

is it compatible with the TPUHalf plugin?
Should we extend clip_gradients there and raise NotImplementedError, what do you think?

[kaushikb11]

is it compatible with the TPUHalf plugin?
Should we extend clip_gradients there and raise NotImplementedError, what do you think?

@awaelchli But it could be the default for the TPU accelerator, right?