DLRM For PyTorch

This repository provides a script and recipe to train the Deep Learning Recommendation Model (DLRM) to achieve state-of-the-art accuracy and is tested and maintained by NVIDIA.

Table Of Contents

Model overview
Setup
- Requirements
Quick Start Guide
Advanced
Performance
- Benchmarking
  - Training performance benchmark
  - Inference performance benchmark
- Results
Release notes
- Changelog
- Known issues

Model overview

The Deep Learning Recommendation Model (DLRM) is a recommendation model designed to make use of both categorical and numerical inputs. It was first described in Deep Learning Recommendation Model for Personalization and Recommendation Systems. This repository provides a reimplementation of the codebase provided originally here. The scripts provided enable you to train DLRM on the Criteo Terabyte Dataset.

Using the scripts provided here, you can efficiently train models that are too large to fit into a single GPU. This is because we use a hybrid-parallel approach, which combines model parallelism for the embedding tables with data parallelism for the Top MLP. This is explained in details in next sections.

This model uses a slightly different preprocessing procedure than the one found in the original implementation. You can find a detailed description of the preprocessing steps in the Dataset guidelines section.

Using DLRM you can train a high-quality general model for providing recommendations.

This model is trained with mixed precision using Tensor Cores on Volta, Turing, and NVIDIA Ampere GPU architectures. Therefore, researchers can get results up to 3.3x faster than training without Tensor Cores while experiencing the benefits of mixed precision training. It is tested against each NGC monthly container release to ensure consistent accuracy and performance over time.

Model architecture

DLRM accepts two types of features: categorical and numerical. For each categorical feature, an embedding table is used to provide dense representation to each unique value. The dense features enter the model and are transformed by a simple neural network referred to as "bottom MLP". This part of the network consists of a series of linear layers with ReLU activations. The output of the bottom MLP and the embedding vectors are then fed into the "dot interaction" operation. The output of "dot interaction" is then concatenated with the features resulting from the bottom MLP and fed into the "top MLP" which is also a series of dense layers with activations. The model outputs a single number which can be interpreted as a likelihood of a certain user clicking an ad.

Figure 1. The architecture of DLRM.

Default configuration

The following features were implemented in this model:

general
- static loss scaling for Tensor Cores (mixed precision) training
- hybrid-parallel multi-GPU training
preprocessing
- dataset preprocessing using Spark 3 on GPUs
- dataset preprocessing using NVTabular on GPUs

Feature support matrix

This model supports the following features:

Feature	DLRM
Automatic mixed precision (AMP)	yes
CUDA Graphs	yes
Hybrid-parallel multi-GPU with all-2-all	yes
Preprocessing on GPU with NVTabular	yes
Preprocessing on GPU with Spark 3	yes

Features

Automatic Mixed Precision (AMP) - enables mixed precision training without any changes to the code-base by performing automatic graph rewrites and loss scaling controlled by an environmental variable.

CUDA Graphs - This feature allows to launch multiple GPU operations through a single CPU operation. The result is a vast reduction in CPU overhead. The benefits are particularly pronounced when training with relatively small batch sizes. The CUDA Graphs feature has been available through a native PyTorch API starting from PyTorch v1.10.

Multi-GPU training with PyTorch distributed - our model uses torch.distributed to implement efficient multi-GPU training with NCCL. For details, see example sources in this repository or see the PyTorch Tutorial.

Preprocessing on GPU with NVTabular - Criteo dataset preprocessing can be conducted using NVTabular. For more information on the framework, see the Announcing the NVIDIA NVTabular Open Beta with Multi-GPU Support and New Data Loaders.

Preprocessing on GPU with Spark 3 - Criteo dataset preprocessing can be conducted using Apache Spark 3.0. For more information on the framework and how to leverage GPU to preprocessing, see the Accelerating Apache Spark 3.0 with GPUs and RAPIDS.

Mixed precision training

Mixed precision is the combined use of different numerical precisions in a computational method. Mixed precision training offers significant computational speedup by performing operations in the half-precision floating-point format while storing minimal information in single-precision to retain as much information as possible in critical parts of the network. Since the introduction of Tensor Cores in Volta, and following with both the Turing and Ampere architectures, significant training speedups are experienced by switching to mixed precision – up to 3.3x overall speedup on the most arithmetically intense model architectures. Using mixed precision training requires two steps:

Porting the model to use the FP16 data type where appropriate.
Adding loss scaling to preserve small gradient values.

The ability to train deep learning networks with lower precision was introduced in the Pascal architecture and first supported in CUDA 8 in the NVIDIA Deep Learning SDK.

For information about:

How to train using mixed precision, see the Mixed Precision Training paper and Training With Mixed Precision documentation.
Techniques used for mixed precision training, see the Mixed-Precision Training of Deep Neural Networks blog.
APEX tools for mixed precision training, see the NVIDIA Apex: Tools for Easy Mixed-Precision Training in PyTorch.

Enabling mixed precision

Mixed precision training is turned off by default. To turn it on issue the --amp flag to the main.py script.

Enabling TF32

TensorFloat-32 (TF32) is the new math mode in NVIDIA A100 GPUs for handling the matrix math also called tensor operations. TF32 running on Tensor Cores in A100 GPUs can provide up to 10x speedups compared to single-precision floating-point math (FP32) on Volta GPUs.

TF32 Tensor Cores can speed up networks using FP32, typically with no loss of accuracy. It is more robust than FP16 for models that require a high dynamic range for weights or activations.

For more information, refer to the TensorFloat-32 in the A100 GPU Accelerates AI Training, HPC up to 20x blog post.

TF32 is supported in the NVIDIA Ampere GPU architecture and is enabled by default.

Hybrid-parallel multi-GPU with all-2-all communication

