  In the last few years, the memory requirements to train state-of-the-art
neural networks have far exceeded the DRAM capacities of modern hardware
accelerators. This has necessitated the development of efficient algorithms to
train these neural networks in parallel on large-scale GPU-based clusters.
Since computation is relatively inexpensive on modern GPUs, designing and
implementing extremely efficient communication in these parallel training
algorithms is critical for extracting the maximum performance. This paper
presents AxoNN, a parallel deep learning framework that exploits asynchrony and
message-driven execution to schedule neural network operations on each GPU,
thereby reducing GPU idle time and maximizing hardware efficiency. By using the
CPU memory as a scratch space for offloading data periodically during training,
AxoNN is able to reduce GPU memory consumption by four times. This allows us to
increase the number of parameters per GPU by four times, thus reducing the
amount of communication and increasing performance by over 13%. When tested
against large transformer models with 12-100 billion parameters on 48-384
NVIDIA Tesla V100 GPUs, AxoNN achieves a per-GPU throughput of 49.4-54.78% of
theoretical peak and reduces the training time by 22-37 days (15-25% speedup)
as compared to the state-of-the-art.
