Novel Weight Prediction Technique to Mitigate Staleness and Consistency Issue for weight updates in DNN Training

DNN training is compute-intensive wherein model is partitioned across GPUs because of large amount of training data and model size. GPU is responsible for the weight updates of assigned model layers and one only GPU is active at a time in parallelism implementation to enable parallel execution. In pipelined parallelism multiple mini-batches are injected into the model concurrently and introduce staleness and consistency issue for weight updates. Since multiple mini- batches are simultaneously processed, later mini-batch could start the training process before its prior mini-batch updates weights. Consistency issue alleviated by ensuring the forward and backward passes of a mini-batch in the same GPU as the same weights through storing multiple versions of weights. However, weight staleness problem still unsolved. In this research article pipelined parallelism addresses unstable learning which leads to worse model accuracy due to weight staleness problem. Proposed mechanism, predicts the future weights in early pipeline stages, so that a mini-batch can adopt the predicted future weights rather than the stale weights to perform its computation. The design is based on the observation that smoothed gradients used in momentum-based optimizers reflect the trend of weight updates and can be leveraged for accurate weight prediction.