Many recommendation models contain very large embedding tables. As a result, the model is often too large to fit onto a single device. This could be easily solved by training in a model-parallel way, using either the CPU or other GPUs as "memory donors". However, this approach is suboptimal as the "memory donor" devices' compute is not utilized. In this repository, we use the model-parallel approach for the bottom part of the model (Embedding Tables + Bottom MLP) while using a usual data parallel approach for the top part of the model (Dot Interaction + Top MLP). This way we can train models much larger than what would normally fit into a single GPU while at the same time making the training faster by using multiple GPUs. We call this approach hybrid-parallel.

The transition from model-parallel to data-parallel in the middle of the neural net needs a specific multi-GPU communication pattern called all-2-all which is available in our PyTorch 21.04-py3 NGC docker container. In the original DLRM whitepaper this has been also referred to as "butterfly shuffle".

In the example shown in this repository we train models of three sizes: "small" (~15 GB), "large" (~82 GB), and "xlarge" (~142 GB). We use the hybrid-parallel approach for the "large" and "xlarge" models, as they do not fit in a single GPU.

Embedding table placement and load balancing

We use the following heuristic for dividing the work between the GPUs:

The Bottom MLP is placed on GPU-0 and no embedding tables are placed on this device.
The tables are sorted from the largest to the smallest
Set max_tables_per_gpu := ceil(number_of_embedding_tables / number_of_available_gpus)
Repeat until all embedding tables have an assigned device:
- Out of all the available GPUs find the one with the largest amount of unallocated memory
- Place the largest unassigned embedding table on this GPU. Raise an exception if it does not fit.
- If the number of embedding tables on this GPU is now equal to max_tables_per_gpu remove this GPU from the list of available GPUs so that no more embedding tables will be placed on this GPU. This ensures the all2all communication is well balanced between all devices.

Preprocessing on GPU

Please refer to the "Preprocessing" section for a detailed description of the Apache Spark 3.0 and NVTabular GPU functionality

BYO dataset functionality overview

This section describes how you can train the DeepLearningExamples RecSys models on your own datasets without changing the model or data loader and with similar performance to the one published in each repository. This can be achieved thanks to Dataset Feature Specification, which describes how the dataset, data loader and model interact with each other during training, inference and evaluation. Dataset Feature Specification has a consistent format across all recommendation models in NVIDIA’s DeepLearningExamples repository, regardless of dataset file type and the data loader, giving you the flexibility to train RecSys models on your own datasets.

Glossary
Dataset Feature Specification
Data Flow in Recommendation Models in DeepLearning examples
Example of Dataset Feature Specification
BYO dataset functionality

Glossary

The Dataset Feature Specification consists of three mandatory and one optional section:

feature_spec provides a base of features that may be referenced in other sections, along with their metadata. Format: dictionary (feature name) => (metadata name => metadata value)

source_spec provides information necessary to extract features from the files that store them. Format: dictionary (mapping name) => (list of chunks)

Mappings are used to represent different versions of the dataset (think: train/validation/test, k-fold splits). A mapping is a list of chunks.
Chunks are subsets of features that are grouped together for saving. For example, some formats may constrain data saved in one file to a single data type. In that case, each data type would correspond to at least one chunk. Another example where this might be used is to reduce file size and enable more parallel loading. Chunk description is a dictionary of three keys:
- type provides information about the format in which the data is stored. Not all formats are supported by all models.
- features is a list of features that are saved in a given chunk. Order of this list may matter: for some formats, it is crucial for assigning read data to the proper feature.
- files is a list of paths to files where the data is saved. For Feature Specification in yaml format, these paths are assumed to be relative to the yaml file’s directory (basename). Order of this list matters: It is assumed that rows 1 to i appear in the first file, rows i+1 to j in the next one, etc.

channel_spec determines how features are used. It is a mapping (channel name) => (list of feature names).

Channels are model specific magic constants. In general, data within a channel is processed using the same logic. Example channels: model output (labels), categorical ids, numerical inputs, user data, and item data.

metadata is a catch-all, wildcard section: If there is some information about the saved dataset that does not fit into the other sections, you can store it here.

Dataset feature specification

Data flow can be described abstractly: Input data consists of a list of rows. Each row has the same number of columns; each column represents a feature. The columns are retrieved from the input files, loaded, aggregated into channels and supplied to the model/training script.

FeatureSpec contains metadata to configure this process and can be divided into three parts:

Specification of how data is organized on disk (source_spec). It describes which feature (from feature_spec) is stored in which file and how files are organized on disk.
Specification of features (feature_spec). Describes a dictionary of features, where key is feature name and values are features’ characteristics such as dtype and other metadata (for example, cardinalities for categorical features)
Specification of model’s inputs and outputs (channel_spec). Describes a dictionary of model’s inputs where keys specify model channel’s names and values specify lists of features to be loaded into that channel. Model’s channels are groups of data streams to which common model logic is applied, for example categorical/continuous data, user/item ids. Required/available channels depend on the model

The FeatureSpec is a common form of description regardless of underlying dataset format, dataset data loader form and model.

Data flow in NVIDIA Deep Learning Examples recommendation models

The typical data flow is as follows:

S.0. Original dataset is downloaded to a specific folder.
S.1. Original dataset is preprocessed into Intermediary Format. For each model, the preprocessing is done differently, using different tools. The Intermediary Format also varies (for example, for DLRM PyTorch, the Intermediary Format is a custom binary one.)
S.2. The Preprocessing Step outputs Intermediary Format with dataset split into training and validation/testing parts along with the Dataset Feature Specification yaml file. Metadata in the preprocessing step is automatically calculated.
S.3. Intermediary Format data together with Dataset Feature Specification are fed into training/evaluation scripts. Data loader reads Intermediary Format and feeds the data into the model according to the description in the Dataset Feature Specification.
S.4. The model is trained and evaluated

Fig.1. Data flow in Recommender models in NVIDIA Deep Learning Examples repository. Channels of the model are drawn in green.

Example of dataset feature specification

As an example, let’s consider a Dataset Feature Specification for a small CSV dataset for some abstract model.

