notebooks/continuous_training.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<img src=\"http://developer.download.nvidia.com/compute/machine-learning/frameworks/nvidia_logo.png\" style=\"width: 90px; float: right;\">\n",
    "\n",
    "# HugeCTR Continuous Training"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Overview\n",
    "The notebook introduces how to use the Embedding Training Cache (ETC) feature in HugeCTR for the continuous training. The ETC feature is designed to handle recommendation models with huge embedding table by the incremental training method, which allows you to train such a model that the model size is much larger than the available GPU memory size.\n",
    "\n",
    "To learn more about the ETC, please check the [Embedding Training Cache](../docs/hugectr_embedding_training_cache.md).\n",
    "\n",
    "To learn how to use the APIs of ETC, please check the [HugeCTR Python Interface](../docs/python_interface.md).\n",
    "\n",
    "## Table of Contents\n",
    "-  [Installation](#1)\n",
    "   * [Get HugeCTR from NGC](#11)\n",
    "   * [Build HugeCTR from Source Code](#12)\n",
    "-  [Continuous Training](#2)\n",
    "   * [Data Preparation](#21)\n",
    "   * [Continuous Training with High-level API](#22)\n",
    "   * [Continuous Training with Low-level API](#23)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id=\"1\"></a>\n",
    "## Installation\n",
    "\n",
    "### 1.1 Get HugeCTR from NGC\n",
    "The continuous training module is preinstalled in the [Merlin Training Container](https://ngc.nvidia.com/catalog/containers/nvidia:merlin:merlin-training): `nvcr.io/nvidia/merlin/merlin-training:22.03`.\n",
    "\n",
    "You can check the existence of required libraries by running the following Python code after launching this container.\n",
    "```bash\n",
    "$ python3 -c \"import hugectr\"\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1.2 Build HugeCTR from Source Code\n",
    "\n",
    "If you want to build HugeCTR from the source code instead of using the NGC container, please refer to the [How to Start Your Development](../../docs/hugectr_contributor_guide.md#how-to-start-your-development)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id=\"2\"></a>\n",
    "## Continuous Training\n",
    "\n",
    "<a id=\"21\"></a>\n",
    "### 2.1 Data Preparation\n",
    "1. Download the Kaggle Criteo dataset using the following command:\n",
    "   ```shell\n",
    "   $ cd ${project_root}/tools\n",
    "   $ wget http://azuremlsampleexperiments.blob.core.windows.net/criteo/day_1.gz\n",
    "   ```\n",
    "   \n",
    "   To preprocess the downloaded Kaggle Criteo dataset, we'll make the following operations: \n",
    "   * Reduce the amounts of data to speed up the preprocessing\n",
    "   * Fill missing values\n",
    "   * Remove the feature values whose occurrences are very rare, etc.\n",
    "\n",
    "2. Preprocessing by Pandas using the following command:\n",
    "   ```shell\n",
    "   $ bash preprocess.sh 1 wdl_data pandas 1 1 100\n",
    "   ```\n",
    "   \n",
    "   Meanings of the command line arguments:\n",
    "   * The 1st argument represents the dataset postfix. It is `1` here since `day_1` is used.\n",
    "   * The 2nd argument `wdl_data` is where the preprocessed data is stored.\n",
    "   * The 3rd argument `pandas` is the processing script going to use, here we choose `pandas`.\n",
    "   * The 4th argument `1` embodies that the normalization is applied to dense features.\n",
    "   * The 5th argument `1` means that the feature crossing is applied.\n",
    "   * The 6th argument `100` means the number of data files in each file list.\n",
    "\n",
    "   For more details about the data preprocessing, please refer to [Preprocess the Criteo Dataset](../samples/criteo#preprocess-the-dataset).\n",
    "   \n",
    "3. Create a soft link of the dataset folder to the path of this notebook using the following command:\n",
    "   ```shell\n",
    "   $ ln -s ${project_root}/tools/wdl_data ${project_root}/notebooks/wdl_data\n",
    "   ```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<a id=\"22\"></a>\n",
    "### 2.2 Continuous Training with High-level API\n",
    "This section gives the code sample of continuous training using Keras-like high-level API. The high-level API encapsulates etct of the complexity for users, making it easy to use and able to handle etct of scenarios in a production environment.\n",
    "\n",
    "Meanwhile, HugeCTR also provides the low-level APIs besides its high-level counterpart to allow you customize the training logic. A code sample using the low-level APIs is provided in the next section.\n",
    "\n",
    "The code sample in this section trains a model from scratch using embedding training cache, gets the incremental model, and saves the trained dense weights and sparse embedding weights. The following steps are required to achieve those logics:\n",
    "\n",
    "1. Create the `solver`, `reader`, `optimizer` and `etc`, 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 overview the model graph.\n",
    "4. Dump the model graph to the JSON file.\n",
    "5. Train the sparse and dense model.\n",
    "6. Set the new training datasets and their corresponding keysets.\n",
    "7. Train the sparse and dense model incrementally.\n",
    "8. Get the incrementally trained embedding table.\n",
    "9. Save the model weights and optimizer states explicitly.\n",
    "\n",
    "Note: `repeat_dataset` should be `False` when using the embedding training cache, while the argument `num_epochs` in `Model::fit` specifies the number of training epochs in this mode."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "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.\"+str(i)+\".txt\" for i in range(2)],\n",
    "                          keyset = [\"wdl_data/file_list.\"+str(i)+\".keyset\" for i in range(2)],\n",
    "                          eval_source = \"wdl_data/file_list.2.txt\",\n",
    "                          check_type = hugectr.Check_t.Sum)\n",
    "optimizer = hugectr.CreateOptimizer(optimizer_type = hugectr.Optimizer_t.Adam)\n",
    "hc_cnfg = hugectr.CreateHMemCache(num_blocks = 2, target_hit_rate = 0.5, max_num_evict = 0)\n",
    "etc = hugectr.CreateETC(ps_types = [hugectr.TrainPSType_t.Staged, hugectr.TrainPSType_t.Cached],\n",
    "                        sparse_models = [\"./wdl_0_sparse_model\", \"./wdl_1_sparse_model\"],\n",
    "                        local_paths = [\"./\"], hmem_cache_configs = [hc_cnfg])\n",
    "model = hugectr.Model(solver, reader, optimizer, etc)\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 = 1000)\n",
    "# Get the updated embedding features in model.fit()\n",
    "# updated_model = model.get_incremental_model()\n",
    "model.set_source(source = [\"wdl_data/file_list.3.txt\", \"wdl_data/file_list.4.txt\"], keyset = [\"wdl_data/file_list.3.keyset\", \"wdl_data/file_list.4.keyset\"], eval_source = \"wdl_data/file_list.5.txt\")\n",
    "model.fit(num_epochs = 1, display = 500, eval_interval = 1000)\n",
    "# Get the updated embedding features in model.fit()\n",
    "updated_model = model.get_incremental_model()\n",
    "model.save_params_to_files(\"wdl_etc\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[HUGECTR][12:36:58][INFO][RANK0]: Empty embedding, trained table will be stored in ./wdl_0_sparse_model\n",
      "[HUGECTR][12:36:58][INFO][RANK0]: Empty embedding, trained table will be stored in ./wdl_1_sparse_model\n",
      "HugeCTR Version: 3.2\n",
      "====================================================Model Init=====================================================\n",
      "[HUGECTR][12:36:58][INFO][RANK0]: Global seed is 3664540043\n",
      "[HUGECTR][12:36:58][INFO][RANK0]: Device to NUMA mapping:\n",
      "  GPU 0 ->  node 0\n",
      "\n",
      "[HUGECTR][12:36:59][WARNING][RANK0]: Peer-to-peer access cannot be fully enabled.\n",
      "[HUGECTR][12:36:59][INFO][RANK0]: Start all2all warmup\n",
      "[HUGECTR][12:36:59][INFO][RANK0]: End all2all warmup\n",
      "[HUGECTR][12:36:59][INFO][RANK0]: Using All-reduce algorithm: NCCL\n",
      "[HUGECTR][12:36:59][INFO][RANK0]: Device 0: Tesla V100-SXM2-32GB\n",
      "[HUGECTR][12:36:59][INFO][RANK0]: num of DataReader workers: 12\n",
      "[HUGECTR][12:36:59][INFO][RANK0]: max_vocabulary_size_per_gpu_=6029312\n",
      "[HUGECTR][12:36:59][INFO][RANK0]: max_vocabulary_size_per_gpu_=5865472\n",
      "[HUGECTR][12:36:59][INFO][RANK0]: Graph analysis to resolve tensor dependency\n",
      "===================================================Model Compile===================================================\n",
      "[HUGECTR][12:37:03][INFO][RANK0]: gpu0 start to init embedding\n",
      "[HUGECTR][12:37:03][INFO][RANK0]: gpu0 init embedding done\n",
      "[HUGECTR][12:37:03][INFO][RANK0]: gpu0 start to init embedding\n",
      "[HUGECTR][12:37:03][INFO][RANK0]: gpu0 init embedding done\n",
      "[HUGECTR][12:37:03][INFO][RANK0]: Enable HMEM-Based Parameter Server\n",
      "[HUGECTR][12:37:03][INFO][RANK0]: ./wdl_0_sparse_model not exist, create and train from scratch\n",
      "[HUGECTR][12:37:03][INFO][RANK0]: Enable HMemCache-Based Parameter Server\n",
      "[HUGECTR][12:37:03][INFO][RANK0]: ./wdl_1_sparse_model/key doesn't exist, created\n",
      "[HUGECTR][12:37:03][INFO][RANK0]: ./wdl_1_sparse_model/emb_vector doesn't exist, created\n",
      "[HUGECTR][12:37:03][INFO][RANK0]: ./wdl_1_sparse_model/Adam.m doesn't exist, created\n",
      "[HUGECTR][12:37:03][INFO][RANK0]: ./wdl_1_sparse_model/Adam.v doesn't exist, created\n",
      "[HUGECTR][12:37:04][INFO][RANK0]: Starting AUC NCCL warm-up\n",
      "[HUGECTR][12:37:04][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",
      "DistributedSlotSparseEmbeddingHash      deep_data                     sparse_embedding1             (None, 26, 16)                \n",
      "Reshape                                 sparse_embedding1             reshape1                      (None, 416)                   \n",
      "Reshape                                 sparse_embedding2             reshape2                      (None, 1)                     \n",
      "Concat                                  reshape1,dense                concat1                       (None, 429)                   \n",
      "InnerProduct                            concat1                       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",
      "InnerProduct                            dropout2                      fc3                           (None, 1)                     \n",
      "Add                                     fc3,reshape2                  add1                          (None, 1)                     \n",
      "BinaryCrossEntropyLoss                  add1,label                    loss                                                        \n",
      "------------------------------------------------------------------------------------------------------------------\n",
      "[HUGECTR][12:37:04][INFO][RANK0]: Save the model graph to wdl.json successfully\n",
      "=====================================================Model Fit=====================================================\n",
      "[HUGECTR][12:37:04][INFO][RANK0]: Use embedding training cache mode with number of training sources: 2, number of epochs: 1\n",
      "[HUGECTR][12:37:04][INFO][RANK0]: Training batchsize: 1024, evaluation batchsize: 1024\n",
      "[HUGECTR][12:37:04][INFO][RANK0]: Evaluation interval: 1000, snapshot interval: 10000\n",
      "[HUGECTR][12:37:04][INFO][RANK0]: Sparse embedding trainable: True, dense network trainable: True\n",
      "[HUGECTR][12:37:04][INFO][RANK0]: Use mixed precision: False, scaler: 1.000000, use cuda graph: True\n",
      "[HUGECTR][12:37:04][INFO][RANK0]: lr: 0.001000, warmup_steps: 1, decay_start: 0, decay_steps: 1, decay_power: 2.000000, end_lr: 0.000000\n",
      "[HUGECTR][12:37:04][INFO][RANK0]: Evaluation source file: wdl_data/file_list.2.txt\n",
      "[HUGECTR][12:37:04][INFO][RANK0]: --------------------Epoch 0, source file: wdl_data/file_list.0.txt--------------------\n",
      "[HUGECTR][12:37:04][INFO][RANK0]: Preparing embedding table for next pass\n",
      "[HUGECTR][12:37:05][INFO][RANK0]: HMEM-Cache PS: Hit rate [load]: 0 %\n",
      "[HUGECTR][12:37:07][INFO][RANK0]: Iter: 500 Time(500 iters): 2.959942s Loss: 0.140601 lr:0.001000\n",
      "[HUGECTR][12:37:10][INFO][RANK0]: Iter: 1000 Time(500 iters): 2.687422s Loss: 0.127723 lr:0.001000\n",
      "[HUGECTR][12:37:15][INFO][RANK0]: Evaluation, AUC: 0.738460\n",
      "[HUGECTR][12:37:15][INFO][RANK0]: Eval Time for 5000 iters: 4.757926s\n",
      "[HUGECTR][12:37:17][INFO][RANK0]: Iter: 1500 Time(500 iters): 7.310160s Loss: 0.152160 lr:0.001000\n",
      "[HUGECTR][12:37:20][INFO][RANK0]: Iter: 2000 Time(500 iters): 2.613197s Loss: 0.124371 lr:0.001000\n",
      "[HUGECTR][12:37:22][INFO][RANK0]: Evaluation, AUC: 0.745345\n",
      "[HUGECTR][12:37:22][INFO][RANK0]: Eval Time for 5000 iters: 1.907179s\n",
      "[HUGECTR][12:37:24][INFO][RANK0]: Iter: 2500 Time(500 iters): 4.343850s Loss: 0.134511 lr:0.001000\n",
      "[HUGECTR][12:37:27][INFO][RANK0]: Iter: 3000 Time(500 iters): 2.505121s Loss: 0.119222 lr:0.001000\n",
      "[HUGECTR][12:37:28][INFO][RANK0]: Evaluation, AUC: 0.751256\n",
      "[HUGECTR][12:37:28][INFO][RANK0]: Eval Time for 5000 iters: 1.900262s\n",
      "[HUGECTR][12:37:31][INFO][RANK0]: Iter: 3500 Time(500 iters): 4.459760s Loss: 0.145278 lr:0.001000\n",
      "[HUGECTR][12:37:34][INFO][RANK0]: Iter: 4000 Time(500 iters): 2.544999s Loss: 0.134373 lr:0.001000\n",
      "[HUGECTR][12:37:35][INFO][RANK0]: Evaluation, AUC: 0.753270\n",
      "[HUGECTR][12:37:35][INFO][RANK0]: Eval Time for 5000 iters: 1.901368s\n",
      "[HUGECTR][12:37:35][INFO][RANK0]: --------------------Epoch 0, source file: wdl_data/file_list.1.txt--------------------\n",
      "[HUGECTR][12:37:35][INFO][RANK0]: Preparing embedding table for next pass\n",
      "[HUGECTR][12:37:37][INFO][RANK0]: HMEM-Cache PS: Hit rate [dump]: 0 %\n",
      "[HUGECTR][12:37:37][INFO][RANK0]: HMEM-Cache PS: Hit rate [load]: 0 %\n",
      "[HUGECTR][12:37:40][INFO][RANK0]: Iter: 4500 Time(500 iters): 6.212693s Loss: 0.131819 lr:0.001000\n",
      "[HUGECTR][12:37:42][INFO][RANK0]: Iter: 5000 Time(500 iters): 2.660587s Loss: 0.117531 lr:0.001000\n",
      "[HUGECTR][12:37:44][INFO][RANK0]: Evaluation, AUC: 0.754530\n",
      "[HUGECTR][12:37:44][INFO][RANK0]: Eval Time for 5000 iters: 1.897969s\n",
      "[HUGECTR][12:37:47][INFO][RANK0]: Iter: 5500 Time(500 iters): 4.340803s Loss: 0.118400 lr:0.001000\n",
      "[HUGECTR][12:37:49][INFO][RANK0]: Iter: 6000 Time(500 iters): 2.497391s Loss: 0.143188 lr:0.001000\n",
      "[HUGECTR][12:37:51][INFO][RANK0]: Evaluation, AUC: 0.755805\n",
      "[HUGECTR][12:37:51][INFO][RANK0]: Eval Time for 5000 iters: 1.904572s\n",
      "[HUGECTR][12:37:54][INFO][RANK0]: Iter: 6500 Time(500 iters): 4.332877s Loss: 0.159262 lr:0.001000\n",
      "[HUGECTR][12:37:56][INFO][RANK0]: Iter: 7000 Time(500 iters): 2.426105s Loss: 0.119848 lr:0.001000\n",
      "[HUGECTR][12:37:58][INFO][RANK0]: Evaluation, AUC: 0.757338\n",
      "[HUGECTR][12:37:58][INFO][RANK0]: Eval Time for 5000 iters: 1.900609s\n",
      "[HUGECTR][12:38:00][INFO][RANK0]: Iter: 7500 Time(500 iters): 4.348594s Loss: 0.139543 lr:0.001000\n",
      "[HUGECTR][12:38:03][INFO][RANK0]: Iter: 8000 Time(500 iters): 2.424926s Loss: 0.109002 lr:0.001000\n",
      "[HUGECTR][12:38:05][INFO][RANK0]: Evaluation, AUC: 0.758067\n",
      "[HUGECTR][12:38:05][INFO][RANK0]: Eval Time for 5000 iters: 1.900712s\n",
      "=====================================================Model Fit=====================================================\n",
      "[HUGECTR][12:38:05][INFO][RANK0]: Use embedding training cache mode with number of training sources: 2, number of epochs: 1\n",
      "[HUGECTR][12:38:05][INFO][RANK0]: Training batchsize: 1024, evaluation batchsize: 1024\n",
      "[HUGECTR][12:38:05][INFO][RANK0]: Evaluation interval: 1000, snapshot interval: 10000\n",
      "[HUGECTR][12:38:05][INFO][RANK0]: Sparse embedding trainable: True, dense network trainable: True\n",
      "[HUGECTR][12:38:05][INFO][RANK0]: Use mixed precision: False, scaler: 1.000000, use cuda graph: True\n",
      "[HUGECTR][12:38: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][12:38:05][INFO][RANK0]: Evaluation source file: wdl_data/file_list.5.txt\n",
      "[HUGECTR][12:38:05][INFO][RANK0]: --------------------Epoch 0, source file: wdl_data/file_list.3.txt--------------------\n",
      "[HUGECTR][12:38:05][INFO][RANK0]: Preparing embedding table for next pass\n",
      "[HUGECTR][12:38:06][INFO][RANK0]: HMEM-Cache PS: Hit rate [dump]: 77.89 %\n",
      "[HUGECTR][12:38:06][INFO][RANK0]: HMEM-Cache PS: Hit rate [load]: 71.22 %\n",
      "[HUGECTR][12:38:08][INFO][RANK0]: Iter: 500 Time(500 iters): 3.652928s Loss: 0.124229 lr:0.001000\n",
      "[HUGECTR][12:38:11][INFO][RANK0]: Iter: 1000 Time(500 iters): 2.455519s Loss: 0.142507 lr:0.001000\n",
      "[HUGECTR][12:38:13][INFO][RANK0]: Evaluation, AUC: 0.757185\n",
      "[HUGECTR][12:38:13][INFO][RANK0]: Eval Time for 5000 iters: 1.909209s\n",
      "[HUGECTR][12:38:15][INFO][RANK0]: Iter: 1500 Time(500 iters): 4.353392s Loss: 0.123939 lr:0.001000\n",
      "[HUGECTR][12:38:18][INFO][RANK0]: Iter: 2000 Time(500 iters): 2.522630s Loss: 0.130625 lr:0.001000\n",
      "[HUGECTR][12:38:21][INFO][RANK0]: Evaluation, AUC: 0.757897\n",
      "[HUGECTR][12:38:21][INFO][RANK0]: Eval Time for 5000 iters: 3.763415s\n",
      "[HUGECTR][12:38:24][INFO][RANK0]: Iter: 2500 Time(500 iters): 6.238394s Loss: 0.138125 lr:0.001000\n",
      "[HUGECTR][12:38:26][INFO][RANK0]: Iter: 3000 Time(500 iters): 2.429449s Loss: 0.126391 lr:0.001000\n",
      "[HUGECTR][12:38:28][INFO][RANK0]: Evaluation, AUC: 0.757794\n",
      "[HUGECTR][12:38:28][INFO][RANK0]: Eval Time for 5000 iters: 1.902641s\n",
      "[HUGECTR][12:38:31][INFO][RANK0]: Iter: 3500 Time(500 iters): 4.398343s Loss: 0.123047 lr:0.001000\n",
      "[HUGECTR][12:38:33][INFO][RANK0]: Iter: 4000 Time(500 iters): 2.420357s Loss: 0.142649 lr:0.001000\n",
      "[HUGECTR][12:38:35][INFO][RANK0]: Evaluation, AUC: 0.760467\n",
      "[HUGECTR][12:38:35][INFO][RANK0]: Eval Time for 5000 iters: 1.899759s\n",
      "[HUGECTR][12:38:35][INFO][RANK0]: --------------------Epoch 0, source file: wdl_data/file_list.4.txt--------------------\n",
      "[HUGECTR][12:38:35][INFO][RANK0]: Preparing embedding table for next pass\n",
      "[HUGECTR][12:38:37][INFO][RANK0]: HMEM-Cache PS: Hit rate [dump]: 64.88 %\n",
      "[HUGECTR][12:38:37][INFO][RANK0]: HMEM-Cache PS: Hit rate [load]: 67.35 %\n",
      "[HUGECTR][12:38:40][INFO][RANK0]: Iter: 4500 Time(500 iters): 6.318908s Loss: 0.165160 lr:0.001000\n",
      "[HUGECTR][12:38:42][INFO][RANK0]: Iter: 5000 Time(500 iters): 2.423369s Loss: 0.112445 lr:0.001000\n",
      "[HUGECTR][12:38:44][INFO][RANK0]: Evaluation, AUC: 0.759795\n",
      "[HUGECTR][12:38:44][INFO][RANK0]: Eval Time for 5000 iters: 1.902252s\n",
      "[HUGECTR][12:38:46][INFO][RANK0]: Iter: 5500 Time(500 iters): 4.329618s Loss: 0.150855 lr:0.001000\n",
      "[HUGECTR][12:38:49][INFO][RANK0]: Iter: 6000 Time(500 iters): 2.422831s Loss: 0.121576 lr:0.001000\n",
      "[HUGECTR][12:38:51][INFO][RANK0]: Evaluation, AUC: 0.760036\n",
      "[HUGECTR][12:38:51][INFO][RANK0]: Eval Time for 5000 iters: 1.896330s\n",
      "[HUGECTR][12:38:53][INFO][RANK0]: Iter: 6500 Time(500 iters): 4.352440s Loss: 0.131191 lr:0.001000\n",
      "[HUGECTR][12:38:56][INFO][RANK0]: Iter: 7000 Time(500 iters): 2.426486s Loss: 0.130866 lr:0.001000\n",
      "[HUGECTR][12:38:57][INFO][RANK0]: Evaluation, AUC: 0.761125\n",
      "[HUGECTR][12:38:57][INFO][RANK0]: Eval Time for 5000 iters: 1.910397s\n",
      "[HUGECTR][12:39:00][INFO][RANK0]: Iter: 7500 Time(500 iters): 4.364026s Loss: 0.096611 lr:0.001000\n",
      "[HUGECTR][12:39:03][INFO][RANK0]: Iter: 8000 Time(500 iters): 2.664058s Loss: 0.142381 lr:0.001000\n",
      "[HUGECTR][12:39:05][INFO][RANK0]: Evaluation, AUC: 0.762636\n",
      "[HUGECTR][12:39:05][INFO][RANK0]: Eval Time for 5000 iters: 1.975668s\n",
      "[HUGECTR][12:39:06][INFO][RANK0]: HMEM-Cache PS: Hit rate [dump]: 64.82 %\n",
      "[HUGECTR][12:39:07][INFO][RANK0]: HMEM-Cache PS: Hit rate [load]: 57.86 %\n",
      "[HUGECTR][12:39:07][INFO][RANK0]: Get updated portion of embedding table [DONE}\n",
      "[HUGECTR][12:39:08][INFO][RANK0]: HMEM-Cache PS: Hit rate [dump]: 64.82 %\n",
      "[HUGECTR][12:39:08][INFO][RANK0]: Updating sparse model in SSD [DONE]\n",
      "[HUGECTR][12:39:10][INFO][RANK0]: Sync blocks from HMEM-Cache to SSD\n",
      " \u001b[38;2;89;255;89m ████████████████████████████████████████▏ \u001b[1m\u001b[31m100.0% \u001b[34m[   2/   2 | 66.7 Hz | 0s<0s]  \u001b[0m\u001b[32m\u001b[0mm\n",
      "[HUGECTR][12:39:10][INFO][RANK0]: Dumping dense weights to file, successful\n",
      "[HUGECTR][12:39:10][INFO][RANK0]: Dumping dense optimizer states to file, successful\n",
      "[HUGECTR][12:39:10][INFO][RANK0]: Dumping untrainable weights to file, successful\n"
     ]
    }
   ],
   "source": [
    "!python3 wdl_train.py"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.3 Continuous Training with Low-level API\n",
    "\n",
    "This section gives the code sample of continuous training using low-level API, which follows the same logics as the code sample in above section.\n",
    "\n",
    "Although the low-level APIs provide fine-grind control to the training logic, we encourage you to use the high-level API if it can satisfy your requirement since the naked data reader and embedding training cache logics are not straightforward and error prone.\n",
    "\n",
    "For more about the low-level API, please refer to [Low-level Training API](../docs/python_interface.md#low-level-training-api) and samples of [Low-level Training](./hugectr_criteo.ipynb)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Writing wdl_etc.py\n"
     ]
    }
   ],
   "source": [
    "%%writefile wdl_etc.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.\"+str(i)+\".txt\" for i in range(2)],\n",
    "                          keyset = [\"wdl_data/file_list.\"+str(i)+\".keyset\" for i in range(2)],\n",
    "                          eval_source = \"wdl_data/file_list.2.txt\",\n",
    "                          check_type = hugectr.Check_t.Sum)\n",
    "optimizer = hugectr.CreateOptimizer(optimizer_type = hugectr.Optimizer_t.Adam)\n",
    "hc_cnfg = hugectr.CreateHMemCache(num_blocks = 2, target_hit_rate = 0.5, max_num_evict = 0)\n",
    "etc = hugectr.CreateETC(ps_types = [hugectr.TrainPSType_t.Staged, hugectr.TrainPSType_t.Cached],\n",
    "                        sparse_models = [\"./wdl_0_sparse_model\", \"./wdl_1_sparse_model\"],\n",
    "                        local_paths = [\"./\"], hmem_cache_configs = [hc_cnfg])\n",
    "model = hugectr.Model(solver, reader, optimizer, etc)\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",
    "etc = model.get_embedding_training_cache()\n",
    "dataset = [(\"wdl_data/file_list.\"+str(i)+\".txt\", \"wdl_data/file_list.\"+str(i)+\".keyset\") for i in range(2)]\n",
    "data_reader_eval.set_source(\"wdl_data/file_list.2.txt\")\n",
    "data_reader_eval_flag = True\n",
    "iteration = 0\n",
    "for file_list, keyset_file in dataset:\n",
    "  data_reader_train.set_source(file_list)\n",
    "  data_reader_train_flag = True\n",
    "  etc.update(keyset_file)\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",
    "  print(\"[HUGECTR][INFO] trained with data in {}\".format(file_list))\n",
    "\n",
    "dataset = [(\"wdl_data/file_list.\"+str(i)+\".txt\", \"wdl_data/file_list.\"+str(i)+\".keyset\") for i in range(3, 5)]\n",
    "for file_list, keyset_file in dataset:\n",
    "  data_reader_train.set_source(file_list)\n",
    "  data_reader_train_flag = True\n",
    "  etc.update(keyset_file)\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",
    "  print(\"[HUGECTR][INFO] trained with data in {}\".format(file_list))\n",
    "incremental_model = model.get_incremental_model()\n",
    "model.save_params_to_files(\"wdl_etc\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[HUGECTR][12:39:44][INFO][RANK0]: Empty embedding, trained table will be stored in ./wdl_0_sparse_model\n",
      "[HUGECTR][12:39:44][INFO][RANK0]: Empty embedding, trained table will be stored in ./wdl_1_sparse_model\n",
      "HugeCTR Version: 3.2\n",
      "====================================================Model Init=====================================================\n",
      "[HUGECTR][12:39:44][INFO][RANK0]: Global seed is 3498697826\n",
      "[HUGECTR][12:39:44][INFO][RANK0]: Device to NUMA mapping:\n",
      "  GPU 0 ->  node 0\n",
      "\n",
      "[HUGECTR][12:39:45][WARNING][RANK0]: Peer-to-peer access cannot be fully enabled.\n",
      "[HUGECTR][12:39:45][INFO][RANK0]: Start all2all warmup\n",
      "[HUGECTR][12:39:45][INFO][RANK0]: End all2all warmup\n",
      "[HUGECTR][12:39:45][INFO][RANK0]: Using All-reduce algorithm: NCCL\n",
      "[HUGECTR][12:39:45][INFO][RANK0]: Device 0: Tesla V100-SXM2-32GB\n",
      "[HUGECTR][12:39:45][INFO][RANK0]: num of DataReader workers: 12\n",
      "[HUGECTR][12:39:45][INFO][RANK0]: max_num_frequent_categories is not specified using default: 1\n",
      "[HUGECTR][12:39:45][INFO][RANK0]: max_num_infrequent_samples is not specified using default: -1\n",
      "[HUGECTR][12:39:45][INFO][RANK0]: p_dup_max is not specified using default: 0.010000\n",
      "[HUGECTR][12:39:45][INFO][RANK0]: max_all_reduce_bandwidth is not specified using default: 130000000000.000000\n",
      "[HUGECTR][12:39:45][INFO][RANK0]: max_all_to_all_bandwidth is not specified using default: 190000000000.000000\n",
      "[HUGECTR][12:39:45][INFO][RANK0]: efficiency_bandwidth_ratio is not specified using default: 1.000000\n",
      "[HUGECTR][12:39:45][INFO][RANK0]: communication_type is not specified using default: IB_NVLink\n",
      "[HUGECTR][12:39:45][INFO][RANK0]: hybrid_embedding_type is not specified using default: Distributed\n",
      "[HUGECTR][12:39:45][INFO][RANK0]: max_vocabulary_size_per_gpu_=6029312\n",
      "[HUGECTR][12:39:45][INFO][RANK0]: max_num_frequent_categories is not specified using default: 1\n",
      "[HUGECTR][12:39:45][INFO][RANK0]: max_num_infrequent_samples is not specified using default: -1\n",
      "[HUGECTR][12:39:45][INFO][RANK0]: p_dup_max is not specified using default: 0.010000\n",
      "[HUGECTR][12:39:45][INFO][RANK0]: max_all_reduce_bandwidth is not specified using default: 130000000000.000000\n",
      "[HUGECTR][12:39:45][INFO][RANK0]: max_all_to_all_bandwidth is not specified using default: 190000000000.000000\n",
      "[HUGECTR][12:39:45][INFO][RANK0]: efficiency_bandwidth_ratio is not specified using default: 1.000000\n",
      "[HUGECTR][12:39:45][INFO][RANK0]: communication_type is not specified using default: IB_NVLink\n",
      "[HUGECTR][12:39:45][INFO][RANK0]: hybrid_embedding_type is not specified using default: Distributed\n",
      "[HUGECTR][12:39:45][INFO][RANK0]: max_vocabulary_size_per_gpu_=5865472\n",
      "[HUGECTR][12:39:45][INFO][RANK0]: Load the model graph from wdl.json successfully\n",
      "[HUGECTR][12:39:45][INFO][RANK0]: Graph analysis to resolve tensor dependency\n",
      "===================================================Model Compile===================================================\n",
      "[HUGECTR][12:39:49][INFO][RANK0]: gpu0 start to init embedding\n",
      "[HUGECTR][12:39:49][INFO][RANK0]: gpu0 init embedding done\n",
      "[HUGECTR][12:39:49][INFO][RANK0]: gpu0 start to init embedding\n",
      "[HUGECTR][12:39:49][INFO][RANK0]: gpu0 init embedding done\n",
      "[HUGECTR][12:39:49][INFO][RANK0]: Enable HMEM-Based Parameter Server\n",
      "[HUGECTR][12:39:49][INFO][RANK0]: ./wdl_0_sparse_model not exist, create and train from scratch\n",
      "[HUGECTR][12:39:49][INFO][RANK0]: Enable HMemCache-Based Parameter Server\n",
      "[HUGECTR][12:39:49][INFO][RANK0]: ./wdl_1_sparse_model/key doesn't exist, created\n",
      "[HUGECTR][12:39:49][INFO][RANK0]: ./wdl_1_sparse_model/emb_vector doesn't exist, created\n",
      "[HUGECTR][12:39:49][INFO][RANK0]: ./wdl_1_sparse_model/Adam.m doesn't exist, created\n",
      "[HUGECTR][12:39:49][INFO][RANK0]: ./wdl_1_sparse_model/Adam.v doesn't exist, created\n",
      "[HUGECTR][12:39:50][INFO][RANK0]: Starting AUC NCCL warm-up\n",
      "[HUGECTR][12:39:50][INFO][RANK0]: Warm-up done\n",
      "[HUGECTR][12:39:50][INFO][RANK0]: Preparing embedding table for next pass\n",
      "[HUGECTR][12:39:50][INFO][RANK0]: HMEM-Cache PS: Hit rate [load]: 0 %\n",
      "[HUGECTR][INFO] iter: 0, metrics: [('AUC', 0.4865134358406067)]\n",
      "[HUGECTR][INFO] iter: 1000, metrics: [('AUC', 0.7405899167060852)]\n",
      "[HUGECTR][INFO] iter: 2000, metrics: [('AUC', 0.7468112707138062)]\n",
      "[HUGECTR][INFO] iter: 3000, metrics: [('AUC', 0.7530832290649414)]\n",
      "[HUGECTR][INFO] trained with data in wdl_data/file_list.0.txt\n",
      "[HUGECTR][12:40:28][INFO][RANK0]: Preparing embedding table for next pass\n",
      "[HUGECTR][12:40:30][INFO][RANK0]: HMEM-Cache PS: Hit rate [dump]: 0 %\n",
      "[HUGECTR][12:40:30][INFO][RANK0]: HMEM-Cache PS: Hit rate [load]: 0 %\n",
      "[HUGECTR][INFO] iter: 4000, metrics: [('AUC', 0.7554274201393127)]\n",
      "[HUGECTR][INFO] iter: 5000, metrics: [('AUC', 0.7563489079475403)]\n",
      "[HUGECTR][INFO] iter: 6000, metrics: [('AUC', 0.7577884197235107)]\n",
      "[HUGECTR][INFO] iter: 7000, metrics: [('AUC', 0.7599539160728455)]\n",
      "[HUGECTR][INFO] trained with data in wdl_data/file_list.1.txt\n",
      "[HUGECTR][12:41:08][INFO][RANK0]: Preparing embedding table for next pass\n",
      "[HUGECTR][12:41:09][INFO][RANK0]: HMEM-Cache PS: Hit rate [dump]: 77.89 %\n",
      "[HUGECTR][12:41:10][INFO][RANK0]: HMEM-Cache PS: Hit rate [load]: 71.22 %\n",
      "[HUGECTR][INFO] iter: 8000, metrics: [('AUC', 0.7602559328079224)]\n",
      "[HUGECTR][INFO] iter: 9000, metrics: [('AUC', 0.7596363425254822)]\n",
      "[HUGECTR][INFO] iter: 10000, metrics: [('AUC', 0.7619153261184692)]\n",
      "[HUGECTR][INFO] iter: 11000, metrics: [('AUC', 0.7607191801071167)]\n",
      "[HUGECTR][INFO] trained with data in wdl_data/file_list.3.txt\n",
      "[HUGECTR][12:41:48][INFO][RANK0]: Preparing embedding table for next pass\n",
      "[HUGECTR][12:41:50][INFO][RANK0]: HMEM-Cache PS: Hit rate [dump]: 64.88 %\n",
      "[HUGECTR][12:41:50][INFO][RANK0]: HMEM-Cache PS: Hit rate [load]: 67.35 %\n",
      "[HUGECTR][INFO] iter: 12000, metrics: [('AUC', 0.763184666633606)]\n",
      "[HUGECTR][INFO] iter: 13000, metrics: [('AUC', 0.7622747421264648)]\n",
      "[HUGECTR][INFO] iter: 14000, metrics: [('AUC', 0.7623080015182495)]\n",
      "[HUGECTR][INFO] iter: 15000, metrics: [('AUC', 0.7622851729393005)]\n",
      "[HUGECTR][INFO] trained with data in wdl_data/file_list.4.txt\n",
      "[HUGECTR][12:42:30][INFO][RANK0]: HMEM-Cache PS: Hit rate [dump]: 64.82 %\n",
      "[HUGECTR][12:42:31][INFO][RANK0]: HMEM-Cache PS: Hit rate [load]: 63.85 %\n",
      "[HUGECTR][12:42:31][INFO][RANK0]: Get updated portion of embedding table [DONE}\n",
      "[HUGECTR][12:42:32][INFO][RANK0]: HMEM-Cache PS: Hit rate [dump]: 64.82 %\n",
      "[HUGECTR][12:42:32][INFO][RANK0]: Updating sparse model in SSD [DONE]\n",
      "[HUGECTR][12:42:34][INFO][RANK0]: Sync blocks from HMEM-Cache to SSD\n",
      " \u001b[38;2;89;255;89m ████████████████████████████████████████▏ \u001b[1m\u001b[31m100.0% \u001b[34m[   2/   2 | 64.6 Hz | 0s<0s]  \u001b[0m\u001b[32m\u001b[0mm\n",
      "[HUGECTR][12:42:34][INFO][RANK0]: Dumping dense weights to file, successful\n",
      "[HUGECTR][12:42:34][INFO][RANK0]: Dumping dense optimizer states to file, successful\n",
      "[HUGECTR][12:42:34][INFO][RANK0]: Dumping untrainable weights to file, successful\n"
     ]
    }
   ],
   "source": [
    "!python3 wdl_etc.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": 4
}