Comparable GPU: Optimizing the BERT Model with AMX Feature

Abstract

BERT is widely used in natural language processing (NLP) tasks in AI field. BERT has wide range of application scenarios. The performance of BERT determines the user experience feeling of application directly. AMX technology is a new feature introduced by Intel CPU, which supports two dimensional vector operations to optimize matrix operations. This paper uses AMX features, combined with optimization techniques such as operator fusion and quantization, to significantly improve the inference performance of BERT. Under the premise of a certain accuracy, compared with NVIDIA’s T4 GPU, in the BF16 small batch size scenario, the performance is improved by 1.2 times; Similarly, the performance of INT8 small batch size scene is 1.48 times higher.