Scaling Distributed Deep Learning Workloads beyond the Memory Capacity with KARMA

Abstract

The dedicated memory of hardware accelerators can be insufficient to store all weights and/or intermediate states of large deep learning models. Although model parallelism is a viable approach to reduce the memory pressure issue, significant modification of the source code and considerations for algorithms are required. An alternative solution is to use out-of-core methods instead of, or in addition to, data parallelism. We propose a performance model based on the concurrency analysis of out-of-core training behavior, and derive a strategy that combines layer swapping and redundant recomputing. We achieve an average of 1.52x speedup in six different models over the state-of-the-art out-of-core methods. We also introduce the first method to solve the challenging problem of out-of-core multi-node training by carefully pipelining gradient exchanges and performing the parameter updates on the host. Our data parallel out-of-core solution can outperform complex hybrid model parallelism in training large models, e.g. Megatron-LM and Turning-NLG.

arXiv commentACM/IEEE Proceedings of the International Conference for High Performance Computing, Networking, Storage and Analysis (SC'20)

References in corpus (8)

• PyTorch: An Imperative Style, High-Performance Deep Learning Library
• Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism
• Training Deep Nets with Sublinear Memory Cost
• Device Placement Optimization with Reinforcement Learning
• Supporting Very Large Models using Automatic Dataflow Graph Partitioning
• Yet Another Accelerated SGD: ResNet-50 Training on ImageNet in 74.7 seconds
• Exploring Hidden Dimensions in Parallelizing Convolutional Neural Networks
• Benanza: Automatic $\mu$Benchmark Generation to Compute "Lower-bound" Latency and Inform Optimizations of Deep Learning Models on GPUs