Deep Learning Recommendation Model (DLRM)

repository·main·Indexed 26 days ago

https://github.com/facebookresearch/dlrm

A Deep Learning Recommendation Model implementation for personalization and recommendation systems. It supports PyTorch and Caffe2 frameworks, handling dense and sparse features via embedding tables and MLP networks. The repository includes support for MLPerf benchmarks, synchronous distributed training using gloo, nccl, or mpi, and integrations with TorchRec, torchx, and torchrun. It provides tools for preprocessing Criteo Kaggle and Criteo Terabyte datasets.

Tokens
3K
Snippets
11
Records
19
Agent score
38%

What's inside DLRM

  1. Materialize synthetic multi-hot dataset

    main

    For MLPerf DLRM v2 benchmarking, you can materialize a synthetic multi-hot dataset from the preprocessed Criteo 1TB dataset.

    Single-process version:

    python materialize_synthetic_multihot_dataset.py \
        --in_memory_binary_criteo_path $PREPROCESSED_CRITEO_1TB_CLICK_LOGS_DATASET_PATH \
        --output_path $MATERIALIZED_DATASET_PATH \
        --num_embeddings_per_feature 40000000,39060,17295,7424,20265,3,7122,1543,63,40000000,3067956,405282,10,2209,11938,155,4,976,14,40000000,40000000,40000000,590152,12973,108,36 \
        --multi_hot_sizes 3,2,1,2,6,1,1,1,1,7,3,8,1,6,9,5,1,1,1,12,100,27,10,3,1,1 \
        --multi_hot_distribution_type uniform
  2. Run a sample DLRM PyTorch training session

    main
    You can run a minimal training session using the dlrm_s_pytorch.py script. Use --mini-batch-size to set the batch size and --data-size to specify the amount of data to generate/use. To inspect the model architecture, weights, and data during execution, add the --debug-mode flag.
  3. Preprocess the Criteo 1TB Click Logs dataset

    main

    To prepare the 1TB Criteo Click Logs dataset for training, you must first download and uncompress the 24 files (days 0-23) using the provided download script, then run the preprocessing script. Note that this process requires 700GB of RAM and typically takes 1-2 days to complete.

    bash ./scripts/process_Criteo_1TB_Click_Logs_dataset.sh \
    ./criteo_1tb/raw_input_dataset_dir \
    ./criteo_1tb/temp_intermediate_files_dir \
    ./criteo_1tb/numpy_contiguous_shuffled_output_dataset_dir
  4. Prepare and run Criteo Terabyte dataset

    main

    To use the Criteo Terabyte Dataset:

    1. Download and unzip raw files day_0.gz through day_23.gz.
    2. Specify the location of the unzipped files using --raw-data-file=<path/day> (the script appends the day number automatically).
    3. The processed .npz files will be stored in <root_dir>/input/.
    4. Use the processed file in future runs with --processed-data-file=<path/*.npz>.
    5. Run the training via the provided benchmark script.
  5. Configure synchronous distributed training

    main

    DLRM supports synchronous distributed training using gloo, nccl, or mpi backends. You can use the torch.distributed.launch utility.

    For a single node with 8 GPUs using the nccl backend on a random dataset, use the following configuration structure:

    Key arguments:

    • --dist-backend: Set to nccl, gloo, or mpi.
    • --arch-embedding-size: String of embedding sizes (e.g., "80000-80000...").
    • --arch-sparse-feature-size: Integer.
    • --arch-mlp-bot: String of layer dimensions (e.g., "128-128-128-128").
    • --arch-mlp-top: String of layer dimensions (e.g., "512-512-512-256-1").
    • --max-ind-range: Maximum index range.
    • --data-generation: Set to random for synthetic data.

    For multi-node training, include launcher-specific arguments:

    • --nnodes: Number of nodes.
    • --node_rank: Rank of the current node.
    • --master_addr: IP address of the master node.
    • --master_port: Port of the master node.
  6. Preprocess Criteo 1TB Click Logs dataset

    main

    To prepare the Criteo 1TB dataset for training, use the provided shell script. This script converts raw TSV files into a format suitable for training.

    Requirements:

    • Requires approximately 700GB of RAM.
    • Takes 1-2 days to run.

    Usage:

    bash ./scripts/process_Criteo_1TB_Click_Logs_dataset.sh \
    ./criteo_1tb/raw_input_dataset_dir \
    ./criteo_1tb/temp_intermediate_files_dir \
    ./criteo_1tb/numpy_contiguous_shuffled_output_dataset_dir
  7. Preprocess Criteo Kaggle dataset to numpy files

    main

    After downloading and uncompressing the Criteo Kaggle dataset, use the npy_preproc_criteo script to convert the data into numpy format.

    python -m torchrec.datasets.scripts.npy_preproc_criteo --input_dir $INPUT_PATH --output_dir $OUTPUT_PATH --dataset_name criteo_kaggle
  8. Run TorchRec DLRM using TorchRun

    main

    You can use torchrun for elastic distributed training. Example for a single node with 2 processes:

    torchrun --nnodes 1 --nproc_per_node 2 --rdzv_backend c10d --rdzv_endpoint localhost --rdzv_id 54321 --role trainer dlrm_main.py

    torchrun --nnodes 1 --nproc_per_node 2 --rdzv_backend c10d --rdzv_endpoint localhost --rdzv_id 54321 --role trainer dlrm_main.py
  9. Setup environment for Criteo Kaggle Display Advertising Challenge

    main

    Before running the Criteo Kaggle benchmark, ensure you have Python >= 3.9 and CUDA >= 12.0. Follow these steps to install the necessary dependencies, including PyTorch nightly, FBGEMM-GPU, and a local build of TorchRec.

    pip install torch --index-url https://download.pytorch.org/whl/nightly/cu126
    pip install fbgemm-gpu --index-url https://download.pytorch.org/whl/nightly/cu126
    git clone https://github.com/pytorch/torchrec.git
    python -m pip install -e torchrec
    pip install -r requirements.txt
  10. Prepare and run Criteo Kaggle dataset

    main

    To use the Criteo Kaggle Display Advertising Challenge Dataset:

    1. Specify the raw train.txt file using --raw-data-file=<path/train.txt>. This triggers preprocessing (categorization, concatenation) which stores results as .npz files in <root_dir>/input/.
    2. For subsequent runs, use the processed file directly with --processed-data-file=<path/*.npz>.
    3. Run the training via the provided benchmark script.