Training Large Neural Networks with Constant Memory using a New Execution Algorithm

By Bharadwaj Pudipeddi, Maral Mesmakhosroshahi, Jinwen Xi, S. Bharadwaj. 2020

75 citations recorded by the source corpus.

Widely popular transformer-based NLP models such as BERT and Turing-NLG have enormous capacity trending to billions of parameters. Current execution methods demand brute-force resources such as HBM devices and high speed interconnectivity for data parallelism. In this paper, we introduce a new relay-style execution technique called L2L (layer-to-layer) where at any given moment, the device memory is primarily populated only with the executing layer(s)'s footprint. The model resides in the DRAM memory attached to either a CPU or an FPGA as an entity we call eager param-server (EPS). To overcome the bandwidth issues of shuttling parameters to and from EPS, the model is executed a layer at a time across many micro-batches instead of the conventional method of minibatches over whole model. L2L is implemented using 16GB V100 devices for BERT-Large running it with a device batch size of up to 256. Our results show 45% reduction in memory and 40% increase in the throughput compared to the state-of-the-art baseline. L2L is also able to fit models up to 50 Billion parameters on a machine with a single 16GB V100 and 512GB CPU memory and without requiring any model partitioning. L2L scales to arbitrary depth allowing researchers to develop on affordable devices which is a big step toward democratizing AI. By running the optimizer in the host EPS, we show a new form of mixed precision for faster throughput and convergence. In addition, the EPS enables dynamic neural architecture approaches by varying layers across iterations. Finally, we also propose and demonstrate a constant memory variation of L2L and we propose future enhancements. This work has been performed on GPUs first, but also targeted towards all high TFLOPS/Watt accelerators.

Subfields: Efficiency & Compression

How researchers see this paper

The researchers whose trusted networks rank it highest, strongest advocate first. Every researcher gets a different order.

In Efficiency & Compression

Citations in the corpus

ZeRO: Memory optimizations Toward Training Trillion Parameter Models (2019, background): “2) CPU Offload: [18], [19] exploit heterogeneous nature of today’s compute nodes, offloading model states to CPU memory through algorithmic design, virtualized memory, respectively.”
Fast Inference of Mixture-of-Experts Language Models with Offloading (2023, methodology): “This technique works particularly well when processing large batches of data, during training Pudipeddi et al. (2020); Ren et al. (2021) or large-batch non-interactive inference Aminabadi et al. (2022); Sheng et al. (2023), where each layer processes a lot of tokens each time the layer is loaded…”
8-bit Optimizers via Block-wise Quantization (2021, background): “They include gradient checkpointing (Chen et al., 2016), reversible residual connections (Gomez et al., 2017), and offloading (Pudipeddi et al., 2020; Rajbhandari et al., 2021).”
Parallel Training of Pre-Trained Models via Chunk-Based Dynamic Memory Management (2021, methodology): “ZeRO-Offload and L2L [31] transfers data among different memory spaces in granularity of tensor.”
Matching Guided Distillation (2020, background): “Whereas this manner can alleviate the issue, it still has two limitations: 1) The adaptation module introduces more parameters (including weights, gradients and optimizer states) into training [27].”