examples/end-to-end-session-based/03-Session-based-Yoochoose-multigpu-training-PyT.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "0fb18ebc",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Copyright 2022 NVIDIA Corporation. All Rights Reserved.\n",
    "#\n",
    "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
    "# you may not use this file except in compliance with the License.\n",
    "# You may obtain a copy of the License at\n",
    "#\n",
    "#     http://www.apache.org/licenses/LICENSE-2.0\n",
    "#\n",
    "# Unless required by applicable law or agreed to in writing, software\n",
    "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
    "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
    "# See the License for the specific language governing permissions and\n",
    "# limitations under the License.\n",
    "# ==============================================================================\n",
    "\n",
    "# Each user is responsible for checking the content of datasets and the\n",
    "# applicable licenses and determining if suitable for the intended use."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "378ff149",
   "metadata": {
    "tags": []
   },
   "source": [
    "<img src=\"https://developer.download.nvidia.com/notebooks/dlsw-notebooks/merlin_transformers4rec_end-to-end-session-based-02-end-to-end-session-based-with-yoochoose-pyt/nvidia_logo.png\" style=\"width: 90px; float: right;\">\n",
    "\n",
    "# Multi-GPU training for session-based recommendations with PyTorch\n",
    "\n",
    "This notebook was prepared by using the latest [merlin-pytorch:22.XX](https://catalog.ngc.nvidia.com/orgs/nvidia/teams/merlin/containers/merlin-pytorch/tags) container."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d8642b00",
   "metadata": {},
   "source": [
    "In the previous two notebooks, we have first created sequential features and saved our processed data frames as parquet files. Then we used these processed parquet files in training a session-based recommendation model with XLNet architecture on a single GPU. We will now expand this exercise to perform a multi-GPU training with the same dataset using PyTorch.\n",
    "\n",
    "There are multiple ways to scale a training pipeline to multiple GPUs: \n",
    "\n",
    "- <b>Model Parallel</b>: If the model is too large to fit on a single GPU, the parameters are distributed over multiple GPUs. This is usually the case for the RecSys domain since the embedding tables can be exceptionally large and memory intensive.\n",
    "- <b>Data Parallel</b>: Every GPU has a copy of all model parameters and runs the forward/backward pass for its batch. Data parallel is useful when you want to speed-up the training/evaluation of data leveraging multiple GPUs in parallel (as typically data won't fit into GPU memory, that is why models are trained on batches).\n",
    "\n",
    "In this example, we demonstrate how to scale a training pipeline to multi-GPU, single node. The goal is to maximize throughput and reduce training time. In that way, models can be trained more frequently and researches can run more experiments in a shorter time duration. \n",
    "\n",
    "This is equivalent to training with a larger batch-size. As we are using more GPUs, we have more computational resources and can achieve higher throughput. In data parallel training, it is often required that all model parameters fit into a single GPU. Every worker (each GPU) has a copy of the model parameters and runs the forward pass on their local batch. The workers synchronize the gradients with each other, which can introduce an overhead.\n",
    "\n",
    "<b>Learning objectives</b> \n",
    "- Scaling training pipeline to multiple GPUs\n",
    "\n",
    "<b>Prerequisites</b>\n",
    "- Run the <b>01-ETL-with-NVTabular.ipynb</b> notebook first to generate the dataset and directories that are also needed by this notebook."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "34bc0fdb",
   "metadata": {},
   "source": [
    "### 1. Creating a multi-gpu training python script "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "78d52653",
   "metadata": {},
   "source": [
    "In this example, we will be using PyTorch, which expects all code including importing libraries, loading data, building the model and training it, to be in a single python script (.py file). We will then spawn torch.distributed.launch, which will distribute the training and evaluation tasks to 2 GPUs with the default <b>DistributedDataParallel</b> configuration.\n",
    "\n",
    "The following cell exports all related code to a <b>pyt_trainer.py</b> file to be created in the same working directory as this notebook. The code is structured as follows:\n",
    "\n",
    "- importing required libraries\n",
    "- specifying and processing command line arguments\n",
    "- specifying the schema file to load and filtering the features that will be used in training\n",
    "- defining the input module\n",
    "- specifying the prediction task\n",
    "- defining the XLNet Transformer architecture configuration\n",
    "- defining the model object and its training arguments\n",
    "- creating a trainer object and running the training loop (over multiple days) with the trainer object\n",
    "\n",
    "All of these steps were covered in the previous two notebooks; feel free to visit the two notebooks to refresh the concepts for each step of the script.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ff36ee4d-6b8a-4c03-b99c-7182c6894fbe",
   "metadata": {},
   "source": [
    "Please note these <b>important points</b> that are relevant to multi-gpu training:\n",
    "- <b>specifying multiple GPUs</b>: PyTorch distributed launch environment will recognize that we have two GPUs, since the `--nproc_per_node` arg of `torch.distributed.launch` takes care of assigning one GPU per process and performs the training loop on multiple GPUs (2 in this case) using different batches of the data in a data-parallel fashion.\n",
    "- <b>data repartitioning</b>: when training on multiple GPUs, data must be re-partitioned into >1 partitions where the number of partitions must be at least equal to the number of GPUs. The torch utility library in Transformers4Rec does this automatically and outputs a UserWarning message. If you would like to avoid this warning message, you may choose to manually re-partition your data files before you launch the training loop or function. See [this document](https://nvidia-merlin.github.io/Transformers4Rec/main/multi_gpu_train.html#distributeddataparallel) for further information on how to do manual re-partitioning.\n",
    "- <b>training and evaluation batch sizes</b>: in the default DistributedDataParallel mode we will be running, keeping the batch size unchanged means each worker will receive the same-size batch despite the fact that you are now using multiple GPUs. If you would like to keep the total batch size constant, you may want to divide the training and evaluation batch sizes by the number of GPUs you are running on, which is expected to reduce time it takes train and evaluate on each batch.\n",
    "\n",
    "Finally, if you have worked with the <b>synthetic</b> data set in notebook 01, remember to change the values 178 and 181 in the next cell accordingly:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "2f2e0a31-4259-4cd3-9618-01a44d2492d0",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "TRAINER_FILE = os.path.join(os.environ.get(\"INPUT_DATA_DIR\", \"/workspace/data\"), \"pyt_trainer.py\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "6c45c899-9c88-4235-8402-03f7f3e40841",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Overwriting /workspace/data/pyt_trainer.py\n"
     ]
    }
   ],
   "source": [
    "%%writefile {TRAINER_FILE}\n",
    "\n",
    "import argparse\n",
    "import os\n",
    "import glob\n",
    "import torch \n",
    "\n",
    "import cupy\n",
    "\n",
    "from transformers4rec import torch as tr\n",
    "from transformers4rec.torch.ranking_metric import NDCGAt, AvgPrecisionAt, RecallAt\n",
    "from transformers4rec.torch.utils.examples_utils import wipe_memory\n",
    "from merlin.schema import Schema\n",
    "from merlin.io import Dataset\n",
    "\n",
    "cupy.cuda.Device(int(os.environ[\"LOCAL_RANK\"])).use()\n",
    "\n",
    "# define arguments that can be passed to this python script\n",
    "parser = argparse.ArgumentParser(description='Hyperparameters for model training')\n",
    "parser.add_argument('--path', type=str, help='Directory with training and validation data')\n",
    "parser.add_argument('--learning-rate', type=float, default=0.0005, help='Learning rate for training')\n",
    "parser.add_argument('--per-device-train-batch-size', type=int, default=64, help='Per device batch size for training')\n",
    "parser.add_argument('--per-device-eval-batch-size', type=int, default=32, help='Per device batch size for evaluation')\n",
    "sh_args = parser.parse_args()\n",
    "\n",
    "# create the schema object by reading the processed train set generated in the previous 01-ETL-with-NVTabular notebook\n",
    "\n",
    "INPUT_DATA_DIR = os.environ.get(\"INPUT_DATA_DIR\", \"/workspace/data\")\n",
    "train = Dataset(os.path.join(INPUT_DATA_DIR, \"processed_nvt/part_0.parquet\"))\n",
    "schema = train.schema\n",
    "\n",
    "# select the subset of features we want to use for training the model by their tags or their names.\n",
    "schema = schema.select_by_name(\n",
    "   ['item_id-list', 'category-list', 'product_recency_days_log_norm-list', 'et_dayofweek_sin-list']\n",
    ")\n",
    "\n",
    "max_sequence_length, d_model = 20, 320\n",
    "# Define input module to process tabular input-features and to prepare masked inputs\n",
    "input_module = tr.TabularSequenceFeatures.from_schema(\n",
    "    schema,\n",
    "    max_sequence_length=max_sequence_length,\n",
    "    continuous_projection=64,\n",
    "    aggregation=\"concat\",\n",
    "    d_output=d_model,\n",
    "    masking=\"mlm\",\n",
    ")\n",
    "\n",
    "# Define Next item prediction-task \n",
    "prediction_task = tr.NextItemPredictionTask(weight_tying=True)\n",
    "\n",
    "# Define the config of the XLNet Transformer architecture\n",
    "transformer_config = tr.XLNetConfig.build(\n",
    "    d_model=d_model, n_head=8, n_layer=2, total_seq_length=max_sequence_length\n",
    ")\n",
    "\n",
    "# Get the end-to-end model \n",
    "model = transformer_config.to_torch_model(input_module, prediction_task)\n",
    "\n",
    "# Set training arguments \n",
    "training_args = tr.trainer.T4RecTrainingArguments(\n",
    "            output_dir=\"./tmp\",\n",
    "            max_sequence_length=20,\n",
    "            data_loader_engine='merlin',\n",
    "            num_train_epochs=10, \n",
    "            dataloader_drop_last=True,\n",
    "            per_device_train_batch_size = sh_args.per_device_train_batch_size,\n",
    "            per_device_eval_batch_size = sh_args.per_device_eval_batch_size,\n",
    "            learning_rate=sh_args.learning_rate,\n",
    "            report_to = [],\n",
    "            logging_steps=200,\n",
    "        )\n",
    "\n",
    "# Instantiate the trainer\n",
    "recsys_trainer = tr.Trainer(\n",
    "    model=model,\n",
    "    args=training_args,\n",
    "    schema=schema,\n",
    "    compute_metrics=True)\n",
    "\n",
    "# Set input and output directories\n",
    "INPUT_DIR=sh_args.path\n",
    "OUTPUT_DIR=sh_args.path\n",
    "\n",
    "import time\n",
    "start = time.time()\n",
    "\n",
    "# main loop for training\n",
    "start_time_window_index = int(os.environ.get(\"START_TIME_INDEX\", \"178\"))\n",
    "final_time_window_index = int(os.environ.get(\"END_TIME_INDEX\", \"181\"))\n",
    "\n",
    "# Iterating over days from 178 to 181\n",
    "for time_index in range(start_time_window_index, final_time_window_index):\n",
    "    # Set data \n",
    "    time_index_train = time_index\n",
    "    time_index_eval = time_index + 1\n",
    "    train_paths = glob.glob(os.path.join(OUTPUT_DIR, f\"{time_index_train}/train.parquet\"))\n",
    "    eval_paths = glob.glob(os.path.join(OUTPUT_DIR, f\"{time_index_eval}/valid.parquet\"))\n",
    "\n",
    "    # Train on day related to time_index \n",
    "    print('*'*20)\n",
    "    print(\"Launch training for day %s are:\" %time_index)\n",
    "    print('*'*20 + '\\n')\n",
    "    recsys_trainer.train_dataset_or_path = train_paths\n",
    "    recsys_trainer.reset_lr_scheduler()\n",
    "    recsys_trainer.train()\n",
    "    recsys_trainer.state.global_step +=1\n",
    "    print('finished')\n",
    "\n",
    "    # Evaluate on the following day\n",
    "    recsys_trainer.eval_dataset_or_path = eval_paths\n",
    "    eval_metrics = recsys_trainer.evaluate(metric_key_prefix='eval')\n",
    "    print('*'*20)\n",
    "    print(\"Eval results for day %s are:\\t\" %time_index_eval)\n",
    "    print('\\n' + '*'*20 + '\\n')\n",
    "    for key in sorted(eval_metrics.keys()):\n",
    "        print(\" %s = %s\" % (key, str(eval_metrics[key]))) \n",
    "    wipe_memory()\n",
    "\n",
    "# export evaluation metrics to a file\n",
    "import json\n",
    "fname = os.path.join(INPUT_DATA_DIR, \"eval_metrics.txt\")\n",
    "f = open(fname, \"w\")\n",
    "f.write(json.dumps(eval_metrics))\n",
    "f.close()\n",
    "\n",
    "end = time.time()\n",
    "print('Total training time:', end-start)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8c8cb68e-a445-439f-a2ae-9c19f3669593",
   "metadata": {},
   "source": [
    "### 2. Executing the multi-gpu training"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c57c0bcd-d9cb-4af5-9cc8-13e9783132d6",
   "metadata": {},
   "source": [
    "You are now ready to execute the python script you created. Run the following shell command which will execute the script and perform data loading, model building and training using 2 GPUs.\n",
    "\n",
    "Note that there are four arguments you can pass to your python script, and only the first one is required:\n",
    "- <b>path</b>: this argument specifies the directory in which to find the multi-day train and validation files to work on\n",
    "- <b>learning rate</b>: you can experiment with different learning rates and see the effect of this hyperparameter when multiple GPUs are used in training (versus 1 GPU). Typically, increasing the learning rate (up to a certain level) as the number of GPUs is increased helps with the accuracy metrics.\n",
    "- <b>per device batch size for training</b>: when using multiple GPUs in DistributedDataParallel mode, you may choose to reduce the batch size in order to keep the total batch size constant. This should help reduce training times.\n",
    "- <b>per device batch size for evaluation</b>: see above"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "c9e83d47-380c-4118-bc29-8bc108163fa0",
   "metadata": {},
   "outputs": [],
   "source": [
    "# If only 1 GPU are available, starts a single process to use that GPU\n",
    "from torch.cuda import device_count\n",
    "num_gpus = device_count()\n",
    "NUM_PROCESSES = min(num_gpus, 2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "1e5678d8-71a9-465f-a986-9ec68cf351ff",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARNING:torch.distributed.run:\n",
      "*****************************************\n",
      "Setting OMP_NUM_THREADS environment variable for each process to be 1 in default, to avoid your system being overloaded, please further tune the variable for optimal performance in your application as needed. \n",
      "*****************************************\n",
      "/usr/local/lib/python3.8/dist-packages/merlin/dtypes/mappings/tf.py:52: UserWarning: Tensorflow dtype mappings did not load successfully due to an error: No module named 'tensorflow'\n",
      "  warn(f\"Tensorflow dtype mappings did not load successfully due to an error: {exc.msg}\")\n",
      "/usr/local/lib/python3.8/dist-packages/merlin/dtypes/mappings/tf.py:52: UserWarning: Tensorflow dtype mappings did not load successfully due to an error: No module named 'tensorflow'\n",
      "  warn(f\"Tensorflow dtype mappings did not load successfully due to an error: {exc.msg}\")\n",
      "Projecting inputs of NextItemPredictionTask to'64' As weight tying requires the input dimension '320' to be equal to the item-id embedding dimension '64'\n",
      "Projecting inputs of NextItemPredictionTask to'64' As weight tying requires the input dimension '320' to be equal to the item-id embedding dimension '64'\n",
      "********************\n",
      "Launch training for day 178 are:\n",
      "********************\n",
      "\n",
      "********************\n",
      "Launch training for day 178 are:\n",
      "********************\n",
      "\n",
      "UserWarning: User is advised to repartition the parquet file before training so npartitions>=global_size. Cudf or pandas can be used for repartitioning eg. pdf.to_parquet('file.parquet',row_group_size=N_ROWS/NPARTITIONS) for pandas or gdf.to_parquet('file.parquet',row_group_size_rows=N_ROWS/NPARTITIONS) for cudf so that npartitions=nr_rows/row_group_size. Also ensure npartitions is divisible by number of GPUs to be used (eg. 2 or 4 partitions, if 2 GPUs will be used).\n",
      "UserWarning: User is advised to repartition the parquet file before training so npartitions>=global_size. Cudf or pandas can be used for repartitioning eg. pdf.to_parquet('file.parquet',row_group_size=N_ROWS/NPARTITIONS) for pandas or gdf.to_parquet('file.parquet',row_group_size_rows=N_ROWS/NPARTITIONS) for cudf so that npartitions=nr_rows/row_group_size. Also ensure npartitions is divisible by number of GPUs to be used (eg. 2 or 4 partitions, if 2 GPUs will be used).\n",
      "***** Running training *****\n",
      "  Num examples = 14080\n",
      "  Num Epochs = 10\n",
      "  Instantaneous batch size per device = 256\n",
      "  Total train batch size (w. parallel, distributed & accumulation) = 512\n",
      "  Gradient Accumulation steps = 1\n",
      "  Total optimization steps = 550\n",
      "{'loss': 7.608, 'learning_rate': 0.0003181818181818182, 'epoch': 3.64}          \n",
      "{'loss': 6.6604, 'learning_rate': 0.00013636363636363637, 'epoch': 7.27}        \n",
      " 91%|█████████████████████████████████████▎   | 500/550 [00:27<00:02, 20.71it/s]Saving model checkpoint to ./tmp/checkpoint-500\n",
      "Trainer.model is not a `PreTrainedModel`, only saving its state dict.\n",
      "100%|████████████████████████████████████████▊| 548/550 [00:29<00:00, 24.82it/s]\n",
      "\n",
      "Training completed. Do not forget to share your model on huggingface.co/models =)\n",
      "\n",
      "\n",
      "finished\n",
      "{'train_runtime': 29.4835, 'train_samples_per_second': 0.339, 'train_steps_per_second': 18.655, 'train_loss': 6.9447163529829545, 'epoch': 10.0}\n",
      "100%|█████████████████████████████████████████| 550/550 [00:29<00:00, 18.66it/s]\n",
      "finished\n",
      "UserWarning: User is advised to repartition the parquet file before training so npartitions>=global_size. Cudf or pandas can be used for repartitioning eg. pdf.to_parquet('file.parquet',row_group_size=N_ROWS/NPARTITIONS) for pandas or gdf.to_parquet('file.parquet',row_group_size_rows=N_ROWS/NPARTITIONS) for cudf so that npartitions=nr_rows/row_group_size. Also ensure npartitions is divisible by number of GPUs to be used (eg. 2 or 4 partitions, if 2 GPUs will be used).\n",
      "UserWarning: User is advised to repartition the parquet file before training so npartitions>=global_size. Cudf or pandas can be used for repartitioning eg. pdf.to_parquet('file.parquet',row_group_size=N_ROWS/NPARTITIONS) for pandas or gdf.to_parquet('file.parquet',row_group_size_rows=N_ROWS/NPARTITIONS) for cudf so that npartitions=nr_rows/row_group_size. Also ensure npartitions is divisible by number of GPUs to be used (eg. 2 or 4 partitions, if 2 GPUs will be used).\n",
      "100%|███████████████████████████████████████████| 10/10 [00:00<00:00, 39.28it/s]********************\n",
      "Eval results for day 179 are:\t\n",
      "\n",
      "********************\n",
      "\n",
      " eval_/loss = 7.320416450500488\n",
      " eval_/next-item/avg_precision_at_10 = 0.07864832133054733\n",
      " eval_/next-item/avg_precision_at_20 = 0.08251794427633286\n",
      " eval_/next-item/ndcg_at_10 = 0.10605569183826447\n",
      " eval_/next-item/ndcg_at_20 = 0.12029790133237839\n",
      " eval_/next-item/recall_at_10 = 0.19570313394069672\n",
      " eval_/next-item/recall_at_20 = 0.2523437440395355\n",
      " eval_runtime = 0.3992\n",
      " eval_samples_per_second = 3206.723\n",
      " eval_steps_per_second = 12.526\n",
      "100%|███████████████████████████████████████████| 10/10 [00:00<00:00, 38.88it/s]\n",
      "********************\n",
      "Eval results for day 179 are:\t\n",
      "\n",
      "********************\n",
      "\n",
      " eval_/loss = 7.320416450500488\n",
      " eval_/next-item/avg_precision_at_10 = 0.07864832133054733\n",
      " eval_/next-item/avg_precision_at_20 = 0.08251794427633286\n",
      " eval_/next-item/ndcg_at_10 = 0.10605569183826447\n",
      " eval_/next-item/ndcg_at_20 = 0.12029790133237839\n",
      " eval_/next-item/recall_at_10 = 0.19570313394069672\n",
      " eval_/next-item/recall_at_20 = 0.2523437440395355\n",
      " eval_runtime = 0.3989\n",
      " eval_samples_per_second = 3209.046\n",
      " eval_steps_per_second = 12.535\n",
      "********************\n",
      "Launch training for day 179 are:\n",
      "********************\n",
      "\n",
      "********************\n",
      "Launch training for day 179 are:\n",
      "********************\n",
      "\n",
      "UserWarning: User is advised to repartition the parquet file before training so npartitions>=global_size. Cudf or pandas can be used for repartitioning eg. pdf.to_parquet('file.parquet',row_group_size=N_ROWS/NPARTITIONS) for pandas or gdf.to_parquet('file.parquet',row_group_size_rows=N_ROWS/NPARTITIONS) for cudf so that npartitions=nr_rows/row_group_size. Also ensure npartitions is divisible by number of GPUs to be used (eg. 2 or 4 partitions, if 2 GPUs will be used).\n",
      "UserWarning: User is advised to repartition the parquet file before training so npartitions>=global_size. Cudf or pandas can be used for repartitioning eg. pdf.to_parquet('file.parquet',row_group_size=N_ROWS/NPARTITIONS) for pandas or gdf.to_parquet('file.parquet',row_group_size_rows=N_ROWS/NPARTITIONS) for cudf so that npartitions=nr_rows/row_group_size. Also ensure npartitions is divisible by number of GPUs to be used (eg. 2 or 4 partitions, if 2 GPUs will be used).\n",
      "***** Running training *****\n",
      "  Num examples = 9984\n",
      "  Num Epochs = 10\n",
      "  Instantaneous batch size per device = 256\n",
      "  Total train batch size (w. parallel, distributed & accumulation) = 512\n",
      "  Gradient Accumulation steps = 1\n",
      "  Total optimization steps = 390\n",
      "{'loss': 6.9038, 'learning_rate': 0.0002435897435897436, 'epoch': 5.13}         \n",
      "100%|█████████████████████████████████████████| 390/390 [00:16<00:00, 25.05it/s]\n",
      "\n",
      "Training completed. Do not forget to share your model on huggingface.co/models =)\n",
      "\n",
      "\n",
      "finished\n",
      "{'train_runtime': 16.8289, 'train_samples_per_second': 0.594, 'train_steps_per_second': 23.174, 'train_loss': 6.640231244991988, 'epoch': 10.0}\n",
      "100%|█████████████████████████████████████████| 390/390 [00:16<00:00, 23.18it/s]\n",
      "finished\n",
      "UserWarning: User is advised to repartition the parquet file before training so npartitions>=global_size. Cudf or pandas can be used for repartitioning eg. pdf.to_parquet('file.parquet',row_group_size=N_ROWS/NPARTITIONS) for pandas or gdf.to_parquet('file.parquet',row_group_size_rows=N_ROWS/NPARTITIONS) for cudf so that npartitions=nr_rows/row_group_size. Also ensure npartitions is divisible by number of GPUs to be used (eg. 2 or 4 partitions, if 2 GPUs will be used).\n",
      "UserWarning: User is advised to repartition the parquet file before training so npartitions>=global_size. Cudf or pandas can be used for repartitioning eg. pdf.to_parquet('file.parquet',row_group_size=N_ROWS/NPARTITIONS) for pandas or gdf.to_parquet('file.parquet',row_group_size_rows=N_ROWS/NPARTITIONS) for cudf so that npartitions=nr_rows/row_group_size. Also ensure npartitions is divisible by number of GPUs to be used (eg. 2 or 4 partitions, if 2 GPUs will be used).\n",
      " 62%|████████████████████████████▏                | 5/8 [00:00<00:00, 46.46it/s]********************\n",
      "Eval results for day 180 are:\t\n",
      "\n",
      "********************\n",
      "\n",
      " eval_/loss = 7.8266496658325195\n",
      " eval_/next-item/avg_precision_at_10 = 0.0614391528069973\n",
      " eval_/next-item/avg_precision_at_20 = 0.0654601976275444\n",
      " eval_/next-item/ndcg_at_10 = 0.08422157913446426\n",
      " eval_/next-item/ndcg_at_20 = 0.09865029156208038\n",
      " eval_/next-item/recall_at_10 = 0.1591796875\n",
      " eval_/next-item/recall_at_20 = 0.2158203125\n",
      " eval_runtime = 0.3334\n",
      " eval_samples_per_second = 3071.067\n",
      " eval_steps_per_second = 11.996\n",
      "100%|█████████████████████████████████████████████| 8/8 [00:00<00:00, 41.08it/s]\n",
      "********************\n",
      "Eval results for day 180 are:\t\n",
      "\n",
      "********************\n",
      "\n",
      " eval_/loss = 7.8266496658325195\n",
      " eval_/next-item/avg_precision_at_10 = 0.0614391528069973\n",
      " eval_/next-item/avg_precision_at_20 = 0.0654601976275444\n",
      " eval_/next-item/ndcg_at_10 = 0.08422157913446426\n",
      " eval_/next-item/ndcg_at_20 = 0.09865029156208038\n",
      " eval_/next-item/recall_at_10 = 0.1591796875\n",
      " eval_/next-item/recall_at_20 = 0.2158203125\n",
      " eval_runtime = 0.3325\n",
      " eval_samples_per_second = 3080.064\n",
      " eval_steps_per_second = 12.031\n",
      "********************\n",
      "Launch training for day 180 are:\n",
      "********************\n",
      "\n",
      "UserWarning: User is advised to repartition the parquet file before training so npartitions>=global_size. Cudf or pandas can be used for repartitioning eg. pdf.to_parquet('file.parquet',row_group_size=N_ROWS/NPARTITIONS) for pandas or gdf.to_parquet('file.parquet',row_group_size_rows=N_ROWS/NPARTITIONS) for cudf so that npartitions=nr_rows/row_group_size. Also ensure npartitions is divisible by number of GPUs to be used (eg. 2 or 4 partitions, if 2 GPUs will be used).\n",
      "********************\n",
      "Launch training for day 180 are:\n",
      "********************\n",
      "\n",
      "UserWarning: User is advised to repartition the parquet file before training so npartitions>=global_size. Cudf or pandas can be used for repartitioning eg. pdf.to_parquet('file.parquet',row_group_size=N_ROWS/NPARTITIONS) for pandas or gdf.to_parquet('file.parquet',row_group_size_rows=N_ROWS/NPARTITIONS) for cudf so that npartitions=nr_rows/row_group_size. Also ensure npartitions is divisible by number of GPUs to be used (eg. 2 or 4 partitions, if 2 GPUs will be used).\n",
      "***** Running training *****\n",
      "  Num examples = 8192\n",
      "  Num Epochs = 10\n",
      "  Instantaneous batch size per device = 256\n",
      "  Total train batch size (w. parallel, distributed & accumulation) = 512\n",
      "  Gradient Accumulation steps = 1\n",
      "  Total optimization steps = 320\n",
      "{'loss': 6.6858, 'learning_rate': 0.0001875, 'epoch': 6.25}                     \n",
      "100%|████████████████████████████████████████▊| 319/320 [00:14<00:00, 24.21it/s]\n",
      "\n",
      "Training completed. Do not forget to share your model on huggingface.co/models =)\n",
      "\n",
      "\n",
      "finished\n",
      "{'train_runtime': 14.4363, 'train_samples_per_second': 0.693, 'train_steps_per_second': 22.166, 'train_loss': 6.4950298309326175, 'epoch': 10.0}\n",
      "100%|█████████████████████████████████████████| 320/320 [00:14<00:00, 22.17it/s]\n",
      "finished\n",
      "UserWarning: User is advised to repartition the parquet file before training so npartitions>=global_size. Cudf or pandas can be used for repartitioning eg. pdf.to_parquet('file.parquet',row_group_size=N_ROWS/NPARTITIONS) for pandas or gdf.to_parquet('file.parquet',row_group_size_rows=N_ROWS/NPARTITIONS) for cudf so that npartitions=nr_rows/row_group_size. Also ensure npartitions is divisible by number of GPUs to be used (eg. 2 or 4 partitions, if 2 GPUs will be used).\n",
      "UserWarning: User is advised to repartition the parquet file before training so npartitions>=global_size. Cudf or pandas can be used for repartitioning eg. pdf.to_parquet('file.parquet',row_group_size=N_ROWS/NPARTITIONS) for pandas or gdf.to_parquet('file.parquet',row_group_size_rows=N_ROWS/NPARTITIONS) for cudf so that npartitions=nr_rows/row_group_size. Also ensure npartitions is divisible by number of GPUs to be used (eg. 2 or 4 partitions, if 2 GPUs will be used).\n",
      "100%|█████████████████████████████████████████████| 4/4 [00:00<00:00, 36.98it/s]********************\n",
      "Eval results for day 181 are:\t\n",
      "\n",
      "********************\n",
      "\n",
      " eval_/loss = 5.795247554779053\n",
      " eval_/next-item/avg_precision_at_10 = 0.12972043454647064\n",
      " eval_/next-item/avg_precision_at_20 = 0.1380162239074707\n",
      " eval_/next-item/ndcg_at_10 = 0.17062446475028992\n",
      " eval_/next-item/ndcg_at_20 = 0.200949028134346\n",
      " eval_/next-item/recall_at_10 = 0.3046875\n",
      " eval_/next-item/recall_at_20 = 0.4248046875\n",
      " eval_runtime = 0.2612\n",
      " eval_samples_per_second = 1960.129\n",
      " eval_steps_per_second = 7.657\n",
      "100%|█████████████████████████████████████████████| 4/4 [00:00<00:00, 34.63it/s]\n",
      "********************\n",
      "Eval results for day 181 are:\t\n",
      "\n",
      "********************\n",
      "\n",
      " eval_/loss = 5.795247554779053\n",
      " eval_/next-item/avg_precision_at_10 = 0.12972043454647064\n",
      " eval_/next-item/avg_precision_at_20 = 0.1380162239074707\n",
      " eval_/next-item/ndcg_at_10 = 0.17062446475028992\n",
      " eval_/next-item/ndcg_at_20 = 0.200949028134346\n",
      " eval_/next-item/recall_at_10 = 0.3046875\n",
      " eval_/next-item/recall_at_20 = 0.4248046875\n",
      " eval_runtime = 0.259\n",
      " eval_samples_per_second = 1976.599\n",
      " eval_steps_per_second = 7.721\n",
      "Total training time: 66.7806625366211\n",
      "Total training time: 66.76544284820557\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "OUTPUT_DIR = os.environ.get(\"OUTPUT_DIR\", \"/workspace/data/preproc_sessions_by_day\")\n",
    "LR = float(os.environ.get(\"LEARNING_RATE\", \"0.0005\"))\n",
    "BATCH_SIZE_TRAIN = int(os.environ.get(\"BATCH_SIZE_TRAIN\", \"256\"))\n",
    "BATCH_SIZE_VALID = int(os.environ.get(\"BATCH_SIZE_VALID\", \"128\"))\n",
    "!python -m torch.distributed.run --nproc_per_node {NUM_PROCESSES} {TRAINER_FILE} --path {OUTPUT_DIR} --learning-rate {LR} --per-device-train-batch-size {BATCH_SIZE_TRAIN} --per-device-eval-batch-size {BATCH_SIZE_VALID}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "847757ca-3c91-4504-b9dc-29127563da01",
   "metadata": {},
   "source": [
    "<b>Congratulations!!!</b> You successfully trained your model using 2 GPUs with a `distributed data parallel` approach. If you choose, you may now go back and experiment with some of the hyperparameters (eg. learning rate, batch sizes, number of GPUs) to collect information on various accuracy metrics as well as total training time, to see what fits best into your workflow."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "68dd532f",
   "metadata": {},
   "source": [
    "## References"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "30f9073e",
   "metadata": {},
   "source": [
    "- Multi-GPU data-parallel training using the Trainer class https://nvidia-merlin.github.io/Transformers4Rec/stable/multi_gpu_train.html\n",
    "\n",
    "- Merlin Transformers4rec: https://github.com/NVIDIA-Merlin/Transformers4Rec\n",
    "\n",
    "- Merlin NVTabular: https://github.com/NVIDIA-Merlin/NVTabular/tree/stable/nvtabular\n",
    "- Merlin Dataloader: https://github.com/NVIDIA-Merlin/dataloader"
   ]
  }
 ],
 "metadata": {
  "interpreter": {
   "hash": "7b543a88d374ac88bf8df97911b380f671b13649694a5b49eb21e60fd27eb479"
  },
  "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.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}