 archived/GPT-J-6B-model-parallel-inference-DJL.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "bc5ab391",
   "metadata": {},
   "source": [
    "# Serve large models on SageMaker with model parallel inference and DJLServing"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "\n",
    "This notebook's CI test result for us-west-2 is as follows. CI test results in other regions can be found at the end of the notebook. \n",
    "\n",
    "![This us-west-2 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://prod.us-west-2.tcx-beacon.docs.aws.dev/sagemaker-nb/us-west-2/advanced_functionality|pytorch_deploy_large_GPT_model|GPT-J-6B-model-parallel-inference-DJL.ipynb)\n",
    "\n",
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "694d9185",
   "metadata": {},
   "source": [
    "> **\u26a0 Notes:** For latest version, see: https://github.com/aws/amazon-sagemaker-examples/blob/main/inference/generativeai/deepspeed/GPT-J-6B_DJLServing_with_PySDK.ipynb"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "98b43ca5",
   "metadata": {},
   "source": [
    "In this notebook, we explore how to host a large language model on SageMaker using model parallelism from DeepSpeed and DJLServing.\n",
    "\n",
    "Language models have recently exploded in both size and popularity. In 2018, BERT-large entered the scene and, with its 340M parameters and novel transformer architecture, set the standard on NLP task accuracy. Within just a few years, state-of-the-art NLP model size has grown by more than 500x with models such as OpenAI\u2019s 175 billion parameter GPT-3 and similarly sized open source Bloom 176B raising the bar on NLP accuracy. This increase in the number of parameters is driven by the simple and empirically-demonstrated positive relationship between model size and accuracy: more is better. With easy access from models zoos such as Hugging Face and improved accuracy in NLP tasks such as classification and text generation, practitioners are increasingly reaching for these large models. However, deploying them can be a challenge because of their size.\n",
    "\n",
    "Model parallelism can help deploy large models that would normally be too large for a single GPU. With model parallelism, we partition and distribute a model across multiple GPUs. Each GPU holds a different part of the model, resolving the memory capacity issue for the largest deep learning models with billions of parameters. This notebook uses tensor parallelism techniques which allow GPUs to work simultaneously on the same layer of a model and achieve low latency inference relative to a pipeline parallel solution.\n",
    "\n",
    "In this notebook, we deploy a PyTorch GPT-J model from Hugging Face with 6 billion parameters across two GPUs on an Amazon SageMaker ml.g5.48xlarge instance. DeepSpeed is used for tensor parallelism inference while DJLServing handles inference requests and the distributed workers. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d6ed354b",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%bash\n",
    "pip install -U pip --quiet\n",
    "pip install -U sagemaker --quiet\n",
    "pip install -U boto3 --quiet\n",
    "\n",
    "pip install -U transformers --quiet"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "81c2bdf4",
   "metadata": {},
   "source": [
    "## Step 1: Get DLC image URL for SageMaker endpoint\n",
    "We get DLC image URL for djl-deepspeed 0.21.0 and set SageMaker settings"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2876d11c",
   "metadata": {},
   "outputs": [],
   "source": [
    "import sagemaker, boto3\n",
    "from sagemaker import image_uris\n",
    "\n",
    "instance_type = \"ml.g5.12xlarge\"\n",
    "\n",
    "role = sagemaker.get_execution_role()  # execution role for the endpoint\n",
    "session = sagemaker.session.Session()  # sagemaker session for interacting with different AWS APIs\n",
    "region = session._region_name\n",
    "bucket = session.default_bucket()  # bucket to house artifacts\n",
    "\n",
    "img_uri = image_uris.retrieve(framework=\"djl-deepspeed\", region=region, version=\"0.21.0\")\n",
    "img_uri"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "47ab31d1",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "s3_location = f\"s3://{bucket}/djl-serving/\"\n",
    "s3_location"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1ac32e96",
   "metadata": {},
   "source": [
    "## Step 2: Create a `model.py` and `serving.properties`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6f4864eb",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%writefile model.py\n",
    "\n",
    "from djl_python import Input, Output\n",
    "import os\n",
    "import deepspeed\n",
    "import torch\n",
    "from transformers import pipeline, AutoModelForCausalLM, AutoTokenizer\n",
    "\n",
    "predictor = None\n",
    "\n",
    "\n",
    "def get_model(properties):\n",
    "    model_name = \"EleutherAI/gpt-j-6B\"\n",
    "    tensor_parallel = properties[\"tensor_parallel_degree\"]\n",
    "    local_rank = int(os.getenv(\"LOCAL_RANK\", \"0\"))\n",
    "    model = AutoModelForCausalLM.from_pretrained(\n",
    "        model_name, revision=\"float32\", torch_dtype=torch.float32\n",
    "    )\n",
    "    tokenizer = AutoTokenizer.from_pretrained(model_name)\n",
    "\n",
    "    model = deepspeed.init_inference(\n",
    "        model,\n",
    "        mp_size=tensor_parallel,\n",
    "        dtype=model.dtype,\n",
    "        replace_method=\"auto\",\n",
    "        replace_with_kernel_inject=True,\n",
    "    )\n",
    "    generator = pipeline(\n",
    "        task=\"text-generation\", model=model, tokenizer=tokenizer, device=local_rank\n",
    "    )\n",
    "    return generator\n",
    "\n",
    "\n",
    "def handle(inputs: Input) -> None:\n",
    "    global predictor\n",
    "    if not predictor:\n",
    "        predictor = get_model(inputs.get_properties())\n",
    "\n",
    "    if inputs.is_empty():\n",
    "        # Model server makes an empty call to warmup the model on startup\n",
    "        return None\n",
    "\n",
    "    data = inputs.get_as_string()\n",
    "    result = predictor(data, do_sample=True, max_new_tokens=256)\n",
    "    return Output().add(result)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f02b6929",
   "metadata": {},
   "source": [
    "### Setup serving.properties\n",
    "\n",
    "User needs to specify DeepSpeed as the engine, as shown below. Use `option.tensor_parallel_degree` to specify number partitions to split the model.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2c5ea96a",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%writefile serving.properties\n",
    "engine = DeepSpeed\n",
    "option.tensor_parallel_degree = 2"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f44488e6",
   "metadata": {},
   "source": [
    "The code below creates the SageMaker model file (`model.tar.gz`) and upload it to S3. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9965dd7c",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%sh\n",
    "if [ -d gpt-j ]; then\n",
    "  rm -d -r gpt-j\n",
    "fi #always start fresh\n",
    "\n",
    "mkdir -p gpt-j\n",
    "mv model.py gpt-j\n",
    "mv serving.properties gpt-j\n",
    "tar -czvf gpt-j.tar.gz gpt-j/\n",
    "#aws s3 cp gpt-j.tar.gz {path}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "db47f969",
   "metadata": {},
   "outputs": [],
   "source": [
    "model_tar_url = sagemaker.s3.S3Uploader.upload(\"gpt-j.tar.gz\", s3_location)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c507e3ef",
   "metadata": {},
   "source": [
    "## Step 3: Create SageMaker endpoint"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "32589338",
   "metadata": {},
   "source": [
    "Now we create our [SageMaker model](https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/sagemaker.html#SageMaker.Client.create_model). Make sure your execution role has access to your model artifacts and ECR image. Please check out our SageMaker Roles [documentation](https://docs.aws.amazon.com/sagemaker/latest/dg/sagemaker-roles.html) for more details. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "026d27d2",
   "metadata": {},
   "outputs": [],
   "source": [
    "from datetime import datetime\n",
    "\n",
    "sm_client = boto3.client(\"sagemaker\")\n",
    "\n",
    "time_stamp = datetime.now().strftime(\"%Y-%m-%d-%H-%M-%S\")\n",
    "model_name = \"gpt-j-\" + time_stamp\n",
    "\n",
    "create_model_response = sm_client.create_model(\n",
    "    ModelName=model_name,\n",
    "    ExecutionRoleArn=role,\n",
    "    PrimaryContainer={\"Image\": img_uri, \"ModelDataUrl\": model_tar_url},\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "22d2fc2b",
   "metadata": {},
   "source": [
    "Now we create an endpoint configuration that SageMaker hosting services uses to deploy models. Note that we configured `ModelDataDownloadTimeoutInSeconds` and `ContainerStartupHealthCheckTimeoutInSeconds` to accommodate the large size of our model. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "84e25dd4",
   "metadata": {},
   "outputs": [],
   "source": [
    "initial_instance_count = 1\n",
    "variant_name = \"AllTraffic\"\n",
    "endpoint_config_name = \"t-j-config-\" + time_stamp\n",
    "\n",
    "production_variants = [\n",
    "    {\n",
    "        \"VariantName\": variant_name,\n",
    "        \"ModelName\": model_name,\n",
    "        \"InitialInstanceCount\": initial_instance_count,\n",
    "        \"InstanceType\": instance_type,\n",
    "        \"ModelDataDownloadTimeoutInSeconds\": 1800,\n",
    "        \"ContainerStartupHealthCheckTimeoutInSeconds\": 3600,\n",
    "    }\n",
    "]\n",
    "\n",
    "endpoint_config = {\n",
    "    \"EndpointConfigName\": endpoint_config_name,\n",
    "    \"ProductionVariants\": production_variants,\n",
    "}\n",
    "\n",
    "ep_conf_res = sm_client.create_endpoint_config(**endpoint_config)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e4b3bc26",
   "metadata": {},
   "source": [
    "We are ready to create an endpoint using the model and the endpoint configuration created from above steps. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "962a1aef",
   "metadata": {},
   "outputs": [],
   "source": [
    "endpoint_name = \"gpt-j\" + time_stamp\n",
    "ep_res = sm_client.create_endpoint(\n",
    "    EndpointName=endpoint_name, EndpointConfigName=endpoint_config_name\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2dc2a85a",
   "metadata": {},
   "source": [
    "The creation of the SageMaker endpoint might take a while. After the endpoint is created, you can test it out using the following code. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5ed7a325",
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "\n",
    "client = boto3.client(\"sagemaker-runtime\")\n",
    "\n",
    "content_type = \"text/plain\"  # The MIME type of the input data in the request body.\n",
    "payload = \"Amazon.com is the best\"  # Payload for inference.\n",
    "response = client.invoke_endpoint(\n",
    "    EndpointName=endpoint_name, ContentType=content_type, Body=payload\n",
    ")\n",
    "print(response[\"Body\"].read())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "92e83c91",
   "metadata": {},
   "source": [
    "## Step 4: Clean up"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a15980a3",
   "metadata": {},
   "outputs": [],
   "source": [
    "sm_client.delete_endpoint(EndpointName=endpoint_name)\n",
    "sm_client.delete_endpoint_config(EndpointConfigName=endpoint_config_name)\n",
    "sm_client.delete_model(ModelName=model_name)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2eff050b",
   "metadata": {},
   "source": [
    "## Conclusion\n",
    "\n",
    "In this notebook, you use tensor parallelism to partition a large language model across multiple GPUs for low latency inference. With tensor parallelism, multiple GPUs work on the same model layer at once allowing for faster inference latency when a low batch size is used. Here, we use open source DeepSpeed as the model parallel library to partition the model and open source Deep Java Library Serving as the model serving solution.\n",
    "\n",
    "As a next step, you can experiment with larger models from Hugging Face such as GPT-NeoX. You can also adjust the tensor parallel degree to see the impact to latency with models of different sizes."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Notebook CI Test Results\n",
    "\n",
    "This notebook was tested in multiple regions. The test results are as follows, except for us-west-2 which is shown at the top of the notebook.\n",
    "\n",
    "![This us-east-1 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://prod.us-west-2.tcx-beacon.docs.aws.dev/sagemaker-nb/us-east-1/advanced_functionality|pytorch_deploy_large_GPT_model|GPT-J-6B-model-parallel-inference-DJL.ipynb)\n",
    "\n",
    "![This us-east-2 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://prod.us-west-2.tcx-beacon.docs.aws.dev/sagemaker-nb/us-east-2/advanced_functionality|pytorch_deploy_large_GPT_model|GPT-J-6B-model-parallel-inference-DJL.ipynb)\n",
    "\n",
    "![This us-west-1 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://prod.us-west-2.tcx-beacon.docs.aws.dev/sagemaker-nb/us-west-1/advanced_functionality|pytorch_deploy_large_GPT_model|GPT-J-6B-model-parallel-inference-DJL.ipynb)\n",
    "\n",
    "![This ca-central-1 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://prod.us-west-2.tcx-beacon.docs.aws.dev/sagemaker-nb/ca-central-1/advanced_functionality|pytorch_deploy_large_GPT_model|GPT-J-6B-model-parallel-inference-DJL.ipynb)\n",
    "\n",
    "![This sa-east-1 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://prod.us-west-2.tcx-beacon.docs.aws.dev/sagemaker-nb/sa-east-1/advanced_functionality|pytorch_deploy_large_GPT_model|GPT-J-6B-model-parallel-inference-DJL.ipynb)\n",
    "\n",
    "![This eu-west-1 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://prod.us-west-2.tcx-beacon.docs.aws.dev/sagemaker-nb/eu-west-1/advanced_functionality|pytorch_deploy_large_GPT_model|GPT-J-6B-model-parallel-inference-DJL.ipynb)\n",
    "\n",
    "![This eu-west-2 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://prod.us-west-2.tcx-beacon.docs.aws.dev/sagemaker-nb/eu-west-2/advanced_functionality|pytorch_deploy_large_GPT_model|GPT-J-6B-model-parallel-inference-DJL.ipynb)\n",
    "\n",
    "![This eu-west-3 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://prod.us-west-2.tcx-beacon.docs.aws.dev/sagemaker-nb/eu-west-3/advanced_functionality|pytorch_deploy_large_GPT_model|GPT-J-6B-model-parallel-inference-DJL.ipynb)\n",
    "\n",
    "![This eu-central-1 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://prod.us-west-2.tcx-beacon.docs.aws.dev/sagemaker-nb/eu-central-1/advanced_functionality|pytorch_deploy_large_GPT_model|GPT-J-6B-model-parallel-inference-DJL.ipynb)\n",
    "\n",
    "![This eu-north-1 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://prod.us-west-2.tcx-beacon.docs.aws.dev/sagemaker-nb/eu-north-1/advanced_functionality|pytorch_deploy_large_GPT_model|GPT-J-6B-model-parallel-inference-DJL.ipynb)\n",
    "\n",
    "![This ap-southeast-1 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://prod.us-west-2.tcx-beacon.docs.aws.dev/sagemaker-nb/ap-southeast-1/advanced_functionality|pytorch_deploy_large_GPT_model|GPT-J-6B-model-parallel-inference-DJL.ipynb)\n",
    "\n",
    "![This ap-southeast-2 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://prod.us-west-2.tcx-beacon.docs.aws.dev/sagemaker-nb/ap-southeast-2/advanced_functionality|pytorch_deploy_large_GPT_model|GPT-J-6B-model-parallel-inference-DJL.ipynb)\n",
    "\n",
    "![This ap-northeast-1 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://prod.us-west-2.tcx-beacon.docs.aws.dev/sagemaker-nb/ap-northeast-1/advanced_functionality|pytorch_deploy_large_GPT_model|GPT-J-6B-model-parallel-inference-DJL.ipynb)\n",
    "\n",
    "![This ap-northeast-2 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://prod.us-west-2.tcx-beacon.docs.aws.dev/sagemaker-nb/ap-northeast-2/advanced_functionality|pytorch_deploy_large_GPT_model|GPT-J-6B-model-parallel-inference-DJL.ipynb)\n",
    "\n",
    "![This ap-south-1 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://prod.us-west-2.tcx-beacon.docs.aws.dev/sagemaker-nb/ap-south-1/advanced_functionality|pytorch_deploy_large_GPT_model|GPT-J-6B-model-parallel-inference-DJL.ipynb)\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "conda_pytorch_p39",
   "language": "python",
   "name": "conda_pytorch_p39"
  },
  "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.15"
  },
  "vscode": {
   "interpreter": {
    "hash": "31f2aee4e71d21fbe5cf8b01ff0e069b9275f58929596ceb00d14d90e3e16cd6"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}