feature_spec:
  user_gender:
    dtype: torch.int8
    cardinality: 3 #M,F,Other
  user_age: #treated as numeric value
    dtype: torch.int8
  user_id:
    dtype: torch.int32
    cardinality: 2655
  item_id:
    dtype: torch.int32
    cardinality: 856
  label:
    dtype: torch.float32

source_spec:
  train:
    - type: csv
      features:
        - user_gender
        - user_age
      files:
        - train_data_0_0.csv
        - train_data_0_1.csv
    - type: csv
      features:
        - user_id
        - item_id
        - label
      files:
        - train_data_1.csv
  test:
    - type: csv
      features:
        - user_id
        - item_id
        - label
        - user_gender
        - user_age
        
      files:
        - test_data.csv

channel_spec:
  numeric_inputs: 
    - user_age
  categorical_user_inputs: 
    - user_gender
    - user_id
  categorical_item_inputs: 
    - item_id
  label_ch: 
    - label

The data contains five features: (user_gender, user_age, user_id, item_id, label). Their data types and necessary metadata are described in the feature specification section.

In the source mapping section, two mappings are provided: one describes the layout of the training data, the other of the testing data. The layout for training data has been chosen arbitrarily to showcase the flexibility. The train mapping consists of two chunks. The first one contains user_gender and user_age, saved as a CSV, and is further broken down into two files. For specifics of the layout, refer to the following example and consult the glossary. The second chunk contains the remaining columns and is saved in a single file. Notice that the order of columns is different in the second chunk - this is alright, as long as the order matches the order in that file (that is, columns in the .csv are also switched)

Let’s break down the train source mapping. The table contains example data color-paired to the files containing it.

The channel spec describes how the data will be consumed. Four streams will be produced and available to the script/model. The feature specification does not specify what happens further: names of these streams are only lookup constants defined by the model/script. Based on this example, we can speculate that the model has three input channels: numeric_inputs, categorical_user_inputs, categorical_item_inputs, and one output channel: label. Feature names are internal to the FeatureSpec and can be freely modified.

BYO dataset functionality

In order to train any Recommendation model in NVIDIA Deep Learning Examples one can follow one of three possible ways:

One delivers already preprocessed dataset in the Intermediary Format supported by data loader used by the training script (different models use different data loaders) together with FeatureSpec yaml file describing at least specification of dataset, features and model channels
One uses a transcoding script
One delivers dataset in non-preprocessed form and uses preprocessing scripts that are a part of the model repository. In order to use already existing preprocessing scripts, the format of the dataset needs to match the one of the original datasets. This way, the FeatureSpec file will be generated automatically, but the user will have the same preprocessing as in the original model repository.

Setup

The following section lists the requirements for training DLRM.

Requirements

This repository contains Dockerfile which extends the PyTorch NGC container and encapsulates some dependencies. Aside from these dependencies, ensure you have the following components:

For more information about how to get started with NGC containers, see the following sections from the NVIDIA GPU Cloud Documentation and the Deep Learning Documentation:

For those unable to use the PyTorch NGC container, to set up the required environment or create your own container, see the versioned NVIDIA Container Support Matrix.

Quick Start Guide

To train your model using mixed or TF32 precision with Tensor Cores or using FP32, perform the following steps using the default parameters of DLRM on the Criteo Terabyte dataset. For the specifics concerning training and inference, see the Advanced section.

Clone the repository.

git clone https://github.com/NVIDIA/DeepLearningExamples
cd DeepLearningExamples/PyTorch/Recommendation/DLRM

Download the dataset.

You can download the data by following the instructions at: http://labs.criteo.com/2013/12/download-terabyte-click-logs/. When you have successfully downloaded it and unpacked it, set the CRITEO_DATASET_PARENT_DIRECTORY to its parent directory:

CRITEO_DATASET_PARENT_DIRECTORY=/raid/dlrm

We recommend to choose the fastest possible file system, otherwise it may lead to an IO bottleneck.

Build DLRM Docker containers

docker build -t nvidia_dlrm_pyt .
docker build -t nvidia_dlrm_preprocessing -f Dockerfile_preprocessing . --build-arg DGX_VERSION=[DGX-2|DGX-A100]

Start an interactive session in the NGC container to run preprocessing. The DLRM PyTorch container can be launched with:

docker run --runtime=nvidia -it --rm --ipc=host  -v ${CRITEO_DATASET_PARENT_DIRECTORY}:/data/dlrm nvidia_dlrm_preprocessing bash

Preprocess the dataset.

Here are a few examples of different preprocessing commands. Out of the box, we support preprocessing on DGX-2 and DGX A100 systems. For the details on how those scripts work and detailed description of dataset types (small FL=15, large FL=3, xlarge FL=2), system requirements, setup instructions for different systems and all the parameters consult the preprocessing section. For an explanation of the FL parameter, see the Dataset Guidelines and Preprocessing sections.

Depending on dataset type (small FL=15, large FL=3, xlarge FL=2) run one of following command:

4.1. Preprocess to small dataset (FL=15) with Spark GPU:

cd /workspace/dlrm/preproc
./prepare_dataset.sh 15 GPU Spark

4.2. Preprocess to large dataset (FL=3) with Spark GPU:

cd /workspace/dlrm/preproc
./prepare_dataset.sh 3 GPU Spark

4.3. Preprocess to xlarge dataset (FL=2) with Spark GPU:

cd /workspace/dlrm/preproc
./prepare_dataset.sh 2 GPU Spark

Start training.

First start the docker container (adding --security-opt seccomp=unconfined option is needed to take the full advantage of processor affinity in multi-GPU training):

docker run --security-opt seccomp=unconfined --runtime=nvidia -it --rm --ipc=host  -v ${PWD}/data:/data nvidia_dlrm_pyt bash

single-GPU:

python -m dlrm.scripts.main --mode train --dataset /data/dlrm/binary_dataset/ --amp --cuda_graphs

multi-GPU for DGX A100:

