notebooks/hugectr_criteo.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "0fe9a708",
   "metadata": {},
   "source": [
    "<img src=\"http://developer.download.nvidia.com/compute/machine-learning/frameworks/nvidia_logo.png\" style=\"width: 90px; float: right;\">\n",
    "\n",
    "# Introduction to the HugeCTR Python Interface"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2f72e368",
   "metadata": {},
   "source": [
    "## Overview\n",
    "\n",
    "HugeCTR version 3.1 introduces an enhanced Python interface\n",
    "The interface supports continuous training and inference with high-level APIs.\n",
    "There are four main improvements.\n",
    "\n",
    "* First, the model graph can be constructed and dumped to a JSON file with Python code and it saves users from writing JSON configuration files.\n",
    "* Second, the API supports the feature of embedding training cache with high-level APIs and extends it further for online training cases.\n",
    "(For learn about continuous training, you can view the [example notebook](./continuous_training.ipynb)).\n",
    "* Third, the freezing method is provided for both sparse embedding and dense network.\n",
    "This method enables transfer learning and fine-tuning for CTR tasks.\n",
    "* Finally, the pre-trained embeddings in other formats can be converted to HugeCTR sparse models and then loaded to facilitate the training process. This is shown in the Load Pre-trained Embeddings section of this notebook.\n",
    "\n",
    "This notebook explains how to access and use the enhanced HugeCTR Python interface.\n",
    "Although the low-level training APIs are still maintained for users who want to have precise control of each training iteration, migrating to the high-level training APIs is strongly recommended.\n",
    "For more details of the usage of the Python API, refer to the [HugeCTR Python Interface](https://nvidia-merlin.github.io/HugeCTR/master/api/python_interface.html) documentation."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "651880c9",
   "metadata": {},
   "source": [
    "## Installation\n",
    "\n",
    "### Get HugeCTR from NGC\n",
    "\n",
    "The HugeCTR Python module is preinstalled in the 22.07 and later [Merlin Training Container](https://catalog.ngc.nvidia.com/orgs/nvidia/teams/merlin/containers/merlin-hugectr): `nvcr.io/nvidia/merlin/merlin-hugectr:22.07`.\n",
    "\n",
    "You can check the existence of the required libraries by running the following Python code after launching this container.\n",
    "\n",
    "```bash\n",
    "$ python3 -c \"import hugectr\"\n",
    "```\n",
    "\n",
    "> If you prefer to build HugeCTR from the source code instead of using the NGC container,\n",
    "> refer to the\n",
    "> [How to Start Your Development](https://nvidia-merlin.github.io/HugeCTR/master/hugectr_contributor_guide.html#how-to-start-your-development)\n",
    "> documentation."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5ec3b0c0",
   "metadata": {},
   "source": [
    "## DCN Model\n",
    "\n",
    "### Download and Preprocess Data\n",
    "\n",
    "1. Download the Criteo dataset using the following command:\n",
    "\n",
    "   ```shell\n",
    "   $ cd ${project-root}/tools\n",
    "   $ wget https://storage.googleapis.com/criteo-cail-datasets/day_1.gz\n",
    "   ```\n",
    "   \n",
    "   In preprocessing, we will further reduce the amounts of data to speedup the preprocessing, fill missing values, remove the feature values whose occurrences are very rare, etc. Here we choose pandas preprocessing method to make the dataset ready for HugeCTR training.\n",
    "\n",
    "2. Preprocessing by Pandas using the following command:\n",
    "\n",
    "   ```shell\n",
    "   $ bash preprocess.sh 1 dcn_data pandas 1 0\n",
    "   ```\n",
    "   \n",
    "   The first argument represents the dataset postfix. It is 1 here since day_1 is used. The second argument dcn_data is where the preprocessed data is stored. The fourth argument (one after pandas) 1 embodies that the normalization is applied to dense features. The last argument 0 means that the feature crossing is not applied.\n",
    "\n",
    "3. Create a soft link to the dataset folder using the following command:\n",
    "\n",
    "   ```shell\n",
    "   $ ln ${project-root}/tools/dcn_data ${project_root}/notebooks/dcn_data\n",
    "   ```\n",
    "   \n",
    "**Note**: It will take a while (dozens of minutes) to preprocess the dataset. Please make sure that it is finished successfully before moving forward to the next section."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "072829a0",
   "metadata": {},
   "source": [
    "### Train from Scratch\n",
    "\n",
    "We can train fom scratch, dump the model graph to a JSON file, and save the model weights and optimizer states by performing the following 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 the JSON file.\n",
    "5. Fit the model, save the model weights and optimizer states implicitly.\n",
    "\n",
    "Please note that the training mode is determined by `repeat_dataset` within `hugectr.CreateSolver`.\n",
    "If it is `True`, the non-epoch mode training is adopted and the maximum iterations should be specified by `max_iter` within `hugectr.Model.fit`.\n",
    "If it is `False`, the epoch-mode training is adopted and the number of epochs should be specified by `num_epochs` within `hugectr.Model.fit`.\n",
    "\n",
    "The optimizer that is used to initialize the model applies to the weights of dense layers, while the optimizer for each sparse embedding layer can be specified independently within `hugectr.SparseEmbedding`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "1aa2c475",
   "metadata": {},
   "outputs": [],
   "source": [
    "import hugectr"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "fe15b878",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Writing dcn_train.py\n"
     ]
    }
   ],
   "source": [
    "%%writefile dcn_train.py\n",
    "import hugectr\n",
    "from mpi4py import MPI\n",
    "solver = hugectr.CreateSolver(max_eval_batches = 1500,\n",
    "                              batchsize_eval = 4096,\n",
    "                              batchsize = 4096,\n",
    "                              lr = 0.001,\n",
    "                              vvgpu = [[0]],\n",
    "                              i64_input_key = False,\n",
    "                              use_mixed_precision = False,\n",
    "                              repeat_dataset = True,\n",
    "                              use_cuda_graph = True)\n",
    "reader = hugectr.DataReaderParams(data_reader_type = hugectr.DataReaderType_t.Norm,\n",
    "                                  source = [\"./dcn_data/file_list.txt\"],\n",
    "                                  eval_source = \"./dcn_data/file_list_test.txt\",\n",
    "                                  check_type = hugectr.Check_t.Sum)\n",
    "optimizer = hugectr.CreateOptimizer(optimizer_type = hugectr.Optimizer_t.Adam)\n",
    "model = hugectr.Model(solver, reader, optimizer)\n",
    "model.add(hugectr.Input(label_dim = 1, label_name = \"label\",\n",
    "                        dense_dim = 13, dense_name = \"dense\",\n",
    "                        data_reader_sparse_param_array = \n",
    "                        [hugectr.DataReaderSparseParam(\"data1\", 2, False, 26)]))\n",
    "model.add(hugectr.SparseEmbedding(embedding_type = hugectr.Embedding_t.DistributedSlotSparseEmbeddingHash, \n",
    "                            workspace_size_per_gpu_in_mb = 264,\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.DenseLayer(layer_type = hugectr.Layer_t.Reshape,\n",
    "                            bottom_names = [\"sparse_embedding1\"],\n",
    "                            top_names = [\"reshape1\"],\n",
    "                            leading_dim=416))\n",
    "model.add(hugectr.DenseLayer(layer_type = hugectr.Layer_t.Concat,\n",
    "                            bottom_names = [\"reshape1\", \"dense\"], top_names = [\"concat1\"]))\n",
    "model.add(hugectr.DenseLayer(layer_type = hugectr.Layer_t.MultiCross,\n",
    "                            bottom_names = [\"concat1\"],\n",
    "                            top_names = [\"multicross1\"],\n",
    "                            num_layers=6))\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.Dropout,\n",
    "                            bottom_names = [\"relu1\"],\n",
    "                            top_names = [\"dropout1\"],\n",
    "                            dropout_rate=0.5))\n",
    "model.add(hugectr.DenseLayer(layer_type = hugectr.Layer_t.InnerProduct,\n",
    "                            bottom_names = [\"dropout1\"],\n",
    "                            top_names = [\"fc2\"],\n",
    "                            num_output=1024))\n",
    "model.add(hugectr.DenseLayer(layer_type = hugectr.Layer_t.ReLU,\n",
    "                            bottom_names = [\"fc2\"],\n",
    "                            top_names = [\"relu2\"]))\n",
    "model.add(hugectr.DenseLayer(layer_type = hugectr.Layer_t.Dropout,\n",
    "                            bottom_names = [\"relu2\"],\n",
    "                            top_names = [\"dropout2\"],\n",
    "                            dropout_rate=0.5))\n",
    "model.add(hugectr.DenseLayer(layer_type = hugectr.Layer_t.Concat,\n",
    "                            bottom_names = [\"dropout2\", \"multicross1\"],\n",
    "                            top_names = [\"concat2\"]))\n",
    "model.add(hugectr.DenseLayer(layer_type = hugectr.Layer_t.InnerProduct,\n",
    "                            bottom_names = [\"concat2\"],\n",
    "                            top_names = [\"fc3\"],\n",
    "                            num_output=1))\n",
    "model.add(hugectr.DenseLayer(layer_type = hugectr.Layer_t.BinaryCrossEntropyLoss,\n",
    "                            bottom_names = [\"fc3\", \"label\"],\n",
    "                            top_names = [\"loss\"]))\n",
    "model.compile()\n",
    "model.summary()\n",
    "model.graph_to_json(graph_config_file = \"dcn.json\")\n",
    "model.fit(max_iter = 1200, display = 500, eval_interval = 100, snapshot = 1000, snapshot_prefix = \"dcn\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "56c99500",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "HugeCTR Version: 3.2\n",
      "====================================================Model Init=====================================================\n",
      "[HUGECTR][03:31:21][INFO][RANK0]: Global seed is 1645340130\n",
      "[HUGECTR][03:31:22][INFO][RANK0]: Device to NUMA mapping:\n",
      "  GPU 0 ->  node 0\n",
      "\n",
      "[HUGECTR][03:31:23][WARNING][RANK0]: Peer-to-peer access cannot be fully enabled.\n",
      "[HUGECTR][03:31:23][INFO][RANK0]: Start all2all warmup\n",
      "[HUGECTR][03:31:23][INFO][RANK0]: End all2all warmup\n",
      "[HUGECTR][03:31:23][INFO][RANK0]: Using All-reduce algorithm: NCCL\n",
      "[HUGECTR][03:31:23][INFO][RANK0]: Device 0: Tesla V100-SXM2-32GB\n",
      "[HUGECTR][03:31:23][INFO][RANK0]: num of DataReader workers: 12\n",
      "[HUGECTR][03:31:23][INFO][RANK0]: max_vocabulary_size_per_gpu_=1441792\n",
      "[HUGECTR][03:31:23][INFO][RANK0]: Graph analysis to resolve tensor dependency\n",
      "[HUGECTR][03:31:23][INFO][RANK0]: Add Slice layer for tensor: concat1, creating 2 copies\n",
      "===================================================Model Compile===================================================\n",
      "[HUGECTR][03:31:35][INFO][RANK0]: gpu0 start to init embedding\n",
      "[HUGECTR][03:31:35][INFO][RANK0]: gpu0 init embedding done\n",
      "[HUGECTR][03:31:35][INFO][RANK0]: Starting AUC NCCL warm-up\n",
      "[HUGECTR][03:31:35][INFO][RANK0]: Warm-up done\n",
      "===================================================Model Summary===================================================\n",
      "label                                   Dense                         Sparse                        \n",
      "label                                   dense                          data1                         \n",
      "(None, 1)                               (None, 13)                              \n",
      "------------------------------------------------------------------------------------------------------------------\n",
      "Layer Type                              Input Name                    Output Name                   Output Shape                  \n",
      "------------------------------------------------------------------------------------------------------------------\n",
      "DistributedSlotSparseEmbeddingHash      data1                         sparse_embedding1             (None, 26, 16)                \n",
      "Reshape                                 sparse_embedding1             reshape1                      (None, 416)                   \n",
      "Concat                                  reshape1,dense                concat1                       (None, 429)                   \n",
      "Slice                                   concat1                       concat1_slice0,concat1_slice1                               \n",
      "MultiCross                              concat1_slice0                multicross1                   (None, 429)                   \n",
      "InnerProduct                            concat1_slice1                fc1                           (None, 1024)                  \n",
      "ReLU                                    fc1                           relu1                         (None, 1024)                  \n",
      "Dropout                                 relu1                         dropout1                      (None, 1024)                  \n",
      "InnerProduct                            dropout1                      fc2                           (None, 1024)                  \n",
      "ReLU                                    fc2                           relu2                         (None, 1024)                  \n",
      "Dropout                                 relu2                         dropout2                      (None, 1024)                  \n",
      "Concat                                  dropout2,multicross1          concat2                       (None, 1453)                  \n",
      "InnerProduct                            concat2                       fc3                           (None, 1)                     \n",
      "BinaryCrossEntropyLoss                  fc3,label                     loss                                                        \n",
      "------------------------------------------------------------------------------------------------------------------\n",
      "[HUGECTR][03:31:35][INFO][RANK0]: Save the model graph to dcn.json successfully\n",
      "=====================================================Model Fit=====================================================\n",
      "[HUGECTR][03:31:35][INFO][RANK0]: Use non-epoch mode with number of iterations: 1200\n",
      "[HUGECTR][03:31:35][INFO][RANK0]: Training batchsize: 4096, evaluation batchsize: 4096\n",
      "[HUGECTR][03:31:35][INFO][RANK0]: Evaluation interval: 100, snapshot interval: 1000\n",
      "[HUGECTR][03:31:35][INFO][RANK0]: Sparse embedding trainable: 1, dense network trainable: 1\n",
      "[HUGECTR][03:31:35][INFO][RANK0]: Use mixed precision: 0, scaler: 1, use cuda graph: -510996182\n",
      "[HUGECTR][03:31:35][INFO][RANK0]: lr: 0.001000, warmup_steps: 1, decay_start: 0, decay_steps: 1, decay_power: 2.000000, end_lr: 0.000000\n",
      "[HUGECTR][03:31:35][INFO][RANK0]: Training source file: ./dcn_data/file_list.txt\n",
      "[HUGECTR][03:31:35][INFO][RANK0]: Evaluation source file: ./dcn_data/file_list_test.txt\n",
      "[HUGECTR][03:31:39][INFO][RANK0]: Evaluation, AUC: 0.717558\n",
      "[HUGECTR][03:31:39][INFO][RANK0]: Eval Time for 1500 iters: 2.957773s\n",
      "[HUGECTR][03:31:43][INFO][RANK0]: Evaluation, AUC: 0.735452\n",
      "[HUGECTR][03:31:43][INFO][RANK0]: Eval Time for 1500 iters: 2.963541s\n",
      "[HUGECTR][03:31:47][INFO][RANK0]: Evaluation, AUC: 0.741079\n",
      "[HUGECTR][03:31:47][INFO][RANK0]: Eval Time for 1500 iters: 2.959102s\n",
      "[HUGECTR][03:31:50][INFO][RANK0]: Evaluation, AUC: 0.745329\n",
      "[HUGECTR][03:31:50][INFO][RANK0]: Eval Time for 1500 iters: 2.964232s\n",
      "[HUGECTR][03:31:51][INFO][RANK0]: Iter: 500 Time(500 iters): 15.479323s Loss: 0.117504 lr:0.001000\n",
      "[HUGECTR][03:31:54][INFO][RANK0]: Evaluation, AUC: 0.749935\n",
      "[HUGECTR][03:31:54][INFO][RANK0]: Eval Time for 1500 iters: 2.961690s\n",
      "[HUGECTR][03:31:58][INFO][RANK0]: Evaluation, AUC: 0.750517\n",
      "[HUGECTR][03:31:58][INFO][RANK0]: Eval Time for 1500 iters: 2.963790s\n",
      "[HUGECTR][03:32:01][INFO][RANK0]: Evaluation, AUC: 0.754112\n",
      "[HUGECTR][03:32:01][INFO][RANK0]: Eval Time for 1500 iters: 2.965818s\n",
      "[HUGECTR][03:32:05][INFO][RANK0]: Evaluation, AUC: 0.755083\n",
      "[HUGECTR][03:32:05][INFO][RANK0]: Eval Time for 1500 iters: 2.962515s\n",
      "[HUGECTR][03:32:09][INFO][RANK0]: Evaluation, AUC: 0.755834\n",
      "[HUGECTR][03:32:09][INFO][RANK0]: Eval Time for 1500 iters: 2.967796s\n",
      "[HUGECTR][03:32:09][INFO][RANK0]: Iter: 1000 Time(500 iters): 18.362356s Loss: 0.154462 lr:0.001000\n",
      "[HUGECTR][03:32:12][INFO][RANK0]: Evaluation, AUC: 0.758410\n",
      "[HUGECTR][03:32:12][INFO][RANK0]: Eval Time for 1500 iters: 2.969008s\n",
      "[HUGECTR][03:32:12][INFO][RANK0]: Rank0: Write hash table to file\n",
      "[HUGECTR][03:32:13][INFO][RANK0]: Dumping sparse weights to files, successful\n",
      "[HUGECTR][03:32:13][INFO][RANK0]: Rank0: Write optimzer state to file\n",
      "[HUGECTR][03:32:13][INFO][RANK0]: Done\n",
      "[HUGECTR][03:32:13][INFO][RANK0]: Rank0: Write optimzer state to file\n",
      "[HUGECTR][03:32:13][INFO][RANK0]: Done\n",
      "[HUGECTR][03:32:15][INFO][RANK0]: Dumping sparse optimzer states to files, successful\n",
      "[HUGECTR][03:32:15][INFO][RANK0]: Dumping dense weights to file, successful\n",
      "[HUGECTR][03:32:15][INFO][RANK0]: Dumping dense optimizer states to file, successful\n",
      "[HUGECTR][03:32:15][INFO][RANK0]: Dumping untrainable weights to file, successful\n",
      "[HUGECTR][03:32:19][INFO][RANK0]: Evaluation, AUC: 0.758818\n",
      "[HUGECTR][03:32:19][INFO][RANK0]: Eval Time for 1500 iters: 2.966687s\n",
      "[HUGECTR][03:32:20][INFO][RANK0]: Finish 1200 iterations with batchsize: 4096 in 44.08s.\n"
     ]
    }
   ],
   "source": [
    "!python3 dcn_train.py"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "650ff4d9",
   "metadata": {},
   "source": [
    "### Continue Training\n",
    "\n",
    "We can continue our training based on the saved model graph, model weights, and optimizer states by performing the following with Python APIs:\n",
    "\n",
    "1. Create the solver, reader and optimizer, then initialize the model.\n",
    "2. Construct the model graph from the saved JSON file.\n",
    "3. Compile the model and have an overview of the model graph.\n",
    "4. Load the model weights and optimizer states.\n",
    "5. Fit the model, save the model weights and optimizer states implicitly."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "763292b8",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "dcn0_opt_sparse_1000.model  dcn_dense_1000.model  dcn_opt_dense_1000.model\r\n",
      "\r\n",
      "dcn0_sparse_1000.model:\r\n",
      "emb_vector  key\r\n"
     ]
    }
   ],
   "source": [
    "!ls *.model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "339fb808",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Writing dcn_continue.py\n"
     ]
    }
   ],
   "source": [
    "%%writefile dcn_continue.py\n",
    "import hugectr\n",
    "from mpi4py import MPI\n",
    "solver = hugectr.CreateSolver(max_eval_batches = 1500,\n",
    "                              batchsize_eval = 4096,\n",
    "                              batchsize = 4096,\n",
    "                              vvgpu = [[0]],\n",
    "                              i64_input_key = False,\n",
    "                              use_mixed_precision = False,\n",
    "                              repeat_dataset = True,\n",
    "                              use_cuda_graph = True)\n",
    "reader = hugectr.DataReaderParams(data_reader_type = hugectr.DataReaderType_t.Norm,\n",
    "                                  source = [\"./dcn_data/file_list.txt\"],\n",
    "                                  eval_source = \"./dcn_data/file_list_test.txt\",\n",
    "                                  check_type = hugectr.Check_t.Sum)\n",
    "optimizer = hugectr.CreateOptimizer(optimizer_type = hugectr.Optimizer_t.Adam)\n",
    "model = hugectr.Model(solver, reader, optimizer)\n",
    "model.construct_from_json(graph_config_file = \"dcn.json\", include_dense_network = True)\n",
    "model.compile()\n",
    "model.load_dense_weights(\"dcn_dense_1000.model\")\n",
    "model.load_sparse_weights([\"dcn0_sparse_1000.model\"])\n",
    "model.load_dense_optimizer_states(\"dcn_opt_dense_1000.model\")\n",
    "model.load_sparse_optimizer_states([\"dcn0_opt_sparse_1000.model\"])\n",
    "model.summary()\n",
    "model.fit(max_iter = 500, display = 50, eval_interval = 100, snapshot = 10000, snapshot_prefix = \"dcn\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "b48ed0a0",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "HugeCTR Version: 3.2\n",
      "====================================================Model Init=====================================================\n",
      "[HUGECTR][03:32:48][INFO][RANK0]: Global seed is 4147354758\n",
      "[HUGECTR][03:32:49][INFO][RANK0]: Device to NUMA mapping:\n",
      "  GPU 0 ->  node 0\n",
      "\n",
      "[HUGECTR][03:32:50][WARNING][RANK0]: Peer-to-peer access cannot be fully enabled.\n",
      "[HUGECTR][03:32:50][INFO][RANK0]: Start all2all warmup\n",
      "[HUGECTR][03:32:50][INFO][RANK0]: End all2all warmup\n",
      "[HUGECTR][03:32:50][INFO][RANK0]: Using All-reduce algorithm: NCCL\n",
      "[HUGECTR][03:32:50][INFO][RANK0]: Device 0: Tesla V100-SXM2-32GB\n",
      "[HUGECTR][03:32:50][INFO][RANK0]: num of DataReader workers: 12\n",
      "[HUGECTR][03:32:50][INFO][RANK0]: max_num_frequent_categories is not specified using default: 1\n",
      "[HUGECTR][03:32:50][INFO][RANK0]: max_num_infrequent_samples is not specified using default: -1\n",
      "[HUGECTR][03:32:50][INFO][RANK0]: p_dup_max is not specified using default: 0.010000\n",
      "[HUGECTR][03:32:50][INFO][RANK0]: max_all_reduce_bandwidth is not specified using default: 130000000000.000000\n",
      "[HUGECTR][03:32:50][INFO][RANK0]: max_all_to_all_bandwidth is not specified using default: 190000000000.000000\n",
      "[HUGECTR][03:32:50][INFO][RANK0]: efficiency_bandwidth_ratio is not specified using default: 1.000000\n",
      "[HUGECTR][03:32:50][INFO][RANK0]: communication_type is not specified using default: IB_NVLink\n",
      "[HUGECTR][03:32:50][INFO][RANK0]: hybrid_embedding_type is not specified using default: Distributed\n",
      "[HUGECTR][03:32:50][INFO][RANK0]: max_vocabulary_size_per_gpu_=1441792\n",
      "[HUGECTR][03:32:50][INFO][RANK0]: Load the model graph from dcn.json successfully\n",
      "[HUGECTR][03:32:50][INFO][RANK0]: Graph analysis to resolve tensor dependency\n",
      "===================================================Model Compile===================================================\n",
      "[HUGECTR][03:33:02][INFO][RANK0]: gpu0 start to init embedding\n",
      "[HUGECTR][03:33:02][INFO][RANK0]: gpu0 init embedding done\n",
      "[HUGECTR][03:33:02][INFO][RANK0]: Starting AUC NCCL warm-up\n",
      "[HUGECTR][03:33:02][INFO][RANK0]: Warm-up done\n",
      "[HUGECTR][03:33:02][INFO][RANK0]: Loading dense model: dcn_dense_1000.model\n",
      "[HUGECTR][03:33:02][INFO][RANK0]: Loading sparse model: dcn0_sparse_1000.model\n",
      "[HUGECTR][03:33:03][INFO][RANK0]: Loading dense opt states: dcn_opt_dense_1000.model\n",
      "[HUGECTR][03:33:03][INFO][RANK0]: Loading sparse optimizer states: dcn0_opt_sparse_1000.model\n",
      "[HUGECTR][03:33:03][INFO][RANK0]: Rank0: Read optimzer state from file\n",
      "[HUGECTR][03:33:04][INFO][RANK0]: Done\n",
      "[HUGECTR][03:33:04][INFO][RANK0]: Rank0: Read optimzer state from file\n",
      "[HUGECTR][03:33:05][INFO][RANK0]: Done\n",
      "===================================================Model Summary===================================================\n",
      "label                                   Dense                         Sparse                        \n",
      "label                                   dense                          data1                         \n",
      "(None, 1)                               (None, 13)                              \n",
      "------------------------------------------------------------------------------------------------------------------\n",
      "Layer Type                              Input Name                    Output Name                   Output Shape                  \n",
      "------------------------------------------------------------------------------------------------------------------\n",
      "DistributedSlotSparseEmbeddingHash      data1                         sparse_embedding1             (None, 26, 16)                \n",
      "Reshape                                 sparse_embedding1             reshape1                      (None, 416)                   \n",
      "Concat                                  reshape1,dense                concat1                       (None, 429)                   \n",
      "Slice                                   concat1                       concat1_slice0,concat1_slice1                               \n",
      "MultiCross                              concat1_slice0                multicross1                   (None, 429)                   \n",
      "InnerProduct                            concat1_slice1                fc1                           (None, 1024)                  \n",
      "ReLU                                    fc1                           relu1                         (None, 1024)                  \n",
      "Dropout                                 relu1                         dropout1                      (None, 1024)                  \n",
      "InnerProduct                            dropout1                      fc2                           (None, 1024)                  \n",
      "ReLU                                    fc2                           relu2                         (None, 1024)                  \n",
      "Dropout                                 relu2                         dropout2                      (None, 1024)                  \n",
      "Concat                                  dropout2,multicross1          concat2                       (None, 1453)                  \n",
      "InnerProduct                            concat2                       fc3                           (None, 1)                     \n",
      "BinaryCrossEntropyLoss                  fc3,label                     loss                                                        \n",
      "------------------------------------------------------------------------------------------------------------------\n",
      "=====================================================Model Fit=====================================================\n",
      "[HUGECTR][03:33:05][INFO][RANK0]: Use non-epoch mode with number of iterations: 500\n",
      "[HUGECTR][03:33:05][INFO][RANK0]: Training batchsize: 4096, evaluation batchsize: 4096\n",
      "[HUGECTR][03:33:05][INFO][RANK0]: Evaluation interval: 100, snapshot interval: 10000\n",
      "[HUGECTR][03:33:05][INFO][RANK0]: Sparse embedding trainable: 1, dense network trainable: 1\n",
      "[HUGECTR][03:33:05][INFO][RANK0]: Use mixed precision: 0, scaler: 1, use cuda graph: 1946517802\n",
      "[HUGECTR][03:33:05][INFO][RANK0]: lr: 0.001000, warmup_steps: 1, decay_start: 0, decay_steps: 1, decay_power: 2.000000, end_lr: 0.000000\n",
      "[HUGECTR][03:33:05][INFO][RANK0]: Training source file: ./dcn_data/file_list.txt\n",
      "[HUGECTR][03:33:05][INFO][RANK0]: Evaluation source file: ./dcn_data/file_list_test.txt\n",
      "[HUGECTR][03:33:06][INFO][RANK0]: Iter: 50 Time(50 iters): 0.451251s Loss: 0.106090 lr:0.001000\n",
      "[HUGECTR][03:33:06][INFO][RANK0]: Iter: 100 Time(50 iters): 0.351385s Loss: 0.128124 lr:0.001000\n",
      "[HUGECTR][03:33:09][INFO][RANK0]: Evaluation, AUC: 0.741880\n",
      "[HUGECTR][03:33:09][INFO][RANK0]: Eval Time for 1500 iters: 2.972658s\n",
      "[HUGECTR][03:33:10][INFO][RANK0]: Iter: 150 Time(50 iters): 3.329078s Loss: 0.128845 lr:0.001000\n",
      "[HUGECTR][03:33:10][INFO][RANK0]: Iter: 200 Time(50 iters): 0.351025s Loss: 0.128085 lr:0.001000\n",
      "[HUGECTR][03:33:13][INFO][RANK0]: Evaluation, AUC: 0.730338\n",
      "[HUGECTR][03:33:13][INFO][RANK0]: Eval Time for 1500 iters: 2.973015s\n",
      "[HUGECTR][03:33:13][INFO][RANK0]: Iter: 250 Time(50 iters): 3.329451s Loss: 0.114888 lr:0.001000\n",
      "[HUGECTR][03:33:14][INFO][RANK0]: Iter: 300 Time(50 iters): 0.350937s Loss: 0.106827 lr:0.001000\n",
      "[HUGECTR][03:33:17][INFO][RANK0]: Evaluation, AUC: 0.728633\n",
      "[HUGECTR][03:33:17][INFO][RANK0]: Eval Time for 1500 iters: 2.972206s\n",
      "[HUGECTR][03:33:17][INFO][RANK0]: Iter: 350 Time(50 iters): 3.328661s Loss: 0.116533 lr:0.001000\n",
      "[HUGECTR][03:33:17][INFO][RANK0]: Iter: 400 Time(50 iters): 0.351332s Loss: 0.110059 lr:0.001000\n",
      "[HUGECTR][03:33:20][INFO][RANK0]: Evaluation, AUC: 0.726629\n",
      "[HUGECTR][03:33:20][INFO][RANK0]: Eval Time for 1500 iters: 2.970395s\n",
      "[HUGECTR][03:33:21][INFO][RANK0]: Iter: 450 Time(50 iters): 3.327333s Loss: 0.117730 lr:0.001000\n",
      "[HUGECTR][03:33:21][INFO][RANK0]: Finish 500 iterations with batchsize: 4096 in 15.53s.\n"
     ]
    }
   ],
   "source": [
    "!python3 dcn_continue.py"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "95eb57e8",
   "metadata": {
    "tags": []
   },
   "source": [
    "### Inference\n",
    "\n",
    "The HugeCTR inference is enabled by `hugectr.inference.InferenceSession.predict` method of InferenceSession.\n",
    "This method requires dense features, embedding columns, and row pointers of slots as the input and gives the prediction result as the output.\n",
    "We need to convert the Criteo data to inference format first."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "0ef1896e",
   "metadata": {},
   "outputs": [],
   "source": [
    "!python3 ../tools/criteo_predict/criteo2predict.py --src_csv_path=dcn_data/val/test.txt --src_config=../tools/criteo_predict/dcn_data.json --dst_path=./dcn_csr.txt --batch_size=1024"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b53ce38b",
   "metadata": {},
   "source": [
    "We can then make inferences based on the saved model graph and model weights by performing the following with Python APIs:\n",
    "\n",
    "1. Configure the inference related parameters.\n",
    "2. Create the inference session.\n",
    "3. Make inference with the `InferenceSession.predict` method. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "f54d51b2",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Writing dcn_inference.py\n"
     ]
    }
   ],
   "source": [
    "%%writefile dcn_inference.py\n",
    "from hugectr.inference import InferenceParams, CreateInferenceSession\n",
    "from mpi4py import MPI\n",
    "\n",
    "def calculate_accuracy(labels, output):\n",
    "    num_samples = len(labels)\n",
    "    flags = [1 if ((labels[i] == 0 and output[i] <= 0.5) or (labels[i] == 1 and output[i] > 0.5)) else 0 for i in range(num_samples)]\n",
    "    correct_samples = sum(flags)\n",
    "    return float(correct_samples)/(float(num_samples)+1e-16)\n",
    "\n",
    "data_file = open(\"dcn_csr.txt\")\n",
    "config_file = \"dcn.json\"\n",
    "labels = [int(item) for item in data_file.readline().split(' ')]\n",
    "dense_features = [float(item) for item in data_file.readline().split(' ') if item!=\"\\n\"]\n",
    "embedding_columns = [int(item) for item in data_file.readline().split(' ')]\n",
    "row_ptrs = [int(item) for item in data_file.readline().split(' ')]\n",
    "\n",
    "# create parameter server, embedding cache and inference session\n",
    "inference_params = InferenceParams(model_name = \"dcn\",\n",
    "                                max_batchsize = 1024,\n",
    "                                hit_rate_threshold = 0.6,\n",
    "                                dense_model_file = \"./dcn_dense_1000.model\",\n",
    "                                sparse_model_files = [\"./dcn0_sparse_1000.model\"],\n",
    "                                device_id = 0,\n",
    "                                use_gpu_embedding_cache = True,\n",
    "                                cache_size_percentage = 0.9,\n",
    "                                i64_input_key = False,\n",
    "                                use_mixed_precision = False)\n",
    "inference_session = CreateInferenceSession(config_file, inference_params)\n",
    "output = inference_session.predict(dense_features, embedding_columns, row_ptrs)\n",
    "accuracy = calculate_accuracy(labels, output)\n",
    "print(\"[HUGECTR][INFO] number samples: {}, accuracy: {}\".format(len(labels), accuracy))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "71e89b8f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[06d04h52m24s][HUGECTR][INFO]: default_emb_vec_value is not specified using default: 0.000000\n",
      "[06d04h52m26s][HUGECTR][INFO]: Global seed is 3956797427\n",
      "[06d04h52m28s][HUGECTR][INFO]: Peer-to-peer access cannot be fully enabled.\n",
      "[06d04h52m28s][HUGECTR][INFO]: Start all2all warmup\n",
      "[06d04h52m28s][HUGECTR][INFO]: End all2all warmup\n",
      "[06d04h52m28s][HUGECTR][INFO]: Use mixed precision: 0\n",
      "[06d04h52m28s][HUGECTR][INFO]: start create embedding for inference\n",
      "[06d04h52m28s][HUGECTR][INFO]: sparse_input name data1\n",
      "[06d04h52m28s][HUGECTR][INFO]: create embedding for inference success\n",
      "[06d04h52m28s][HUGECTR][INFO]: Inference stage skip BinaryCrossEntropyLoss layer, replaced by Sigmoid layer\n",
      "[HUGECTR][INFO] number samples: 1024, accuracy: 0.96875\n"
     ]
    }
   ],
   "source": [
    "!python3 dcn_inference.py"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ce4cb459",
   "metadata": {},
   "source": [
    "## Wide and Deep Model\n",
    "\n",
    "### Download and Preprocess Data\n",
    "\n",
    "1. Download the Criteo dataset using the following command:\n",
    "\n",
    "   ```shell\n",
    "   $ cd ${project_root}/tools\n",
    "   $ wget https://storage.googleapis.com/criteo-cail-datasets/day_1.gz\n",
    "   ```\n",
    "   \n",
    "   In preprocessing, we will further reduce the amounts of data to speedup the preprocessing, fill missing values, remove the feature values whose occurrences are very rare, etc. Here we choose pandas preprocessing method to make the dataset ready for HugeCTR training.\n",
    "\n",
    "2. Preprocessing by Pandas using the following command:\n",
    "\n",
    "   ```shell\n",
    "   $ bash preprocess.sh 1 wdl_data pandas 1 1 100\n",
    "   ```\n",
    "   \n",
    "   The first argument represents the dataset postfix. It is 1 here since day_1 is used. The second argument wdl_data is where the preprocessed data is stored. The fourth argument (one after pandas) 1 embodies that the normalization is applied to dense features. The fifth argument 1 means that the feature crossing is applied. The last argument 100 means the number of data files in each file list.\n",
    "   \n",
    "3. Create a soft link to the dataset folder using the following command:\n",
    "\n",
    "   ```shell\n",
    "   $ ln ${project_root}/tools/wdl_data ${project_root}/notebooks/wdl_data\n",
    "   ```\n",
    "   \n",
    "**Note**: It will take a while (dozens of minutes) to preprocess the dataset. Please make sure that it is finished successfully before moving forward to the next section."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a1aebf8d",
   "metadata": {},
   "source": [
    "### Train from Scratch\n",
    "\n",
    "We can train fom scratch, dump the model graph to a JSON file, and save the model weights and optimizer states by performing the same steps that we followed with the DCN Model."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "18c12473",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Writing wdl_train.py\n"
     ]
    }
   ],
   "source": [
    "%%writefile wdl_train.py\n",
    "import hugectr\n",
    "from mpi4py import MPI\n",
    "solver = hugectr.CreateSolver(max_eval_batches = 5000,\n",
    "                              batchsize_eval = 1024,\n",
    "                              batchsize = 1024,\n",
    "                              lr = 0.001,\n",
    "                              vvgpu = [[0]],\n",
    "                              i64_input_key = False,\n",
    "                              use_mixed_precision = False,\n",
    "                              repeat_dataset = False,\n",
    "                              use_cuda_graph = True)\n",
    "reader = hugectr.DataReaderParams(data_reader_type = hugectr.DataReaderType_t.Norm,\n",
    "                          source = [\"wdl_data/file_list.0.txt\"],\n",
    "                          eval_source = \"wdl_data/file_list.1.txt\",\n",
    "                          check_type = hugectr.Check_t.Sum)\n",
    "optimizer = hugectr.CreateOptimizer(optimizer_type = hugectr.Optimizer_t.Adam)\n",
    "model = hugectr.Model(solver, reader, optimizer)\n",
    "model.add(hugectr.Input(label_dim = 1, label_name = \"label\",\n",
    "                        dense_dim = 13, dense_name = \"dense\",\n",
    "                        data_reader_sparse_param_array = \n",
    "                        [hugectr.DataReaderSparseParam(\"wide_data\", 30, True, 1),\n",
    "                        hugectr.DataReaderSparseParam(\"deep_data\", 2, False, 26)]))\n",
    "model.add(hugectr.SparseEmbedding(embedding_type = hugectr.Embedding_t.DistributedSlotSparseEmbeddingHash, \n",
    "                            workspace_size_per_gpu_in_mb = 69,\n",
    "                            embedding_vec_size = 1,\n",
    "                            combiner = \"sum\",\n",
    "                            sparse_embedding_name = \"sparse_embedding2\",\n",
    "                            bottom_name = \"wide_data\",\n",
    "                            optimizer = optimizer))\n",
    "model.add(hugectr.SparseEmbedding(embedding_type = hugectr.Embedding_t.DistributedSlotSparseEmbeddingHash, \n",
    "                            workspace_size_per_gpu_in_mb = 1074,\n",
    "                            embedding_vec_size = 16,\n",
    "                            combiner = \"sum\",\n",
    "                            sparse_embedding_name = \"sparse_embedding1\",\n",
    "                            bottom_name = \"deep_data\",\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=416))\n",
    "model.add(hugectr.DenseLayer(layer_type = hugectr.Layer_t.Reshape,\n",
    "                            bottom_names = [\"sparse_embedding2\"],\n",
    "                            top_names = [\"reshape2\"],\n",
    "                            leading_dim=1))\n",
    "model.add(hugectr.DenseLayer(layer_type = hugectr.Layer_t.Concat,\n",
    "                            bottom_names = [\"reshape1\", \"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.Dropout,\n",
    "                            bottom_names = [\"relu1\"],\n",
    "                            top_names = [\"dropout1\"],\n",
    "                            dropout_rate=0.5))\n",
    "model.add(hugectr.DenseLayer(layer_type = hugectr.Layer_t.InnerProduct,\n",
    "                            bottom_names = [\"dropout1\"],\n",
    "                            top_names = [\"fc2\"],\n",
    "                            num_output=1024))\n",
    "model.add(hugectr.DenseLayer(layer_type = hugectr.Layer_t.ReLU,\n",
    "                            bottom_names = [\"fc2\"],\n",
    "                            top_names = [\"relu2\"]))\n",
    "model.add(hugectr.DenseLayer(layer_type = hugectr.Layer_t.Dropout,\n",
    "                            bottom_names = [\"relu2\"],\n",
    "                            top_names = [\"dropout2\"],\n",
    "                            dropout_rate=0.5))\n",
    "model.add(hugectr.DenseLayer(layer_type = hugectr.Layer_t.InnerProduct,\n",
    "                            bottom_names = [\"dropout2\"],\n",
    "                            top_names = [\"fc3\"],\n",
    "                            num_output=1))\n",
    "model.add(hugectr.DenseLayer(layer_type = hugectr.Layer_t.Add,\n",
    "                            bottom_names = [\"fc3\", \"reshape2\"],\n",
    "                            top_names = [\"add1\"]))\n",
    "model.add(hugectr.DenseLayer(layer_type = hugectr.Layer_t.BinaryCrossEntropyLoss,\n",
    "                            bottom_names = [\"add1\", \"label\"],\n",
    "                            top_names = [\"loss\"]))\n",
    "model.compile()\n",
    "model.summary()\n",
    "model.graph_to_json(graph_config_file = \"wdl.json\")\n",
    "model.fit(num_epochs = 1, display = 500, eval_interval = 500, snapshot = 4000, snapshot_prefix = \"wdl\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "6a54ff32",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "HugeCTR Version: 3.2\n",
      "====================================================Model Init=====================================================\n",
      "[HUGECTR][07:13:04][INFO][RANK0]: Global seed is 1910256490\n",
      "[HUGECTR][07:13:04][INFO][RANK0]: Device to NUMA mapping:\n",
      "  GPU 0 ->  node 0\n",
      "\n",
      "[HUGECTR][07:13:06][WARNING][RANK0]: Peer-to-peer access cannot be fully enabled.\n",
      "[HUGECTR][07:13:06][INFO][RANK0]: Start all2all warmup\n",
      "[HUGECTR][07:13:06][INFO][RANK0]: End all2all warmup\n",
      "[HUGECTR][07:13:06][INFO][RANK0]: Using All-reduce algorithm: NCCL\n",
      "[HUGECTR][07:13:06][INFO][RANK0]: Device 0: Tesla V100-SXM2-16GB\n",
      "[HUGECTR][07:13:06][INFO][RANK0]: num of DataReader workers: 12\n",
      "[HUGECTR][07:13:06][INFO][RANK0]: max_vocabulary_size_per_gpu_=6029312\n",
      "[HUGECTR][07:13:06][INFO][RANK0]: max_vocabulary_size_per_gpu_=5865472\n",
      "[HUGECTR][07:13:06][INFO][RANK0]: Graph analysis to resolve tensor dependency\n",
      "===================================================Model Compile===================================================\n",
      "[HUGECTR][07:13:09][INFO][RANK0]: gpu0 start to init embedding\n",
      "[HUGECTR][07:13:09][INFO][RANK0]: gpu0 init embedding done\n",
      "[HUGECTR][07:13:09][INFO][RANK0]: gpu0 start to init embedding\n",
      "[HUGECTR][07:13:09][INFO][RANK0]: gpu0 init embedding done\n",
      "[HUGECTR][07:13:09][INFO][RANK0]: Starting AUC NCCL warm-up\n",
      "[HUGECTR][07:13:09][INFO][RANK0]: Warm-up done\n",
      "===================================================Model Summary===================================================\n",
      "label                                   Dense                         Sparse                        \n",
      "label                                   dense                          wide_data,deep_data           \n",
      "(None, 1)                               (None, 13)                              \n",
      "——————————————————————————————————————————————————————————————————————————————————————————————————————————————————\n",
      "Layer Type                              Input Name                    Output Name                   Output Shape                  \n",
      "——————————————————————————————————————————————————————————————————————————————————————————————————————————————————\n",
      "DistributedSlotSparseEmbeddingHash      wide_data                     sparse_embedding2             (None, 1, 1)                  \n",
      "------------------------------------------------------------------------------------------------------------------\n",
      "DistributedSlotSparseEmbeddingHash      deep_data                     sparse_embedding1             (None, 26, 16)                \n",
      "------------------------------------------------------------------------------------------------------------------\n",
      "Reshape                                 sparse_embedding1             reshape1                      (None, 416)                   \n",
      "------------------------------------------------------------------------------------------------------------------\n",
      "Reshape                                 sparse_embedding2             reshape2                      (None, 1)                     \n",
      "------------------------------------------------------------------------------------------------------------------\n",
      "Concat                                  reshape1                      concat1                       (None, 429)                   \n",
      "                                        dense                                                                                     \n",
      "------------------------------------------------------------------------------------------------------------------\n",
      "InnerProduct                            concat1                       fc1                           (None, 1024)                  \n",
      "------------------------------------------------------------------------------------------------------------------\n",
      "ReLU                                    fc1                           relu1                         (None, 1024)                  \n",
      "------------------------------------------------------------------------------------------------------------------\n",
      "Dropout                                 relu1                         dropout1                      (None, 1024)                  \n",
      "------------------------------------------------------------------------------------------------------------------\n",
      "InnerProduct                            dropout1                      fc2                           (None, 1024)                  \n",
      "------------------------------------------------------------------------------------------------------------------\n",
      "ReLU                                    fc2                           relu2                         (None, 1024)                  \n",
      "------------------------------------------------------------------------------------------------------------------\n",
      "Dropout                                 relu2                         dropout2                      (None, 1024)                  \n",
      "------------------------------------------------------------------------------------------------------------------\n",
      "InnerProduct                            dropout2                      fc3                           (None, 1)                     \n",
      "------------------------------------------------------------------------------------------------------------------\n",
      "Add                                     fc3                           add1                          (None, 1)                     \n",
      "                                        reshape2                                                                                  \n",
      "------------------------------------------------------------------------------------------------------------------\n",
      "BinaryCrossEntropyLoss                  add1                          loss                                                        \n",
      "                                        label                                                                                     \n",
      "------------------------------------------------------------------------------------------------------------------\n",
      "[HUGECTR][07:13:09][INFO][RANK0]: Save the model graph to wdl.json successfully\n",
      "=====================================================Model Fit=====================================================\n",
      "[HUGECTR][07:13:09][INFO][RANK0]: Use epoch mode with number of epochs: 1\n",
      "[HUGECTR][07:13:09][INFO][RANK0]: Training batchsize: 1024, evaluation batchsize: 1024\n",
      "[HUGECTR][07:13:09][INFO][RANK0]: Evaluation interval: 500, snapshot interval: 4000\n",
      "[HUGECTR][07:13:09][INFO][RANK0]: Dense network trainable: True\n",
      "[HUGECTR][07:13:09][INFO][RANK0]: Sparse embedding sparse_embedding1 trainable: True\n",
      "[HUGECTR][07:13:09][INFO][RANK0]: Sparse embedding sparse_embedding2 trainable: True\n",
      "[HUGECTR][07:13:09][INFO][RANK0]: Use mixed precision: False, scaler: 1.000000, use cuda graph: True\n",
      "[HUGECTR][07:13:09][INFO][RANK0]: lr: 0.001000, warmup_steps: 1, end_lr: 0.000000\n",
      "[HUGECTR][07:13:09][INFO][RANK0]: decay_start: 0, decay_steps: 1, decay_power: 2.000000\n",
      "[HUGECTR][07:13:09][INFO][RANK0]: Training source file: wdl_data/file_list.0.txt\n",
      "[HUGECTR][07:13:09][INFO][RANK0]: Evaluation source file: wdl_data/file_list.1.txt\n",
      "[HUGECTR][07:13:09][INFO][RANK0]: -----------------------------------Epoch 0-----------------------------------\n",
      "[HUGECTR][07:13:12][INFO][RANK0]: Iter: 500 Time(500 iters): 2.633535s Loss: 0.146090 lr:0.001000\n",
      "[HUGECTR][07:13:24][INFO][RANK0]: Evaluation, AUC: 0.729127\n",
      "[HUGECTR][07:13:24][INFO][RANK0]: Eval Time for 5000 iters: 12.329920s\n",
      "[HUGECTR][07:13:27][INFO][RANK0]: Iter: 1000 Time(500 iters): 14.758588s Loss: 0.133614 lr:0.001000\n",
      "[HUGECTR][07:13:29][INFO][RANK0]: Evaluation, AUC: 0.739283\n",
      "[HUGECTR][07:13:29][INFO][RANK0]: Eval Time for 5000 iters: 1.892544s\n",
      "[HUGECTR][07:13:31][INFO][RANK0]: Iter: 1500 Time(500 iters): 4.309439s Loss: 0.145478 lr:0.001000\n",
      "[HUGECTR][07:13:33][INFO][RANK0]: Evaluation, AUC: 0.744546\n",
      "[HUGECTR][07:13:33][INFO][RANK0]: Eval Time for 5000 iters: 1.888380s\n",
      "[HUGECTR][07:13:35][INFO][RANK0]: Iter: 2000 Time(500 iters): 4.322345s Loss: 0.142099 lr:0.001000\n",
      "[HUGECTR][07:13:37][INFO][RANK0]: Evaluation, AUC: 0.748392\n",
      "[HUGECTR][07:13:37][INFO][RANK0]: Eval Time for 5000 iters: 1.894575s\n",
      "[HUGECTR][07:13:40][INFO][RANK0]: Iter: 2500 Time(500 iters): 4.354580s Loss: 0.167694 lr:0.001000\n",
      "[HUGECTR][07:13:42][INFO][RANK0]: Evaluation, AUC: 0.748089\n",
      "[HUGECTR][07:13:42][INFO][RANK0]: Eval Time for 5000 iters: 1.853501s\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[HUGECTR][07:13:44][INFO][RANK0]: Iter: 3000 Time(500 iters): 4.269730s Loss: 0.124279 lr:0.001000\n",
      "[HUGECTR][07:13:46][INFO][RANK0]: Evaluation, AUC: 0.753290\n",
      "[HUGECTR][07:13:46][INFO][RANK0]: Eval Time for 5000 iters: 1.906998s\n",
      "[HUGECTR][07:13:48][INFO][RANK0]: Iter: 3500 Time(500 iters): 4.328614s Loss: 0.114806 lr:0.001000\n",
      "[HUGECTR][07:13:50][INFO][RANK0]: Evaluation, AUC: 0.755007\n",
      "[HUGECTR][07:13:50][INFO][RANK0]: Eval Time for 5000 iters: 1.897527s\n",
      "[HUGECTR][07:13:53][INFO][RANK0]: Iter: 4000 Time(500 iters): 4.308839s Loss: 0.128652 lr:0.001000\n",
      "[HUGECTR][07:13:55][INFO][RANK0]: Evaluation, AUC: 0.756323\n",
      "[HUGECTR][07:13:55][INFO][RANK0]: Eval Time for 5000 iters: 1.849973s\n",
      "[HUGECTR][07:13:55][INFO][RANK0]: Rank0: Write hash table to file\n",
      "[HUGECTR][07:13:55][INFO][RANK0]: Rank0: Write hash table to file\n",
      "[HUGECTR][07:13:55][INFO][RANK0]: Dumping sparse weights to files, successful\n",
      "[HUGECTR][07:13:56][INFO][RANK0]: Rank0: Write optimzer state to file\n",
      "[HUGECTR][07:13:56][INFO][RANK0]: Done\n",
      "[HUGECTR][07:13:56][INFO][RANK0]: Rank0: Write optimzer state to file\n",
      "[HUGECTR][07:13:56][INFO][RANK0]: Done\n",
      "[HUGECTR][07:13:56][INFO][RANK0]: Rank0: Write optimzer state to file\n",
      "[HUGECTR][07:13:56][INFO][RANK0]: Done\n",
      "[HUGECTR][07:13:57][INFO][RANK0]: Rank0: Write optimzer state to file\n",
      "[HUGECTR][07:13:57][INFO][RANK0]: Done\n",
      "[HUGECTR][07:14:04][INFO][RANK0]: Dumping sparse optimzer states to files, successful\n",
      "[HUGECTR][07:14:04][INFO][RANK0]: Dumping dense weights to file, successful\n",
      "[HUGECTR][07:14:04][INFO][RANK0]: Dumping dense optimizer states to file, successful\n",
      "[HUGECTR][07:14:04][INFO][RANK0]: Dumping untrainable weights to file, successful\n",
      "[HUGECTR][07:14:04][INFO][RANK0]: Finish 1 epochs 4001 global iterations with batchsize 1024 in 54.53s.\n"
     ]
    }
   ],
   "source": [
    "!python3 wdl_train.py"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "79f88be7",
   "metadata": {},
   "source": [
    "### Fine-tuning\n",
    "\n",
    "We can only load the sparse embedding layers, their corresponding weights, and then construct a new dense network. The dense weights will be trained first and the sparse weights will be fine-tuned later. We can achieve this by performing the following with Python APIs:\n",
    "\n",
    "1. Create the solver, reader and optimizer, then initialize the model.\n",
    "2. Load the sparse embedding layers from the saved JSON file.\n",
    "3. Add the dense layers on top of the loaded model graph.\n",
    "4. Compile the model and have an overview of the model graph.\n",
    "5. Load the sparse weights and freeze the sparse embedding layers.\n",
    "6. Train the dense weights.\n",
    "7. Unfreeze the sparse embedding layers and freeze the dense layers, reset the learning rate scheduler with a small rate.\n",
    "8. Fine-tune the sparse weights."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "60908e93",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Overwriting wdl_fine_tune.py\n"
     ]
    }
   ],
   "source": [
    "%%writefile wdl_fine_tune.py\n",
    "import hugectr\n",
    "from mpi4py import MPI\n",
    "solver = hugectr.CreateSolver(max_eval_batches = 5000,\n",
    "                              batchsize_eval = 1024,\n",
    "                              batchsize = 1024,\n",
    "                              vvgpu = [[0]],\n",
    "                              i64_input_key = False,\n",
    "                              use_mixed_precision = False,\n",
    "                              repeat_dataset = False,\n",
    "                              use_cuda_graph = True)\n",
    "reader = hugectr.DataReaderParams(data_reader_type = hugectr.DataReaderType_t.Norm,\n",
    "                          source = [\"wdl_data/file_list.2.txt\"],\n",
    "                          eval_source = \"wdl_data/file_list.3.txt\",\n",
    "                          check_type = hugectr.Check_t.Sum)\n",
    "optimizer = hugectr.CreateOptimizer(optimizer_type = hugectr.Optimizer_t.Adam)\n",
    "model = hugectr.Model(solver, reader, optimizer)\n",
    "model.construct_from_json(graph_config_file = \"wdl.json\", include_dense_network = False)\n",
    "model.add(hugectr.DenseLayer(layer_type = hugectr.Layer_t.Reshape,\n",
    "                            bottom_names = [\"sparse_embedding1\"],\n",
    "                            top_names = [\"reshape1\"],\n",
    "                            leading_dim=416))\n",
    "model.add(hugectr.DenseLayer(layer_type = hugectr.Layer_t.Reshape,\n",
    "                            bottom_names = [\"sparse_embedding2\"],\n",
    "                            top_names = [\"reshape2\"],\n",
    "                            leading_dim=1))\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.Dropout,\n",
    "                            bottom_names = [\"relu1\"],\n",
    "                            top_names = [\"dropout1\"],\n",
    "                            dropout_rate=0.5))\n",
    "model.add(hugectr.DenseLayer(layer_type = hugectr.Layer_t.InnerProduct,\n",
    "                            bottom_names = [\"dropout1\"],\n",
    "                            top_names = [\"fc2\"],\n",
    "                            num_output=1))\n",
    "model.add(hugectr.DenseLayer(layer_type = hugectr.Layer_t.BinaryCrossEntropyLoss,\n",
    "                            bottom_names = [\"fc2\", \"label\"],\n",
    "                            top_names = [\"loss\"]))\n",
    "model.compile()\n",
    "model.summary()\n",
    "model.load_sparse_weights([\"wdl0_sparse_4000.model\", \"wdl1_sparse_4000.model\"])\n",
    "model.freeze_embedding()\n",
    "model.fit(num_epochs = 1, display = 500, eval_interval = 1000, snapshot = 100000, snapshot_prefix = \"wdl\")\n",
    "model.unfreeze_embedding()\n",
    "model.freeze_dense()\n",
    "model.reset_learning_rate_scheduler(base_lr = 0.0001)\n",
    "model.fit(num_epochs = 2, display = 500, eval_interval = 1000, snapshot = 100000, snapshot_prefix = \"wdl\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "b6af6081",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "HugeCTR Version: 3.2\n",
      "====================================================Model Init=====================================================\n",
      "[HUGECTR][07:29:56][INFO][RANK0]: Global seed is 2136095432\n",
      "[HUGECTR][07:29:56][INFO][RANK0]: Device to NUMA mapping:\n",
      "  GPU 0 ->  node 0\n",
      "\n",
      "[HUGECTR][07:29:58][WARNING][RANK0]: Peer-to-peer access cannot be fully enabled.\n",
      "[HUGECTR][07:29:58][INFO][RANK0]: Start all2all warmup\n",
      "[HUGECTR][07:29:58][INFO][RANK0]: End all2all warmup\n",
      "[HUGECTR][07:29:58][INFO][RANK0]: Using All-reduce algorithm: NCCL\n",
      "[HUGECTR][07:29:58][INFO][RANK0]: Device 0: Tesla V100-SXM2-16GB\n",
      "[HUGECTR][07:29:58][INFO][RANK0]: num of DataReader workers: 12\n",
      "[HUGECTR][07:29:58][INFO][RANK0]: max_num_frequent_categories is not specified using default: 1\n",
      "[HUGECTR][07:29:58][INFO][RANK0]: max_num_infrequent_samples is not specified using default: -1\n",
      "[HUGECTR][07:29:58][INFO][RANK0]: p_dup_max is not specified using default: 0.010000\n",
      "[HUGECTR][07:29:58][INFO][RANK0]: max_all_reduce_bandwidth is not specified using default: 130000000000.000000\n",
      "[HUGECTR][07:29:58][INFO][RANK0]: max_all_to_all_bandwidth is not specified using default: 190000000000.000000\n",
      "[HUGECTR][07:29:58][INFO][RANK0]: efficiency_bandwidth_ratio is not specified using default: 1.000000\n",
      "[HUGECTR][07:29:58][INFO][RANK0]: communication_type is not specified using default: IB_NVLink\n",
      "[HUGECTR][07:29:58][INFO][RANK0]: hybrid_embedding_type is not specified using default: Distributed\n",
      "[HUGECTR][07:29:58][INFO][RANK0]: max_vocabulary_size_per_gpu_=6029312\n",
      "[HUGECTR][07:29:58][INFO][RANK0]: max_num_frequent_categories is not specified using default: 1\n",
      "[HUGECTR][07:29:58][INFO][RANK0]: max_num_infrequent_samples is not specified using default: -1\n",
      "[HUGECTR][07:29:58][INFO][RANK0]: p_dup_max is not specified using default: 0.010000\n",
      "[HUGECTR][07:29:58][INFO][RANK0]: max_all_reduce_bandwidth is not specified using default: 130000000000.000000\n",
      "[HUGECTR][07:29:58][INFO][RANK0]: max_all_to_all_bandwidth is not specified using default: 190000000000.000000\n",
      "[HUGECTR][07:29:58][INFO][RANK0]: efficiency_bandwidth_ratio is not specified using default: 1.000000\n",
      "[HUGECTR][07:29:58][INFO][RANK0]: communication_type is not specified using default: IB_NVLink\n",
      "[HUGECTR][07:29:58][INFO][RANK0]: hybrid_embedding_type is not specified using default: Distributed\n",
      "[HUGECTR][07:29:58][INFO][RANK0]: max_vocabulary_size_per_gpu_=5865472\n",
      "[HUGECTR][07:29:58][INFO][RANK0]: Load the model graph from wdl.json successfully\n",
      "[HUGECTR][07:29:58][INFO][RANK0]: Graph analysis to resolve tensor dependency\n",
      "===================================================Model Compile===================================================\n",
      "[HUGECTR][07:30:00][INFO][RANK0]: gpu0 start to init embedding\n",
      "[HUGECTR][07:30:00][INFO][RANK0]: gpu0 init embedding done\n",
      "[HUGECTR][07:30:00][INFO][RANK0]: gpu0 start to init embedding\n",
      "[HUGECTR][07:30:00][INFO][RANK0]: gpu0 init embedding done\n",
      "[HUGECTR][07:30:00][INFO][RANK0]: Starting AUC NCCL warm-up\n",
      "[HUGECTR][07:30:00][INFO][RANK0]: Warm-up done\n",
      "===================================================Model Summary===================================================\n",
      "label                                   Dense                         Sparse                        \n",
      "label                                   dense                          wide_data,deep_data           \n",
      "(None, 1)                               (None, 13)                              \n",
      "——————————————————————————————————————————————————————————————————————————————————————————————————————————————————\n",
      "Layer Type                              Input Name                    Output Name                   Output Shape                  \n",
      "——————————————————————————————————————————————————————————————————————————————————————————————————————————————————\n",
      "DistributedSlotSparseEmbeddingHash      wide_data                     sparse_embedding2             (None, 1, 1)                  \n",
      "------------------------------------------------------------------------------------------------------------------\n",
      "DistributedSlotSparseEmbeddingHash      deep_data                     sparse_embedding1             (None, 26, 16)                \n",
      "------------------------------------------------------------------------------------------------------------------\n",
      "Reshape                                 sparse_embedding1             reshape1                      (None, 416)                   \n",
      "------------------------------------------------------------------------------------------------------------------\n",
      "Reshape                                 sparse_embedding2             reshape2                      (None, 1)                     \n",
      "------------------------------------------------------------------------------------------------------------------\n",
      "Concat                                  reshape1                      concat1                       (None, 430)                   \n",
      "                                        reshape2                                                                                  \n",
      "                                        dense                                                                                     \n",
      "------------------------------------------------------------------------------------------------------------------\n",
      "InnerProduct                            concat1                       fc1                           (None, 1024)                  \n",
      "------------------------------------------------------------------------------------------------------------------\n",
      "ReLU                                    fc1                           relu1                         (None, 1024)                  \n",
      "------------------------------------------------------------------------------------------------------------------\n",
      "Dropout                                 relu1                         dropout1                      (None, 1024)                  \n",
      "------------------------------------------------------------------------------------------------------------------\n",
      "InnerProduct                            dropout1                      fc2                           (None, 1)                     \n",
      "------------------------------------------------------------------------------------------------------------------\n",
      "BinaryCrossEntropyLoss                  fc2                           loss                                                        \n",
      "                                        label                                                                                     \n",
      "------------------------------------------------------------------------------------------------------------------\n",
      "[HUGECTR][07:30:00][INFO][RANK0]: Loading sparse model: wdl0_sparse_4000.model\n",
      "[HUGECTR][07:30:00][INFO][RANK0]: Loading sparse model: wdl1_sparse_4000.model\n",
      "=====================================================Model Fit=====================================================\n",
      "[HUGECTR][07:30:00][INFO][RANK0]: Use epoch mode with number of epochs: 1\n",
      "[HUGECTR][07:30:00][INFO][RANK0]: Training batchsize: 1024, evaluation batchsize: 1024\n",
      "[HUGECTR][07:30:00][INFO][RANK0]: Evaluation interval: 1000, snapshot interval: 100000\n",
      "[HUGECTR][07:30:00][INFO][RANK0]: Dense network trainable: True\n",
      "[HUGECTR][07:30:00][INFO][RANK0]: Sparse embedding sparse_embedding1 trainable: False\n",
      "[HUGECTR][07:30:00][INFO][RANK0]: Sparse embedding sparse_embedding2 trainable: False\n",
      "[HUGECTR][07:30:00][INFO][RANK0]: Use mixed precision: False, scaler: 1.000000, use cuda graph: True\n",
      "[HUGECTR][07:30:00][INFO][RANK0]: lr: 0.001000, warmup_steps: 1, end_lr: 0.000000\n",
      "[HUGECTR][07:30:00][INFO][RANK0]: decay_start: 0, decay_steps: 1, decay_power: 2.000000\n",
      "[HUGECTR][07:30:00][INFO][RANK0]: Training source file: wdl_data/file_list.2.txt\n",
      "[HUGECTR][07:30:00][INFO][RANK0]: Evaluation source file: wdl_data/file_list.3.txt\n",
      "[HUGECTR][07:30:00][INFO][RANK0]: -----------------------------------Epoch 0-----------------------------------\n",
      "[HUGECTR][07:30:02][INFO][RANK0]: Iter: 500 Time(500 iters): 2.288018s Loss: 0.115263 lr:0.001000\n",
      "[HUGECTR][07:30:04][INFO][RANK0]: Iter: 1000 Time(500 iters): 2.084800s Loss: 0.130941 lr:0.001000\n",
      "[HUGECTR][07:30:06][INFO][RANK0]: Evaluation, AUC: 0.753592\n",
      "[HUGECTR][07:30:06][INFO][RANK0]: Eval Time for 5000 iters: 1.233550s\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[HUGECTR][07:30:08][INFO][RANK0]: Iter: 1500 Time(500 iters): 3.320545s Loss: 0.160203 lr:0.001000\n",
      "[HUGECTR][07:30:10][INFO][RANK0]: Iter: 2000 Time(500 iters): 2.083907s Loss: 0.133159 lr:0.001000\n",
      "[HUGECTR][07:30:11][INFO][RANK0]: Evaluation, AUC: 0.757654\n",
      "[HUGECTR][07:30:11][INFO][RANK0]: Eval Time for 5000 iters: 1.257166s\n",
      "[HUGECTR][07:30:13][INFO][RANK0]: Iter: 2500 Time(500 iters): 3.344821s Loss: 0.114668 lr:0.001000\n",
      "[HUGECTR][07:30:15][INFO][RANK0]: Iter: 3000 Time(500 iters): 2.085232s Loss: 0.131622 lr:0.001000\n",
      "[HUGECTR][07:30:17][INFO][RANK0]: Evaluation, AUC: 0.759316\n",
      "[HUGECTR][07:30:17][INFO][RANK0]: Eval Time for 5000 iters: 1.307634s\n",
      "[HUGECTR][07:30:19][INFO][RANK0]: Iter: 3500 Time(500 iters): 3.395008s Loss: 0.140864 lr:0.001000\n",
      "[HUGECTR][07:30:21][INFO][RANK0]: Iter: 4000 Time(500 iters): 2.080470s Loss: 0.132377 lr:0.001000\n",
      "[HUGECTR][07:30:22][INFO][RANK0]: Evaluation, AUC: 0.759804\n",
      "[HUGECTR][07:30:22][INFO][RANK0]: Eval Time for 5000 iters: 1.176526s\n",
      "[HUGECTR][07:30:22][INFO][RANK0]: Finish 1 epochs 4001 global iterations with batchsize 1024 in 21.89s.\n",
      "=====================================================Model Fit=====================================================\n",
      "[HUGECTR][07:30:22][INFO][RANK0]: Use epoch mode with number of epochs: 2\n",
      "[HUGECTR][07:30:22][INFO][RANK0]: Training batchsize: 1024, evaluation batchsize: 1024\n",
      "[HUGECTR][07:30:22][INFO][RANK0]: Evaluation interval: 1000, snapshot interval: 100000\n",
      "[HUGECTR][07:30:22][INFO][RANK0]: Dense network trainable: False\n",
      "[HUGECTR][07:30:22][INFO][RANK0]: Sparse embedding sparse_embedding1 trainable: True\n",
      "[HUGECTR][07:30:22][INFO][RANK0]: Sparse embedding sparse_embedding2 trainable: True\n",
      "[HUGECTR][07:30:22][INFO][RANK0]: Use mixed precision: False, scaler: 1.000000, use cuda graph: True\n",
      "[HUGECTR][07:30:22][INFO][RANK0]: lr: 0.001000, warmup_steps: 1, end_lr: 0.000000\n",
      "[HUGECTR][07:30:22][INFO][RANK0]: decay_start: 0, decay_steps: 1, decay_power: 2.000000\n",
      "[HUGECTR][07:30:22][INFO][RANK0]: Training source file: wdl_data/file_list.2.txt\n",
      "[HUGECTR][07:30:22][INFO][RANK0]: Evaluation source file: wdl_data/file_list.3.txt\n",
      "[HUGECTR][07:30:22][INFO][RANK0]: -----------------------------------Epoch 0-----------------------------------\n",
      "[HUGECTR][07:30:24][INFO][RANK0]: Iter: 500 Time(500 iters): 2.143974s Loss: 0.113414 lr:0.000100\n",
      "[HUGECTR][07:30:26][INFO][RANK0]: Iter: 1000 Time(500 iters): 2.082424s Loss: 0.128542 lr:0.000100\n",
      "[HUGECTR][07:30:27][INFO][RANK0]: Evaluation, AUC: 0.761524\n",
      "[HUGECTR][07:30:27][INFO][RANK0]: Eval Time for 5000 iters: 1.205501s\n",
      "[HUGECTR][07:30:30][INFO][RANK0]: Iter: 1500 Time(500 iters): 3.291612s Loss: 0.161557 lr:0.000100\n",
      "[HUGECTR][07:30:32][INFO][RANK0]: Iter: 2000 Time(500 iters): 2.083802s Loss: 0.131485 lr:0.000100\n",
      "[HUGECTR][07:30:33][INFO][RANK0]: Evaluation, AUC: 0.762616\n",
      "[HUGECTR][07:30:33][INFO][RANK0]: Eval Time for 5000 iters: 1.170735s\n",
      "[HUGECTR][07:30:35][INFO][RANK0]: Iter: 2500 Time(500 iters): 3.260273s Loss: 0.111285 lr:0.000100\n",
      "[HUGECTR][07:30:37][INFO][RANK0]: Iter: 3000 Time(500 iters): 2.086440s Loss: 0.128462 lr:0.000100\n",
      "[HUGECTR][07:30:38][INFO][RANK0]: Evaluation, AUC: 0.763377\n",
      "[HUGECTR][07:30:38][INFO][RANK0]: Eval Time for 5000 iters: 1.226101s\n",
      "[HUGECTR][07:30:40][INFO][RANK0]: Iter: 3500 Time(500 iters): 3.316375s Loss: 0.140509 lr:0.000100\n",
      "[HUGECTR][07:30:42][INFO][RANK0]: Iter: 4000 Time(500 iters): 2.082592s Loss: 0.127911 lr:0.000100\n",
      "[HUGECTR][07:30:44][INFO][RANK0]: Evaluation, AUC: 0.763914\n",
      "[HUGECTR][07:30:44][INFO][RANK0]: Eval Time for 5000 iters: 1.270215s\n",
      "[HUGECTR][07:30:44][INFO][RANK0]: -----------------------------------Epoch 1-----------------------------------\n",
      "[HUGECTR][07:30:46][INFO][RANK0]: Iter: 4500 Time(500 iters): 3.370336s Loss: 0.146994 lr:0.000100\n",
      "[HUGECTR][07:30:48][INFO][RANK0]: Iter: 5000 Time(500 iters): 2.087686s Loss: 0.110219 lr:0.000100\n",
      "[HUGECTR][07:30:49][INFO][RANK0]: Evaluation, AUC: 0.764316\n",
      "[HUGECTR][07:30:49][INFO][RANK0]: Eval Time for 5000 iters: 1.142174s\n",
      "[HUGECTR][07:30:51][INFO][RANK0]: Iter: 5500 Time(500 iters): 3.232587s Loss: 0.144252 lr:0.000100\n",
      "[HUGECTR][07:30:53][INFO][RANK0]: Iter: 6000 Time(500 iters): 2.087027s Loss: 0.122446 lr:0.000100\n",
      "[HUGECTR][07:30:54][INFO][RANK0]: Evaluation, AUC: 0.764680\n",
      "[HUGECTR][07:30:54][INFO][RANK0]: Eval Time for 5000 iters: 1.234112s\n",
      "[HUGECTR][07:30:56][INFO][RANK0]: Iter: 6500 Time(500 iters): 3.325564s Loss: 0.098065 lr:0.000100\n",
      "[HUGECTR][07:30:59][INFO][RANK0]: Iter: 7000 Time(500 iters): 2.087908s Loss: 0.132715 lr:0.000100\n",
      "[HUGECTR][07:31:00][INFO][RANK0]: Evaluation, AUC: 0.764872\n",
      "[HUGECTR][07:31:00][INFO][RANK0]: Eval Time for 5000 iters: 1.179473s\n",
      "[HUGECTR][07:31:02][INFO][RANK0]: Iter: 7500 Time(500 iters): 3.268473s Loss: 0.132111 lr:0.000100\n",
      "[HUGECTR][07:31:04][INFO][RANK0]: Iter: 8000 Time(500 iters): 2.083339s Loss: 0.126090 lr:0.000100\n",
      "[HUGECTR][07:31:05][INFO][RANK0]: Evaluation, AUC: 0.764933\n",
      "[HUGECTR][07:31:05][INFO][RANK0]: Eval Time for 5000 iters: 1.272813s\n",
      "[HUGECTR][07:31:05][INFO][RANK0]: Finish 2 epochs 8002 global iterations with batchsize 1024 in 43.22s.\n"
     ]
    }
   ],
   "source": [
    "!python3 wdl_fine_tune.py"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8847365b",
   "metadata": {},
   "source": [
    "### Load Pre-trained Embeddings\n",
    "\n",
    "If you have the pre-trained embeddings in other formats, you can convert them to the HugeCTR sparse models and then load them to facilitate the training process. For the sake of simplicity and generality, we represent the pretrained embeddings with the dictionary of randomly initialized numpy arrays, of which the keys indicate the embedding keys and the array values embody the embedding values. It is worth mentioning that there are two embedding tables for the Wide&Deep model, and here we only load the pre-trained embeddings for one table and freeze the corresponding embedding layer."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "ec6b1785",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Overwriting wdl_load_pretrained.py\n"
     ]
    }
   ],
   "source": [
    "%%writefile wdl_load_pretrained.py\n",
    "import hugectr\n",
    "from mpi4py import MPI\n",
    "import numpy as np\n",
    "import os\n",
    "import struct\n",
    "solver = hugectr.CreateSolver(max_eval_batches = 5000,\n",
    "                              batchsize_eval = 1024,\n",
    "                              batchsize = 1024,\n",
    "                              vvgpu = [[0]],\n",
    "                              i64_input_key = False,\n",
    "                              use_mixed_precision = False,\n",
    "                              repeat_dataset = False,\n",
    "                              use_cuda_graph = True)\n",
    "reader = hugectr.DataReaderParams(data_reader_type = hugectr.DataReaderType_t.Norm,\n",
    "                          source = [\"wdl_data/file_list.0.txt\"],\n",
    "                          eval_source = \"wdl_data/file_list.1.txt\",\n",
    "                          check_type = hugectr.Check_t.Sum)\n",
    "optimizer = hugectr.CreateOptimizer(optimizer_type = hugectr.Optimizer_t.Adam)\n",
    "model = hugectr.Model(solver, reader, optimizer)\n",
    "model.construct_from_json(graph_config_file = \"wdl.json\", include_dense_network = True)\n",
    "model.compile()\n",
    "model.summary()\n",
    "\n",
    "def convert_pretrained_embeddings_to_sparse_model(pre_trained_sparse_embeddings, hugectr_sparse_model, embedding_vec_size):\n",
    "    os.system(\"mkdir -p {}\".format(hugectr_sparse_model))\n",
    "    with open(\"{}/key\".format(hugectr_sparse_model), 'wb') as key_file, \\\n",
    "        open(\"{}/emb_vector\".format(hugectr_sparse_model), 'wb') as vec_file:\n",
    "      for key in pre_trained_sparse_embeddings:\n",
    "        vec = pre_trained_sparse_embeddings[key]\n",
    "        key_struct = struct.pack('q', key)\n",
    "        vec_struct = struct.pack(str(embedding_vec_size) + \"f\", *vec)\n",
    "        key_file.write(key_struct)\n",
    "        vec_file.write(vec_struct)\n",
    "\n",
    "# Convert the pretrained embeddings\n",
    "pretrained_embeddings = dict()\n",
    "hugectr_sparse_model = \"wdl1_pretrained.model\"\n",
    "embedding_vec_size = 16\n",
    "key_range = (0, 100000)\n",
    "for key in range(key_range[0], key_range[1]):\n",
    "    pretrained_embeddings[key] = np.random.randn(embedding_vec_size).astype(np.float32)\n",
    "convert_pretrained_embeddings_to_sparse_model(pretrained_embeddings, hugectr_sparse_model, embedding_vec_size)\n",
    "print(\"Successfully convert pretrained embeddings to {}\".format(hugectr_sparse_model))\n",
    "\n",
    "# Load the pretrained sparse models\n",
    "model.load_sparse_weights({\"sparse_embedding1\": hugectr_sparse_model})\n",
    "model.freeze_embedding(\"sparse_embedding1\")\n",
    "model.fit(num_epochs = 1, display = 500, eval_interval = 1000, snapshot = 100000, snapshot_prefix = \"wdl\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "e1efe377",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "HugeCTR Version: 3.2\n",
      "====================================================Model Init=====================================================\n",
      "[HUGECTR][07:31:36][INFO][RANK0]: Global seed is 3369591795\n",
      "[HUGECTR][07:31:36][INFO][RANK0]: Device to NUMA mapping:\n",
      "  GPU 0 ->  node 0\n",
      "\n",
      "[HUGECTR][07:31:38][WARNING][RANK0]: Peer-to-peer access cannot be fully enabled.\n",
      "[HUGECTR][07:31:38][INFO][RANK0]: Start all2all warmup\n",
      "[HUGECTR][07:31:38][INFO][RANK0]: End all2all warmup\n",
      "[HUGECTR][07:31:38][INFO][RANK0]: Using All-reduce algorithm: NCCL\n",
      "[HUGECTR][07:31:38][INFO][RANK0]: Device 0: Tesla V100-SXM2-16GB\n",
      "[HUGECTR][07:31:38][INFO][RANK0]: num of DataReader workers: 12\n",
      "[HUGECTR][07:31:38][INFO][RANK0]: max_num_frequent_categories is not specified using default: 1\n",
      "[HUGECTR][07:31:38][INFO][RANK0]: max_num_infrequent_samples is not specified using default: -1\n",
      "[HUGECTR][07:31:38][INFO][RANK0]: p_dup_max is not specified using default: 0.010000\n",
      "[HUGECTR][07:31:38][INFO][RANK0]: max_all_reduce_bandwidth is not specified using default: 130000000000.000000\n",
      "[HUGECTR][07:31:38][INFO][RANK0]: max_all_to_all_bandwidth is not specified using default: 190000000000.000000\n",
      "[HUGECTR][07:31:38][INFO][RANK0]: efficiency_bandwidth_ratio is not specified using default: 1.000000\n",
      "[HUGECTR][07:31:38][INFO][RANK0]: communication_type is not specified using default: IB_NVLink\n",
      "[HUGECTR][07:31:38][INFO][RANK0]: hybrid_embedding_type is not specified using default: Distributed\n",
      "[HUGECTR][07:31:38][INFO][RANK0]: max_vocabulary_size_per_gpu_=6029312\n",
      "[HUGECTR][07:31:38][INFO][RANK0]: max_num_frequent_categories is not specified using default: 1\n",
      "[HUGECTR][07:31:38][INFO][RANK0]: max_num_infrequent_samples is not specified using default: -1\n",
      "[HUGECTR][07:31:38][INFO][RANK0]: p_dup_max is not specified using default: 0.010000\n",
      "[HUGECTR][07:31:38][INFO][RANK0]: max_all_reduce_bandwidth is not specified using default: 130000000000.000000\n",
      "[HUGECTR][07:31:38][INFO][RANK0]: max_all_to_all_bandwidth is not specified using default: 190000000000.000000\n",
      "[HUGECTR][07:31:38][INFO][RANK0]: efficiency_bandwidth_ratio is not specified using default: 1.000000\n",
      "[HUGECTR][07:31:38][INFO][RANK0]: communication_type is not specified using default: IB_NVLink\n",
      "[HUGECTR][07:31:38][INFO][RANK0]: hybrid_embedding_type is not specified using default: Distributed\n",
      "[HUGECTR][07:31:38][INFO][RANK0]: max_vocabulary_size_per_gpu_=5865472\n",
      "[HUGECTR][07:31:38][INFO][RANK0]: Load the model graph from wdl.json successfully\n",
      "[HUGECTR][07:31:38][INFO][RANK0]: Graph analysis to resolve tensor dependency\n",
      "===================================================Model Compile===================================================\n",
      "[HUGECTR][07:31:42][INFO][RANK0]: gpu0 start to init embedding\n",
      "[HUGECTR][07:31:42][INFO][RANK0]: gpu0 init embedding done\n",
      "[HUGECTR][07:31:42][INFO][RANK0]: gpu0 start to init embedding\n",
      "[HUGECTR][07:31:42][INFO][RANK0]: gpu0 init embedding done\n",
      "[HUGECTR][07:31:42][INFO][RANK0]: Starting AUC NCCL warm-up\n",
      "[HUGECTR][07:31:42][INFO][RANK0]: Warm-up done\n",
      "===================================================Model Summary===================================================\n",
      "label                                   Dense                         Sparse                        \n",
      "label                                   dense                          wide_data,deep_data           \n",
      "(None, 1)                               (None, 13)                              \n",
      "——————————————————————————————————————————————————————————————————————————————————————————————————————————————————\n",
      "Layer Type                              Input Name                    Output Name                   Output Shape                  \n",
      "——————————————————————————————————————————————————————————————————————————————————————————————————————————————————\n",
      "DistributedSlotSparseEmbeddingHash      wide_data                     sparse_embedding2             (None, 1, 1)                  \n",
      "------------------------------------------------------------------------------------------------------------------\n",
      "DistributedSlotSparseEmbeddingHash      deep_data                     sparse_embedding1             (None, 26, 16)                \n",
      "------------------------------------------------------------------------------------------------------------------\n",
      "Reshape                                 sparse_embedding1             reshape1                      (None, 416)                   \n",
      "------------------------------------------------------------------------------------------------------------------\n",
      "Reshape                                 sparse_embedding2             reshape2                      (None, 1)                     \n",
      "------------------------------------------------------------------------------------------------------------------\n",
      "Concat                                  reshape1                      concat1                       (None, 429)                   \n",
      "                                        dense                                                                                     \n",
      "------------------------------------------------------------------------------------------------------------------\n",
      "InnerProduct                            concat1                       fc1                           (None, 1024)                  \n",
      "------------------------------------------------------------------------------------------------------------------\n",
      "ReLU                                    fc1                           relu1                         (None, 1024)                  \n",
      "------------------------------------------------------------------------------------------------------------------\n",
      "Dropout                                 relu1                         dropout1                      (None, 1024)                  \n",
      "------------------------------------------------------------------------------------------------------------------\n",
      "InnerProduct                            dropout1                      fc2                           (None, 1024)                  \n",
      "------------------------------------------------------------------------------------------------------------------\n",
      "ReLU                                    fc2                           relu2                         (None, 1024)                  \n",
      "------------------------------------------------------------------------------------------------------------------\n",
      "Dropout                                 relu2                         dropout2                      (None, 1024)                  \n",
      "------------------------------------------------------------------------------------------------------------------\n",
      "InnerProduct                            dropout2                      fc3                           (None, 1)                     \n",
      "------------------------------------------------------------------------------------------------------------------\n",
      "Add                                     fc3                           add1                          (None, 1)                     \n",
      "                                        reshape2                                                                                  \n",
      "------------------------------------------------------------------------------------------------------------------\n",
      "BinaryCrossEntropyLoss                  add1                          loss                                                        \n",
      "                                        label                                                                                     \n",
      "------------------------------------------------------------------------------------------------------------------\n",
      "Successfully convert pretrained embeddings to wdl1_pretrained.model\n",
      "[HUGECTR][07:31:42][INFO][RANK0]: Loading sparse model: wdl1_pretrained.model\n",
      "=====================================================Model Fit=====================================================\n",
      "[HUGECTR][07:31:42][INFO][RANK0]: Use epoch mode with number of epochs: 1\n",
      "[HUGECTR][07:31:42][INFO][RANK0]: Training batchsize: 1024, evaluation batchsize: 1024\n",
      "[HUGECTR][07:31:42][INFO][RANK0]: Evaluation interval: 1000, snapshot interval: 100000\n",
      "[HUGECTR][07:31:42][INFO][RANK0]: Dense network trainable: True\n",
      "[HUGECTR][07:31:42][INFO][RANK0]: Sparse embedding sparse_embedding1 trainable: False\n",
      "[HUGECTR][07:31:42][INFO][RANK0]: Sparse embedding sparse_embedding2 trainable: True\n",
      "[HUGECTR][07:31:42][INFO][RANK0]: Use mixed precision: False, scaler: 1.000000, use cuda graph: True\n",
      "[HUGECTR][07:31:42][INFO][RANK0]: lr: 0.001000, warmup_steps: 1, end_lr: 0.000000\n",
      "[HUGECTR][07:31:42][INFO][RANK0]: decay_start: 0, decay_steps: 1, decay_power: 2.000000\n",
      "[HUGECTR][07:31:42][INFO][RANK0]: Training source file: wdl_data/file_list.0.txt\n",
      "[HUGECTR][07:31:42][INFO][RANK0]: Evaluation source file: wdl_data/file_list.1.txt\n",
      "[HUGECTR][07:31:42][INFO][RANK0]: -----------------------------------Epoch 0-----------------------------------\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[HUGECTR][07:31:45][INFO][RANK0]: Iter: 500 Time(500 iters): 2.578954s Loss: 0.144762 lr:0.001000\n",
      "[HUGECTR][07:31:47][INFO][RANK0]: Iter: 1000 Time(500 iters): 2.417656s Loss: 0.136326 lr:0.001000\n",
      "[HUGECTR][07:31:49][INFO][RANK0]: Evaluation, AUC: 0.713149\n",
      "[HUGECTR][07:31:49][INFO][RANK0]: Eval Time for 5000 iters: 1.900026s\n",
      "[HUGECTR][07:31:52][INFO][RANK0]: Iter: 1500 Time(500 iters): 4.323596s Loss: 0.148682 lr:0.001000\n",
      "[HUGECTR][07:31:54][INFO][RANK0]: Iter: 2000 Time(500 iters): 2.416977s Loss: 0.145738 lr:0.001000\n",
      "[HUGECTR][07:31:56][INFO][RANK0]: Evaluation, AUC: 0.725260\n",
      "[HUGECTR][07:31:56][INFO][RANK0]: Eval Time for 5000 iters: 1.876535s\n",
      "[HUGECTR][07:31:58][INFO][RANK0]: Iter: 2500 Time(500 iters): 4.297001s Loss: 0.168649 lr:0.001000\n",
      "[HUGECTR][07:32:01][INFO][RANK0]: Iter: 3000 Time(500 iters): 2.418015s Loss: 0.134682 lr:0.001000\n",
      "[HUGECTR][07:32:03][INFO][RANK0]: Evaluation, AUC: 0.732183\n",
      "[HUGECTR][07:32:03][INFO][RANK0]: Eval Time for 5000 iters: 1.877291s\n",
      "[HUGECTR][07:32:05][INFO][RANK0]: Iter: 3500 Time(500 iters): 4.296748s Loss: 0.117909 lr:0.001000\n",
      "[HUGECTR][07:32:08][INFO][RANK0]: Iter: 4000 Time(500 iters): 2.411790s Loss: 0.133109 lr:0.001000\n",
      "[HUGECTR][07:32:09][INFO][RANK0]: Evaluation, AUC: 0.736392\n",
      "[HUGECTR][07:32:09][INFO][RANK0]: Eval Time for 5000 iters: 1.901939s\n",
      "[HUGECTR][07:32:09][INFO][RANK0]: Finish 1 epochs 4001 global iterations with batchsize 1024 in 27.09s.\n"
     ]
    }
   ],
   "source": [
    "!python3 wdl_load_pretrained.py"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c1378815",
   "metadata": {},
   "source": [
    "### Low-level Training\n",
    "\n",
    "The low-level training APIs are maintained in the enhanced HugeCTR Python interface. If you want to have precise control of each training iteration and each evaluation step, you may find it helpful to use these APIs. Since the data reader behavior is different in epoch mode and non-epoch mode, we should pay attention to how to tweak the data reader when using low-level training.\n",
    "We will demonstrate how to write the low-level training scripts for non-epoch mode, epoch mode, and embedding training cache mode."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "ecc97b71",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Writing wdl_non_epoch.py\n"
     ]
    }
   ],
   "source": [
    "%%writefile wdl_non_epoch.py\n",
    "import hugectr\n",
    "from mpi4py import MPI\n",
    "solver = hugectr.CreateSolver(max_eval_batches = 5000,\n",
    "                              batchsize_eval = 1024,\n",
    "                              batchsize = 1024,\n",
    "                              vvgpu = [[0]],\n",
    "                              i64_input_key = False,\n",
    "                              use_mixed_precision = False,\n",
    "                              repeat_dataset = True,\n",
    "                              use_cuda_graph = True)\n",
    "reader = hugectr.DataReaderParams(data_reader_type = hugectr.DataReaderType_t.Norm,\n",
    "                          source = [\"wdl_data/file_list.0.txt\"],\n",
    "                          eval_source = \"wdl_data/file_list.1.txt\",\n",
    "                          check_type = hugectr.Check_t.Sum)\n",
    "optimizer = hugectr.CreateOptimizer(optimizer_type = hugectr.Optimizer_t.Adam)\n",
    "model = hugectr.Model(solver, reader, optimizer)\n",
    "model.construct_from_json(graph_config_file = \"wdl.json\", include_dense_network = True)\n",
    "model.compile()\n",
    "model.start_data_reading()\n",
    "lr_sch = model.get_learning_rate_scheduler()\n",
    "max_iter = 2000\n",
    "for i in range(max_iter):\n",
    "    lr = lr_sch.get_next()\n",
    "    model.set_learning_rate(lr)\n",
    "    model.train()\n",
    "    if (i%100 == 0):\n",
    "        loss = model.get_current_loss()\n",
    "        print(\"[HUGECTR][INFO] iter: {}; loss: {}\".format(i, loss))\n",
    "    if (i%1000 == 0 and i != 0):\n",
    "        for _ in range(solver.max_eval_batches):\n",
    "            model.eval()\n",
    "        metrics = model.get_eval_metrics()\n",
    "        print(\"[HUGECTR][INFO] iter: {}, {}\".format(i, metrics))\n",
    "model.save_params_to_files(\"./\", max_iter)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "9ca7d711",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "HugeCTR Version: 3.2.0\n",
      "====================================================Model Init=====================================================\n",
      "[28d09h36m32s][HUGECTR][INFO]: Global seed is 3898093135\n",
      "[28d09h36m33s][HUGECTR][INFO]: Device to NUMA mapping:\n",
      "  GPU 0 ->  node 0\n",
      "\n",
      "[28d09h36m35s][HUGECTR][INFO]: Peer-to-peer access cannot be fully enabled.\n",
      "[28d09h36m35s][HUGECTR][INFO]: Start all2all warmup\n",
      "[28d09h36m35s][HUGECTR][INFO]: End all2all warmup\n",
      "[28d09h36m35s][HUGECTR][INFO]: Using All-reduce algorithm NCCL\n",
      "Device 0: Tesla V100-SXM2-16GB\n",
      "[28d09h36m35s][HUGECTR][INFO]: num of DataReader workers: 12\n",
      "[28d09h36m35s][HUGECTR][INFO]: max_num_frequent_categories is not specified using default: 1\n",
      "[28d09h36m35s][HUGECTR][INFO]: max_num_infrequent_samples is not specified using default: -1\n",
      "[28d09h36m35s][HUGECTR][INFO]: p_dup_max is not specified using default: 0.010000\n",
      "[28d09h36m35s][HUGECTR][INFO]: max_all_reduce_bandwidth is not specified using default: 130000000000.000000\n",
      "[28d09h36m35s][HUGECTR][INFO]: max_all_to_all_bandwidth is not specified using default: 190000000000.000000\n",
      "[28d09h36m35s][HUGECTR][INFO]: efficiency_bandwidth_ratio is not specified using default: 1.000000\n",
      "[28d09h36m35s][HUGECTR][INFO]: communication_type is not specified using default: IB_NVLink\n",
      "[28d09h36m35s][HUGECTR][INFO]: hybrid_embedding_type is not specified using default: Distributed\n",
      "[28d09h36m35s][HUGECTR][INFO]: max_vocabulary_size_per_gpu_=6029312\n",
      "[28d09h36m35s][HUGECTR][INFO]: max_num_frequent_categories is not specified using default: 1\n",
      "[28d09h36m35s][HUGECTR][INFO]: max_num_infrequent_samples is not specified using default: -1\n",
      "[28d09h36m35s][HUGECTR][INFO]: p_dup_max is not specified using default: 0.010000\n",
      "[28d09h36m35s][HUGECTR][INFO]: max_all_reduce_bandwidth is not specified using default: 130000000000.000000\n",
      "[28d09h36m35s][HUGECTR][INFO]: max_all_to_all_bandwidth is not specified using default: 190000000000.000000\n",
      "[28d09h36m35s][HUGECTR][INFO]: efficiency_bandwidth_ratio is not specified using default: 1.000000\n",
      "[28d09h36m35s][HUGECTR][INFO]: communication_type is not specified using default: IB_NVLink\n",
      "[28d09h36m35s][HUGECTR][INFO]: hybrid_embedding_type is not specified using default: Distributed\n",
      "[28d09h36m35s][HUGECTR][INFO]: max_vocabulary_size_per_gpu_=5865472\n",
      "[28d09h36m35s][HUGECTR][INFO]: Load the model graph from wdl.json, successful\n",
      "===================================================Model Compile===================================================\n",
      "[28d09h36m38s][HUGECTR][INFO]: gpu0 start to init embedding\n",
      "[28d09h36m38s][HUGECTR][INFO]: gpu0 init embedding done\n",
      "[28d09h36m38s][HUGECTR][INFO]: gpu0 start to init embedding\n",
      "[28d09h36m38s][HUGECTR][INFO]: gpu0 init embedding done\n",
      "[28d09h36m38s][HUGECTR][INFO]: Starting AUC NCCL warm-up\n",
      "[28d09h36m38s][HUGECTR][INFO]: Warm-up done\n",
      "[HUGECTR][INFO] iter: 0; loss: 1.0029206275939941\n",
      "[HUGECTR][INFO] iter: 100; loss: 0.12538853287696838\n",
      "[HUGECTR][INFO] iter: 200; loss: 0.10476257652044296\n",
      "[HUGECTR][INFO] iter: 300; loss: 0.1463421732187271\n",
      "[HUGECTR][INFO] iter: 400; loss: 0.1541304737329483\n",
      "[HUGECTR][INFO] iter: 500; loss: 0.14912495017051697\n",
      "[HUGECTR][INFO] iter: 600; loss: 0.12571805715560913\n",
      "[HUGECTR][INFO] iter: 700; loss: 0.13279415667057037\n",
      "[HUGECTR][INFO] iter: 800; loss: 0.13649113476276398\n",
      "[HUGECTR][INFO] iter: 900; loss: 0.1288434863090515\n",
      "[HUGECTR][INFO] iter: 1000; loss: 0.13555476069450378\n",
      "[HUGECTR][INFO] iter: 1000, [('AUC', 0.7378068566322327)]\n",
      "[HUGECTR][INFO] iter: 1100; loss: 0.15259310603141785\n",
      "[HUGECTR][INFO] iter: 1200; loss: 0.15795981884002686\n",
      "[HUGECTR][INFO] iter: 1300; loss: 0.13971731066703796\n",
      "[HUGECTR][INFO] iter: 1400; loss: 0.14082138240337372\n",
      "[HUGECTR][INFO] iter: 1500; loss: 0.14722011983394623\n",
      "[HUGECTR][INFO] iter: 1600; loss: 0.13573814928531647\n",
      "[HUGECTR][INFO] iter: 1700; loss: 0.12339376658201218\n",
      "[HUGECTR][INFO] iter: 1800; loss: 0.15557655692100525\n",
      "[HUGECTR][INFO] iter: 1900; loss: 0.1399267613887787\n",
      "[28d09h36m55s][HUGECTR][INFO]: Rank0: Write hash table to file\n",
      "[28d09h36m56s][HUGECTR][INFO]: Rank0: Write hash table to file\n",
      "[28d09h36m56s][HUGECTR][INFO]: Dumping sparse weights to files, successful\n",
      "[28d09h36m56s][HUGECTR][INFO]: Rank0: Write optimzer state to file\n",
      "[28d09h36m56s][HUGECTR][INFO]: Done\n",
      "[28d09h36m56s][HUGECTR][INFO]: Rank0: Write optimzer state to file\n",
      "[28d09h36m56s][HUGECTR][INFO]: Done\n",
      "[28d09h36m57s][HUGECTR][INFO]: Rank0: Write optimzer state to file\n",
      "[28d09h36m57s][HUGECTR][INFO]: Done\n",
      "[28d09h36m58s][HUGECTR][INFO]: Rank0: Write optimzer state to file\n",
      "[28d09h36m58s][HUGECTR][INFO]: Done\n",
      "[28d09h37m05s][HUGECTR][INFO]: Dumping sparse optimzer states to files, successful\n",
      "[28d09h37m05s][HUGECTR][INFO]: Dumping dense weights to file, successful\n",
      "[28d09h37m05s][HUGECTR][INFO]: Dumping dense optimizer states to file, successful\n",
      "[28d09h37m05s][HUGECTR][INFO]: Dumping untrainable weights to file, successful\n"
     ]
    }
   ],
   "source": [
    "!python3 wdl_non_epoch.py"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "c218a76b",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Writing wdl_epoch.py\n"
     ]
    }
   ],
   "source": [
    "%%writefile wdl_epoch.py\n",
    "import hugectr\n",
    "from mpi4py import MPI\n",
    "solver = hugectr.CreateSolver(max_eval_batches = 5000,\n",
    "                              batchsize_eval = 1024,\n",
    "                              batchsize = 1024,\n",
    "                              vvgpu = [[0]],\n",
    "                              i64_input_key = False,\n",
    "                              use_mixed_precision = False,\n",
    "                              repeat_dataset = False,\n",
    "                              use_cuda_graph = True)\n",
    "reader = hugectr.DataReaderParams(data_reader_type = hugectr.DataReaderType_t.Norm,\n",
    "                          source = [\"wdl_data/file_list.0.txt\"],\n",
    "                          eval_source = \"wdl_data/file_list.1.txt\",\n",
    "                          check_type = hugectr.Check_t.Sum)\n",
    "optimizer = hugectr.CreateOptimizer(optimizer_type = hugectr.Optimizer_t.Adam)\n",
    "model = hugectr.Model(solver, reader, optimizer)\n",
    "model.construct_from_json(graph_config_file = \"wdl.json\", include_dense_network = True)\n",
    "model.compile()\n",
    "lr_sch = model.get_learning_rate_scheduler()\n",
    "data_reader_train = model.get_data_reader_train()\n",
    "data_reader_eval = model.get_data_reader_eval()\n",
    "data_reader_eval.set_source()\n",
    "data_reader_eval_flag = True\n",
    "iteration = 0\n",
    "for epoch in range(2):\n",
    "  print(\"[HUGECTR][INFO] epoch: \", epoch)\n",
    "  data_reader_train.set_source()\n",
    "  data_reader_train_flag = True\n",
    "  while True:\n",
    "    lr = lr_sch.get_next()\n",
    "    model.set_learning_rate(lr)\n",
    "    data_reader_train_flag = model.train()\n",
    "    if not data_reader_train_flag:\n",
    "      break\n",
    "    if iteration % 1000 == 0:\n",
    "      batches = 0\n",
    "      while data_reader_eval_flag:\n",
    "        if batches >= solver.max_eval_batches:\n",
    "          break\n",
    "        data_reader_eval_flag = model.eval()\n",
    "        batches += 1\n",
    "      if not data_reader_eval_flag:\n",
    "        data_reader_eval.set_source()\n",
    "        data_reader_eval_flag = True\n",
    "      metrics = model.get_eval_metrics()\n",
    "      print(\"[HUGECTR][INFO] iter: {}, metrics: {}\".format(iteration, metrics))\n",
    "    iteration += 1\n",
    "model.save_params_to_files(\"./\", iteration)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "746b9fa0",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "HugeCTR Version: 3.2.0\n",
      "====================================================Model Init=====================================================\n",
      "[28d09h37m13s][HUGECTR][INFO]: Global seed is 2582215374\n",
      "[28d09h37m14s][HUGECTR][INFO]: Device to NUMA mapping:\n",
      "  GPU 0 ->  node 0\n",
      "\n",
      "[28d09h37m16s][HUGECTR][INFO]: Peer-to-peer access cannot be fully enabled.\n",
      "[28d09h37m16s][HUGECTR][INFO]: Start all2all warmup\n",
      "[28d09h37m16s][HUGECTR][INFO]: End all2all warmup\n",
      "[28d09h37m16s][HUGECTR][INFO]: Using All-reduce algorithm NCCL\n",
      "Device 0: Tesla V100-SXM2-16GB\n",
      "[28d09h37m16s][HUGECTR][INFO]: num of DataReader workers: 12\n",
      "[28d09h37m16s][HUGECTR][INFO]: max_num_frequent_categories is not specified using default: 1\n",
      "[28d09h37m16s][HUGECTR][INFO]: max_num_infrequent_samples is not specified using default: -1\n",
      "[28d09h37m16s][HUGECTR][INFO]: p_dup_max is not specified using default: 0.010000\n",
      "[28d09h37m16s][HUGECTR][INFO]: max_all_reduce_bandwidth is not specified using default: 130000000000.000000\n",
      "[28d09h37m16s][HUGECTR][INFO]: max_all_to_all_bandwidth is not specified using default: 190000000000.000000\n",
      "[28d09h37m16s][HUGECTR][INFO]: efficiency_bandwidth_ratio is not specified using default: 1.000000\n",
      "[28d09h37m16s][HUGECTR][INFO]: communication_type is not specified using default: IB_NVLink\n",
      "[28d09h37m16s][HUGECTR][INFO]: hybrid_embedding_type is not specified using default: Distributed\n",
      "[28d09h37m16s][HUGECTR][INFO]: max_vocabulary_size_per_gpu_=6029312\n",
      "[28d09h37m16s][HUGECTR][INFO]: max_num_frequent_categories is not specified using default: 1\n",
      "[28d09h37m16s][HUGECTR][INFO]: max_num_infrequent_samples is not specified using default: -1\n",
      "[28d09h37m16s][HUGECTR][INFO]: p_dup_max is not specified using default: 0.010000\n",
      "[28d09h37m16s][HUGECTR][INFO]: max_all_reduce_bandwidth is not specified using default: 130000000000.000000\n",
      "[28d09h37m16s][HUGECTR][INFO]: max_all_to_all_bandwidth is not specified using default: 190000000000.000000\n",
      "[28d09h37m16s][HUGECTR][INFO]: efficiency_bandwidth_ratio is not specified using default: 1.000000\n",
      "[28d09h37m16s][HUGECTR][INFO]: communication_type is not specified using default: IB_NVLink\n",
      "[28d09h37m16s][HUGECTR][INFO]: hybrid_embedding_type is not specified using default: Distributed\n",
      "[28d09h37m16s][HUGECTR][INFO]: max_vocabulary_size_per_gpu_=5865472\n",
      "[28d09h37m16s][HUGECTR][INFO]: Load the model graph from wdl.json, successful\n",
      "===================================================Model Compile===================================================\n",
      "[28d09h37m19s][HUGECTR][INFO]: gpu0 start to init embedding\n",
      "[28d09h37m19s][HUGECTR][INFO]: gpu0 init embedding done\n",
      "[28d09h37m19s][HUGECTR][INFO]: gpu0 start to init embedding\n",
      "[28d09h37m19s][HUGECTR][INFO]: gpu0 init embedding done\n",
      "[28d09h37m19s][HUGECTR][INFO]: Starting AUC NCCL warm-up\n",
      "[28d09h37m19s][HUGECTR][INFO]: Warm-up done\n",
      "[HUGECTR][INFO] epoch:  0\n",
      "[HUGECTR][INFO] iter: 0, metrics: [('AUC', 0.5298923850059509)]\n",
      "[HUGECTR][INFO] iter: 1000, metrics: [('AUC', 0.7407246828079224)]\n",
      "[HUGECTR][INFO] iter: 2000, metrics: [('AUC', 0.7498546242713928)]\n",
      "[HUGECTR][INFO] iter: 3000, metrics: [('AUC', 0.7546358704566956)]\n",
      "[HUGECTR][INFO] epoch:  1\n",
      "[HUGECTR][INFO] iter: 4000, metrics: [('AUC', 0.7573446035385132)]\n",
      "[HUGECTR][INFO] iter: 5000, metrics: [('AUC', 0.7208017706871033)]\n",
      "[HUGECTR][INFO] iter: 6000, metrics: [('AUC', 0.7227433323860168)]\n",
      "[HUGECTR][INFO] iter: 7000, metrics: [('AUC', 0.7216600775718689)]\n",
      "[28d09h38m42s][HUGECTR][INFO]: Rank0: Write hash table to file\n",
      "[28d09h38m42s][HUGECTR][INFO]: Rank0: Write hash table to file\n",
      "[28d09h38m43s][HUGECTR][INFO]: Dumping sparse weights to files, successful\n",
      "[28d09h38m43s][HUGECTR][INFO]: Rank0: Write optimzer state to file\n",
      "[28d09h38m43s][HUGECTR][INFO]: Done\n",
      "[28d09h38m43s][HUGECTR][INFO]: Rank0: Write optimzer state to file\n",
      "[28d09h38m43s][HUGECTR][INFO]: Done\n",
      "[28d09h38m44s][HUGECTR][INFO]: Rank0: Write optimzer state to file\n",
      "[28d09h38m44s][HUGECTR][INFO]: Done\n",
      "[28d09h38m44s][HUGECTR][INFO]: Rank0: Write optimzer state to file\n",
      "[28d09h38m45s][HUGECTR][INFO]: Done\n",
      "[28d09h38m52s][HUGECTR][INFO]: Dumping sparse optimzer states to files, successful\n",
      "[28d09h38m52s][HUGECTR][INFO]: Dumping dense weights to file, successful\n",
      "[28d09h38m52s][HUGECTR][INFO]: Dumping dense optimizer states to file, successful\n",
      "[28d09h38m52s][HUGECTR][INFO]: Dumping untrainable weights to file, successful\n"
     ]
    }
   ],
   "source": [
    "!python3 wdl_epoch.py"
   ]
  }
 ],
 "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
}