Mixed Precision Training

[matthew-z]

It seems that currently there is no fp16 support in AllenNLP, but I think it is quite straightforward to do in a clean way with Apex.

Advantages of FP16:

1. Larger batch size, especially useful for large embedding models like ELMO and BERT
2. Faster training: the latest GPUs (P100, V100, RTX) can achieve 60-80% speed up
3. Many experiments have showed that FP16 does not hurt accuracy for CV models.

Challenge:
Some functions in AllenNLP are not FP16 stable, such as masked_softmax

[matt-gardner]

Thanks! For now, this is pretty low on our priority list, but contributions would be welcome here.

[matt-peters]

Would love to see fp16 support in allennlp!

[Hyperparticle]

I played around with apex/amp and found that integrating it into allennlp is straightforward, with a measurable performance boost on a large CRF tagger model (RTX GPU).

After cloning and installing apex, you only need to modify allennlp/training/trainer.py (or create your own trainer).

Near the top after the imports, import amp and initialize it.

from apex import amp
amp_handle = amp.init()

Replace the self.optimizer = optimizer assignment, wrapping amp in the __init__() function.

self.optimizer = amp_handle.wrap_optimizer(optimizer)

Finally, replace loss.backward() in _train_epoch() with scaled loss.

with self.optimizer.scale_loss(loss) as scaled_loss:
    scaled_loss.backward()

I tested this with a larger version of the crf_tagger.json tutorial config on conll 2003 using cuda 10.0 and pytorch 1.0.0. With such a small model, performance was actually slower with fp16 enabled. However, once I boosted the embedding/RNN size to something like 256/4096, I could see ~17% faster training times per epoch on a Titan RTX (over 10 epochs). It's not close to 2x like I hoped, but there definitely may be optimizations in order that can improve this. I also notice a small regularization effect on this model in that fp32 training overfits early on where fp16 does not.

During training, I see log messages like "Gradient overflow, skipping update" being emitted from apex. I'm not sure how to suppress them without interfering with allennlp's logging.

As @matthew-z mentioned, some allennlp functions are not stable, and I've not investigated these.

I'm not sure what additional work is necessary to submit a pull request. Thoughts?

[matt-gardner]

Awesome, thanks for trying this out! @joelgrus is currently thinking about how to make the Trainer easier to modify, and this is another thing to keep in mind for that, I think. He can probably give better advice than me on what exactly to put into a PR.

For functions that aren't stable, this is because of hard-coded values that were made with fp32 in mind, right? Like 1e-13 in masked_softmax and 1e-45 in masked_log_softmax? Are there any other issues? I think before we can integrate anything into the library, we'd need to make sure that you won't silently get incorrect behavior when using these functions. Crashing is ok - you tried to use something we don't officially support, and it crashed, you can write your own functions to work around it. But doing the wrong thing without warning is no good.

[joelgrus]

I'm still not 100% sure what I'm going to do about the trainer. In the meantime, if you wanted to do a PR, I would probably do the following (once you've done the correctness checking that mattg pointed out):

(1) give the current trainer a private method

def _backpropagate(self, loss: torch.Tensor) -> None:
    """
    This is pulled out as a method because subclasses might want to do something
    special (e.g. mixed-precision training).
    """
    loss.backward()

(2) replace the call to loss.backward() in _train_epoch with a call to that method

(3) create a subclass ApexTrainer(Trainer) that just overrides the constructor (so it can wrap the optimizer) and also overrides the new _backpropagate function.

Like I said, I'm not quite sure yet how I'm going to clean up the Trainer code, but I'm pretty sure I could work with this.

[matthew-z]

I have tested masked_softmax and masked_log_softmax, and their current implementations will not work in pure fp16 mode, but in apex.amp mode they will be converted into fp32 automatically, so we don't need to worry them.

[matthew-z]

@Hyperparticle

I guess there is no need to wrap the optimizer, just

from apex import amp
amp_handler = amp.init()

with amp_handler.scale_loss(loss, self.optimizer) as scaled_loss:
    scaled_loss.backward()

[Hyperparticle]

@matthew-z

Calling amp_handler.scale_loss() is an alternative to wrapping the optimizer, but this assumes there is only one optimizer in the code. This assumption holds currently, but in the future, you may want multiple optimizers or backward passes (in which case, you must wrap the optimizer). Just something to consider.

[Hyperparticle]

Also, I ran a similar experiment with amp on/off with a Transformer-based tagger. For some reason, I'm getting ~5% slower training times on RTX with mixed precision enabled (even with larger models) despite Transformer being most suitable for tensor cores. Tested this on a vanilla PyTorch Transformer tagger as well, and I'm seeing something similar. There's got to be some bottleneck I'm missing, as fairseq claims a nearly 5x speedup over plain fp32.

[matthew-z]

@Hyperparticle
In my tests, apex can achieve 1.8x speed for BERT-large model, but nearly no speed up for RNN models.

fp16 provides about 200% ~ 250% speed over fp32 in very simple matrix multiplication and CNN with V100 in pytorch. I guess some special CUDA code required to achieve more speed up (5x sounds too surprising to me)

BTW, RTX game card's tensor cores are much lower than titan or V100 intentionally (otherwise people will not buy titan).

[scarecrow1123]

Wondering if there's a PR on this. The solution given by @matthew-z using the new amp API should work seamlessly with the current setup I reckon. Regarding multiple optimizer and the use of scale_loss, the current Trainer class anyways support only a single optimizer and hence the suggested solution should work very well. It does work for me in my local setup and would like to know if there are any other PRs for this. Thanks!

[matthew-z]

There was an unfinished PR #2467...

[jiacheng-xu]

any update?