python -m torch.distributed.launch --no_python --use_env --nproc_per_node 8 \
          bash  -c './bind.sh --cpu=dgxa100_ccx.sh --mem=dgxa100_ccx.sh python -m dlrm.scripts.main \
          --dataset /data/dlrm/binary_dataset/ --seed 0 --epochs 1 --amp --cuda_graphs'

multi-GPU for DGX-1 and DGX-2:

python -m torch.distributed.launch --no_python --use_env --nproc_per_node 8 \
          bash  -c './bind.sh  --cpu=exclusive -- python -m dlrm.scripts.main \
          --dataset /data/dlrm/binary_dataset/ --seed 0 --epochs 1 --amp --cuda_graphs'

Start validation/evaluation. If you want to run validation or evaluation, you can either:

use the checkpoint obtained from the training commands above, or
download the pretrained checkpoint from NGC.

In order to download the checkpoint from NGC, visit ngc.nvidia.com website and browse the available models. Download the checkpoint files and unzip them to some path, for example, to $CRITEO_DATASET_PARENT_DIRECTORY/checkpoints/. The checkpoint requires around 15GB of disk space.

Commands:

single-GPU:

python -m dlrm.scripts.main --mode test --dataset /data/dlrm/binary_dataset/ --load_checkpoint_path `$CRITEO_DATASET_PARENT_DIRECTORY/checkpoints/checkpoint`

multi-GPU for DGX A100:

python -m torch.distributed.launch --no_python --use_env --nproc_per_node 8 \
          bash  -c './bind.sh --cpu=dgxa100_ccx.sh --mem=dgxa100_ccx.sh python -m dlrm.scripts.main \
          --dataset /data/dlrm/binary_dataset/ --seed 0 --epochs 1 --amp --load_checkpoint_path `$CRITEO_DATASET_PARENT_DIRECTORY/checkpoints/checkpoint`'

multi-GPU for DGX-1 and DGX-2:

python -m torch.distributed.launch --no_python --use_env --nproc_per_node 8 \
          bash  -c './bind.sh  --cpu=exclusive -- python -m dlrm.scripts.main \
          --dataset /data/dlrm/binary_dataset/ --seed 0 --epochs 1 --amp --load_checkpoint_path `$CRITEO_DATASET_PARENT_DIRECTORY/checkpoints/checkpoint`'

Advanced

The following sections provide greater details of the dataset, running training and inference, and the training results.

Scripts and sample code

The dlrm/scripts/main.py script provides an entry point to most of the functionality. Using different command-line flags allows you to run training, validation, and benchmark both training and inference on real or synthetic data.

Utilities related to loading the data reside in the data directory.

Command-line options

The dlrm/scripts/main.py script supports a number of command-line flags. You can get the descriptions of those by running python -m dlrm.scripts.main --help.

The following example output is printed when running the model:

Epoch:[0/1] [200/128028]  eta: 1:28:44  loss: 0.1782  step_time: 0.041657  lr: 0.8794
Epoch:[0/1] [400/128028]  eta: 1:25:15  loss: 0.1403  step_time: 0.038504  lr: 1.7544
Epoch:[0/1] [600/128028]  eta: 1:23:56  loss: 0.1384  step_time: 0.038422  lr: 2.6294
Epoch:[0/1] [800/128028]  eta: 1:23:13  loss: 0.1370  step_time: 0.038421  lr: 3.5044
Epoch:[0/1] [1000/128028]  eta: 1:22:45  loss: 0.1362  step_time: 0.038464  lr: 4.3794
Epoch:[0/1] [1200/128028]  eta: 1:22:24  loss: 0.1346  step_time: 0.038455  lr: 5.2544
Epoch:[0/1] [1400/128028]  eta: 1:22:07  loss: 0.1339  step_time: 0.038459  lr: 6.1294
Epoch:[0/1] [1600/128028]  eta: 1:21:52  loss: 0.1320  step_time: 0.038481  lr: 7.0044
Epoch:[0/1] [1800/128028]  eta: 1:21:39  loss: 0.1315  step_time: 0.038482  lr: 7.8794
Epoch:[0/1] [2000/128028]  eta: 1:21:27  loss: 0.1304  step_time: 0.038466  lr: 8.7544
Epoch:[0/1] [2200/128028]  eta: 1:21:15  loss: 0.1305  step_time: 0.038430  lr: 9.6294

Getting the data

This example uses the Criteo Terabyte Dataset. The first 23 days are used as the training set. The last day is split in half. The first part, referred to as "test", is used for validating training results. The second one, referred to as "validation", is unused.

Dataset guidelines

The preprocessing steps applied to the raw data include:

Replacing the missing values with 0
Replacing the categorical values that exist fewer than FL times with a special value (FL value is called a frequency threshold or a frequency limit)
Converting the hash values to consecutive integers
Adding 3 to all the numerical features so that all of them are greater or equal to 1
Taking a natural logarithm of all numerical features

BYO dataset

This implementation supports using other datasets thanks to BYO dataset functionality. The BYO dataset functionality allows users to plug in their dataset in a common fashion for all Recommender models that support this functionality. Using BYO dataset functionality, the user does not have to modify the source code of the model thanks to the Feature Specification file. For general information on how BYO dataset works, refer to the BYO dataset overview section.

There are three ways to plug in user's dataset:

1. Provide an unprocessed dataset in a format matching the one used by Criteo 1TB, then use Criteo 1TB's preprocessing. Feature Specification file is then generated automatically.

The required format of the user's dataset is:

The data should be split into text files. Each line of those text files should contain a single training example. An example should consist of multiple fields separated by tabulators:

The first field is the label – 1 for a positive example and 0 for negative.
The next N tokens should contain the numerical features separated by tabs.
The next M tokens should contain the hashed categorical features separated by tabs.

The correct dataset files together with the Feature Specification yaml file will be generated automatically by preprocessing script.

For an example of using this process, refer to the Quick Start Guide

