Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

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

Bharadwaj Pudipeddi, Maral Mesmakhosroshahi, Jinwen Xi, Sujeeth Bharadwaj

TL;DR

The paper tackles the memory bottlenecks of training giant transformer models by introducing L2L, a layer-to-layer execution paradigm where the executing layer resides on-device while the rest of the model and optimizer sit in a host-based eager param-server (EPS). By combining graph reduction to the EPS, per-layer micro-batching (inner looping), and cross mixed precision, L2L achieves constant device memory, enabling depths and widths far beyond conventional limits and even fitting 50B-parameter models on a single 16GB GPU with CPU memory. Empirical results on GLUE show comparable accuracy to baselines with substantial memory savings (≈45%) and throughput gains (≈40%), with near-linear multi-GPU scaling and the ability to push model depth without partitioning. The approach promises democratization of AI by lowering hardware requirements and lays out an extended version, L2Lp, for further parallelization and scalability.

Abstract

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.

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

TL;DR

The paper tackles the memory bottlenecks of training giant transformer models by introducing L2L, a layer-to-layer execution paradigm where the executing layer resides on-device while the rest of the model and optimizer sit in a host-based eager param-server (EPS). By combining graph reduction to the EPS, per-layer micro-batching (inner looping), and cross mixed precision, L2L achieves constant device memory, enabling depths and widths far beyond conventional limits and even fitting 50B-parameter models on a single 16GB GPU with CPU memory. Empirical results on GLUE show comparable accuracy to baselines with substantial memory savings (≈45%) and throughput gains (≈40%), with near-linear multi-GPU scaling and the ability to push model depth without partitioning. The approach promises democratization of AI by lowering hardware requirements and lays out an extended version, L2Lp, for further parallelization and scalability.

Abstract

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.

Paper Structure

This paper contains 20 sections, 6 equations, 5 figures, 7 tables, 2 algorithms.

Table of Contents

  1. Introduction
  2. Related Work
  3. L2L Contribution
  4. Proposed Algorithm
  5. L2L Graph Reduction
  6. Inner Looping in L2L
  7. Computational Complexity Analysis on Minimizing Transfer Overhead
  8. Cross Mixed Precision
  9. Experimental Results
  10. Experimental Data and Setups
  11. Memory and Throughput Test
  12. Increasing L2L Throughput with Inner Looping
  13. L2L Scaling From Single GPU to Multi-GPU
  14. Going Beyond BERT
  15. Constant Memory by Depth of the Model
  16. ...and 5 more sections

Figures (5)

  • Figure 1: L2L execution with EPS compared with conventional execution.
  • Figure 2: Model graph residing in device for L2L compared with conventional execution.
  • Figure 3: L2L execution with micro-batch looping compared with the conventional data parallelism.
  • Figure 4: Validation curves for finetuning the GLUE tasks for $3$ epochs.
  • Figure 5: L2Lp's Execution Process.