Potential solution to weird loss scaling behaviour

Hi there,

I'm currently working on improving the numerical stability of mixed precision training by proposing a new loss scaling method. It is called adaptive loss scaling, which calculates a proper loss scale for each gradient tensor through statistical properties. You may find the arxiv paper that describes this approach here: https://arxiv.org/abs/1910.12385, and the codebase: https://github.com/kumasento/ada-loss

Adaptive Loss Scaling

Existing loss scaling methods either decide the loss scale statically or dynamically, for all gradients. It is highly possible that the scale value doesn't fit for all gradients and all iterations.

Our approach calculates the loss scale for each gradient on-the-fly. Since the purpose of loss scaling is to reduce the rate of underflow, we use a statistical model to estimate that rate by given mean and variance of weights and activation gradients, and we aim to find a scale that can reduce the underflow rate down below a threshold. More details can be found in the paper.

The additional functionality we should implement besides the standard mixed precision training includes:

1. Calculating mean and variance of weight and gradient
2. Gradient tensors should carry their own loss scale, in order to perform unscaling before gradient update
3. The statistical model for calculating loss scale.

Migration to apex

I've already got some promising results from our current codebase. However, since this codebase is in chainer, we've met some difficulties to further improve our results:

1. There aren't much benchmarks that I can test with written in chainer, while PyTorch has a lot. I'm quite interested in seeing how our approach perform on NLP models.
2. There are potential performance issue with chainer, especially when we intend to perform mean and variance calculation on-the-fly. PyTorch seems to be ok, and at least we can revise its C++ core.

Therefore, we are interested in migrating this method to PyTorch/apex.

Help Needed

I'm currently going through the apex codebase and figuring out how our adaptive loss scaling method can be integrated. Meanwhile, if anyone is interested in this technique, please don't hesitate to let me know :). We're very keen to find any failure case from the current loss scaling method to see whether ours can make improvement.

@mcarilli @ptrblck your comments will be highly appreciated :)

Thanks!
Ruizhe

Start by reviewing Apex's existing loss-scaling implementation alongside the linked adaptive loss scaling paper and reference code. Define how mean and variance calculation, per-gradient loss scales, and unscaling would integrate with PyTorch/Apex, then validate the approach against mixed-precision and NLP benchmarks; done means the method is integrated and its numerical benefits are measured.