2. Provide a CSV containing preprocessed data and a simplified Feature Specification yaml file, then transcode the data with `transcode.py` script

This option should be used if the user has their own CSV file with a preprocessed dataset they want to train on.

The required format of the user's dataset is:

CSV files containing the data, already split into train and test sets.
Feature Specification yaml file describing the layout of the CSV data

For an example of a feature specification file, refer to the tests/transcoding folder.

The CSV containing the data:

should be already split into train and test
should contain no header
should contain one column per feature, in the order specified by the list of features for that chunk in the source_spec section of the feature specification file
categorical features should be non-negative integers in the range [0,cardinality-1] if cardinality is specified

The Feature Specification yaml file:

needs to describe the layout of data in CSV files
may contain information about cardinalities. However, if set to auto, they will be inferred from the data by the transcoding script.

Refer to tests/transcoding/small_csv.yaml for an example of the yaml Feature Specification.

The following example shows how to use this way of plugging user's dataset:

Prepare your data and save the path:

DATASET_PARENT_DIRECTORY=/raid/dlrm

Build the DLRM image with:

docker build -t nvidia_dlrm_pyt .

Launch the container with:

docker run --runtime=nvidia -it --rm --ipc=host  -v ${DATASET_PARENT_DIRECTORY}:/data nvidia_dlrm_preprocessing bash

If you are just testing the process, you can create synthetic csv data:

python -m dlrm.scripts.gen_csv --feature_spec_in tests/transcoding/small_csv.yaml

Convert the data:

mkdir /data/conversion_output
python -m dlrm.scripts.transcode --input /data --output /data/converted

You may need to tune the --chunk_size parameter. Higher values speed up the conversion but require more RAM.

This will convert the data from /data and save the output in /data/converted. A feature specification file describing the new data will be automatically generated.

To run the training on 1 GPU:

python -m dlrm.scripts.main --mode train --dataset /data/converted --amp --cuda_graphs

multi-GPU for DGX A100:

python -m torch.distributed.launch --no_python --use_env --nproc_per_node 8 \
          bash  -c './bind.sh --cpu=dgxa100_ccx.sh --mem=dgxa100_ccx.sh python -m dlrm.scripts.main \
          --dataset /data/converted --seed 0 --epochs 1 --amp --cuda_graphs'

multi-GPU for DGX-1 and DGX-2:

python -m torch.distributed.launch --no_python --use_env --nproc_per_node 8 \
          bash  -c './bind.sh  --cpu=exclusive -- python -m dlrm.scripts.main \
          --dataset /data/converted --seed 0 --epochs 1 --amp --cuda_graphs'

3. Provide a fully preprocessed dataset, saved in split binary files, and a Feature Specification yaml file

This is the option to choose if you want full control over preprocessing and/or want to preprocess data directly to the target format.

Your final output will need to contain a Feature Specification yaml describing data and file layout. For an example feature specification file, refer to tests/feature_specs/criteo_f15.yaml

For details, refer to the BYO dataset overview section.

Channel definitions and requirements

This model defines three channels:

categorical, accepting an arbitrary number of features
numerical, accepting an arbitrary number of features
label, accepting a single feature

The training script expects two mappings:

train
test

For performance reasons:

The only supported dataset type is split binary
Splitting chunks into multiple files is not supported.
Each categorical feature has to be provided in a separate chunk
All numerical features have to be provided in a single chunk
All numerical features have to appear in the same order in channel_spec and source_spec
Only integer types are supported for categorical features
Only float16 is supported for numerical features

BYO dataset constraints for the model

There are the following constraints of BYO dataset functionality for this model:

The performance of the model depends on the dataset size. Generally, the model should scale better for datasets containing more data points. For a smaller dataset, you might experience slower performance than the one reported for Criteo
Using other datasets might require tuning some hyperparameters (for example, learning rate, beta1 and beta2) to reach desired accuracy.
The optimized cuda interaction kernels for FP16 and TF32 assume that the number of categorical variables is smaller than WARP_SIZE=32 and embedding size is <=128

Preprocessing

The preprocessing scripts provided in this repository support running both on CPU and GPU using NVtabular (GPU only) and Apache Spark 3.0.

Please note that the preprocessing will require about 4TB of disk storage.

The syntax for the preprocessing script is as follows:

cd /workspace/dlrm/preproc
./prepare_dataset.sh <frequency_threshold> <GPU|CPU> <NVTabular|Spark>

For the Criteo Terabyte dataset, we recommend a frequency threshold of FL=3(when using A100 40GB or V100 32 GB) or FL=2(when using A100 80GB) if you intend to run the hybrid-parallel mode on multiple GPUs. If you want to make the model fit into a single NVIDIA Tesla V100-32GB, you can set FL=15.

The first argument means the frequency threshold to apply to the categorical variables. For a frequency threshold FL, the categorical values that occur less often than FL will be replaced with one special value for each category. Thus, a larger value of FL will require smaller embedding tables and will substantially reduce the overall size of the model.

The second argument is the hardware to use (either GPU or CPU).

The third arguments is a framework to use (either NVTabular or Spark). In case of choosing a CPU preprocessing this argument is omitted as it only Apache Spark is supported on CPU.

The preprocessing scripts make use of the following environment variables to configure the data directory paths:

download_dir – this directory should contain the original Criteo Terabyte CSV files
spark_output_path – directory to which the parquet data will be written
conversion_intermediate_dir – directory used for storing intermediate data used to convert from parquet to train-ready format
final_output_dir – directory to store the final results of the preprocessing which can then be used to train DLRM

In the final_output_dir will be three subdirectories created: train, test, validation, and one json file – model_size.json – containing a maximal index of each category. The train is the train dataset transformed from day_0 to day_22. The test is the test dataset transformed from the prior half of day_23. The validation is the dataset transformed from the latter half of day_23.

