examples/nvidia_cuda_q/3_distributed_statevector_simulations.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "021cfd04-df68-414d-a864-48f62fc8ddfb",
   "metadata": {},
   "source": [
    "# Distributed state vector simulations on multiple GPUs (advanced)\n",
    "\n",
    "In the notebook \"2_parallel_simulations.ipynb\", you learned how to use CUDA-Q and Braket Hybrid Jobs to parallelize the simulation of a batch of observables and circuits over multiple GPUs, where each GPU simulates a single QPU. For workloads with larger qubit counts, however, it may be necessary to distribute a single state vector simulation across multiple GPUs, so that multiple GPUs together simulate a single QPU.\n",
    "\n",
    "In this notebook, you will learn how to use CUDA-Q and Braket Hybrid Jobs to tackle this."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "32b46659-6dcc-4900-a13a-e971f8bf0590",
   "metadata": {},
   "source": [
    "We start with necessary imports that are used in the examples below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8738f65f-969c-4b58-96f8-69bbc1bad5e1",
   "metadata": {},
   "outputs": [],
   "source": [
    "from braket.jobs import hybrid_job\n",
    "from braket.jobs.config import InstanceConfig"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1331136a-a369-4eef-8bfe-252c79103a3e",
   "metadata": {},
   "source": [
    "Next, we need to create and upload a container which contains both CUDA-Q and the underlying CUDA support required for distributing our computation across multiple GPUs. Note: this container image will be different than the one used in the previous notebooks illustrating more basic CUDA-Q scenarios.\n",
    "\n",
    "To do this, we need to run the commands in the cell below. (For more information about what these commands are doing, please see the detailed documentation in \"0_hello_cudaq_jobs.ipynb\". The difference here is that we specify the dockerfile `Dockerfile.mgpu` in order to ensure full support for this advanced scenario.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f552c738",
   "metadata": {},
   "outputs": [],
   "source": [
    "!chmod +x container/container_build_and_push.sh\n",
    "!container/container_build_and_push.sh cudaq-mgpu-job us-west-2 Dockerfile.mgpu"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fc46e446",
   "metadata": {},
   "source": [
    "Now we prepare the URI of the container image. Fill the proper value of `aws_account_id`, `region_name` and `container_image_name` in the cell below. For example, with the shell command above, `region_name=\"us-west-2\"` and `container_image_name=\"cudaq-mgpu-job\"`. The cell below prints out the image URI. When you use a container image to run a job, it ensures that your code is run in the same environment every time. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "25fdc720-3143-411a-8bef-f9623369b516",
   "metadata": {},
   "outputs": [],
   "source": [
    "aws_account_id = \"<aws-account-id>\"\n",
    "region_name = \"<region-name>\"\n",
    "container_image_name = \"<container-image-name>\"\n",
    "\n",
    "image_uri = f\"{aws_account_id}.dkr.ecr.{region_name}.amazonaws.com/{container_image_name}:latest\"\n",
    "print(image_uri)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a844344d-0978-4b11-8fe4-66b387e80c72",
   "metadata": {},
   "source": [
    "## Distributed state vector simulations\n",
    "Now that we have the container image URI, we are ready to run our workload. The `nvidia` target with `mgpu` option supports distributing state vector simulations to multiple GPUs. This enables GPU simulations for circuits with higher qubit count, to up to 34 qubits. The example below shows how to submit a job with the `mgpu` option."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d38b73de-6a7e-45b7-b64b-afec60c0a6c3",
   "metadata": {},
   "outputs": [],
   "source": [
    "@hybrid_job(\n",
    "    device=\"local:nvidia/nvidia-mgpu\",\n",
    "    instance_config=InstanceConfig(instanceType=\"ml.p3.8xlarge\", instanceCount=1),\n",
    "    image_uri=image_uri,\n",
    ")\n",
    "def distributed_gpu_job(\n",
    "    n_qubits,\n",
    "    n_shots,\n",
    "    sagemaker_mpi_enabled=True,\n",
    "):\n",
    "    import cudaq\n",
    "\n",
    "    # Define target\n",
    "    cudaq.set_target(\"nvidia\", option=\"mgpu\")\n",
    "    print(\"CUDA-Q backend: \", cudaq.get_target())\n",
    "    print(\"num_available_gpus: \", cudaq.num_available_gpus())\n",
    "\n",
    "    # Initialize MPI and view the MPI properties\n",
    "    cudaq.mpi.initialize()\n",
    "    rank = cudaq.mpi.rank()\n",
    "\n",
    "    # Define circuit and observables\n",
    "    @cudaq.kernel\n",
    "    def ghz():\n",
    "        qubits = cudaq.qvector(n_qubits)\n",
    "        h(qubits[0])\n",
    "        for q in range(1, n_qubits):\n",
    "            cx(qubits[0], qubits[q])\n",
    "\n",
    "    hamiltonian = cudaq.SpinOperator.random(n_qubits, 1)\n",
    "\n",
    "    # Parallelize circuit simulation\n",
    "    result = cudaq.observe(ghz, hamiltonian, shots_count=n_shots)\n",
    "\n",
    "    # End the MPI interface\n",
    "    cudaq.mpi.finalize()\n",
    "\n",
    "    if rank == 0:\n",
    "        return {\"expectation\": result.expectation()}\n",
    "\n",
    "\n",
    "n_qubits = 25\n",
    "n_shots = 1000\n",
    "distributed_job = distributed_gpu_job(n_qubits, n_shots)\n",
    "print(\"Job ARN: \", distributed_job.arn)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a054c24e",
   "metadata": {},
   "outputs": [],
   "source": [
    "distributed_job_result = distributed_job.result()\n",
    "print(f\"result: {distributed_job_result['expectation']}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "18004fe8-24a0-4316-ab0a-a4e0aac6ba1e",
   "metadata": {},
   "source": [
    "## Summary\n",
    "This notebook shows you how to distribute a single state vector simulation across multiple GPUs, so that multiple GPUs together simulate a single QPU. If you have workloads with a qubit count that is too large to simulate on a single GPU, you can use this technique to make these large workloads feasible."
   ]
  }
 ],
 "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.10.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}