[RFC] Gradient clipping hooks in the LightningModule

[carmocca]

🚀 Feature

Add clipping hooks to the LightningModule

Motivation

It's currently very difficult to change the clipping logic

Pitch

class LightningModule:
    def clip_gradients(self, optimizer, optimizer_idx):
        ...

The default implementation would be the same as we currently provide, where the trainer's clipping flags are used.

Maybe those would be deprecated in favor of LightningModule properties.

class LightningOptimizer
    def step(closure=closure)
        if closure is None:
            closure = do_nothing_closure
        def wrapper_closure()
            closure()
            self._trainer.call_hook("clip_gradients", self.optimizer)
        self.optimizer.step(closure=wrapper_closure)

Need to evaluate the limitations since clipping is currently tied to plugins

Additional context

This would fix #5096, #6123 (comment), #5671, #5982, and allow easily implementing new clipping techniques without having to merge them into Lightning

cc: @rohitgr7 who has been pushing for this for a while

[tchaton]

@SeanNaren @awaelchli Any thoughts ?

[rohitgr7]

Thanks @carmocca for bringing this up ☺️. This is on my TODO list actually but I kept pushing it further due to refactors. It would be great to have this since it's a part of optimization just like optimizer_step.

[edenlightning]

@rohitgr7 firendly ping?

[ericharper]

Being able to customize gradient clipping is still really important to the NeMo team. I hope these hooks can be added soon.