Gradient clipping when using Lite

[Dekakhrone]

Hello everybody! As I understood from the code, it is impossible to clip gradients when using Lite without some deep modification. Am I right, or am I missing something?

Ahead of the suggestion to use the LightningModule I will say that I am writing deeply custom code that is incompatible with its limitations. However Lite is seemed to me as a solution that would take over all the work with the devices.

[jsdtwry]

Similar question! I use lightning lite and "torch.nn.utils.clip_grad_norm_" within training for clip gradients. But it seems abnormal and usually cause "NAN" loss when adding mixed-precision training. It seems not like lightning, there is no grad scaler performed automatically. Does anyone meet these problems? @Dekakhrone Have you solved this problem?

[Dekakhrone]

Not yet, I had to take my mind off solving this problem. I will write a solution as soon as I find it.

[rohitgr7]

the ideal way is to use torch.nn.utils.clip_grad_norm_ but when it comes to different strategies like deepspeed which handles gradient_clipping internally, this is a limitation of lite that you need to configure that differently, but this is not the case with normal pytorchlightning code using LightningModules and Trainer.

But it seems abnormal and usually cause "NAN" loss when adding mixed-precision training. It seems not like lightning, there is no grad scaler performed automatically.

loss is scaled automatically when you do self.backward

[jsdtwry]

loss is scaled automatically when you do self.backward

What about gradient? We know when use gradient clipping under mixed-precision training, we need do some grad unscale like

optimizer.zero_grad()
with autocast():
    output = model(input)
    loss = loss_fn(output, target)
scaler.scale(loss).backward()

scaler.unscale_(optimizer)
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm)
scaler.step(optimizer)
scaler.update()

Will this cause some problems with lightning-lite Multi-card training? Indeed we found some abnormalities in training. My optimizer is wrappered and use self.backward(scaler.scale(loss)) to unscale.
And what's your suggestion on gradient clip in mixed-precision multi-card training.

[KavyaBPadmesh]

similar issue, I am trying to apply gradient clipping to certain parameter groups by customizing the configure_gradient_clipping trainer hook. Since clip_gradients is handled internally by the lightning plugin, is there a way to apply clipping to few parameter groups?