Deep Learning Recommendation Model (DLRM)

June 4, 2020

---

Deep Learning Recommendation Model is a Deep Neural Network for personalization and recommendation developed by Facebook. Due to its importance for large internet companies, it is one of the benchmarks in MLPerf. It is a simple network with 26 Embedding layers, 8 Linear layers and a Batched Matrix Multiply for calculating pairwise interactions. However, a single training step (forward and backward propagation) invokes about 800 GPU CUDA kernels (depending on the batch size). The best way to understand a DL network and GPU performance is to understand every single CUDA kernel i.e. which layer of the network invoked the kernel, with what arguments (tensor shapes and datatypes) and in which direction (forward propagation or backward propagation).

In this blog post, I will categorize every kernel used in the training of DLRM. All the information in the tables below was obtaining using Nvidia's PyTorch Profiler, PyProf, on a Turing T4 GPU. The information below is only a subset of what is provided by PyProf. The code and instructions for obtaining a detailed profile are here. Note that different GPUs will have slightly different kernel names e.g. volta_* as opposed to turing_*.

Model Architecture

Model Parameters

The parameters in the profiled code are as follows. These parameters were obtained from Nvidia Deep Learning Examples. The only constraint is that the output of the bottom MLP should be equal to the embedding size. The vocabulary sizes of the categorical embeddings are for a synthetic dataset and not the real Criteo Terabyte Dataset. For the purpose of this article, I modified the shape of the categorical embeddings from [100000]*26 to [100000 + 100*x for x in range(26)], so that its easier to differentiate and understand the kernels associated with each of the 26 embeddings, both in forward and backward propagation. With the Criteo dataset, the categorical embeddings tables become really big. With an embedding size of 128, the embedding tables have ~ 24 Billion parameters ~ 96 GB, and it is not possible to store them on a single GPU.

Bottom MLP = [512, 256, 128]
Top MLP = [1024, 1024, 512, 256, 1]
Numerical Features (N) = 13
Categorical Features (M) = 26
Categorical Embeddings = [100000 + 100*x for x in range(26)]
Embedding Size = 128
Batch Size = 32768

GPU Kernels

The tables below show the GPU kernels invoked in 1 training step. For every GPU kernel we show the direction (fprop, bprop), name of the layer, name of the operation, and the input tensor shapes / matrix dimensions for the operation. PyProf provides a lot of additional information for every GPU kernel e.g. grid dimensions, block dimensions, silicon time, datatypes, flops, bytes, tensor core usage and so on.

Bottom MLP (fprop)

Kernels 1 through 11 correspond to the forward propagation through the 3 Linear layers in the Bottom MLP. Kernels 2, 3 and 4 correspond to bias, linear and relu operations of the first Linear layer respectively. Kernels 5-7, 8-10 correspond to the second and third Linear layers respectively. Kernels 1 and 11 correspond to data type conversions from fp32 to fp16 and back.

Embedding (fprop)

Kernel 12 corresponds to a data type cast operation on the index tensor. Kernels 13 through 38 correspond to the forward propagation through the 26 Embedding layers. You can notice the size of the index tensor (batch size) and the size of the embedding tables (vocab x embedding size).

Pairwise Interaction (fprop)

Kernels 39 through 44 correspond to the pairwise interaction layer. Pairwise interaction is implemented as a dot product of every vector with all other vectors. If there are $N$ vectors (of the same length), then we will have ${N \choose 2} = \frac{N\times(N-1)}{2}$ interactions (scalar outputs). This can be implemented by creating a matrix $A$ containing the $N$ vectors, calculating $A.A^T$ and then taking the lower (or upper) traingular matrix, since the output is symmetric. In some sense, this is similar to attention. When this operation is done for multiple batches (32768 in this example), it is implemented as a Batched Matrix Multiplication (BMM).

Kernel 39 concatenates the outputs of the 26 embedding lookups and the Bottom MLP, resulting in 27 vectors. Kernel 40 computes the pairwise interaction of 27 vectors, each of length 128 using a BMM operation. Kernel 41 extract the lower (or upper) triangular matrix, which gives us ${27 \choose 2} = 351$ values.

Concat (fprop)

Kernels 42 through 44 concatenate the output of the Pairwise Interaction layer with the output of the Bottom MLP and pad it with 1 element before feeding it to the Top MLP.

Top MLP (fprop)

Kernels 45 through 60 correspond to the forward propagation through the 5 Linear layers in the Top MLP. Kernels 46, 47 and 48 correspond to bias, linear and relu operations of the first Linear layer respectively. Kernels 49-51, 52-54, 55-57, 58-59 correspond to the second, third, fourth and fifth Linear layers respectively. Kernels 45 and 60 correspond to data type conversions from fp32 to fp16 and back.

Loss (fprop)

Kernels 61 through 75 compute the BCE loss.

Gradient

In this profile, we are not accumulating gradients, and hence kernels 76 through 117 zero out the buffers used for storing the gradients during back propagation. Note that there are separate buffers for storing the gradients of the weights, biases and embeddings. Kernels 76-81, 82-91, 92-117 correspond to the Bottom MLP, Top MLP and the Embedding layers respectively.

Loss Scaling

Kernels 118 through 122 most likely correspond to loss scaling and the beginning of back propagation.

Loss (bprop)

Kernels 123 through 126 correspond to the backward propagation through the Loss layer.

Top MLP (bprop)

Kernels 127 through 149 correspond to the backward propagation through the Top MLP layer. Kernels 128-130 calculate the data gradient and weight gradient through the last Linear layer. Kernel 131 calculates the bias gradient. Likewise kernels 132-136, 137-140, 141-144, 145-148 correspond to the backward propagation through the fourth, third, second, and the first Linear layer of the Top MLP respectively.

Pairwise Interaction (bprop)

Kernels 150 through 165 correspond to the backward propagation through the pairwise interaction layer. Kernels 150-161 probably correspond to the lower (or upper) triangular matrix indexing. Kernels 162 and 163 correspond to the two gradients through the BMM operation. The gradients correspond to the same input and therefore kernels 164 and 165 are probably adding the gradients.

Embedding (bprop)

Kernels 166 through 711 correspond to the backward propagation through the 26 Embedding layers.

Bottom MLP (bprop)

Kernels 712 through 726 correspond to the backward propagation through the Bottom MLP layer. Kernel 713 does the backward propagation through the ReLU layer, 714-716 calculate the data gradient and weight gradient through the last Linear layer. Kernel 717 calculates the bias gradient. Likewise kernels 718-722 and 723-726 correspond to the backward propagation through the second and first Linear layer of the Bottom MLP respectively.

At the end of back propagation, we have all the weight gradients. Kernels 727 through 788 correspond to reversing the loss scaling and applying the SGD optimizer on all parameters.