[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.

[rohitgr7]

yes, working on it.. will update soon.

[ananthsub]

@rohitgr7 @carmocca Some of the issues we ought to resolve here:

• What's the expected API on the LightningModule?
• How do we reconcile the LightningModule that the user sees vs the automatically parallelized model the trainer sees? Not all gradient clipping techniques are compatible with techniques like DeepSpeed or FSDP. For example: https://github.com/PyTorchLightning/pytorch-lightning/blob/c7451b3ccf742b0e8971332caf2e041ceabd9fe8/pytorch_lightning/plugins/precision/fully_sharded_native_amp.py#L28-L46
• What happens to the trainer flags for gradient clip value or gradient clip algorithm? How does someone know if those flags are being used or not?
• @awaelchli asked why not implement this in on_after_backward ?

Using self.trainer.accelerator as part of the default implementation in the LightningModule is too fragile: someone will override the whole method with a custom implementation, and it'll fail when using different trainer flags. At the same time, it's not good if we have 2 different places where gradient clipping is implemented: once in the plugin, and once in the LIghtningModule.

It also goes against the principle of #7315

[tchaton]

• on_after_backward

Hey @ananthsub,

1 ) I believe we could pass the clip_val and gradient_clip_algorithm as extra arguments to the function. However, users might ignore both if they have a custom implementation.

2 ) DeepSpeed and FSDP don't support gradient clipping as the gradients are concatenated. I don't think we should block the features entirely for 2 plugins which are non mature for this feature.

3. The PrecisionPlugin should raise a MisConfiguration if the values are being provided and the plugins doesn't support it. If the users override the clip_gradient, we would just put a warning.

4. on_after_backward is actually the wrong place to implement this. This won't work with accumulated_grad_batches as it is called after each backward, it should be done in on_before_optimizer_step hook. The fact we are confused ourself where to apply gradient clipping makes me think users will be too. A dedicated hook seems simpler

5. Should we do loop -> lightning module -> trainer -> accelerator or loop -> accelerator -> lightning module. I would prefer the first one as it is makes maintenance simpler and the code more minimalist. But don't have a strong opinion there.

Best,
T.C

[carmocca]

+1 to passing the existing arguments to the hook so that they can be used for a default implementation

[rohitgr7]

+1 to passing the existing arguments to the hook so that they can be used for a default implementation

1. so are you suggesting keeping both trainer arguments add adding arguments in the new clip_gradients hook? I'd prefer not to since it might confuse users a little bit.. like having the same hparams at 2 different locations.. but that's just a personal opinion.

2. regarding loop -> lightning module -> trainer -> accelerator part I think we have the same workflow for optimizer_step so this looks ok to me.

[rohitgr7]

ok seems like keeping trainer arguments is a good option and we can pass them to this hook directly as defaults.

for now I can think of 2 ways for api design.

def clip_gradients(self, opt, opt_idx, clip_val, clip_algo):
    # if someone wants to use the internal implementation, use super
   if opt_idx == 0:
     super().clip_gradients(opt, clip_val, clip_algo)
   else:
      # implement your own
   

on this one, pointed out by @awaelchli that setting the values in the hooks signature and then calling super() is not something we have anywhere in Lightning so it will be very unfamiliar to users.

2. or we can put the current implementation inside clip_gradients and make configure_gradient_clipping available to users

def configure_gradient_clipping(self, optimizer, optimizer_idx, clip_val, clip_algo):
    if optimizer_idx == 0:
        self.clip_gradients(optimizer, clip_val, clip_algo). # lightning will handle this
    else:
        # implement your own

but open to suggestions :)

[tchaton]

Hey @rohitgr7,

Interesting. It is a great point !

I like the option 2.

@ananthsub @carmocca Any thoughts ?

Best,
T.C