examples/hybrid_jobs/5_Parallelize_training_for_QML/Parallelize_training_for_QML.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "1360239b",
   "metadata": {},
   "source": [
    "# Parallelize training for Quantum machine learning"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9ea4b44b",
   "metadata": {},
   "source": [
    "Quantum machine learning (QML) is a special type of hybrid quantum-classical workload. Like classical machine learning (ML), there is usually a parameterized model, a dataset and a loss function. When training the model with the dataset, the parameters of the model are updated to minimize the loss function. In QML, the model contains one or many quantum circuits. The model may or may not also include classical neural nets. A loss function is usually defined for a single data point. Say a dataset $D$ has $N_D$ data points, $d_1$, $d_2$, ..., $d_{N_D}$. The losses associated with the data points are $L(d_1)$, $L(d_2)$, ..., $L(d_{N_D})$, where $L$ is the loss function. Without invoking any advanced feature, the algorithm script would compute these losses in serial, and then average them to be the total loss for gradient computations. This procedure is time consuming, especially when there are hundreds of data points. \n",
    "\n",
    "## Data parallelism\n",
    "The loss from one data point is independent of the other data points. The order of the loss evaluations therefore does not need to follow a specific order. They can even be evaluated <i>in parallel!</i> Losses and gradients of variational parameters associated with different data points can be evaluated on different GPUs at the same time. This is known as data parallelism. In this notebook, we will learn to use [SageMaker's distributed data parallel library](https://docs.aws.amazon.com/sagemaker/latest/dg/data-parallel.html) in Braket Jobs to accelerate the training of your quantum model. We go through examples to show you how to parallelize trainings across multiple GPUs in an instance, and even multiple GPUs over multiple instances! "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "77bf8c78",
   "metadata": {},
   "source": [
    "## Binary Classification of Sonar dataset\n",
    "Let's use a binary classification of the [Sonar dataset](https://archive.ics.uci.edu/ml/datasets/Connectionist+Bench+%28Sonar%2C+Mines+vs.+Rocks%29) as the QML example. The Sonar dataset contains 208 data points each with 60 features that are collected from sonar signals bouncing off materials. Each data point is either labeled as \"M\" for mines or \"R\" for rocks. Our QML model consists of an input layer, a quantum circuit and an output layer. The input and output layer are classical dense layers. The dimension of classical input layer is $60\\times N$, where $N$ is the number of qubits in the quantum circuit. The result of the input layer is encoded into the quantum circuit using [angle embedding](https://pennylane.readthedocs.io/en/stable/code/api/pennylane.AngleEmbedding.html). After the angle embedding, the quantum circuit has the same structure as figure 4 of [this paper](https://arxiv.org/abs/1804.00633). It is a generic circuit ansatz that has two parametrized [strongly entangling layers](https://pennylane.readthedocs.io/en/user-docs-refactor/code/pennylane.templates.layers.StronglyEntanglingLayer.html) and a single parametrized rotation gate at the first qubit. The measurement is only performed at the first qubit. Using the concept of classical ML, the dimension of the quantum circuit layer is $N\\times1$. The classical output layer has dimension $1\\times1$ which takes the measurement of the quantum circuit and outputs a real number. The loss function is the [margin loss function](https://pytorch.org/docs/stable/generated/torch.nn.MarginRankingLoss.html) between the model output and the label. See [model_def.py](source_script/model_def.py) and [quantum_circuit.py](source_script/quantum_circuit.py) for more detail about the model and the quantum circuits."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "036ae723",
   "metadata": {},
   "outputs": [],
   "source": [
    "input_file_path = \"data/sonar.all-data\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "19c06a90",
   "metadata": {},
   "source": [
    "## Training with single GPU\n",
    "\n",
    "Let's start by running the hybrid job with `lightning.gpu` simulator on a single GPU in a `ml.p3.2xlarge` instance which has one Nvidia V100 GPU. The algorithm script to train our quantum model is [train_single.py](source_script/train_single.py). In the algorithm script, we use PennyLane with PyTorch as our framework, which are both included in Braket's pre-configured PyTorch container. As faster demonstration, we only use a portion of the dataset (64 data points) instead of the full dataset. You can experiment with more data by setting `ndata` in the hyperparameters to a higher number, up to 208, the size of dataset. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "d4a17777",
   "metadata": {},
   "outputs": [],
   "source": [
    "from braket.jobs.config import InstanceConfig\n",
    "from braket.aws import AwsSession\n",
    "from braket.jobs.image_uris import Framework, retrieve_image\n",
    "\n",
    "instance_config = InstanceConfig(instanceType='ml.p3.2xlarge')\n",
    "\n",
    "hyperparameters={\"nwires\": 10, \n",
    "                 \"ndata\": 64, \n",
    "                 \"batch_size\": 64, \n",
    "                 \"epochs\": 5, \n",
    "                 \"gamma\": 0.99, \n",
    "                 \"lr\": 0.1,\n",
    "                 \"seed\": 42,\n",
    "                }\n",
    "\n",
    "input_file_path = \"data/sonar.all-data\"\n",
    "\n",
    "image_uri = retrieve_image(Framework.PL_PYTORCH, AwsSession().region)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f77a2e35",
   "metadata": {},
   "source": [
    "We submit our hybrid job after setting up instance configuration, hyperparameters and the hybrid job container image."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ab8a1598",
   "metadata": {},
   "source": [
    "**Note:** The following cell may be unable to complete with the default resource limits. You may contact [AWS Support](https://support.console.aws.amazon.com/support/home#/case/create?issueType=service-limit-increase) to increase the limits on your account."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "a14487eb",
   "metadata": {},
   "outputs": [],
   "source": [
    "import time\n",
    "from braket.aws import AwsQuantumJob\n",
    "\n",
    "job = AwsQuantumJob.create(\n",
    "    device=\"local:pennylane/lightning.gpu\",\n",
    "    source_module=\"qml_script\",\n",
    "    entry_point=\"qml_script.train_single\",\n",
    "    job_name=\"qml-single-\" + str(int(time.time())),\n",
    "    hyperparameters=hyperparameters,\n",
    "    input_data={\"input-data\": input_file_path},\n",
    "    instance_config=instance_config,\n",
    "    image_uri=image_uri,\n",
    "    wait_until_complete=False,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "863cb6b8",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'last loss': 0.09291739761829376}\n"
     ]
    }
   ],
   "source": [
    "# This cell should take about 7 minutes\n",
    "print(job.result())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bbdcb978",
   "metadata": {},
   "source": [
    "## Modify your algorithm script for data parallelism <a class=\"anchor\" id=\"modify\"></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0801d236",
   "metadata": {},
   "source": [
    "[PyTorch](https://pytorch.org/tutorials/intermediate/ddp_tutorial.html) has built-in features for data parallelism. With SageMaker's distributed data parallel library, Braket Hybrid Jobs makes it easier for you to leverage data parallelism to accelerate your training. To use data parallelism, you need to slightly modify your algorithm script. As an example, we modify the algorithm script [train_single.py](source_script/train_single.py) to [train_dp.py](source_script/train_dp.py). Let's go through the changes."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "af6e0428",
   "metadata": {},
   "source": [
    "First, we import the `smdistributed` package which does most of the heavy lifting for distributing your workloads across multiple GPUs and/or multiple instances. This package is pre-configured in the Braket PyTorch and TensorFlow containers. The `DDP` class from `smdistributed` converts the quantum model into a data parallelizable model. The `dist` module  tell our algorithm script the total number of GPUs for the training (`world_size`), and the `rank` and `local_rank` of a GPU. `rank` is the absolute index of a GPU across all instances, while `local_rank` is the index of a GPU within an instance. For example, if there are four instances each with eight GPUs allocated for the training, the `rank` ranges from 0 to 31 and the `local_rank` ranges from 0 to 7.\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "99921aa9",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "import smdistributed.dataparallel.torch.distributed as dist\n",
      "from smdistributed.dataparallel.torch.parallel.distributed import DistributedDataParallel as DDP\n",
      "    dp_info = {\n",
      "        \"world_size\": dist.get_world_size(),\n",
      "        \"rank\": dist.get_rank(),\n",
      "        \"local_rank\": dist.get_local_rank(),\n",
      "    }\n",
      "    batch_size //= dp_info[\"world_size\"] // 8\n",
      "    batch_size = max(batch_size, 1)\n",
      "    print(\"dp_info: \", dp_info)\n"
     ]
    }
   ],
   "source": [
    "!sed -n 22,23p qml_script/train_dp.py\n",
    "!sed -n 52,60p qml_script/train_dp.py"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d2abe53b",
   "metadata": {},
   "source": [
    "Next, we define a `DistributedSampler` according to the `world_size` and `rank`, and pass it into the data loader. This sampler avoids GPUs accessing the same slice of a dataset. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "bc7d8572",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "    train_sampler = torch.utils.data.distributed.DistributedSampler(\n",
      "        train_dataset, \n",
      "        num_replicas=dp_info[\"world_size\"], \n",
      "        rank=dp_info[\"rank\"]\n",
      "    )\n",
      "    train_loader = torch.utils.data.DataLoader(\n",
      "        train_dataset,\n",
      "        batch_size=batch_size,\n",
      "        shuffle=False,\n",
      "        num_workers=0,\n",
      "        pin_memory=True,\n",
      "        sampler=train_sampler,        \n",
      "    )\n"
     ]
    }
   ],
   "source": [
    "!sed -n 67,79p qml_script/train_dp.py"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5f5c75ce",
   "metadata": {},
   "source": [
    "Next, we use the `DDP` class to make our quantum model parallelizable. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "2a9eb130",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "    model = DressedQNN(qc_dev).to(device)\n",
      "    model = DDP(model)\n",
      "    torch.cuda.set_device(dp_info[\"local_rank\"])\n",
      "    model.cuda(dp_info[\"local_rank\"])\n"
     ]
    }
   ],
   "source": [
    "!sed -n 91,94p qml_script/train_dp.py"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f1c16d74",
   "metadata": {},
   "source": [
    "The above are the programming change you need to make to use data parallelism. In QML, we often want to save models, save results and print training progress. If each GPU executes the saving and printing command, the log would be flooded with the repeated information, and the model and results would be overwriting each other. To avoid this, we only save and print from the GPU that has `rank` 0. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "d06eda95",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "    if dp_info[\"rank\"] == 0:    \n",
      "        print(\"Training Finished!!\")\n",
      "        torch.save(model.state_dict(), f\"{output_dir}/test_local.pt\")\n",
      "        save_job_result({\"last loss\": float(loss_before.detach().cpu())})\n"
     ]
    }
   ],
   "source": [
    "!sed -n 118,121p qml_script/train_dp.py"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0bb63156",
   "metadata": {},
   "source": [
    "## Training with multiple GPUs in single instance\n",
    "With the modified algorithm, we can now submit our hybrid job with data parallelism. Amazon Braket Hybrid Jobs supports `ml.p3.16xlarge` for SageMaker distributed data parallel library. Be sure to choose this instance type from the list and configure it through the `InstanceConfig` argument in Jobs. See this [documentation](https://aws.amazon.com/ec2/instance-types/?tag=local002-20) for the specification, and see the [Amazon Braket pricing page](https://aws.amazon.com/braket/pricing/?tag=local002-20) for the cost of the instance type.\n",
    "\n",
    "For the SageMaker distributed data parallel library to know that data parallelism is enabled, we set the `distribution` argument to be `\"data_parallel\"` when creating a hybrid job. This argument triggers Braket Hybrid Jobs to add two additional hyperparameters, `\"sagemaker_distributed_dataparallel_enabled\"` setting to `\"true\"` and `\"sagemaker_instance_type\"` setting to the instance type we are using. These two hyperparameters are used by the `smdistributed` package at runtime. You do not need to explicitly call them in the algorithm script. Keep in mind that data parallelism only works correctly when you modify your algorithm script according to the [previous section](#modify). If the data parallelism option is enabled in the hyperparameters without a correctly modified algorithm script, the hybrid job may throw errors, or each GPU may repeat the same workload without data parallelism.\n",
    "\n",
    "<i>Warning: The p3.16xlarge instance has higher cost per minute. The cell below may incur charge up to $4. Run this cell only if you are comfortable with the charge. </i>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "62316645",
   "metadata": {},
   "outputs": [],
   "source": [
    "from braket.jobs.config import InstanceConfig\n",
    "from braket.aws import AwsSession\n",
    "from braket.jobs.image_uris import Framework, retrieve_image\n",
    "\n",
    "instance_config = InstanceConfig(instanceType='ml.p3.16xlarge')\n",
    "\n",
    "hyperparameters={\"nwires\": 10, \n",
    "                 \"ndata\": 64, \n",
    "                 \"batch_size\": 64, \n",
    "                 \"epochs\": 5, \n",
    "                 \"gamma\": 0.99, \n",
    "                 \"lr\": 0.1,\n",
    "                 \"seed\": 42,\n",
    "                }\n",
    "\n",
    "input_file_path = \"data/sonar.all-data\"\n",
    "\n",
    "image_uri = retrieve_image(Framework.PL_PYTORCH, AwsSession().region)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e707aeed",
   "metadata": {},
   "source": [
    "With the instance type and data parallelism configured, we can now submit our hybrid job. There are 8 GPUs in a `ml.p3.16xlarge` instance. When the instance spins up, the workload is distributed across the 8 GPUs."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6b9e9c31",
   "metadata": {},
   "source": [
    "**Note:** The following cell may be unable to complete with the default resource limits. You may contact [AWS Support](https://support.console.aws.amazon.com/support/home#/case/create?issueType=service-limit-increase) to increase the limits on your account."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "59d1a03e",
   "metadata": {},
   "outputs": [],
   "source": [
    "import time\n",
    "from braket.aws import AwsQuantumJob\n",
    "\n",
    "job = AwsQuantumJob.create(\n",
    "    device=\"local:pennylane/lightning.gpu\",\n",
    "    source_module=\"qml_script\",\n",
    "    entry_point=\"qml_script.train_dp\",\n",
    "    job_name=\"qml-dp1x-\" + str(int(time.time())),\n",
    "    hyperparameters=hyperparameters,\n",
    "    input_data={\"input-data\": input_file_path},\n",
    "    instance_config=instance_config,\n",
    "    distribution=\"data_parallel\",\n",
    "    image_uri=image_uri,\n",
    "    wait_until_complete=False,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "d261eba8",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'last loss': 0.1675658121705055}\n"
     ]
    }
   ],
   "source": [
    "# This cell should take about 7 minutes\n",
    "print(job.result())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5b89a298",
   "metadata": {},
   "source": [
    "## Training with multiple GPUs across multiple instances"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2fd12e6d",
   "metadata": {},
   "source": [
    "We are not limited to parallelizing the workload inside a single instance. We can perform distributed training by parallelizing our workload across multiple instances. With the algorithm script modified and the data parallelism enabled, we can perform distributed data parallelism by setting instance count larger than 1. To configure instance count, we use the `instanceCount` argument in `InstanceConfig`. The SageMaker distributed library we included in our algorithm script coordinates the multiple instances and conducts the distributed training for us. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "09e245d8",
   "metadata": {},
   "outputs": [],
   "source": [
    "instance_config = InstanceConfig(instanceType='ml.p3.16xlarge', instanceCount=2)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "687b4b67",
   "metadata": {},
   "source": [
    "Be mindful that, when using multiple instances, each instance incurs charge based on how long you use it. In distributed data parallelism, when you set `instanceCount=2`, two instances are allocated to run your hybrid job. SageMaker distributed library managed the instances that they start and end at the same time. If your workload takes 200 seconds, you will be billed for 200 seconds for each instance used, which adds to 400 seconds in total.\n",
    "\n",
    "<i> Warning: The p3.16xlarge instance has higher cost per minute. The cell below may incur charge up to $8. Run this cell only if you are comfortable with the charge. </i>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "c3da81b2",
   "metadata": {},
   "outputs": [],
   "source": [
    "from braket.jobs.config import InstanceConfig\n",
    "from braket.aws import AwsSession\n",
    "from braket.jobs.image_uris import Framework, retrieve_image\n",
    "\n",
    "hyperparameters={\"nwires\": 10, \n",
    "                 \"ndata\": 64, \n",
    "                 \"batch_size\": 64, \n",
    "                 \"epochs\": 5, \n",
    "                 \"gamma\": 0.99, \n",
    "                 \"lr\": 0.1,\n",
    "                 \"seed\": 42,\n",
    "                }\n",
    "\n",
    "input_file_path = \"data/sonar.all-data\"\n",
    "\n",
    "image_uri = retrieve_image(Framework.PL_PYTORCH, AwsSession().region)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "20e83e99",
   "metadata": {},
   "source": [
    "Now we can submit our hybrid job with distributed data parallelism!"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0842096f",
   "metadata": {},
   "source": [
    "**Note:** The following cell may be unable to complete with the default resource limits. You may contact [AWS Support](https://support.console.aws.amazon.com/support/home#/case/create?issueType=service-limit-increase) to increase the limits on your account."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "fe83fd35",
   "metadata": {},
   "outputs": [],
   "source": [
    "import time\n",
    "from braket.aws import AwsQuantumJob\n",
    "\n",
    "job = AwsQuantumJob.create(\n",
    "    device=\"local:pennylane/lightning.gpu\",\n",
    "    source_module=\"qml_script\",\n",
    "    entry_point=\"qml_script.train_dp\",\n",
    "    job_name=\"qml-dp2x-\" + str(int(time.time())),\n",
    "    hyperparameters=hyperparameters,\n",
    "    input_data={\"input-data\": input_file_path},\n",
    "    instance_config=instance_config,\n",
    "    distribution=\"data_parallel\",\n",
    "    image_uri=image_uri,\n",
    "    wait_until_complete=False,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "ed0ead17",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'last loss': 0.26321490332484243}\n"
     ]
    }
   ],
   "source": [
    "# This cell should take about 7 minutes\n",
    "print(job.result())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "54dc54f4",
   "metadata": {},
   "source": [
    "## Summary\n",
    "In this notebook, we show you how to use SageMaker distributed library to parallelize quantum machine learning workloads. To learn more about distributed training, you can read the [Amazon Braket documentation](https://docs.aws.amazon.com/braket/latest/developerguide/braket-jobs.html) and [Amazon SageMaker Distributed Training Libraries](https://docs.aws.amazon.com/sagemaker/latest/dg/distributed-training.html?tag=local002-20)."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "conda_braket",
   "language": "python",
   "name": "conda_braket"
  },
  "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.7.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}