The model is tested on 3 datasets resulting from Criteo dataset preprocessing: small (Freqency threshold = 15), large (Freqency threshold = 3) and xlarge (Freqency threshold = 2). Each dataset occupies approx 370GB of disk space. Table below presents information on the supercomputer and GPU count that are needed to train model on particular dataset.

Dataset	GPU VRAM consumption*	Model checkpoint size*	FL setting	DGX A100 40GB, 1GPU	DGX A100 40GB, 8GPU	DGX A100 80GB, 1GPU	DGX A100 80GB, 8GPU	DGX-1** or DGX-2, 1 GPU	DGX-1** or DGX-2, 8GPU	DGX-2, 16GPU
small (FL=15)	20.5	15.0	15	Yes	Yes	Yes	Yes	Yes	Yes	Yes
large (FL=3)	132.3	81.9	3	NA	Yes	NA	Yes	NA	Yes	Yes
xlarge (FL=2)	198.8	141.3	2	NA	NA	NA	Yes	NA	NA	NA

*with default embedding dimension setting **DGX-1 V100 32GB

NVTabular

NVTabular preprocessing is calibrated to run on DGX A100 and DGX-2 AI systems. However, it should be possible to change the values of ALL_DS_MEM_FRAC, TRAIN_DS_MEM_FRAC, TEST_DS_MEM_FRAC, VALID_DS_MEM_FRAC in preproc/preproc_NVTabular.py, so that they'll work on also on other hardware platforms such as DGX-1 or a custom one.

Spark

The script spark_data_utils.py is a PySpark application, which is used to preprocess the Criteo Terabyte Dataset. In the Docker image, we have installed Spark 3.0.1, which will start a standalone cluster of Spark. The scripts run_spark_cpu.sh and run_spark_gpu.sh start Spark, then run several PySpark jobs with spark_data_utils.py.

Note that the Spark job requires about 3TB disk space used for data shuffling.

Spark preprocessing is calibrated to run on DGX A100 and DGX-2 AI systems. However, it should be possible to change the values in preproc/DGX-2_config.sh or preproc/DGX-A100_config.sh so that they'll work on also on other hardware platforms such as DGX-1 or a custom one.

Training process

The main training script resides in dlrm/scripts/main.py. Once the training is completed, it stores the checkpoint in the path specified by --save_checkpoint_path and a JSON training log in --log_path. The quality of the predictions generated by the model is measured by the ROC AUC metric. The speed of training and inference is measured by throughput i.e., the number of samples processed per second. We use mixed precision training with static loss scaling for the bottom and top MLPs while embedding tables are stored in FP32 format.

Inference process

This section describes inference with PyTorch in Python. If you're interested in inference using the Triton Inference Server, refer to triton/README.md file.

Two modes for inference are currently supported by the dlrm/scripts/main.py script:

Inference benchmark – this mode will measure and print out throughput and latency numbers for multiple batch sizes. You can activate it by passing the --mode inference_benchmark command line flag. The batch sizes to be tested can be set with the --inference_benchmark_batch_sizes command-line argument.
Test-only – this mode can be used to run a full validation on a checkpoint to measure ROC AUC. You can enable it by passing --mode test.

Deploying DLRM Using NVIDIA Triton Inference Server

The NVIDIA Triton Inference Server provides a cloud inferencing solution optimized for NVIDIA GPUs. The server provides an inference service via an HTTP or gRPC endpoint, allowing remote clients to request inferencing for any model being managed by the server. More information on how to perform inference using NVIDIA Triton Inference Server can be found in triton/README.md.

Performance

The performance measurements in this document were conducted at the time of publication and may not reflect the performance achieved from NVIDIA’s latest software release. For the most up-to-date performance measurements, go to NVIDIA Data Center Deep Learning Product Performance.

Benchmarking

The following section shows how to run benchmarks measuring the model performance in training and inference modes.

Training performance benchmark

To benchmark the training performance on a specific batch size, please follow the instructions in the Quick Start Guide. You can also add the --max_steps 1000 --benchmark_warmup_steps 500 if you want to get a reliable throughput measurement without running the entire training.

You can create a synthetic dataset by running python -m dlrm.scripts.prepare_synthetic_dataset --synthetic_dataset_dir /tmp/dlrm_synthetic_data if you haven't yet downloaded the dataset.

Inference performance benchmark

To benchmark the inference performance on a specific batch size, run:

python -m dlrm.scripts.main --mode inference_benchmark --dataset /data

You can also create a synthetic dataset by running python -m dlrm.scripts.prepare_synthetic_dataset --synthetic_dataset_dir /tmp/dlrm_synthetic_data if you haven't yet downloaded the dataset.

Results

The following sections provide details on how we achieved our performance and accuracy in training and inference.

We used three model size variants to show memory scalability in a multi-GPU setup:

Model variant	Frequency threshold	Model size
small	15	15 GB
large	3	82 GB
xlarge	2	142 GB

Training accuracy results

Training accuracy: NVIDIA DGX A100 (8x A100 80GB)

Our results were obtained by running dlrm/scripts/main.py script as described in the Quick Start Guide in the DLRM Docker container using NVIDIA A100 80GB GPUs.

GPUs	Model size	Batch size / GPU	Accuracy (AUC) - TF32	Accuracy (AUC) - mixed precision	Time to train - TF32]	Time to train - mixed precision	Time to train speedup (TF32 to mixed precision)
8	large	8k	0.802509	0.802528	0:06:27	0:04:36	1.40217
1	small	64k	0.802537	0.802521	0:24:26	0:17:47	1.37395

Training accuracy: NVIDIA DGX-1 (8x V100 32GB)

Our results were obtained by running dlrm/scripts/main.py script as described in the Quick Start Guide in the DLRM Docker container using NVIDIA V100 32GB GPUs.

GPUs	Model size	Batch size / GPU	Accuracy (AUC) - FP32	Accuracy (AUC) - mixed precision	Time to train - FP32]	Time to train - mixed precision	Time to train speedup (FP32 to mixed precision)
8	large	8k	0.802568	0.802562	0:28:24	0:11:45	2.41702
1	small	64k	0.802784	0.802723	1:58:10	0:38:17	3.08663

