Chapter06/3_SDP_finetuning_pytorch_models.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Using SageMaker Data Parallel\n",
    "\n",
    "in this example we will learn how to use proprietary SageMaker Distributed Data Parallel library (\"SDDP\"). The SDDP library provides a proprietary implementation of data parallelism with native integration with other SageMaker capabilities. SDDP is packaged in SageMaker DL containers and supports both the TensorFlow 2 and PyTorch frameworks.\n",
    "\n",
    "As a training task, we use the same binary classification CV as in the PyTorch DDP sample. Since SDDP is natively supported by SageMaker, we don’t need to develop any custom launcher utilities. As a result, we will have to make minimal changes in training script and job configuration to enabled SDDP.\n",
    "\n",
    "We start with imports and data preparations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sagemaker\n",
    "from sagemaker import get_execution_role\n",
    "\n",
    "sagemaker_session = sagemaker.Session()\n",
    "role = get_execution_role()\n",
    "\n",
    "bucket = sagemaker_session.default_bucket()\n",
    "prefix = 'sagemaker/sm-dataparallel-distribution-options'\n",
    "print('Bucket:\\n{}'.format(bucket))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Data preparation was already done in Chapter06/2_distributed_training_PyTorch.ipynb\n",
    "# If you skipped it, then run following code below\n",
    "\n",
    "! wget https://download.pytorch.org/tutorial/hymenoptera_data.zip\n",
    "! unzip hymenoptera_data.zip\n",
    "data_url = sagemaker_session.upload_data(path=\"./hymenoptera_data\", key_prefix=\"hymenoptera_data\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Running SDDP Training Job\n",
    "\n",
    "To run SDDP-enabled training job, we need to do minor modification to our training script and training job configuration. Let's review them in details.\n",
    "\n",
    "### Modifying Training Script\n",
    "\n",
    "SDDP library starting version 1.4.0 is an integrated PyTorch DDP package that we used in the previous example as a specific backend option. This significantly reduces the changes needed to use SDDP. In fact, if you already have a DDP-enabled training script, you will only need to add an import of the torch_sddp package and use the smddp communication backend when initializing the process group, as follows:\n",
    "\n",
    "```python\n",
    "import smdistributed.dataparallel.torch.torch_smddp\n",
    "import torch.distributed as dist\n",
    "dist.init_process_group(backend='smddp')\n",
    "```\n",
    "\n",
    "Keep in mind that SDDP v1.4 is only available with the latest PyTorch v10 DL containers. For earlier versions, the SDDP API is slightly different. For more details, please refer to the official API documentation [here](https://sagemaker.readthedocs.io/en/stable/api/training/distributed.html#the-sagemaker-distributed-data-parallel-library).\n",
    "\n",
    "Execute the cell below to review training script in full."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "! pygmentize 3_sources/train_sm_dp.py"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Training Job Configuration\n",
    "\n",
    "Starting the SDDP job requires you to provide a `distribution` object with the configuration of data parallelism. In `distribution` we specify that we need to run `dataparallel` job type. You can also provide additional MPI configuration in `custom_mpi_options` parameter.\n",
    "\n",
    "```python\n",
    "distribution = {\n",
    "    \"smdistributed\": {\n",
    "        \"dataparallel\": {\n",
    "            \"enabled\": True,\n",
    "            \"custom_mpi_options\": \"-verbose -x NCCL_DEBUG=VERSION\"\n",
    "} }\n",
    "}\n",
    "```\n",
    "\n",
    "\n",
    "Another thing to keep in mind is that SDDP is only available for a limited set of multi-GPU instance types: `ml.p3.16xlarge`, `ml.p3dn.24xlarge`, and `ml.p4d.24xlarge`. Execute the cell below to start SDDP training job."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sagemaker.pytorch import PyTorch\n",
    "\n",
    "instance_type = 'ml.p3.2xlarge'\n",
    "instance_count = 2\n",
    "\n",
    "distribution = { \n",
    "    \"smdistributed\": { \n",
    "        \"dataparallel\": {\n",
    "            \"enabled\": True, \n",
    "            \"custom_mpi_options\": \"-verbose -x NCCL_DEBUG=VERSION\"\n",
    "        }\n",
    "    }\n",
    "}\n",
    "\n",
    "sm_dp_estimator = PyTorch(\n",
    "          entry_point=\"train.py\", # Pick your train script\n",
    "          source_dir='3_sources',\n",
    "          role=role,\n",
    "          instance_type=instance_type,\n",
    "          sagemaker_session=sagemaker_session,\n",
    "          framework_version='1.6.0',\n",
    "          py_version='py36',\n",
    "          instance_count=1,\n",
    "          hyperparameters={\n",
    "              \"batch-size\":64,\n",
    "              \"epochs\":20,\n",
    "              \"model-name\":\"squeezenet\",\n",
    "              \"num-classes\": 2,\n",
    "              \"feature-extract\":True,\n",
    "              \"sync-s3-path\":f\"s3://{bucket}/distributed-training/output\"\n",
    "          },\n",
    "          disable_profiler=True,\n",
    "          debugger_hook_config=False,\n",
    "          distribution=distribution,\n",
    "          base_job_name=\"SM-DP\",\n",
    "      )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sm_dp_estimator.fit(inputs={\"train\":f\"{data_url}/train\", \"val\":f\"{data_url}/val\"})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Summary\n",
    "\n",
    "In this example, we learn how wiht minimal modification you can use SDDP library to run distributed data parallel jobs."
   ]
  }
 ],
 "metadata": {
  "interpreter": {
   "hash": "cae6ef5e525c6d5a8daa33565a4e32326fcdb22bb4405c41032726ef6ebbb77e"
  },
  "kernelspec": {
   "display_name": "conda_pytorch_p36",
   "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.9.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}