Efficient DNN Training Using Vectorized Block-Scaled GeMMs with Adaptive Block Shapes

Abstract

Reduced precision datatypes have become essential to the efficient training and deployment of Deep Neural Networks (DNNs). A recent development in the field has been the emergence of block-scaled datatypes: tensor representation formats derived from floating-point, that share a common exponent across multiple elements. While these formats are being broadly adopted and optimized for by DNN-specific inference accelerators, the potential benefits for training workloads on general-purpose vector processors has yet to be thoroughly explored. This work proposes implementations of block-scaled general matrix multiplications (GeMM) for DNN training at the edge using commercially available vector instruction sets (ARM SVE). Using these implementations, we highlight an accuracy-speed trade-off involving the shape of shared exponent blocks—vectors or squares. We exploit this result to optimize the training of fully connected networks by adapting the shared exponent block shapes during training. We also explore deployments on scalable hardware vector lengths utilizing different degrees of data-level parallelism. We demonstrate more efficient DNN training with our block-scaled datatypes compared to standard IEEE 32-bit floating point (FP32) formats through effective hardware-software cooptimizations.