[feat] Multiscale grad scaler

🚀 Feature

The GradScaler as supported in Torch AMP handles an optimal scaling factor in between fp32 and fp16 to catch possible underflow and overflows, as explained here.

Implement a multiscale grad scaler, either per param or or per param group, while trying to stick as close as possible to the current GradScaler API.

Motivation

One limitation of the current GradScaler is that there's a single scaling factor for the whole model, which means that this will only work if a single "window" is enough across the model for all the gradients. Empirically, it seems that this is a limitation for some very deep or hard to initialize models, which means that users fall back to not using AMP in that case (and loose the Tensor cores benefits on a V100 for instance).

Pitch

Enable Torch AMP for everyone, with a one-stop shop.

Alternatives

Not doing that, meaning that many people cannot use Torch AMP and fall back to fp32.

Additional context

Discussed internally at FB, got good feedback from a couple of users.