notebooks/continuous_training.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cdfec37b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Copyright 2021 NVIDIA Corporation. All Rights Reserved.\n",
    "#\n",
    "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
    "# you may not use this file except in compliance with the License.\n",
    "# You may obtain a copy of the License at\n",
    "#\n",
    "#     http://www.apache.org/licenses/LICENSE-2.0\n",
    "#\n",
    "# Unless required by applicable law or agreed to in writing, software\n",
    "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
    "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
    "# See the License for the specific language governing permissions and\n",
    "# limitations under the License.\n",
    "# ==============================================================================\n",
    "\n",
    "# Each user is responsible for checking the content of datasets and the\n",
    "# applicable licenses and determining if suitable for the intended use."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<img src=\"http://developer.download.nvidia.com/notebooks/dlsw-notebooks/merlin_hugectr_continuous-training/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, see the [Embedding Training Cache](https://nvidia-merlin.github.io/HugeCTR/master/hugectr_core_features.html#embedding-training-cache) documentation.\n",
    "\n",
    "To learn how to use the APIs of ETC, see the [HugeCTR Python Interface](https://nvidia-merlin.github.io/HugeCTR/master/api/python_interface.html) documentation."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setup\n",
    "\n",
    "To setup the environment, refer to [HugeCTR Example Notebooks](../notebooks) and follow the instructions there before running the following."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Continuous Training\n",
    "\n",
    "\n",
    "To download and prepare the dataset we will be doing the following steps. At the end of this cell, we provide the shell commands you can run on the terminal to get the data ready for this notebook.\n",
    "\n",
    "**Note**: If you already have the data downloaded, then skip to the preprocessing step (2). If preprocessing is also done, skip to creating the softlink between the processed data to the `notebooks/` directory (3).\n",
    "\n",
    "\n",
    "### Data Preparation\n",
    "\n",
    "\n",
    "1. Download the Criteo dataset\n",
    "\n",
    "To preprocess the downloaded Kaggle Criteo dataset, we'll make the following operations: \n",
    "\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",
    "\n",
    "2. Preprocessing by Pandas:\n",
    "   \n",
    "   Meanings of the command line arguments:\n",
    "\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 the \"Preprocess the Criteo Dataset\" section of the README in the [samples/criteo](https://github.com/NVIDIA-Merlin/HugeCTR/tree/master/samples/criteo) directory of the repository on GitHub.\n",
    "   \n",
    "3. Create a soft link of the dataset folder to the path of this notebook\n",
    "\n",
    "#### Run the following commands on the terminal to prepare the data for this notebook\n",
    "\n",
    "\n",
    "```shell\n",
    "export project_root=/home/hugectr # set this to the directory where hugectr is downloaded\n",
    "cd ${project_root}/tools\n",
    "# Step 1\n",
    "wget https://storage.googleapis.com/criteo-cail-datasets/day_0.gz\n",
    "#Step 2\n",
    "bash preprocess.sh 0 wdl_data pandas 1 1 10\n",
    "#Step 3\n",
    "ln -s ${project_root}/tools/wdl_data ${project_root}/notebooks/wdl_data\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Continuous Training with High-level API\n",
    "\n",
    "This section gives the code sample of continuous training using a Keras-like high-level API. The high-level API encapsulates much of the complexity for users, making it easy to use and able to handle many of the scenarios in a production environment.\n",
    "\n",
    "Meanwhile, in addition to a high-level API, HugeCTR also provides low-level APIs that enable 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 the 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": 9,
   "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",
    "                          num_workers=8,\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": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[HCTR][08:03:26.675][INFO][RK0][main]: Empty embedding, trained table will be stored in ./wdl_0_sparse_model\n",
      "[HCTR][08:03:26.675][INFO][RK0][main]: Empty embedding, trained table will be stored in ./wdl_1_sparse_model\n",
      "HugeCTR Version: 4.1\n",
      "====================================================Model Init=====================================================\n",
      "[HCTR][08:03:26.676][WARNING][RK0][main]: The model name is not specified when creating the solver.\n",
      "[HCTR][08:03:26.676][INFO][RK0][main]: Global seed is 1441282772\n",
      "[HCTR][08:03:26.678][INFO][RK0][main]: Device to NUMA mapping:\n",
      "  GPU 0 ->  node 0\n",
      "[HCTR][08:03:28.494][WARNING][RK0][main]: Peer-to-peer access cannot be fully enabled.\n",
      "[HCTR][08:03:28.494][INFO][RK0][main]: Start all2all warmup\n",
      "[HCTR][08:03:28.495][INFO][RK0][main]: End all2all warmup\n",
      "[HCTR][08:03:28.495][INFO][RK0][main]: Using All-reduce algorithm: NCCL\n",
      "[HCTR][08:03:28.496][INFO][RK0][main]: Device 0: Tesla V100-SXM2-32GB\n",
      "[HCTR][08:03:28.496][INFO][RK0][main]: num of DataReader workers for train: 8\n",
      "[HCTR][08:03:28.496][INFO][RK0][main]: num of DataReader workers for eval: 8\n",
      "[HCTR][08:03:28.505][INFO][RK0][main]: max_vocabulary_size_per_gpu_=6029312\n",
      "[HCTR][08:03:28.507][INFO][RK0][main]: max_vocabulary_size_per_gpu_=5865472\n",
      "[HCTR][08:03:28.511][INFO][RK0][main]: Graph analysis to resolve tensor dependency\n",
      "===================================================Model Compile===================================================\n",
      "[HCTR][08:03:32.482][INFO][RK0][main]: gpu0 start to init embedding\n",
      "[HCTR][08:03:32.482][INFO][RK0][main]: gpu0 init embedding done\n",
      "[HCTR][08:03:32.482][INFO][RK0][main]: gpu0 start to init embedding\n",
      "[HCTR][08:03:32.484][INFO][RK0][main]: gpu0 init embedding done\n",
      "[HCTR][08:03:32.484][INFO][RK0][main]: Enable HMEM-Based Parameter Server\n",
      "[HCTR][08:03:32.484][INFO][RK0][main]: ./wdl_0_sparse_model not exist, create and train from scratch\n",
      "[HCTR][08:03:32.492][INFO][RK0][main]: Enable HMemCache-Based Parameter Server\n",
      "[HCTR][08:03:32.492][INFO][RK0][main]: ./wdl_1_sparse_model/key doesn't exist, created\n",
      "[HCTR][08:03:32.495][INFO][RK0][main]: ./wdl_1_sparse_model/emb_vector doesn't exist, created\n",
      "[HCTR][08:03:32.498][INFO][RK0][main]: ./wdl_1_sparse_model/Adam.m doesn't exist, created\n",
      "[HCTR][08:03:32.502][INFO][RK0][main]: ./wdl_1_sparse_model/Adam.v doesn't exist, created\n",
      "[HCTR][08:03:33.843][INFO][RK0][main]: Starting AUC NCCL warm-up\n",
      "[HCTR][08:03:33.847][INFO][RK0][main]: Warm-up done\n",
      "===================================================Model Summary===================================================\n",
      "[HCTR][08:03:33.847][INFO][RK0][main]: Model structure on each GPU\n",
      "Label                                   Dense                         Sparse                        \n",
      "label                                   dense                          wide_data,deep_data           \n",
      "(1024,1)                                (1024,13)                               \n",
      "——————————————————————————————————————————————————————————————————————————————————————————————————————————————————\n",
      "Layer Type                              Input Name                    Output Name                   Output Shape                  \n",
      "——————————————————————————————————————————————————————————————————————————————————————————————————————————————————\n",
      "DistributedSlotSparseEmbeddingHash      wide_data                     sparse_embedding2             (1024,1,1)                    \n",
      "------------------------------------------------------------------------------------------------------------------\n",
      "DistributedSlotSparseEmbeddingHash      deep_data                     sparse_embedding1             (1024,26,16)                  \n",
      "------------------------------------------------------------------------------------------------------------------\n",
      "Reshape                                 sparse_embedding1             reshape1                      (1024,416)                    \n",
      "------------------------------------------------------------------------------------------------------------------\n",
      "Reshape                                 sparse_embedding2             reshape2                      (1024,1)                      \n",
      "------------------------------------------------------------------------------------------------------------------\n",
      "Concat                                  reshape1                      concat1                       (1024,429)                    \n",
      "                                        dense                                                                                     \n",
      "------------------------------------------------------------------------------------------------------------------\n",
      "InnerProduct                            concat1                       fc1                           (1024,1024)                   \n",
      "------------------------------------------------------------------------------------------------------------------\n",
      "ReLU                                    fc1                           relu1                         (1024,1024)                   \n",
      "------------------------------------------------------------------------------------------------------------------\n",
      "Dropout                                 relu1                         dropout1                      (1024,1024)                   \n",
      "------------------------------------------------------------------------------------------------------------------\n",
      "InnerProduct                            dropout1                      fc2                           (1024,1024)                   \n",
      "------------------------------------------------------------------------------------------------------------------\n",
      "ReLU                                    fc2                           relu2                         (1024,1024)                   \n",
      "------------------------------------------------------------------------------------------------------------------\n",
      "Dropout                                 relu2                         dropout2                      (1024,1024)                   \n",
      "------------------------------------------------------------------------------------------------------------------\n",
      "InnerProduct                            dropout2                      fc3                           (1024,1)                      \n",
      "------------------------------------------------------------------------------------------------------------------\n",
      "Add                                     fc3                           add1                          (1024,1)                      \n",
      "                                        reshape2                                                                                  \n",
      "------------------------------------------------------------------------------------------------------------------\n",
      "BinaryCrossEntropyLoss                  add1                          loss                                                        \n",
      "                                        label                                                                                     \n",
      "------------------------------------------------------------------------------------------------------------------\n",
      "[HCTR][08:03:33.857][INFO][RK0][main]: Save the model graph to wdl.json successfully\n",
      "=====================================================Model Fit=====================================================\n",
      "[HCTR][08:03:33.857][INFO][RK0][main]: Use embedding training cache mode with number of training sources: 2, number of epochs: 1\n",
      "[HCTR][08:03:33.857][INFO][RK0][main]: Training batchsize: 1024, evaluation batchsize: 1024\n",
      "[HCTR][08:03:33.857][INFO][RK0][main]: Evaluation interval: 1000, snapshot interval: 10000\n",
      "[HCTR][08:03:33.857][INFO][RK0][main]: Dense network trainable: True\n",
      "[HCTR][08:03:33.857][INFO][RK0][main]: Sparse embedding sparse_embedding1 trainable: True\n",
      "[HCTR][08:03:33.857][INFO][RK0][main]: Sparse embedding sparse_embedding2 trainable: True\n",
      "[HCTR][08:03:33.857][INFO][RK0][main]: Use mixed precision: False, scaler: 1.000000, use cuda graph: True\n",
      "[HCTR][08:03:33.857][INFO][RK0][main]: lr: 0.001000, warmup_steps: 1, end_lr: 0.000000\n",
      "[HCTR][08:03:33.857][INFO][RK0][main]: decay_start: 0, decay_steps: 1, decay_power: 2.000000\n",
      "[HCTR][08:03:33.858][INFO][RK0][main]: Evaluation source file: wdl_data/file_list.2.txt\n",
      "[HCTR][08:03:33.858][INFO][RK0][main]: --------------------Epoch 0, source file: wdl_data/file_list.0.txt--------------------\n",
      "[HCTR][08:03:33.860][INFO][RK0][main]: Preparing embedding table for next pass\n",
      "[HCTR][08:03:33.952][INFO][RK0][main]: HMEM-Cache PS: Hit rate [load]: 0 %\n",
      "[HCTR][08:03:36.359][INFO][RK0][main]: --------------------Epoch 0, source file: wdl_data/file_list.1.txt--------------------\n",
      "[HCTR][08:03:36.360][INFO][RK0][main]: Preparing embedding table for next pass\n",
      "[HCTR][08:03:37.255][INFO][RK0][main]: HMEM-Cache PS: Hit rate [dump]: 0 %\n",
      "[HCTR][08:03:37.355][INFO][RK0][main]: HMEM-Cache PS: Hit rate [load]: 0 %\n",
      "[HCTR][08:03:37.964][INFO][RK0][main]: Iter: 500 Time(500 iters): 4.10376s Loss: 0.108516 lr:0.001\n",
      "=====================================================Model Fit=====================================================\n",
      "[HCTR][08:03:39.695][INFO][RK0][main]: Use embedding training cache mode with number of training sources: 2, number of epochs: 1\n",
      "[HCTR][08:03:39.695][INFO][RK0][main]: Training batchsize: 1024, evaluation batchsize: 1024\n",
      "[HCTR][08:03:39.695][INFO][RK0][main]: Evaluation interval: 1000, snapshot interval: 10000\n",
      "[HCTR][08:03:39.695][INFO][RK0][main]: Dense network trainable: True\n",
      "[HCTR][08:03:39.695][INFO][RK0][main]: Sparse embedding sparse_embedding1 trainable: True\n",
      "[HCTR][08:03:39.695][INFO][RK0][main]: Sparse embedding sparse_embedding2 trainable: True\n",
      "[HCTR][08:03:39.695][INFO][RK0][main]: Use mixed precision: False, scaler: 1.000000, use cuda graph: True\n",
      "[HCTR][08:03:39.695][INFO][RK0][main]: lr: 0.001000, warmup_steps: 1, end_lr: 0.000000\n",
      "[HCTR][08:03:39.695][INFO][RK0][main]: decay_start: 0, decay_steps: 1, decay_power: 2.000000\n",
      "[HCTR][08:03:39.695][INFO][RK0][main]: Evaluation source file: wdl_data/file_list.5.txt\n",
      "[HCTR][08:03:39.695][INFO][RK0][main]: --------------------Epoch 0, source file: wdl_data/file_list.3.txt--------------------\n",
      "[HCTR][08:03:39.696][INFO][RK0][main]: Preparing embedding table for next pass\n",
      "[HCTR][08:03:40.424][INFO][RK0][main]: HMEM-Cache PS: Hit rate [dump]: 78.9 %\n",
      "[HCTR][08:03:40.501][INFO][RK0][main]: HMEM-Cache PS: Hit rate [load]: 72.53 %\n",
      "[HCTR][08:03:42.840][INFO][RK0][main]: --------------------Epoch 0, source file: wdl_data/file_list.4.txt--------------------\n",
      "[HCTR][08:03:42.841][INFO][RK0][main]: Preparing embedding table for next pass\n",
      "[HCTR][08:03:43.696][INFO][RK0][main]: HMEM-Cache PS: Hit rate [dump]: 66.2 %\n",
      "[HCTR][08:03:43.767][INFO][RK0][main]: HMEM-Cache PS: Hit rate [load]: 68.85 %\n",
      "[HCTR][08:03:44.382][INFO][RK0][main]: Iter: 500 Time(500 iters): 4.68363s Loss: 0.103496 lr:0.001\n",
      "[HCTR][08:03:46.952][INFO][RK0][main]: HMEM-Cache PS: Hit rate [dump]: 66.14 %\n",
      "[HCTR][08:03:47.217][INFO][RK0][main]: HMEM-Cache PS: Hit rate [load]: 58.85 %\n",
      "[HCTR][08:03:47.222][INFO][RK0][main]: Get updated portion of embedding table [DONE}\n",
      "[HCTR][08:03:48.314][INFO][RK0][main]: HMEM-Cache PS: Hit rate [dump]: 66.14 %\n",
      "[HCTR][08:03:48.318][INFO][RK0][main]: Updating sparse model in SSD\n",
      "[HCTR][08:03:48.544][INFO][RK0][main]: Done!\n",
      "[HCTR][08:03:48.544][INFO][RK0][main]: Sync blocks from HMEM-Cache to SSD\n",
      " \u001b[38;2;89;255;89m ████████████████████████████████████████▏ \u001b[1m\u001b[31m100.0% \u001b[34m[   2/   2 | 13.3 Hz | 0s<0s]  \u001b[0m\u001b[32m\u001b[0mm\n",
      "[HCTR][08:03:48.695][INFO][RK0][main]: Using Local file system backend.\n",
      "[HCTR][08:03:48.710][INFO][RK0][main]: Dumping dense weights to file, successful\n",
      "[HCTR][08:03:48.715][INFO][RK0][main]: Using Local file system backend.\n",
      "[HCTR][08:03:48.744][INFO][RK0][main]: Dumping dense optimizer states to file, successful\n"
     ]
    }
   ],
   "source": [
    "!python3 wdl_train.py"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Continuous Training with the Low-level API\n",
    "\n",
    "This section gives the code sample for continuous training using the low-level API.\n",
    "The program logic is the same as the preceding code sample.\n",
    "\n",
    "Although the low-level APIs provide fine-grained control of the training logic, we encourage you to use the high-level API if it can satisfy your requirements because 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](https://nvidia-merlin.github.io/HugeCTR/master/api/python_interface.html#low-level-training-api) and samples of [Low-level Training](./hugectr_criteo.ipynb)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "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",
    "                          num_workers = 10,\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": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[HCTR][08:03:56.005][INFO][RK0][main]: Use existing embedding: ./wdl_0_sparse_model\n",
      "[HCTR][08:03:56.005][INFO][RK0][main]: Use existing embedding: ./wdl_1_sparse_model\n",
      "HugeCTR Version: 4.1\n",
      "====================================================Model Init=====================================================\n",
      "[HCTR][08:03:56.005][WARNING][RK0][main]: The model name is not specified when creating the solver.\n",
      "[HCTR][08:03:56.005][INFO][RK0][main]: Global seed is 3575646790\n",
      "[HCTR][08:03:56.008][INFO][RK0][main]: Device to NUMA mapping:\n",
      "  GPU 0 ->  node 0\n",
      "[HCTR][08:03:57.823][WARNING][RK0][main]: Peer-to-peer access cannot be fully enabled.\n",
      "[HCTR][08:03:57.823][INFO][RK0][main]: Start all2all warmup\n",
      "[HCTR][08:03:57.823][INFO][RK0][main]: End all2all warmup\n",
      "[HCTR][08:03:57.824][INFO][RK0][main]: Using All-reduce algorithm: NCCL\n",
      "[HCTR][08:03:57.825][INFO][RK0][main]: Device 0: Tesla V100-SXM2-32GB\n",
      "[HCTR][08:03:57.825][INFO][RK0][main]: num of DataReader workers for train: 10\n",
      "[HCTR][08:03:57.825][INFO][RK0][main]: num of DataReader workers for eval: 10\n",
      "[HCTR][08:03:57.836][WARNING][RK0][main]: Embedding vector size(1) is not a multiple of 32, which may affect the GPU resource utilization.\n",
      "[HCTR][08:03:57.836][INFO][RK0][main]: max_num_frequent_categories is not specified using default: 1\n",
      "[HCTR][08:03:57.836][INFO][RK0][main]: max_num_infrequent_samples is not specified using default: -1\n",
      "[HCTR][08:03:57.836][INFO][RK0][main]: p_dup_max is not specified using default: 0.01\n",
      "[HCTR][08:03:57.836][INFO][RK0][main]: max_all_reduce_bandwidth is not specified using default: 1.3e+11\n",
      "[HCTR][08:03:57.836][INFO][RK0][main]: max_all_to_all_bandwidth is not specified using default: 1.9e+11\n",
      "[HCTR][08:03:57.836][INFO][RK0][main]: efficiency_bandwidth_ratio is not specified using default: 1\n",
      "[HCTR][08:03:57.836][INFO][RK0][main]: communication_type is not specified using default: IB_NVLink\n",
      "[HCTR][08:03:57.836][INFO][RK0][main]: hybrid_embedding_type is not specified using default: Distributed\n",
      "[HCTR][08:03:57.836][INFO][RK0][main]: max_vocabulary_size_per_gpu_=6029312\n",
      "[HCTR][08:03:57.838][WARNING][RK0][main]: Embedding vector size(16) is not a multiple of 32, which may affect the GPU resource utilization.\n",
      "[HCTR][08:03:57.838][INFO][RK0][main]: max_num_frequent_categories is not specified using default: 1\n",
      "[HCTR][08:03:57.838][INFO][RK0][main]: max_num_infrequent_samples is not specified using default: -1\n",
      "[HCTR][08:03:57.838][INFO][RK0][main]: p_dup_max is not specified using default: 0.01\n",
      "[HCTR][08:03:57.838][INFO][RK0][main]: max_all_reduce_bandwidth is not specified using default: 1.3e+11\n",
      "[HCTR][08:03:57.838][INFO][RK0][main]: max_all_to_all_bandwidth is not specified using default: 1.9e+11\n",
      "[HCTR][08:03:57.838][INFO][RK0][main]: efficiency_bandwidth_ratio is not specified using default: 1\n",
      "[HCTR][08:03:57.838][INFO][RK0][main]: communication_type is not specified using default: IB_NVLink\n",
      "[HCTR][08:03:57.838][INFO][RK0][main]: hybrid_embedding_type is not specified using default: Distributed\n",
      "[HCTR][08:03:57.838][INFO][RK0][main]: max_vocabulary_size_per_gpu_=5865472\n",
      "[HCTR][08:03:57.842][INFO][RK0][main]: Load the model graph from wdl.json successfully\n",
      "[HCTR][08:03:57.842][INFO][RK0][main]: Graph analysis to resolve tensor dependency\n",
      "===================================================Model Compile===================================================\n",
      "[HCTR][08:04:01.810][INFO][RK0][main]: gpu0 start to init embedding\n",
      "[HCTR][08:04:01.811][INFO][RK0][main]: gpu0 init embedding done\n",
      "[HCTR][08:04:01.811][INFO][RK0][main]: gpu0 start to init embedding\n",
      "[HCTR][08:04:01.813][INFO][RK0][main]: gpu0 init embedding done\n",
      "[HCTR][08:04:01.813][INFO][RK0][main]: Enable HMEM-Based Parameter Server\n",
      "[HCTR][08:04:01.948][INFO][RK0][main]: Enable HMemCache-Based Parameter Server\n",
      "[HCTR][08:04:03.308][INFO][RK0][main]: Starting AUC NCCL warm-up\n",
      "[HCTR][08:04:03.312][INFO][RK0][main]: Warm-up done\n",
      "[HCTR][08:04:03.315][INFO][RK0][main]: Preparing embedding table for next pass\n",
      "[HCTR][08:04:03.454][INFO][RK0][main]: HMEM-Cache PS: Hit rate [load]: 0 %\n",
      "[HUGECTR][INFO] iter: 0, metrics: [('AUC', 0.6509891152381897)]\n",
      "[HUGECTR][INFO] trained with data in wdl_data/file_list.0.txt\n",
      "[HCTR][08:04:06.386][INFO][RK0][main]: Preparing embedding table for next pass\n",
      "[HCTR][08:04:06.865][INFO][RK0][main]: HMEM-Cache PS: Hit rate [dump]: 100 %\n",
      "[HCTR][08:04:06.966][INFO][RK0][main]: HMEM-Cache PS: Hit rate [load]: 78.9 %\n",
      "[HUGECTR][INFO] trained with data in wdl_data/file_list.1.txt\n",
      "[HCTR][08:04:09.307][INFO][RK0][main]: Preparing embedding table for next pass\n",
      "[HCTR][08:04:10.058][INFO][RK0][main]: HMEM-Cache PS: Hit rate [dump]: 100 %\n",
      "[HCTR][08:04:10.133][INFO][RK0][main]: HMEM-Cache PS: Hit rate [load]: 91.27 %\n",
      "[HUGECTR][INFO] iter: 1000, metrics: [('AUC', 0.716963529586792)]\n",
      "[HUGECTR][INFO] trained with data in wdl_data/file_list.3.txt\n",
      "[HCTR][08:04:12.976][INFO][RK0][main]: Preparing embedding table for next pass\n",
      "[HCTR][08:04:13.647][INFO][RK0][main]: HMEM-Cache PS: Hit rate [dump]: 91.27 %\n",
      "[HCTR][08:04:13.724][INFO][RK0][main]: HMEM-Cache PS: Hit rate [load]: 91.25 %\n",
      "[HUGECTR][INFO] trained with data in wdl_data/file_list.4.txt\n",
      "[HCTR][08:04:16.497][INFO][RK0][main]: HMEM-Cache PS: Hit rate [dump]: 91.25 %\n",
      "[HCTR][08:04:16.674][INFO][RK0][main]: HMEM-Cache PS: Hit rate [load]: 100 %\n",
      "[HCTR][08:04:16.680][INFO][RK0][main]: Get updated portion of embedding table [DONE}\n",
      "[HCTR][08:04:17.705][INFO][RK0][main]: HMEM-Cache PS: Hit rate [dump]: 91.25 %\n",
      "[HCTR][08:04:17.705][INFO][RK0][main]: Updating sparse model in SSD\n",
      "[HCTR][08:04:17.732][INFO][RK0][main]: Done!\n",
      "[HCTR][08:04:17.732][INFO][RK0][main]: Sync blocks from HMEM-Cache to SSD\n",
      " \u001b[38;2;89;255;89m ████████████████████████████████████████▏ \u001b[1m\u001b[31m100.0% \u001b[34m[   2/   2 | 11.9 Hz | 0s<0s]  \u001b[0m\u001b[32m\u001b[0mm\n",
      "[HCTR][08:04:17.902][INFO][RK0][main]: Using Local file system backend.\n",
      "[HCTR][08:04:17.917][INFO][RK0][main]: Dumping dense weights to file, successful\n",
      "[HCTR][08:04:17.919][INFO][RK0][main]: Using Local file system backend.\n",
      "[HCTR][08:04:17.954][INFO][RK0][main]: Dumping dense optimizer states to file, successful\n"
     ]
    }
   ],
   "source": [
    "!python3 wdl_etc.py"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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
}