notebooks/multi_gpu_offline_inference.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cdfec37b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Copyright 2021 NVIDIA Corporation. All Rights Reserved.\n",
    "#\n",
    "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
    "# you may not use this file except in compliance with the License.\n",
    "# You may obtain a copy of the License at\n",
    "#\n",
    "#     http://www.apache.org/licenses/LICENSE-2.0\n",
    "#\n",
    "# Unless required by applicable law or agreed to in writing, software\n",
    "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
    "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
    "# See the License for the specific language governing permissions and\n",
    "# limitations under the License.\n",
    "# ==============================================================================\n",
    "\n",
    "# Each user is responsible for checking the content of datasets and the\n",
    "# applicable licenses and determining if suitable for the intended use."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f75c838b",
   "metadata": {},
   "source": [
    "<img src=\"http://developer.download.nvidia.com/notebooks/dlsw-notebooks/merlin_hugectr_multi-gpu-offline-inference/nvidia_logo.png\" style=\"width: 90px; float: right;\">\n",
    "\n",
    "# Multi-GPU Offline Inference"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ea4ae25b",
   "metadata": {},
   "source": [
    "## Overview\n",
    "\n",
    "In HugeCTR version 3.4.1, we provide Python APIs to perform multi-GPU offline inference.\n",
    "This work leverages the [HugeCTR Hierarchical Parameter Server](https://nvidia-merlin.github.io/HugeCTR/master/hugectr_core_features.html#hierarchical-parameter-server) and enables concurrent execution on multiple devices.\n",
    "The `Norm` or `Parquet` dataset format is currently supported by multi-GPU offline inference.\n",
    "\n",
    "This notebook explains how to perform multi-GPU offline inference with the HugeCTR Python APIs.\n",
    "For more details about the API, see the [HugeCTR Python Interface](https://nvidia-merlin.github.io/HugeCTR/master/api/python_interface.html#inference-api) documentation."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "902f3ef1",
   "metadata": {},
   "source": [
    "## Setup\n",
    "\n",
    "To setup the environment, refer to [HugeCTR Example Notebooks](../notebooks) and follow the instructions there before running the following."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "240b78ac",
   "metadata": {},
   "source": [
    "## Data Generation\n",
    "\n",
    "HugeCTR provides a tool to generate synthetic datasets.\n",
    "The [Data Generator](https://nvidia-merlin.github.io/HugeCTR/master/api/python_interface.html#data-generator-api) class is capable of generating datasets in different formats and with different distributions.\n",
    "We will generate multi-hot Parquet datasets with a power-law distribution for this notebook:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "db37ef07",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[HCTR][08:59:54.134][INFO][RK0][main]: Generate Parquet dataset\n",
      "[HCTR][08:59:54.134][INFO][RK0][main]: train data folder: ./multi_hot_parquet, eval data folder: ./multi_hot_parquet, slot_size_array: 10000, 10000, 10000, nnz array: 2, 1, 3, #files for train: 32, #files for eval: 8, #samples per file: 40960, Use power law distribution: 1, alpha of power law: 1.3\n",
      "[HCTR][08:59:54.136][INFO][RK0][main]: ./multi_hot_parquet exist\n",
      "[HCTR][08:59:54.140][INFO][RK0][main]: ./multi_hot_parquet/train/gen_0.parquet\n",
      "[HCTR][08:59:55.615][INFO][RK0][main]: ./multi_hot_parquet/train/gen_1.parquet\n",
      "[HCTR][08:59:55.850][INFO][RK0][main]: ./multi_hot_parquet/train/gen_2.parquet\n",
      "[HCTR][08:59:56.078][INFO][RK0][main]: ./multi_hot_parquet/train/gen_3.parquet\n",
      "[HCTR][08:59:56.311][INFO][RK0][main]: ./multi_hot_parquet/train/gen_4.parquet\n",
      "[HCTR][08:59:56.534][INFO][RK0][main]: ./multi_hot_parquet/train/gen_5.parquet\n",
      "[HCTR][08:59:56.770][INFO][RK0][main]: ./multi_hot_parquet/train/gen_6.parquet\n",
      "[HCTR][08:59:56.959][INFO][RK0][main]: ./multi_hot_parquet/train/gen_7.parquet\n",
      "[HCTR][08:59:57.152][INFO][RK0][main]: ./multi_hot_parquet/train/gen_8.parquet\n",
      "[HCTR][08:59:57.309][INFO][RK0][main]: ./multi_hot_parquet/train/gen_9.parquet\n",
      "[HCTR][08:59:57.496][INFO][RK0][main]: ./multi_hot_parquet/train/gen_10.parquet\n",
      "[HCTR][08:59:57.671][INFO][RK0][main]: ./multi_hot_parquet/train/gen_11.parquet\n",
      "[HCTR][08:59:57.879][INFO][RK0][main]: ./multi_hot_parquet/train/gen_12.parquet\n",
      "[HCTR][08:59:58.069][INFO][RK0][main]: ./multi_hot_parquet/train/gen_13.parquet\n",
      "[HCTR][08:59:58.240][INFO][RK0][main]: ./multi_hot_parquet/train/gen_14.parquet\n",
      "[HCTR][08:59:58.423][INFO][RK0][main]: ./multi_hot_parquet/train/gen_15.parquet\n",
      "[HCTR][08:59:58.619][INFO][RK0][main]: ./multi_hot_parquet/train/gen_16.parquet\n",
      "[HCTR][08:59:58.833][INFO][RK0][main]: ./multi_hot_parquet/train/gen_17.parquet\n",
      "[HCTR][08:59:59.017][INFO][RK0][main]: ./multi_hot_parquet/train/gen_18.parquet\n",
      "[HCTR][08:59:59.176][INFO][RK0][main]: ./multi_hot_parquet/train/gen_19.parquet\n",
      "[HCTR][08:59:59.358][INFO][RK0][main]: ./multi_hot_parquet/train/gen_20.parquet\n",
      "[HCTR][08:59:59.527][INFO][RK0][main]: ./multi_hot_parquet/train/gen_21.parquet\n",
      "[HCTR][08:59:59.722][INFO][RK0][main]: ./multi_hot_parquet/train/gen_22.parquet\n",
      "[HCTR][08:59:59.939][INFO][RK0][main]: ./multi_hot_parquet/train/gen_23.parquet\n",
      "[HCTR][09:00:00.107][INFO][RK0][main]: ./multi_hot_parquet/train/gen_24.parquet\n",
      "[HCTR][09:00:00.294][INFO][RK0][main]: ./multi_hot_parquet/train/gen_25.parquet\n",
      "[HCTR][09:00:00.509][INFO][RK0][main]: ./multi_hot_parquet/train/gen_26.parquet\n",
      "[HCTR][09:00:00.695][INFO][RK0][main]: ./multi_hot_parquet/train/gen_27.parquet\n",
      "[HCTR][09:00:00.955][INFO][RK0][main]: ./multi_hot_parquet/train/gen_28.parquet\n",
      "[HCTR][09:00:01.190][INFO][RK0][main]: ./multi_hot_parquet/train/gen_29.parquet\n",
      "[HCTR][09:00:01.365][INFO][RK0][main]: ./multi_hot_parquet/train/gen_30.parquet\n",
      "[HCTR][09:00:01.509][INFO][RK0][main]: ./multi_hot_parquet/train/gen_31.parquet\n",
      "[HCTR][09:00:01.698][INFO][RK0][main]: ./multi_hot_parquet/file_list.txt done!\n",
      "[HCTR][09:00:01.708][INFO][RK0][main]: ./multi_hot_parquet/val/gen_0.parquet\n",
      "[HCTR][09:00:01.895][INFO][RK0][main]: ./multi_hot_parquet/val/gen_1.parquet\n",
      "[HCTR][09:00:02.062][INFO][RK0][main]: ./multi_hot_parquet/val/gen_2.parquet\n",
      "[HCTR][09:00:02.255][INFO][RK0][main]: ./multi_hot_parquet/val/gen_3.parquet\n",
      "[HCTR][09:00:02.472][INFO][RK0][main]: ./multi_hot_parquet/val/gen_4.parquet\n",
      "[HCTR][09:00:02.665][INFO][RK0][main]: ./multi_hot_parquet/val/gen_5.parquet\n",
      "[HCTR][09:00:02.888][INFO][RK0][main]: ./multi_hot_parquet/val/gen_6.parquet\n",
      "[HCTR][09:00:03.110][INFO][RK0][main]: ./multi_hot_parquet/val/gen_7.parquet\n",
      "[HCTR][09:00:03.303][INFO][RK0][main]: ./multi_hot_parquet/file_list_test.txt done!\n"
     ]
    }
   ],
   "source": [
    "import hugectr\n",
    "from hugectr.tools import DataGeneratorParams, DataGenerator\n",
    "\n",
    "data_generator_params = DataGeneratorParams(\n",
    "  format = hugectr.DataReaderType_t.Parquet,\n",
    "  label_dim = 2,\n",
    "  dense_dim = 2,\n",
    "  num_slot = 3,\n",
    "  i64_input_key = True,\n",
    "  nnz_array = [2, 1, 3],\n",
    "  source = \"./multi_hot_parquet/file_list.txt\",\n",
    "  eval_source = \"./multi_hot_parquet/file_list_test.txt\",\n",
    "  slot_size_array = [10000, 10000, 10000],\n",
    "  check_type = hugectr.Check_t.Non,\n",
    "  dist_type = hugectr.Distribution_t.PowerLaw,\n",
    "  power_law_type = hugectr.PowerLaw_t.Short,\n",
    "  num_files = 32,\n",
    "  eval_num_files = 8)\n",
    "data_generator = DataGenerator(data_generator_params)\n",
    "data_generator.generate()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "19109028",
   "metadata": {},
   "source": [
    "## Train from Scratch\n",
    "\n",
    "We can train fom scratch by performing the following steps with Python APIs:\n",
    "\n",
    "1. Create the solver, reader and optimizer, then initialize the model.\n",
    "2. Construct the model graph by adding input, sparse embedding and dense layers in order.\n",
    "3. Compile the model and have an overview of the model graph.\n",
    "4. Dump the model graph to a JSON file.\n",
    "5. Fit the model, save the model weights and optimizer states implicitly.\n",
    "6. Dump one batch of evaluation results to files."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "e2b0d9d6",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Overwriting multi_hot_train.py\n"
     ]
    }
   ],
   "source": [
    "%%writefile multi_hot_train.py\n",
    "import hugectr\n",
    "from mpi4py import MPI\n",
    "solver = hugectr.CreateSolver(model_name = \"multi_hot\",\n",
    "                              max_eval_batches = 1,\n",
    "                              batchsize_eval = 131072,\n",
    "                              batchsize = 16384,\n",
    "                              lr = 0.001,\n",
    "                              vvgpu = [[0]],\n",
    "                              i64_input_key = True,\n",
    "                              repeat_dataset = True,\n",
    "                              use_cuda_graph = True)\n",
    "reader = hugectr.DataReaderParams(data_reader_type = hugectr.DataReaderType_t.Parquet,\n",
    "                                  source = [\"./multi_hot_parquet/file_list.txt\"],\n",
    "                                  eval_source = \"./multi_hot_parquet/file_list_test.txt\",\n",
    "                                  check_type = hugectr.Check_t.Non,\n",
    "                                  slot_size_array = [10000, 10000, 10000])\n",
    "optimizer = hugectr.CreateOptimizer(optimizer_type = hugectr.Optimizer_t.Adam)\n",
    "model = hugectr.Model(solver, reader, optimizer)\n",
    "model.add(hugectr.Input(label_dim = 2, label_name = \"label\",\n",
    "                        dense_dim = 2, dense_name = \"dense\",\n",
    "                        data_reader_sparse_param_array = \n",
    "                        [hugectr.DataReaderSparseParam(\"data1\", [2, 1], False, 2),\n",
    "                        hugectr.DataReaderSparseParam(\"data2\", 3, False, 1),]))\n",
    "model.add(hugectr.SparseEmbedding(embedding_type = hugectr.Embedding_t.DistributedSlotSparseEmbeddingHash, \n",
    "                            workspace_size_per_gpu_in_mb = 4,\n",
    "                            embedding_vec_size = 16,\n",
    "                            combiner = \"sum\",\n",
    "                            sparse_embedding_name = \"sparse_embedding1\",\n",
    "                            bottom_name = \"data1\",\n",
    "                            optimizer = optimizer))\n",
    "model.add(hugectr.SparseEmbedding(embedding_type = hugectr.Embedding_t.DistributedSlotSparseEmbeddingHash, \n",
    "                            workspace_size_per_gpu_in_mb = 2,\n",
    "                            embedding_vec_size = 16,\n",
    "                            combiner = \"sum\",\n",
    "                            sparse_embedding_name = \"sparse_embedding2\",\n",
    "                            bottom_name = \"data2\",\n",
    "                            optimizer = optimizer))\n",
    "model.add(hugectr.DenseLayer(layer_type = hugectr.Layer_t.Reshape,\n",
    "                            bottom_names = [\"sparse_embedding1\"],\n",
    "                            top_names = [\"reshape1\"],\n",
    "                            leading_dim=32))                            \n",
    "model.add(hugectr.DenseLayer(layer_type = hugectr.Layer_t.Reshape,\n",
    "                            bottom_names = [\"sparse_embedding2\"],\n",
    "                            top_names = [\"reshape2\"],\n",
    "                            leading_dim=16))                            \n",
    "model.add(hugectr.DenseLayer(layer_type = hugectr.Layer_t.Concat,\n",
    "                            bottom_names = [\"reshape1\", \"reshape2\", \"dense\"], top_names = [\"concat1\"]))\n",
    "model.add(hugectr.DenseLayer(layer_type = hugectr.Layer_t.InnerProduct,\n",
    "                            bottom_names = [\"concat1\"],\n",
    "                            top_names = [\"fc1\"],\n",
    "                            num_output=1024))\n",
    "model.add(hugectr.DenseLayer(layer_type = hugectr.Layer_t.ReLU,\n",
    "                            bottom_names = [\"fc1\"],\n",
    "                            top_names = [\"relu1\"]))\n",
    "model.add(hugectr.DenseLayer(layer_type = hugectr.Layer_t.InnerProduct,\n",
    "                            bottom_names = [\"relu1\"],\n",
    "                            top_names = [\"fc2\"],\n",
    "                            num_output=2))\n",
    "model.add(hugectr.DenseLayer(layer_type = hugectr.Layer_t.MultiCrossEntropyLoss,\n",
    "                            bottom_names = [\"fc2\", \"label\"],\n",
    "                            top_names = [\"loss\"],\n",
    "                            target_weight_vec = [0.5, 0.5]))\n",
    "model.compile()\n",
    "model.summary()\n",
    "model.graph_to_json(\"multi_hot.json\")\n",
    "model.fit(max_iter = 1100, display = 200, eval_interval = 1000, snapshot = 1000, snapshot_prefix = \"multi_hot\")\n",
    "model.export_predictions(\"multi_hot_pred_\" + str(1000), \"multi_hot_label_\" + str(1000))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "d0f29350",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "HugeCTR Version: 3.7\n",
      "====================================================Model Init=====================================================\n",
      "[HCTR][09:00:10.032][INFO][RK0][main]: Initialize model: multi_hot\n",
      "[HCTR][09:00:10.032][INFO][RK0][main]: Global seed is 69819197\n",
      "[HCTR][09:00:10.135][INFO][RK0][main]: Device to NUMA mapping:\n",
      "  GPU 0 ->  node 0\n",
      "[HCTR][09:00:11.978][WARNING][RK0][main]: Peer-to-peer access cannot be fully enabled.\n",
      "[HCTR][09:00:11.978][INFO][RK0][main]: Start all2all warmup\n",
      "[HCTR][09:00:11.978][INFO][RK0][main]: End all2all warmup\n",
      "[HCTR][09:00:11.979][INFO][RK0][main]: Using All-reduce algorithm: NCCL\n",
      "[HCTR][09:00:11.980][INFO][RK0][main]: Device 0: Tesla V100-SXM2-32GB\n",
      "[HCTR][09:00:11.985][INFO][RK0][main]: num of DataReader workers for train: 1\n",
      "[HCTR][09:00:11.985][INFO][RK0][main]: num of DataReader workers for eval: 1\n",
      "[HCTR][09:00:12.176][INFO][RK0][main]: Vocabulary size: 30000\n",
      "[HCTR][09:00:12.177][INFO][RK0][main]: max_vocabulary_size_per_gpu_=21845\n",
      "[HCTR][09:00:12.179][INFO][RK0][main]: max_vocabulary_size_per_gpu_=10922\n",
      "[HCTR][09:00:12.181][INFO][RK0][main]: Graph analysis to resolve tensor dependency\n",
      "===================================================Model Compile===================================================\n",
      "[HCTR][09:00:43.965][INFO][RK0][main]: gpu0 start to init embedding\n",
      "[HCTR][09:00:43.965][INFO][RK0][main]: gpu0 init embedding done\n",
      "[HCTR][09:00:43.965][INFO][RK0][main]: gpu0 start to init embedding\n",
      "[HCTR][09:00:43.965][INFO][RK0][main]: gpu0 init embedding done\n",
      "[HCTR][09:00:43.969][INFO][RK0][main]: Starting AUC NCCL warm-up\n",
      "[HCTR][09:00:43.972][INFO][RK0][main]: Warm-up done\n",
      "===================================================Model Summary===================================================\n",
      "[HCTR][09:00:43.972][INFO][RK0][main]: label                                   Dense                         Sparse                        \n",
      "label                                   dense                          data1,data2                   \n",
      "(None, 2)                               (None, 2)                               \n",
      "——————————————————————————————————————————————————————————————————————————————————————————————————————————————————\n",
      "Layer Type                              Input Name                    Output Name                   Output Shape                  \n",
      "——————————————————————————————————————————————————————————————————————————————————————————————————————————————————\n",
      "DistributedSlotSparseEmbeddingHash      data1                         sparse_embedding1             (None, 2, 16)                 \n",
      "------------------------------------------------------------------------------------------------------------------\n",
      "DistributedSlotSparseEmbeddingHash      data2                         sparse_embedding2             (None, 1, 16)                 \n",
      "------------------------------------------------------------------------------------------------------------------\n",
      "Reshape                                 sparse_embedding1             reshape1                      (None, 32)                    \n",
      "------------------------------------------------------------------------------------------------------------------\n",
      "Reshape                                 sparse_embedding2             reshape2                      (None, 16)                    \n",
      "------------------------------------------------------------------------------------------------------------------\n",
      "Concat                                  reshape1                      concat1                       (None, 50)                    \n",
      "                                        reshape2                                                                                  \n",
      "                                        dense                                                                                     \n",
      "------------------------------------------------------------------------------------------------------------------\n",
      "InnerProduct                            concat1                       fc1                           (None, 1024)                  \n",
      "------------------------------------------------------------------------------------------------------------------\n",
      "ReLU                                    fc1                           relu1                         (None, 1024)                  \n",
      "------------------------------------------------------------------------------------------------------------------\n",
      "InnerProduct                            relu1                         fc2                           (None, 2)                     \n",
      "------------------------------------------------------------------------------------------------------------------\n",
      "MultiCrossEntropyLoss                   fc2                           loss                                                        \n",
      "                                        label                                                                                     \n",
      "------------------------------------------------------------------------------------------------------------------\n",
      "[HCTR][09:00:43.977][INFO][RK0][main]: Save the model graph to multi_hot.json successfully\n",
      "=====================================================Model Fit=====================================================\n",
      "[HCTR][09:00:43.977][INFO][RK0][main]: Use non-epoch mode with number of iterations: 1100\n",
      "[HCTR][09:00:43.977][INFO][RK0][main]: Training batchsize: 16384, evaluation batchsize: 131072\n",
      "[HCTR][09:00:43.977][INFO][RK0][main]: Evaluation interval: 1000, snapshot interval: 1000\n",
      "[HCTR][09:00:43.977][INFO][RK0][main]: Dense network trainable: True\n",
      "[HCTR][09:00:43.977][INFO][RK0][main]: Sparse embedding sparse_embedding1 trainable: True\n",
      "[HCTR][09:00:43.977][INFO][RK0][main]: Sparse embedding sparse_embedding2 trainable: True\n",
      "[HCTR][09:00:43.977][INFO][RK0][main]: Use mixed precision: False, scaler: 1.000000, use cuda graph: True\n",
      "[HCTR][09:00:43.977][INFO][RK0][main]: lr: 0.001000, warmup_steps: 1, end_lr: 0.000000\n",
      "[HCTR][09:00:43.977][INFO][RK0][main]: decay_start: 0, decay_steps: 1, decay_power: 2.000000\n",
      "[HCTR][09:00:43.977][INFO][RK0][main]: Training source file: ./multi_hot_parquet/file_list.txt\n",
      "[HCTR][09:00:43.977][INFO][RK0][main]: Evaluation source file: ./multi_hot_parquet/file_list_test.txt\n",
      "[HCTR][09:00:46.346][INFO][RK0][main]: Iter: 200 Time(200 iters): 2.36888s Loss: 0.346413 lr:0.001\n",
      "[HCTR][09:00:48.421][INFO][RK0][main]: Iter: 400 Time(200 iters): 2.07362s Loss: 0.345891 lr:0.001\n",
      "[HCTR][09:00:50.519][INFO][RK0][main]: Iter: 600 Time(200 iters): 2.09809s Loss: 0.345239 lr:0.001\n",
      "[HCTR][09:00:52.586][INFO][RK0][main]: Iter: 800 Time(200 iters): 2.06616s Loss: 0.344346 lr:0.001\n",
      "[HCTR][09:00:54.656][INFO][RK0][main]: Iter: 1000 Time(200 iters): 2.0697s Loss: 0.343731 lr:0.001\n",
      "[HCTR][09:00:54.686][INFO][RK0][main]: Evaluation, AUC: 0.499013\n",
      "[HCTR][09:00:54.686][INFO][RK0][main]: Eval Time for 1 iters: 0.006811s\n",
      "[HCTR][09:00:54.692][INFO][RK0][main]: Rank0: Write hash table to file\n",
      "[HCTR][09:00:54.830][INFO][RK0][main]: Rank0: Write hash table to file\n",
      "[HCTR][09:00:54.848][INFO][RK0][main]: Dumping sparse weights to files, successful\n",
      "[HCTR][09:00:54.851][INFO][RK0][main]: Rank0: Write optimzer state to file\n",
      "[HCTR][09:00:54.852][INFO][RK0][main]: Done\n",
      "[HCTR][09:00:54.852][INFO][RK0][main]: Rank0: Write optimzer state to file\n",
      "[HCTR][09:00:54.853][INFO][RK0][main]: Done\n",
      "[HCTR][09:00:54.886][INFO][RK0][main]: Rank0: Write optimzer state to file\n",
      "[HCTR][09:00:54.887][INFO][RK0][main]: Done\n",
      "[HCTR][09:00:54.887][INFO][RK0][main]: Rank0: Write optimzer state to file\n",
      "[HCTR][09:00:54.888][INFO][RK0][main]: Done\n",
      "[HCTR][09:00:54.904][INFO][RK0][main]: Dumping sparse optimzer states to files, successful\n",
      "[HCTR][09:00:54.906][INFO][RK0][main]: Dumping dense weights to file, successful\n",
      "[HCTR][09:00:54.909][INFO][RK0][main]: Dumping dense optimizer states to file, successful\n",
      "[HCTR][09:00:55.915][INFO][RK0][main]: Finish 1100 iterations with batchsize: 16384 in 11.94s.\n"
     ]
    }
   ],
   "source": [
    "!python3 multi_hot_train.py"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0bb3b86c",
   "metadata": {},
   "source": [
    "### Multi-GPU Offline Inference\n",
    "\n",
    "We can demonstrate multi-GPU offline inference by performing the following steps with Python APIs:\n",
    "\n",
    "1. Configure the inference hyperparameters.\n",
    "2. Initialize the inference model. The model is a collection of inference sessions deployed on multiple devices.\n",
    "3. Make an inference from the evaluation dataset.\n",
    "4. Check the correctness of the inference by comparing it with the dumped evaluation results.\n",
    "\n",
    "**Note**: The `max_batchsize` configured within `InferenceParams` is the global batch size.\n",
    "The value for `max_batchsize` should be divisible by the number of deployed devices.\n",
    "The numpy array returned by `InferenceModel.predict` is of the shape `(max_batchsize * num_batches, label_dim)`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "8e25d216",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[HCTR][09:01:06.069][WARNING][RK0][main]: default_value_for_each_table.size() is not equal to the number of embedding tables\n",
      "[HCTR][09:01:06.072][INFO][RK0][main]: Global seed is 3072588155\n",
      "[HCTR][09:01:06.222][INFO][RK0][main]: Device to NUMA mapping:\n",
      "  GPU 0 ->  node 0\n",
      "  GPU 1 ->  node 0\n",
      "  GPU 2 ->  node 0\n",
      "  GPU 3 ->  node 0\n",
      "  GPU 4 ->  node 1\n",
      "  GPU 5 ->  node 1\n",
      "  GPU 6 ->  node 1\n",
      "  GPU 7 ->  node 1\n",
      "[HCTR][09:01:23.761][WARNING][RK0][main]: Peer-to-peer access cannot be fully enabled.\n",
      "[HCTR][09:01:23.763][INFO][RK0][main]: Start all2all warmup\n",
      "[HCTR][09:01:23.996][INFO][RK0][main]: End all2all warmup\n",
      "[HCTR][09:01:24.013][INFO][RK0][main]: default_emb_vec_value is not specified using default: 0\n",
      "[HCTR][09:01:24.013][INFO][RK0][main]: default_emb_vec_value is not specified using default: 0\n",
      "[HCTR][09:01:24.013][INFO][RK0][main]: Creating HashMap CPU database backend...\n",
      "[HCTR][09:01:24.013][INFO][RK0][main]: Volatile DB: initial cache rate = 1\n",
      "[HCTR][09:01:24.013][INFO][RK0][main]: Volatile DB: cache missed embeddings = 0\n",
      "[HCTR][09:01:24.347][INFO][RK0][main]: Table: hps_et.multi_hot.sparse_embedding1; cached 19849 / 19849 embeddings in volatile database (PreallocatedHashMapBackend); load: 19849 / 18446744073709551615 (0.00%).\n",
      "[HCTR][09:01:24.622][INFO][RK0][main]: Table: hps_et.multi_hot.sparse_embedding2; cached 9996 / 9996 embeddings in volatile database (PreallocatedHashMapBackend); load: 9996 / 18446744073709551615 (0.00%).\n",
      "[HCTR][09:01:24.622][DEBUG][RK0][main]: Real-time subscribers created!\n",
      "[HCTR][09:01:24.622][INFO][RK0][main]: Create embedding cache in device 0.\n",
      "[HCTR][09:01:24.628][INFO][RK0][main]: Use GPU embedding cache: True, cache size percentage: 0.500000\n",
      "[HCTR][09:01:24.628][INFO][RK0][main]: Configured cache hit rate threshold: 1.000000\n",
      "[HCTR][09:01:24.628][INFO][RK0][main]: The size of thread pool: 80\n",
      "[HCTR][09:01:24.628][INFO][RK0][main]: The size of worker memory pool: 2\n",
      "[HCTR][09:01:24.628][INFO][RK0][main]: The size of refresh memory pool: 1\n",
      "[HCTR][09:01:24.641][INFO][RK0][main]: Create embedding cache in device 1.\n",
      "[HCTR][09:01:24.646][INFO][RK0][main]: Use GPU embedding cache: True, cache size percentage: 0.500000\n",
      "[HCTR][09:01:24.646][INFO][RK0][main]: Configured cache hit rate threshold: 1.000000\n",
      "[HCTR][09:01:24.646][INFO][RK0][main]: The size of thread pool: 80\n",
      "[HCTR][09:01:24.646][INFO][RK0][main]: The size of worker memory pool: 2\n",
      "[HCTR][09:01:24.646][INFO][RK0][main]: The size of refresh memory pool: 1\n",
      "[HCTR][09:01:24.647][INFO][RK0][main]: Create embedding cache in device 2.\n",
      "[HCTR][09:01:24.652][INFO][RK0][main]: Use GPU embedding cache: True, cache size percentage: 0.500000\n",
      "[HCTR][09:01:24.652][INFO][RK0][main]: Configured cache hit rate threshold: 1.000000\n",
      "[HCTR][09:01:24.652][INFO][RK0][main]: The size of thread pool: 80\n",
      "[HCTR][09:01:24.652][INFO][RK0][main]: The size of worker memory pool: 2\n",
      "[HCTR][09:01:24.652][INFO][RK0][main]: The size of refresh memory pool: 1\n",
      "[HCTR][09:01:24.654][INFO][RK0][main]: Create embedding cache in device 3.\n",
      "[HCTR][09:01:24.659][INFO][RK0][main]: Use GPU embedding cache: True, cache size percentage: 0.500000\n",
      "[HCTR][09:01:24.659][INFO][RK0][main]: Configured cache hit rate threshold: 1.000000\n",
      "[HCTR][09:01:24.659][INFO][RK0][main]: The size of thread pool: 80\n",
      "[HCTR][09:01:24.659][INFO][RK0][main]: The size of worker memory pool: 2\n",
      "[HCTR][09:01:24.659][INFO][RK0][main]: The size of refresh memory pool: 1\n",
      "[HCTR][09:01:24.662][INFO][RK0][main]: Create embedding cache in device 4.\n",
      "[HCTR][09:01:24.667][INFO][RK0][main]: Use GPU embedding cache: True, cache size percentage: 0.500000\n",
      "[HCTR][09:01:24.667][INFO][RK0][main]: Configured cache hit rate threshold: 1.000000\n",
      "[HCTR][09:01:24.667][INFO][RK0][main]: The size of thread pool: 80\n",
      "[HCTR][09:01:24.667][INFO][RK0][main]: The size of worker memory pool: 2\n",
      "[HCTR][09:01:24.667][INFO][RK0][main]: The size of refresh memory pool: 1\n",
      "[HCTR][09:01:24.669][INFO][RK0][main]: Create embedding cache in device 5.\n",
      "[HCTR][09:01:24.675][INFO][RK0][main]: Use GPU embedding cache: True, cache size percentage: 0.500000\n",
      "[HCTR][09:01:24.675][INFO][RK0][main]: Configured cache hit rate threshold: 1.000000\n",
      "[HCTR][09:01:24.675][INFO][RK0][main]: The size of thread pool: 80\n",
      "[HCTR][09:01:24.675][INFO][RK0][main]: The size of worker memory pool: 2\n",
      "[HCTR][09:01:24.675][INFO][RK0][main]: The size of refresh memory pool: 1\n",
      "[HCTR][09:01:24.679][INFO][RK0][main]: Create embedding cache in device 6.\n",
      "[HCTR][09:01:24.683][INFO][RK0][main]: Use GPU embedding cache: True, cache size percentage: 0.500000\n",
      "[HCTR][09:01:24.683][INFO][RK0][main]: Configured cache hit rate threshold: 1.000000\n",
      "[HCTR][09:01:24.683][INFO][RK0][main]: The size of thread pool: 80\n",
      "[HCTR][09:01:24.683][INFO][RK0][main]: The size of worker memory pool: 2\n",
      "[HCTR][09:01:24.683][INFO][RK0][main]: The size of refresh memory pool: 1\n",
      "[HCTR][09:01:24.685][INFO][RK0][main]: Create embedding cache in device 7.\n",
      "[HCTR][09:01:24.688][INFO][RK0][main]: Use GPU embedding cache: True, cache size percentage: 0.500000\n",
      "[HCTR][09:01:24.688][INFO][RK0][main]: Configured cache hit rate threshold: 1.000000\n",
      "[HCTR][09:01:24.688][INFO][RK0][main]: The size of thread pool: 80\n",
      "[HCTR][09:01:24.688][INFO][RK0][main]: The size of worker memory pool: 2\n",
      "[HCTR][09:01:24.688][INFO][RK0][main]: The size of refresh memory pool: 1\n",
      "[HCTR][09:01:24.768][INFO][RK0][main]: Create inference session on device: 0\n",
      "[HCTR][09:01:24.768][INFO][RK0][main]: Model name: multi_hot\n",
      "[HCTR][09:01:24.768][INFO][RK0][main]: Use mixed precision: False\n",
      "[HCTR][09:01:24.768][INFO][RK0][main]: Use cuda graph: True\n",
      "[HCTR][09:01:24.768][INFO][RK0][main]: Max batchsize: 2048\n",
      "[HCTR][09:01:24.768][INFO][RK0][main]: Use I64 input key: True\n",
      "[HCTR][09:01:24.768][INFO][RK0][main]: start create embedding for inference\n",
      "[HCTR][09:01:24.768][INFO][RK0][main]: sparse_input name data1\n",
      "[HCTR][09:01:24.768][INFO][RK0][main]: sparse_input name data2\n",
      "[HCTR][09:01:24.768][INFO][RK0][main]: create embedding for inference success\n",
      "[HCTR][09:01:24.768][INFO][RK0][main]: Inference stage skip MultiCrossEntropyLoss layer, replaced by Sigmoid layer\n",
      "[HCTR][09:01:25.520][INFO][RK0][main]: Create inference session on device: 1\n",
      "[HCTR][09:01:25.520][INFO][RK0][main]: Model name: multi_hot\n",
      "[HCTR][09:01:25.520][INFO][RK0][main]: Use mixed precision: False\n",
      "[HCTR][09:01:25.520][INFO][RK0][main]: Use cuda graph: True\n",
      "[HCTR][09:01:25.520][INFO][RK0][main]: Max batchsize: 2048\n",
      "[HCTR][09:01:25.520][INFO][RK0][main]: Use I64 input key: True\n",
      "[HCTR][09:01:25.520][INFO][RK0][main]: start create embedding for inference\n",
      "[HCTR][09:01:25.520][INFO][RK0][main]: sparse_input name data1\n",
      "[HCTR][09:01:25.520][INFO][RK0][main]: sparse_input name data2\n",
      "[HCTR][09:01:25.520][INFO][RK0][main]: create embedding for inference success\n",
      "[HCTR][09:01:25.520][INFO][RK0][main]: Inference stage skip MultiCrossEntropyLoss layer, replaced by Sigmoid layer\n",
      "[HCTR][09:01:26.275][INFO][RK0][main]: Create inference session on device: 2\n",
      "[HCTR][09:01:26.275][INFO][RK0][main]: Model name: multi_hot\n",
      "[HCTR][09:01:26.275][INFO][RK0][main]: Use mixed precision: False\n",
      "[HCTR][09:01:26.275][INFO][RK0][main]: Use cuda graph: True\n",
      "[HCTR][09:01:26.275][INFO][RK0][main]: Max batchsize: 2048\n",
      "[HCTR][09:01:26.275][INFO][RK0][main]: Use I64 input key: True\n",
      "[HCTR][09:01:26.275][INFO][RK0][main]: start create embedding for inference\n",
      "[HCTR][09:01:26.275][INFO][RK0][main]: sparse_input name data1\n",
      "[HCTR][09:01:26.275][INFO][RK0][main]: sparse_input name data2\n",
      "[HCTR][09:01:26.275][INFO][RK0][main]: create embedding for inference success\n",
      "[HCTR][09:01:26.275][INFO][RK0][main]: Inference stage skip MultiCrossEntropyLoss layer, replaced by Sigmoid layer\n",
      "[HCTR][09:01:27.035][INFO][RK0][main]: Create inference session on device: 3\n",
      "[HCTR][09:01:27.035][INFO][RK0][main]: Model name: multi_hot\n",
      "[HCTR][09:01:27.035][INFO][RK0][main]: Use mixed precision: False\n",
      "[HCTR][09:01:27.035][INFO][RK0][main]: Use cuda graph: True\n",
      "[HCTR][09:01:27.035][INFO][RK0][main]: Max batchsize: 2048\n",
      "[HCTR][09:01:27.035][INFO][RK0][main]: Use I64 input key: True\n",
      "[HCTR][09:01:27.035][INFO][RK0][main]: start create embedding for inference\n",
      "[HCTR][09:01:27.035][INFO][RK0][main]: sparse_input name data1\n",
      "[HCTR][09:01:27.035][INFO][RK0][main]: sparse_input name data2\n",
      "[HCTR][09:01:27.035][INFO][RK0][main]: create embedding for inference success\n",
      "[HCTR][09:01:27.035][INFO][RK0][main]: Inference stage skip MultiCrossEntropyLoss layer, replaced by Sigmoid layer\n",
      "[HCTR][09:01:27.781][INFO][RK0][main]: Create inference session on device: 4\n",
      "[HCTR][09:01:27.781][INFO][RK0][main]: Model name: multi_hot\n",
      "[HCTR][09:01:27.781][INFO][RK0][main]: Use mixed precision: False\n",
      "[HCTR][09:01:27.781][INFO][RK0][main]: Use cuda graph: True\n",
      "[HCTR][09:01:27.781][INFO][RK0][main]: Max batchsize: 2048\n",
      "[HCTR][09:01:27.781][INFO][RK0][main]: Use I64 input key: True\n",
      "[HCTR][09:01:27.781][INFO][RK0][main]: start create embedding for inference\n",
      "[HCTR][09:01:27.781][INFO][RK0][main]: sparse_input name data1\n",
      "[HCTR][09:01:27.781][INFO][RK0][main]: sparse_input name data2\n",
      "[HCTR][09:01:27.781][INFO][RK0][main]: create embedding for inference success\n",
      "[HCTR][09:01:27.781][INFO][RK0][main]: Inference stage skip MultiCrossEntropyLoss layer, replaced by Sigmoid layer\n",
      "[HCTR][09:01:28.534][INFO][RK0][main]: Create inference session on device: 5\n",
      "[HCTR][09:01:28.534][INFO][RK0][main]: Model name: multi_hot\n",
      "[HCTR][09:01:28.534][INFO][RK0][main]: Use mixed precision: False\n",
      "[HCTR][09:01:28.534][INFO][RK0][main]: Use cuda graph: True\n",
      "[HCTR][09:01:28.534][INFO][RK0][main]: Max batchsize: 2048\n",
      "[HCTR][09:01:28.534][INFO][RK0][main]: Use I64 input key: True\n",
      "[HCTR][09:01:28.534][INFO][RK0][main]: start create embedding for inference\n",
      "[HCTR][09:01:28.534][INFO][RK0][main]: sparse_input name data1\n",
      "[HCTR][09:01:28.534][INFO][RK0][main]: sparse_input name data2\n",
      "[HCTR][09:01:28.534][INFO][RK0][main]: create embedding for inference success\n",
      "[HCTR][09:01:28.534][INFO][RK0][main]: Inference stage skip MultiCrossEntropyLoss layer, replaced by Sigmoid layer\n",
      "[HCTR][09:01:29.291][INFO][RK0][main]: Create inference session on device: 6\n",
      "[HCTR][09:01:29.291][INFO][RK0][main]: Model name: multi_hot\n",
      "[HCTR][09:01:29.291][INFO][RK0][main]: Use mixed precision: False\n",
      "[HCTR][09:01:29.291][INFO][RK0][main]: Use cuda graph: True\n",
      "[HCTR][09:01:29.291][INFO][RK0][main]: Max batchsize: 2048\n",
      "[HCTR][09:01:29.291][INFO][RK0][main]: Use I64 input key: True\n",
      "[HCTR][09:01:29.291][INFO][RK0][main]: start create embedding for inference\n",
      "[HCTR][09:01:29.291][INFO][RK0][main]: sparse_input name data1\n",
      "[HCTR][09:01:29.291][INFO][RK0][main]: sparse_input name data2\n",
      "[HCTR][09:01:29.291][INFO][RK0][main]: create embedding for inference success\n",
      "[HCTR][09:01:29.291][INFO][RK0][main]: Inference stage skip MultiCrossEntropyLoss layer, replaced by Sigmoid layer\n",
      "[HCTR][09:01:30.037][INFO][RK0][main]: Create inference session on device: 7\n",
      "[HCTR][09:01:30.037][INFO][RK0][main]: Model name: multi_hot\n",
      "[HCTR][09:01:30.037][INFO][RK0][main]: Use mixed precision: False\n",
      "[HCTR][09:01:30.037][INFO][RK0][main]: Use cuda graph: True\n",
      "[HCTR][09:01:30.037][INFO][RK0][main]: Max batchsize: 2048\n",
      "[HCTR][09:01:30.037][INFO][RK0][main]: Use I64 input key: True\n",
      "[HCTR][09:01:30.038][INFO][RK0][main]: start create embedding for inference\n",
      "[HCTR][09:01:30.038][INFO][RK0][main]: sparse_input name data1\n",
      "[HCTR][09:01:30.038][INFO][RK0][main]: sparse_input name data2\n",
      "[HCTR][09:01:30.038][INFO][RK0][main]: create embedding for inference success\n",
      "[HCTR][09:01:30.038][INFO][RK0][main]: Inference stage skip MultiCrossEntropyLoss layer, replaced by Sigmoid layer\n",
      "[HCTR][09:01:30.807][INFO][RK0][main]: Create inference data reader on 8 GPU(s)\n",
      "[HCTR][09:01:30.807][INFO][RK0][main]: num of DataReader workers: 8\n",
      "[HCTR][09:01:30.915][INFO][RK0][main]: Vocabulary size: 30000\n",
      "\n",
      "[INFO] Inference time for 8 batches: 0.182527\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "pred:  [[0.51329887 0.4888402 ]\n",
      " [0.55268604 0.62567735]\n",
      " [0.48302165 0.5015869 ]\n",
      " ...\n",
      " [0.52275413 0.46319592]\n",
      " [0.46984023 0.5436093 ]\n",
      " [0.48216432 0.48920953]]\n",
      "grount_truth:  [0.513299 0.48884  0.552686 ... 0.543609 0.482164 0.48921 ]\n",
      "mse:  8.482603947165404e-14\n"
     ]
    }
   ],
   "source": [
    "import hugectr\n",
    "from hugectr.inference import InferenceModel, InferenceParams\n",
    "import numpy as np\n",
    "from mpi4py import MPI\n",
    "\n",
    "model_config = \"multi_hot.json\"\n",
    "inference_params = InferenceParams(\n",
    "    model_name = \"multi_hot\",\n",
    "    max_batchsize = 16384,\n",
    "    hit_rate_threshold = 1.0,\n",
    "    dense_model_file = \"multi_hot_dense_1000.model\",\n",
    "    sparse_model_files = [\"multi_hot0_sparse_1000.model\", \"multi_hot1_sparse_1000.model\"],\n",
    "    deployed_devices = [0, 1, 2, 3, 4, 5, 6, 7],\n",
    "    use_gpu_embedding_cache = True,\n",
    "    cache_size_percentage = 0.5,\n",
    "    i64_input_key = True\n",
    ")\n",
    "inference_model = InferenceModel(model_config, inference_params)\n",
    "pred = inference_model.predict(\n",
    "    8,\n",
    "    \"./multi_hot_parquet/file_list_test.txt\",\n",
    "    hugectr.DataReaderType_t.Parquet,\n",
    "    hugectr.Check_t.Non,\n",
    "    [10000, 10000, 10000]\n",
    ")\n",
    "grount_truth = np.loadtxt(\"multi_hot_pred_1000\")\n",
    "print(\"pred: \", pred)\n",
    "print(\"grount_truth: \", grount_truth)\n",
    "diff = pred.flatten()-grount_truth\n",
    "mse = np.mean(diff*diff)\n",
    "print(\"mse: \", mse)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}