Background

During the inference of a large model with hundreds of millions of parameters, the model cannot be loaded to one card or one machine due to the large weight or input data. In this case, the inference computing task needs to be distributed to multiple cards or computers. Currently, common distributed parallelism policies include data parallelism, tensor parallelism (TP), and pipeline parallelism. Currently, the foundation model uses transformers and their variants as the basic structure. A transformer layer consists of a self attention block and a multilayer perceptron (MLP) layer, as shown in Figure 1. One tensor parallelism policy is to perform tensor parallelism on the self attention block and the MLP block at a transformer layer.

Figure 1 Tensor parallelism

Taking the MLP layer as an example, for the two GEMM (general matrix multiplication) calculations of the MLP, right matrices A and B, A is parallelized by column slicing, and after obtaining the GELU layer output, B is parallelized by row slicing, and finally, an Allreduce operation is performed before Dropout.

Figure 2 Tensor parallelism of MLP

For a large model with a large number of transformer layers, frequent communication operations make communication latency a key factor in distributed inference. Introduction to ATB Communication Operators introduces four types of communication operators in the ATB, integrates the high-performance collective communication library and small-packet communication mechanism, and provides single-server multi-device and multi-server multi-device collective communication primitives to implement distributed inference. In addition, the ATB provides parallel communication operators, improving the inference performance.

Parent topic: Distributed Inference