clip_gradient with clip_grad_value

[dhkim0225]

🚀 Feature

Same issue with #4927
The current clip_gradient uses clip_grad_norm; can we add clip_grad_value?

https://github.com/PyTorchLightning/pytorch-lightning/blob/f2e99d617f05ec65fded81ccc6d0d59807c47573/pytorch_lightning/plugins/native_amp.py#L63-L65

[tchaton]

Hey @dhkim0225,

Sounds like a great feature. Would you mind making a PR to enable it ?

Best regards,
T.C

[dhkim0225]

@tchaton I'm really glad to do that !

Renaming the clip_grad_value variable to clip_grad_norm and,
adding a new clip_grad_value feature seems most intuitive,
but this is BC-breaking.

Do you have any good ideas?

Sincerely,
Anthony Kim.

[tchaton]

Hey @dhkim0225,

I just checked. You can pass those arguments directly to the Trainer.
trainer = Trainer(gradient_clip_val=10, track_grad_norm=2) and they will be used in NativeAMP clip_gradient.

It will use self.trainer.gradient_clip_val.

    def clip_gradients(self, optimizer, clip_val=None):
        # use the trainer's clip val if none passed
        grad_clip_val = self.trainer.gradient_clip_val
        if clip_val is not None:
            grad_clip_val = clip_val
        grad_clip_val = float(grad_clip_val)

        if grad_clip_val <= 0:
            return
        self._clip_gradients(optimizer, grad_clip_val)

    def _clip_gradients(self, optimizer: Optimizer, grad_clip_val: Union[float, int], norm_type: float = 2.0):
        if self.trainer.amp_backend:
            self.trainer.precision_connector.backend.clip_gradients(grad_clip_val, optimizer, norm_type)
        else:
            model = self.trainer.get_model()
            torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=grad_clip_val, norm_type=norm_type)

Closing this issue as the parameter is already exposed in the Trainer. Feel free to re-open it if you think I missed something.

Best,
T.C