(A) NAND flash-based on-die processing architecture for stochastic optimization of large scale DNNs

Abstract

Training deep neural network (DNN) models is a resource-intensive, iterative process. For this reason, nowadays, complex optimizers like Adam are widely adopted as it increases the speed and efficiency of training. These optimizers, however, employ additional variables and raise the memory demand 2x to 3x of model parameters itself, worsening the memory capacity bottleneck when training large scale DNN models. Moreover, as the size of DNN models is projected to grow even further, it is not practical to assume that the future models will fit in accelerator memory. This has triggered various efforts to offload models to flash-based storage. However, when the model, especially the optimizer is offloaded to flash, the limited I/O bandwidth becomes a bottleneck, severely slowing down the overall training process. To this end, a solid-state drive (SSD) system with on-die processing (ODP) flash memory architecture is presented for gradient descent-based DNN models. Proposed scheme accelerates, or enables the training of large scale models by processing optimization stage in the storage device, specifically inside the flash dies. ODP capability of proposed architecture eliminates the heavy data movement over external interconnect and internal flash channels. In addition, proposed scheme involves customized flash translation layer (FTL) to support efficient memory management. Overall, it achieves, on average, a 2.8x speedup and a 3.6x improved energy efficiency in the weight update stage over baseline SSD offloading.