Training accuracy plots

Models trained with FP32, TF32, and Automatic Mixed Precision (AMP) achieve similar accuracy.

The plot represents ROC AUC metric as a function of steps (step is single batch) during training for default precision (FP32 for Volta architecture (DGX-1) and TF32 for Ampere GPU architecture (DGX-A100)), and AMP for all three datasets. All other parameters of training are default.

Figure 1. Training stability for a FL3 dataset: distribution of ROC AUC across different configurations. 'All configurations' refer to the distribution of ROC AUC for cartesian product of architecture, training precision.

Figure 2. Training stability for a FL15 dataset: distribution of ROC AUC across different configurations. 'All configurations' refer to the distribution of ROC AUC for cartesian product of architecture, training precision.

Training stability test

Training of the model is stable for multiple configurations achieving a standard deviation of 10e-4. The model achieves similar ROC AUC scores for A100 and V100, training precisions. It was trained for one epoch (roughly 4 billion samples, 64014 batches), starting from 10 different initial random seeds for each setup. The training was performed in the pytorch:21.10-py3 NGC container with and without mixed precision enabled. The provided charts and numbers consider single and multi GPU training. After training, the models were evaluated on the test set. The following plots compare distributions of ROC AUC on the test set.

Figure 3. Training stability for a FL3 dataset: distribution of ROC AUC across different configurations. 'All configurations' refer to the distribution of ROC AUC for cartesian product of architecture, training precision.

Figure 4. Training stability for a FL15 dataset: distribution of ROC AUC across different configurations. 'All configurations' refer to the distribution of ROC AUC for cartesian product of architecture, training precision.

Impact of mixed precision on training accuracy

The accuracy of training, measured with ROC AUC on the test set after the final epoch metric was not impacted by enabling mixed precision. The obtained results were statistically similar. The similarity was measured according to the following procedure:

The model was trained 10 times for default settings (FP32 or TF32 for Volta and Ampere architecture respectively) and 10 times for AMP. After the last epoch, the accuracy score ROC AUC was calculated on the test set.

Distributions for two hardware configurations (A100, V100) for 2 datasets are presented below.

Figure 5. Impact of AMP on ROC AUC distribution for A100 and V100 GPUs for single- and multi-gpu training on a dataset with a frequency threshold of 3.

Figure 6. Impact of AMP on ROC AUC distribution for A100 and V100 GPUs for single- and multi-gpu training on a dataset with a frequency threshold of 15.

Distribution of AUC ROC for single precision training (TF32 for A100, FP32 for Volta) and AMP training were compared in terms of mean, variance and Kolmogorov–Smirnov test to state statistical difference between single precision and AMP results. Refer to the expandable table below.

Full tabular data for AMP influence on AUC ROC

Hardware	Dataset	GPUs	mean AUC ROC for full precision	std AUC ROC for full precision	mean AUC ROC for AMP	std AUC ROC for AMP	KS test value: statictics, p-value
DGX A100	FL3	8	0.802681	0.000073	0.802646	0.000063	('0.400', '0.418')
DGX-2	FL3	16	0.802614	0.000073	0.802623	0.000122	('0.267', '0.787')
Sample size was set to 10 experiments for each training setup.

Training performance results

We used throughput in items processed per second as the performance metric.

Training performance: NVIDIA DGX A100 (8x A100 80GB)

Our results were obtained by running the following commands:

for single-GPU setup:

python -m dlrm.scripts.main --dataset /data --amp --cuda_graphs

for multi-GPU setup:

python -m torch.distributed.launch --no_python --use_env --nproc_per_node 8 \
          bash  -c './bind.sh --cpu=dgxa100_ccx.sh --mem=dgxa100_ccx.sh python -m dlrm.scripts.main \
          --dataset /data --amp --cuda_graphs'

in the DLRM Docker container on NVIDIA DGX A100 (8x A100 80GB) GPUs. Performance numbers (in records of data per second) were averaged over an entire training epoch.

GPUs	Model size	Batch size / GPU	Throughput - TF32	Throughput - mixed precision	Throughput speedup (TF32 to mixed precision)
8	large	8k	11,400,000	16,500,000	1.447
1	small	64k	2,880,000	4,020,000	1.396

To achieve these same results, follow the steps in the Quick Start Guide.

Training performance: NVIDIA DGX-1 (8x V100 32GB)

Our results were obtained by running the following commands:

for single-GPU:

python -m dlrm.scripts.main --mode train --dataset /data --amp --cuda_graphs

for multi-GPU :

python -m torch.distributed.launch --no_python --use_env --nproc_per_node 8 \
          bash  -c './bind.sh  --cpu=exclusive -- python -m dlrm.scripts.main \
          --dataset /data --amp --cuda_graphs'

in the DLRM Docker container on NVIDIA DGX-1 with (8x V100 32GB) GPUs. Performance numbers (in records of data per second) were averaged over an entire training epoch.

GPUs	Model size	Batch size / GPU	Throughput - FP32	Throughput - mixed precision	Throughput speedup (FP32 to mixed precision)
8	large	8k	2,880,000	6,920,000	2.403
1	small	64k	672,000	2,090,000	3.110

To achieve these same results, follow the steps in the Quick Start Guide.

Training performance: NVIDIA DGX-2 (16x V100 32GB)

Our results were obtained by running the following commands:

for single-GPU:

python -m dlrm.scripts.main --dataset /data --amp --cuda_graphs

for multi-GPU:

python -m torch.distributed.launch --no_python --use_env --nproc_per_node [8/16] \
          bash  -c './bind.sh  --cpu=exclusive -- python -m dlrm.scripts.main \
          --dataset /data --amp --cuda_graphs'

in the DLRM Docker container on NVIDIA DGX-2 with (16x V100 32GB) GPUs. Performance numbers (in records of data per second) were averaged over an entire training epoch.

