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

Scripts and sample code

These are the important modules in this repository:

• main.py - The main entrypoint script for training, evaluating, and benchmarking.
• model.py - Defines the DLRM model and some auxiliary functions used to train it.
• dataloader.py - Handles defining the dataset objects based on command-line flags.
• datasets.py - Defines the TfRawBinaryDataset class responsible for storing and loading the training data.
• slurm_multinode.sh - Example batch script for multi-node training on SLURM clusters.
• lr_scheduler.py - Defines a TensorFlow learning rate scheduler that supports both learning rate warmup and polynomial decay.
• embedding.py - Implementations of the embedding layers.
• interaction.py - Implementation of the dot-interaction layer using TensorFlow operations.
• tensorflow-dot-based-interact - A directory with a set of custom CUDA kernels. They provide fast implementations of the dot-interaction operation for various precisions and hardware platforms.
• utils.py - General utilities, such as a timer used for taking performance measurements.

Parameters

The table below lists the most important command-line parameters of the main.py script.

Command-line options

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

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 is used as a validation set and the second set is used as a hold-out test set.

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 15 times with a special value.
• Converting the hash values to consecutive integers.
• Adding 2 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:

DATASET_PARENT_DIRECTORY=/raid/dlrm

docker build -t nvidia_dlrm_tf .

docker run --cap-add SYS_NICE --runtime=nvidia -it --rm --ipc=host  -v ${DATASET_PARENT_DIRECTORY}/data:/data nvidia_dlrm_tf bash

python gen_csv.py --feature_spec_in tests/transcoding/small_csv.yaml

mkdir /data/conversion_output
cp tests/transcoding/small_csv.yaml /data/feature_spec.yaml
python transcode.py --input /data --output /data/converted

horovodrun -np 1 -H localhost:1 --mpi-args=--oversubscribe numactl --interleave=all -- python -u main.py --dataset_path /data/converted --amp --xla

horovodrun -np 8 -H localhost:8 --mpi-args=--oversubscribe numactl --interleave=all -- python -u main.py --dataset_path /data/converted --amp --xla

horovodrun -np 8 -H localhost:8 --mpi-args=--oversubscribe numactl --interleave=all -- python -u main.py --dataset_path /data/converted --amp --xla

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:

1. 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
2. Using other datasets might require tuning some hyperparameters (for example, learning rate, beta1 and beta2) to reach desired accuracy.
3. 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

Preprocess with Spark

The preprocessing scripts provided in this repository support running both on CPU and on DGX-2 using Apache Spark 3.0. It should be possible to change the values in preproc/dgx2_config.sh so that they'll work on other hardware platforms such as DGX-1.

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

The syntax for the preprocessing script is as follows:

cd preproc
./prepare_dataset.sh <DGX2|CPU> <frequency_threshold>

The first argument is the hardware platform to use (either DGX-2 or pure-CPU). The second argument means the frequency threshold to apply to the categorical variables. For a frequency threshold T, the categorical values that occur less often than T will be replaced with a special embedding. Thus, a larger value of T will require smaller embedding tables and will substantially reduce the overall size of the model.

For the Criteo Terabyte dataset we recommend a frequency threshold of T=3 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 T=15.

The preprocessing scripts makes 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

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 runs several PySpark jobs with spark_data_utils.py, for example: generates the dictionary

• transforms the train dataset

• transforms the test dataset

• transforms the validation dataset

  Change the variables in the run-spark.sh script according to your environment. Configure the paths.

export SPARK_LOCAL_DIRS=/data/spark-tmp
export INPUT_PATH=/data/criteo
export OUTPUT_PATH=/data/output

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

Where: SPARK_LOCAL_DIRS is the path where Spark uses to write shuffle data. INPUT_PATH is the path of the Criteo Terabyte Dataset, including uncompressed files like day_0, day_1… OUTPUT_PATH is where the script writes the output data. It will generate the following subdirectories of models, train, test, and validation.

• The model is the dictionary folder.
• 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.

Configure the resources which Spark will use.

export TOTAL_CORES=80
export TOTAL_MEMORY=800

Where: TOTAL_CORES is the total CPU cores you want Spark to use.

TOTAL_MEMORY is the total memory Spark will use.

Configure frequency limit.

USE_FREQUENCY_LIMIT=15

The frequency limit is used to filter out the categorical values which appear less than n times in the whole dataset, and make them be 0. Change this variable to 1 to enable it. The default frequency limit is 15 in the script. You also can change the number as you want by changing the line of OPTS="--frequency_limit 8".

Training process

The main training script resides in main.py. The speed of training 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.