[GitHub] [incubator-mxnet] ptrendx commented on issue #14173: [WIP] MXNet AMP (automatic mixed precision)

[GitBox <gi...@apache.org>]

ptrendx commented on issue #14173: [WIP] MXNet AMP (automatic mixed precision)
URL: https://github.com/apache/incubator-mxnet/pull/14173#issuecomment-475703714
 
 
   @pengzhao-intel The goal for AMP is to enable users to get the benefits of mixed precision training (so mixed fp16/fp32 dtypes) with as little code changes as possible. NVIDIA had a set of guidelines for a long time how to do mixed precision training, but it turned out that implementing them is too hard for typical users. That is the reason of making AMP (automatic mixed precision) tool that would automatically apply the guidelines to the code to make mixed precision training as easy as possible.
   
   To give an example (I will push the tutorial to this PR soon as well), the typical training loop in Gluon looks like this:
   ```
   net = get_network()
   trainer = mx.gluon.Trainer(...)
   for data in dataloader:
       with autograd.record(True):
           out = net(data)
           l = loss(out)
           autograd.backward(l)
        trainer.step()
   ```
   With AMP, the same loop would look like this:
   ```
   amp.init()
   net = get_network()
   trainer = mx.gluon.Trainer(...)
   amp.init_trainer(trainer)
   for data in dataloader:
       with autograd.record(True):
           out = net(data)
           l = loss(out)
           with amp.scale_loss(l, trainer) as scaled_loss:
               autograd.backward(scaled_loss)
        trainer.step()
   ```
   
   The methodology of AMP integration is pretty simple:
    - we wrap the functions in `mx.sym` and `mx.ndarray` (as well as the `Symbol` and `NDArray` class members) to insert appropriate casts
    - since we don't have dtype information at this point for the symbolic graph we insert a special `amp_cast` and `amp_multicast` functions there (for cast to fp16/fp32 and casting to widest type respectively) that mark themselves as as invisible as possible (inplace as well as inplaceidentity) in order to be able to just exit early if the dtypes on both ends of the op end up being the same and input shares the memory with the output.
    - we give utilities to scale the loss (by default dynamically, but I will also add ability to do static loss scaling, since dynamic loss scaling is slightly costly)
   
   We tested all (except for the newest pose estimation) of the networks from GluonCV and for all of them we saw between few to over 100% speedups running on V100, with typical speedup being over 50%, depending on what the relative cost of convolution/gemms are compared to the rest of the network.
   
   For nearly all of the networks the procedure was basically applying the 3 lines I showed earlier. The only problem was with detection networks that GluonCV first runs on the CPU to generate anchors and lack of FP16 GEMM on CPU makes it impossible (@pengzhao-intel it would be great if you could close that gap). For those networks I needed to change the script to generate anchors on the GPU and then copy them to the CPU for Gluon data loader to consume.

----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
 
For queries about this service, please contact Infrastructure at:
users@infra.apache.org


With regards,
Apache Git Services