tensorrt_0.7.1_aot_chatglm3_6B_32K.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "6c518e46-1464-4cc9-83ca-d992364ae438",
   "metadata": {},
   "source": [
    "# AOT 方式使用 Tensorrt-LMI 部署 Chatglm3-6b-32"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2b5ed753-95bb-411e-a132-38f4002531f4",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "%pip install sagemaker --upgrade  --quiet"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "925d01e1-1153-4005-9dcf-5cca9a3f3f2d",
   "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": "0fa2eb33-d9aa-4249-9063-05903414f5c7",
   "metadata": {},
   "source": [
    "## Pull 转换模型时需要的镜像"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "75da62b6-454f-4a51-ae0e-c26d7b970435",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "!aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/s0w3f1p2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a5837f0c-df31-4a04-90a7-c0c10d4e57e0",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "!docker pull public.ecr.aws/s0w3f1p2/tensorrt-lmi-xq:v1"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8ba2e611-604f-45f4-9547-0272eae2e939",
   "metadata": {
    "tags": []
   },
   "source": [
    "## 转换模型"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cbbb0977-5c9d-42be-b1c0-b08a3dcbbe25",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "%%writefile serving.properties\n",
    "option.model_id=THUDM/chatglm3-6b-32k\n",
    "option.tensor_parallel_degree=2\n",
    "option.max_rolling_batch_size=64\n",
    "option.dtype=fp16\n",
    "option.chatglm_model_version=chatglm3_6b_32k\n",
    "option.trust_remote_code=True"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d5c72d7a-c8a0-4dd7-83f6-569f5cd353eb",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "!cat serving.properties"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d49b3e4b-b5af-4035-9acc-68817f5f47bb",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "file_name = \"chatglm3_6b_32k_aot\"\n",
    "MODEL_REPO_DIR=f\"output/{file_name}\"\n",
    "!mkdir -p $MODEL_REPO_DIR\n",
    "!mv serving.properties $MODEL_REPO_DIR"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8ecf8687-64a1-46e0-8605-dc2488f3a45b",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "bucket=\"llm-trt\"\n",
    "s3_model_prefix = f\"lmi/{file_name}\"\n",
    "s3url=f\"s3://{bucket}/{s3_model_prefix}\"\n",
    "s3url"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "904fc642-8cf0-475d-85d3-ba201aee13e1",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "! readlink -f $MODEL_REPO_DIR"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "de059920",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from pathlib import Path\n",
    "current_path = Path.cwd()\n",
    "trt_model_path = Path(current_path, MODEL_REPO_DIR)\n",
    "print(trt_model_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aa5c2061-96ac-4f8f-b252-d5935c2f9b4f",
   "metadata": {
    "collapsed": true,
    "jupyter": {
     "outputs_hidden": true
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "!docker run --runtime=nvidia --gpus all --shm-size 12gb \\\n",
    "-v $trt_model_path:/tmp/trtllm \\\n",
    "-p 8080:8080 \\\n",
    "public.ecr.aws/s0w3f1p2/tensorrt-lmi-xq:v1 /opt/djl/partition/trt_llm_partition.py \\\n",
    "--properties_dir /tmp/trtllm \\\n",
    "--trt_llm_model_repo /tmp/trtllm \\\n",
    "--tensor_parallel_degree 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c57ec1d0",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "!mv ./model_tokenizer/chatglm3_6B_32K_tokenizer_config.json $trt_model_path/tokenizer_config.json"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cadddcce-81fa-42a7-8d4e-e05afa7d14c4",
   "metadata": {},
   "source": [
    "## 上传转换后的模型到 S3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a616c75f-bc72-4490-839d-98c0e35f1fd7",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "!aws s3 sync $trt_model_path $s3url"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c8de9e3b-5f6b-4122-9112-56147b6e764a",
   "metadata": {},
   "source": [
    "## 创建，上传，部署所需配置文件到S3\n",
    " - 修改 serving.properties 中 model_id 为上传的 S3 模型地址\n",
    " - 根据自己的输入输出，写 model.py 文件"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ca46be6a-52b4-44ff-aab8-cf736d7de4a2",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "with open('serving.properties', 'w') as f:\n",
    "    f.write(f\"option.model_id={s3url}\\n\")\n",
    "    f.write(\"option.tensor_parallel_degree=2\\n\")\n",
    "    f.write(\"option.max_rolling_batch_size=64\\n\")\n",
    "    f.write(\"option.dtype=fp16\\n\")\n",
    "    f.write(\"option.chatglm_model_version=chatglm3_6b_32k\\n\")\n",
    "    f.write(\"option.trust_remote_code=True\\n\")\n",
    "    f.write(\"option.output_formatter=jsonlines\\n\")\n",
    "    f.write(\"option.rolling_batch=trtllm\\n\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fd2cb0af",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "!cat serving.properties"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9b6d70a2-b5dc-4b98-9627-6993c1dba267",
   "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": "code",
   "execution_count": null,
   "id": "744579de-f8fa-450a-89a1-329e84e067bc",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "s3_code_prefix = f\"large-model-lmi/{file_name}\"\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}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6a15ffed-5c5b-4809-b4de-512aad3291a4",
   "metadata": {
    "tags": []
   },
   "source": [
    "## 设置部署使用的镜像初始化model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "df140c25-c20c-4dd0-b306-286fc57ce52b",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "image_uri = image_uris.retrieve(\n",
    "        framework=\"djl-tensorrtllm\",\n",
    "        region=sess.boto_session.region_name,\n",
    "        version=\"0.26.0\"\n",
    "    )\n",
    "\n",
    "model = Model(image_uri=image_uri, model_data=code_artifact, role=role)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "411ac2af-54c7-44c2-8ee7-3239526570aa",
   "metadata": {},
   "source": [
    "## 开始部署"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b56e2d51-2800-4524-a484-1d7cacdb2bf6",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "instance_type = \"ml.g5.12xlarge\"\n",
    "endpoint_name = sagemaker.utils.name_from_base(\"chatglm3-6b-32k-aot\")\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",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4af75b82-8ee3-4330-b85e-e3b5361ae142",
   "metadata": {},
   "source": [
    "## 测试"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "32e86063-b6d3-42c2-9da9-f43d4a6b5ae2",
   "metadata": {
    "collapsed": true,
    "jupyter": {
     "outputs_hidden": true
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "role = \"user\"\n",
    "input_text = \"世界上第二高的山峰是哪座\"\n",
    "\n",
    "response = predictor.predict(\n",
    "    {\"inputs\": f\"<|{role}|>{input_text}<|assistant|>\\n\", \n",
    "     \"parameters\": {\n",
    "             \"max_new_tokens\":128,\n",
    "             \"temperature\": 0.8,\n",
    "             \"do_sample\": True,\n",
    "             \"top_p\": 0.8,\n",
    "             \"return_log_probs\": False\n",
    "         }\n",
    "    }\n",
    ")\n",
    "print(type(response))\n",
    "text = str(response, 'utf-8')\n",
    "print(text)\n",
    "# print(json.loads(text)[\"generated_text\"].strip())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a99776b0-e8be-4b49-82d9-f35726925efc",
   "metadata": {
    "tags": []
   },
   "source": [
    "## 删除部署的endpoint以及对应的cogfug"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "28f70111-21c4-42c7-835c-259a7b7918ba",
   "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.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}