GPUs	Model size	Batch size / GPU	Throughput - FP32	Throughput - mixed precision	Throughput speedup (FP32 to mixed precision)
16	large	4k	4,740,000	10,800,000	2.278
8	large	8k	3,330,000	7,930,000	2.381
1	small	64k	717,000	2,250,000	3.138

To achieve these same results, follow the steps in the Quick Start Guide.

Inference performance results

Inference performance: NVIDIA A100 (1x A100 80GB)

Our results were obtained by running the --inference_benchmark mode in the DLRM Docker container on on the NVIDIA A100 (1x A100 80GB) GPU.

	Mixed Precision				TF32
	CUDA Graphs ON		CUDA Graphs OFF		CUDA Graphs ON		CUDA Graphs OFF
Batch size	Throughput Avg	Latency Avg	Throughput Avg	Latency Avg	Throughput Avg	Latency Avg	Throughput Avg	Latency Avg
32768	14,796,024	0.00221	14,369,047	0.00228	8,832,225	0.00371	8,637,000	0.00379
16384	14,217,340	0.00115	13,673,623	0.00120	8,540,191	0.00192	8,386,694	0.00195
8192	12,769,583	0.00064	11,336,204	0.00072	7,658,459	0.00107	7,463,740	0.00110
4096	10,556,140	0.00039	8,203,285	0.00050	6,777,965	0.00060	6,142,076	0.00067
2048	8,415,889	0.00024	4,785,479	0.00043	5,214,990	0.00039	4,365,954	0.00047
1024	5,045,754	0.00020	2,357,953	0.00043	3,854,504	0.00027	2,615,601	0.00039
512	3,168,261	0.00016	1,190,989	0.00043	2,441,310	0.00021	1,332,944	0.00038
256	1,711,749	0.00015	542,310	0.00047	1,365,320	0.00019	592,034	0.00043
128	889,777	0.00014	274,223	0.00047	790,984	0.00016	300,908	0.00043
64	459,728	0.00014	136,180	0.00047	416,463	0.00015	150,382	0.00043
32	222,386	0.00014	70,107	0.00046	174,163	0.00018	75,768	0.00042
16	117,386	0.00014	34,983	0.00046	108,992	0.00015	38,369	0.00042
8	59,200	0.00014	18,852	0.00042	55,661	0.00014	19,440	0.00041
4	29,609	0.00014	8,505	0.00047	27,957	0.00014	10,206	0.00039
2	14,066	0.00014	4,610	0.00043	13,010	0.00015	5,229	0.00038

To achieve these same results, follow the steps in the Quick Start Guide.

Inference performance: NVIDIA DGX-1 (1x V100 32GB)

Our results were obtained by running the --inference_benchmark mode in the DLRM Docker container on NVIDIA DGX-1 with (1x V100 32GB) GPU.

	Mixed Precision				FP32
	CUDA Graphs ON		CUDA Graphs OFF		CUDA Graphs ON		CUDA Graphs OFF
Batch size	Throughput Avg	Latency Avg	Throughput Avg	Latency Avg	Throughput Avg	Latency Avg	Throughput Avg	Latency Avg
32768	6,716,240	0.00488	6,792,739	0.00482	1,809,345	0.01811	1,802,851	0.01818
16384	6,543,544	0.00250	6,520,519	0.00251	1,754,713	0.00934	1,745,214	0.00939
8192	6,215,194	0.00132	6,074,446	0.00135	1,669,188	0.00491	1,656,393	0.00495
4096	5,230,443	0.00078	4,901,451	0.00084	1,586,666	0.00258	1,574,068	0.00260
2048	4,261,124	0.00048	3,523,239	0.00058	1,462,006	0.00140	1,416,985	0.00145
1024	3,306,724	0.00031	2,047,274	0.00050	1,277,860	0.00080	1,161,032	0.00088
512	2,049,382	0.00025	1,005,919	0.00051	1,016,186	0.00050	841,732	0.00061
256	1,149,997	0.00022	511,102	0.00050	726,349	0.00035	485,162	0.00053
128	663,048	0.00019	264,015	0.00048	493,878	0.00026	238,936	0.00054
64	359,505	0.00018	132,913	0.00048	295,273	0.00022	124,120	0.00052
32	175,465	0.00018	64,287	0.00050	157,629	0.00020	63,919	0.00050
16	99,207	0.00016	31,062	0.00052	83,019	0.00019	34,660	0.00046
8	52,532	0.00015	16,492	0.00049	43,289	0.00018	17,893	0.00045
4	27,626	0.00014	8,391	0.00048	22,692	0.00018	8,923	0.00045
2	13,791	0.00015	4,146	0.00048	11,747	0.00017	4,487	0.00045

To achieve these same results, follow the steps in the Quick Start Guide.

Release notes

Changelog

October 2021

Added support for CUDA Graphs
Switched to PyTorch native AMP for mixed precision training
Unified the single-GPU and multi-GPU training scripts
Added support for BYO dataset
Updated performance results
Updated container version

June 2021

Updated container version
Updated performance results

March 2021

Added NVTabular as a new preprocessing option
Added a new dataset - xlarge, which uses a frequency threshold of 2
Introduced a new GPU - A100 80GB, and its performance results
Updated Spark preprocessing
Added Adam as an optional optimizer for embedding and MLPs, for multi-GPU training
Improved README

August 2020

Preprocessing with Spark 3 on GPU
Multiple performance optimizations
Automatic placement and load balancing of embedding tables
Improved README

June 2020

Updated performance tables to include A100 results and multi-GPU setup
Multi-GPU optimizations

May 2020

Performance optimizations

April 2020

Initial release

Known issues

Adam optimizer performance is not optimized.
For some seeds, the model's loss can become NaN due to aggressive learning rate schedule.
Custom dot interaction kernels for FP16 and TF32 assume that embedding size <= 128 and number of categorical variables < 32. Pass --interaction_op dot to use the slower native operation in those cases.