vllm_deploy_qwen1.5_72B_chat_AWQ.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "71a329f0",
   "metadata": {},
   "source": [
    "# vLLM Qwen1.5-72B-Chat-AWQ vLLM deployment guide\n",
    "In this tutorial, you will use LMI container from DLC to SageMaker and run inference with it.\n",
    "\n",
    "Please make sure the following permission granted before running the notebook:\n",
    "\n",
    "- S3 bucket push access\n",
    "- SageMaker access\n",
    "\n",
    "## Step 1: Let's bump up SageMaker and import stuff"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "67fa3208",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "%pip install sagemaker --upgrade  --quiet\n",
    "%pip install boto3==1.34.101"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ec9ac353",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import boto3\n",
    "import sagemaker\n",
    "from sagemaker import Model, image_uris, serializers, deserializers\n",
    "\n",
    "role = sagemaker.get_execution_role()  # execution role for the endpoint\n",
    "sess = sagemaker.session.Session()  # sagemaker session for interacting with different AWS APIs\n",
    "region = sess._region_name  # region name of the current SageMaker Studio environment\n",
    "account_id = sess.account_id()  # account_id of the current SageMaker Studio environment"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "81deac79",
   "metadata": {},
   "source": [
    "## Step 2: Start preparing model artifacts\n",
    "In LMI contianer, we expect some artifacts to help setting up the model\n",
    "- serving.properties (required): Defines the model server settings\n",
    "- model.py (optional): A python file to define the core inference logic\n",
    "- requirements.txt (optional): Any additional pip wheel need to install"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b011bf5f",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "%%writefile serving.properties\n",
    "engine=Python\n",
    "option.model_id=Qwen/Qwen1.5-72B-Chat-AWQ\n",
    "option.task=text-generation\n",
    "option.trust_remote_code=true\n",
    "option.tensor_parallel_degree=4\n",
    "option.rolling_batch=vllm\n",
    "option.quantize=awq\n",
    "option.dtype=fp16\n",
    "option.max_model_len=10272"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b0142973",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "%%sh\n",
    "mkdir mymodel\n",
    "mv serving.properties mymodel/\n",
    "tar czvf mymodel.tar.gz mymodel/\n",
    "rm -rf mymodel"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2e58cf33",
   "metadata": {},
   "source": [
    "## Step 3: Start building SageMaker endpoint\n",
    "In this step, we will build SageMaker endpoint from scratch"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4d955679",
   "metadata": {},
   "source": [
    "### Getting the container image URI\n",
    "\n",
    "[Large Model Inference available DLC](https://github.com/aws/deep-learning-containers/blob/master/available_images.md#large-model-inference-containers)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7a174b36",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "image_uri = image_uris.retrieve(\n",
    "        framework=\"djl-deepspeed\",\n",
    "        region=sess.boto_session.region_name,\n",
    "        version=\"0.27.0\"\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "11601839",
   "metadata": {},
   "source": [
    "### Upload artifact on S3 and create SageMaker model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "38b1e5ca",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "s3_code_prefix = \"large-model-lmi/code-qwen1.5-72B\"\n",
    "bucket = sess.default_bucket()  # bucket to house artifacts\n",
    "code_artifact = sess.upload_data(\"mymodel.tar.gz\", bucket, s3_code_prefix)\n",
    "print(f\"S3 Code or Model tar ball uploaded to --- > {code_artifact}\")\n",
    "\n",
    "model = Model(image_uri=image_uri, model_data=code_artifact, role=role)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "004f39f6",
   "metadata": {},
   "source": [
    "### 4.2 Create SageMaker endpoint\n",
    "\n",
    "You need to specify the instance to use and endpoint names"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8e0e61cd",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "instance_type = \"ml.g5.12xlarge\"\n",
    "endpoint_name = sagemaker.utils.name_from_base(\"lmi-model-qwen1-5-72B\")\n",
    "\n",
    "model.deploy(initial_instance_count=1,\n",
    "             instance_type=instance_type,\n",
    "             endpoint_name=endpoint_name,\n",
    "             # container_startup_health_check_timeout=3600\n",
    "            )\n",
    "\n",
    "# our requests and responses will be in json format so we specify the serializer and the deserializer\n",
    "predictor = sagemaker.Predictor(\n",
    "    endpoint_name=endpoint_name,\n",
    "    sagemaker_session=sess,\n",
    "    serializer=serializers.JSONSerializer(),\n",
    "    deserializer=sagemaker.deserializers.JSONDeserializer(),\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bb63ee65",
   "metadata": {},
   "source": [
    "## Step 5: Test and benchmark the inference"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4ee11ce2-c488-42c5-89eb-df453878c245",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "!pip install transformers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aaae61f7-b669-4a02-97a9-85085db6632c",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from transformers import AutoTokenizer\n",
    "\n",
    "MODEL_DIR = \"Qwen/Qwen1.5-72B-Chat-AWQ\"\n",
    "# model = AutoModelForCausalLM.from_pretrained(MODEL_DIR, torch_dtype=\"auto\")\n",
    "tokenizer = AutoTokenizer.from_pretrained(MODEL_DIR, use_fast=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2bcef095",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "prompt = \"世界上第二高峰是哪座？\"\n",
    "messages = [\n",
    "    {\"role\": \"system\", \"content\": \"You are a helpful assistant.\"},\n",
    "    {\"role\": \"user\", \"content\": prompt}\n",
    "]\n",
    "inputs = tokenizer.apply_chat_template(\n",
    "    messages,\n",
    "    tokenize=False,\n",
    "    add_generation_prompt=True\n",
    ")\n",
    "parameters = {\n",
    "        \"max_new_tokens\":1024, \n",
    "        \"do_sample\": True,\n",
    "        \"stop_token_ids\":[151645,151643],\n",
    "        \"repetition_penalty\": 1.05,\n",
    "        \"temperature\": 0.7,\n",
    "        \"top_p\": 0.8,\n",
    "        \"top_k\": 20\n",
    "    }\n",
    "response = predictor.predict(\n",
    "    {\"inputs\": inputs, \"parameters\": parameters}\n",
    ")\n",
    "# text = str(response, \"utf-8\")\n",
    "print(response)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "95732e49-db07-4de2-ae1d-8ed96e95f056",
   "metadata": {},
   "source": [
    "# Streaming"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "11c67435-fc81-42ae-85c5-608a9eb0e366",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import json\n",
    "import boto3\n",
    "\n",
    "smr_client = boto3.client(\"sagemaker-runtime\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2cae82c8-c390-4d70-ac18-f1d71c852727",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import io\n",
    "import json\n",
    "\n",
    "class TokenIterator:\n",
    "    def __init__(self, stream):\n",
    "        self.byte_iterator = iter(stream)\n",
    "        self.buffer = io.BytesIO()\n",
    "        self.read_pos = 0\n",
    "\n",
    "    def __iter__(self):\n",
    "        return self\n",
    "\n",
    "    def __next__(self):\n",
    "        while True:\n",
    "            self.buffer.seek(self.read_pos)\n",
    "            line = self.buffer.readline()\n",
    "            \n",
    "            # print(line)\n",
    "            if line and line[-1] == ord(\"\\n\"):\n",
    "                self.read_pos += len(line)\n",
    "                full_line = line[:-1].decode(\"utf-8\")\n",
    "                # print(full_line)\n",
    "                line_data = json.loads(full_line.lstrip(\"data:\").rstrip(\"/n\"))\n",
    "                return line_data[\"token\"].get(\"text\", \"\")\n",
    "            chunk = next(self.byte_iterator)\n",
    "            self.buffer.seek(0, io.SEEK_END)\n",
    "            self.buffer.write(chunk[\"PayloadPart\"][\"Bytes\"])\n",
    "        \n",
    "def get_realtime_response_stream(sagemaker_runtime, endpoint_name, payload):\n",
    "    response_stream = sagemaker_runtime.invoke_endpoint_with_response_stream(\n",
    "        EndpointName=endpoint_name,\n",
    "        Body=json.dumps(payload),\n",
    "        ContentType=\"application/json\",\n",
    "        CustomAttributes='accept_eula=false'\n",
    "    )\n",
    "    return response_stream"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c16918d8-35d9-4b56-ad50-73f7ab858a1e",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "prompt = \"世界上第二高峰是哪座？\"\n",
    "messages = [\n",
    "    {\"role\": \"system\", \"content\": \"You are a helpful assistant.\"},\n",
    "    {\"role\": \"user\", \"content\": prompt}\n",
    "]\n",
    "inputs = tokenizer.apply_chat_template(\n",
    "    messages,\n",
    "    tokenize=False,\n",
    "    add_generation_prompt=True\n",
    ")\n",
    "\n",
    "parameters = {\n",
    "        \"max_new_tokens\":1024, \n",
    "        \"do_sample\": True,\n",
    "        \"stop_token_ids\":[151645,151643],\n",
    "        \"repetition_penalty\": 1.05,\n",
    "        \"temperature\": 0.7,\n",
    "        \"top_p\": 0.8,\n",
    "        \"top_k\": 20,\n",
    "    }\n",
    "\n",
    "payload = {\n",
    "    \"inputs\":  inputs,\n",
    "    \"parameters\": parameters,\n",
    "    \"stream\": True ## <-- to have response stream.\n",
    "}\n",
    "response_stream = get_realtime_response_stream(smr_client, endpoint_name, payload)\n",
    "# print_response_stream(response_stream)\n",
    "for token in TokenIterator(response_stream[\"Body\"]):\n",
    "    # pass\n",
    "    print(token, end=\"\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c1cd9042",
   "metadata": {},
   "source": [
    "## Clean up the environment"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3d674b41",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "sess.delete_endpoint(endpoint_name)\n",
    "sess.delete_endpoint_config(endpoint_name)\n",
    "model.delete_model()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "conda_pytorch_p310",
   "language": "python",
   "name": "conda_pytorch_p310"
  },
  "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.14"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}