Configurable gradient clipping to avoid exploding gradients

[forklady42]

In an early training run on the original grid data, I ran into NaNs for losses.

After adding gradient logging locally, I discovered some major spikes in the gradients. This indicates that exploding gradients are likely the cause of the NaNs.

Indeed, after adding gradient clipping to avoid the large spikes, I ran 10 epochs on the original gridded data without any NaNs for losses. This will be another hyperparameter for us to tune.

I suspect smarter weight initialization, learning rate adjustments, and revisiting normalization will help stabilize the gradients, but in the meantime, gradient clipping will prevent them from spiraling out of control.

Note: I'm basing this PR on hananol/setup-ptlightning since this configuration change is being passed into the Lightning trainer. However, no need to block merging #26 on this PR. I can rebase onto main once the PyTorch Lightning refactor is merged.

[hanaol]

@forklady42, what gradient_clip_value did you use in your runs? I’ll add it to the configuration file and include it in this PR.

[forklady42]

I used 1.0, same as the default here, but honestly, I think that's more aggressive than necessary. Let's start with 20.0. That will head off any particularly large gradients without being excessively restrictive.

I'll go ahead and add it to the MP config.