{ "cells": [ { "cell_type": "markdown", "id": "f24a09e9-262f-4943-84a2-b60b6c24c389", "metadata": {}, "source": [ "# RWKV SageMaker Inference\n", "\n", "RWKV を SageMaker でデプロイするサンプルコード。\n", "\n", "例として、日本語が 10% 学習に使用されている RWKV Raven 7B を利用します。必要に応じてモデルを変更してご利用ください。" ] }, { "cell_type": "code", "execution_count": null, "id": "0428530b-8f42-4ab7-83ce-c3bdfa96a51b", "metadata": {}, "outputs": [], "source": [ "!pip install \"sagemaker>=2.143.0\" -U" ] }, { "cell_type": "code", "execution_count": null, "id": "3f94c4e8-de95-4657-97e0-dc213415bee5", "metadata": { "tags": [] }, "outputs": [], "source": [ "import sagemaker, boto3, json\n", "from sagemaker import get_execution_role\n", "from sagemaker.pytorch.model import PyTorchModel\n", "from sagemaker.huggingface import HuggingFace\n", "\n", "role = get_execution_role()\n", "region = boto3.Session().region_name\n", "sess = sagemaker.Session()\n", "bucket = sess.default_bucket()\n", "\n", "sagemaker.__version__" ] }, { "cell_type": "markdown", "id": "cde108aa-cb07-44f8-aaa3-fac6873664e8", "metadata": {}, "source": [ "## Prepare Model" ] }, { "cell_type": "code", "execution_count": null, "id": "67cb8f80-813f-4aa8-8283-4bc0f2962609", "metadata": { "tags": [] }, "outputs": [], "source": [ "# Download Tokenizer to data directory\n", "!curl https://raw.githubusercontent.com/BlinkDL/RWKV-LM/main/RWKV-v4/20B_tokenizer.json --create-dirs -o data/20B_tokenizer.json" ] }, { "cell_type": "code", "execution_count": null, "id": "e6883643-3322-47da-abe3-a570afdada87", "metadata": { "tags": [] }, "outputs": [], "source": [ "!rm scripts/code/base_models/*" ] }, { "cell_type": "code", "execution_count": null, "id": "a8d46ecf-6753-4f7c-8496-c87353c7c260", "metadata": { "tags": [] }, "outputs": [], "source": [ "base_model_name = \"rwkv-4-raven\"\n", "base_model_file = \"RWKV-4-Raven-7B-v10-Eng89%25-Jpn10%25-Other1%25-20230420-ctx4096.pth\"\n", "n_layer = 32\n", "n_embd = 4096\n", "\n", "base_model_url = f\"https://huggingface.co/BlinkDL/{base_model_name.lower()}/resolve/main/{base_model_file}\"\n", "base_model_path = f\"scripts/code/base_models/{base_model_file}\"" ] }, { "cell_type": "markdown", "id": "9882e8b3-4b8a-4700-8233-8f7e504c6b8d", "metadata": {}, "source": [ "## Package and Upload Model" ] }, { "cell_type": "code", "execution_count": null, "id": "a08d3c5f-2537-433b-a7c9-e2da3480a5f4", "metadata": { "scrolled": true, "tags": [] }, "outputs": [], "source": [ "!rm -rf scripts/model && mkdir scripts/model\n", "!cp data/20B_tokenizer.json scripts/model\n", "%cd scripts\n", "!tar -czvf ../package.tar.gz *\n", "%cd -" ] }, { "cell_type": "code", "execution_count": null, "id": "11d0ae99-8cfe-43f5-8677-863e91bce7a0", "metadata": { "tags": [] }, "outputs": [], "source": [ "model_path = sess.upload_data('package.tar.gz', bucket=bucket, key_prefix=f\"RWKV\")\n", "model_path" ] }, { "cell_type": "markdown", "id": "6432b637-a883-4d94-ba28-54fc1be73408", "metadata": {}, "source": [ "## Deploy Model\n", "\n", "RWKV をデプロイします。\n", "\n", "RWKV Raven 7B の場合、`ml.g5.2xlarge` での GPU Memory 使用率は 66% でした。" ] }, { "cell_type": "code", "execution_count": null, "id": "a20f1592-7ff0-45a2-bbb7-3f9528a11a98", "metadata": { "tags": [] }, "outputs": [], "source": [ "from sagemaker.async_inference import AsyncInferenceConfig\n", "from sagemaker.serializers import JSONSerializer\n", "\n", "endpoint_name = \"RWKV\"\n", "\n", "huggingface_model = PyTorchModel(\n", " model_data=model_path,\n", " framework_version=\"1.13\",\n", " py_version='py39',\n", " role=role,\n", " name=endpoint_name,\n", " env={\n", " \"model_params\": json.dumps({\n", " \"model_url\": base_model_url,\n", " \"model_path\": f\"/tmp/{base_model_file}\",\n", " \"tokenizer_path\": \"model/20B_tokenizer.json\",\n", " \"strategy\": \"cuda bf16\",\n", " }),\n", " \"SAGEMAKER_MODEL_SERVER_TIMEOUT\": \"3600\"\n", " }\n", ")\n", "\n", "# deploy model to SageMaker Inference\n", "predictor = huggingface_model.deploy(\n", " initial_instance_count=1,\n", " instance_type='ml.g5.2xlarge',\n", " endpoint_name=endpoint_name,\n", " serializer=JSONSerializer(),\n", " # async_inference_config=AsyncInferenceConfig()\n", ")" ] }, { "cell_type": "markdown", "id": "d07b0901-2675-49e5-8ccd-968f3ca2820e", "metadata": {}, "source": [ "## Run Inference" ] }, { "cell_type": "code", "execution_count": null, "id": "b6eb4904-4b77-4cd0-bbe9-4d4685bc018e", "metadata": { "tags": [] }, "outputs": [], "source": [ "import sagemaker\n", "from sagemaker.predictor import Predictor\n", "from sagemaker.predictor_async import AsyncPredictor\n", "from sagemaker.serializers import JSONSerializer\n", "from sagemaker.deserializers import NumpyDeserializer, JSONDeserializer\n", "\n", "endpoint_name = \"RWKV\"\n", "\n", "# predictor_client = AsyncPredictor(\n", "# predictor=Predictor(\n", "# endpoint_name=endpoint_name,\n", "# sagemaker_session=sagemaker.Session(),\n", "# serializer=JSONSerializer(),\n", "# deserializer=JSONDeserializer()\n", "# ),\n", "# name=endpoint_name\n", "# )\n", "predictor_client = Predictor(\n", " endpoint_name=endpoint_name,\n", " sagemaker_session=sagemaker.Session(),\n", " serializer=JSONSerializer(),\n", " deserializer=JSONDeserializer()\n", ")\n", "data = {\n", " \"instruction\": \"ヴァージン・オーストラリアはいつから運航を開始したのですか？\",\n", " \"input\": \"\"\"ヴァージン・オーストラリア航空（Virgin Australia Airlines Pty Ltd）の商号で、オーストラリアを拠点とする航空会社です。ヴァージン・ブランドを使用する航空会社の中で、保有機材数では最大の航空会社である。2000年8月31日にヴァージン・ブルーとして、2機の航空機で単一路線で運航を開始した[3]。2001年9月のアンセット・オーストラリアの破綻後、突然オーストラリア国内市場の大手航空会社としての地位を確立した。その後、ブリスベン、メルボルン、シドニーをハブとして、オーストラリア国内の32都市に直接乗り入れるまでに成長した[4]。\"\"\",\n", " \"token_count\": 200,\n", " \"alpha_frequency\": 0.25, \n", " \"alpha_presence\": 0.25, \n", " \"token_ban\": [],\n", " \"token_stop\": [0],\n", " \"chunk_len\": 256\n", "}\n", "response = predictor_client.predict(\n", " data=data\n", ")\n", "print(response)" ] }, { "cell_type": "code", "execution_count": null, "id": "e4a87f87-cad2-49e5-9a6d-693b7063eebc", "metadata": { "tags": [] }, "outputs": [], "source": [ "import boto3\n", "import json\n", "\n", "endpoint_name = \"RWKV\"\n", "sagemaker_client = boto3.client('sagemaker-runtime')\n", "\n", "data = {\n", " \"instruction\": \"ヴァージン・オーストラリアはいつから運航を開始したのですか？\",\n", " \"input\": \"\"\"ヴァージン・オーストラリア航空（Virgin Australia Airlines Pty Ltd）の商号で、オーストラリアを拠点とする航空会社です。ヴァージン・ブランドを使用する航空会社の中で、保有機材数では最大の航空会社である。2000年8月31日にヴァージン・ブルーとして、2機の航空機で単一路線で運航を開始した[3]。2001年9月のアンセット・オーストラリアの破綻後、突然オーストラリア国内市場の大手航空会社としての地位を確立した。その後、ブリスベン、メルボルン、シドニーをハブとして、オーストラリア国内の32都市に直接乗り入れるまでに成長した[4]。\"\"\",\n", " \"token_count\": 300,\n", " \"alpha_frequency\": 0.25, \n", " \"alpha_presence\": 0.25, \n", " \"token_ban\": [],\n", " \"token_stop\": [0],\n", " \"chunk_len\": 256\n", "}\n", "\n", "response = sagemaker_client.invoke_endpoint(\n", " EndpointName=endpoint_name,\n", " ContentType='application/json',\n", " Accept='application/json',\n", " Body=json.dumps(data)\n", ")\n", "\n", "result = json.loads(response['Body'].read())\n", "print(result)" ] }, { "cell_type": "markdown", "id": "709f337d-b552-4bbf-b609-f415fa101ddd", "metadata": {}, "source": [ "## Benchmark Speed" ] }, { "cell_type": "code", "execution_count": null, "id": "0b1256a8-85b0-4717-a59a-c1b2f14e9256", "metadata": { "tags": [] }, "outputs": [], "source": [ "%timeit response = predictor_client.predict(data=data)" ] }, { "cell_type": "markdown", "id": "219ebbfd-3b41-4305-95af-8d35c013969b", "metadata": {}, "source": [ "## Delete Endpoint" ] }, { "cell_type": "code", "execution_count": null, "id": "1cd9be37-f988-4394-b920-61d1cfeaf066", "metadata": { "tags": [] }, "outputs": [], "source": [ "predictor.delete_model()\n", "predictor.delete_endpoint()" ] }, { "cell_type": "code", "execution_count": null, "id": "dca3e9ef-793e-4d30-8900-6396432d49d4", "metadata": {}, "outputs": [], "source": [] } ], "metadata": { "availableInstances": [ { "_defaultOrder": 0, "_isFastLaunch": true, "category": "General purpose", "gpuNum": 0, "hideHardwareSpecs": false, "memoryGiB": 4, "name": "ml.t3.medium", "vcpuNum": 2 }, { "_defaultOrder": 1, "_isFastLaunch": false, "category": "General purpose", "gpuNum": 0, "hideHardwareSpecs": false, "memoryGiB": 8, "name": "ml.t3.large", "vcpuNum": 2 }, { "_defaultOrder": 2, "_isFastLaunch": false, "category": "General purpose", "gpuNum": 0, "hideHardwareSpecs": false, "memoryGiB": 16, "name": "ml.t3.xlarge", "vcpuNum": 4 }, { "_defaultOrder": 3, "_isFastLaunch": false, "category": "General purpose", "gpuNum": 0, "hideHardwareSpecs": false, "memoryGiB": 32, "name": "ml.t3.2xlarge", "vcpuNum": 8 }, { "_defaultOrder": 4, "_isFastLaunch": true, "category": "General purpose", "gpuNum": 0, "hideHardwareSpecs": false, "memoryGiB": 8, "name": "ml.m5.large", "vcpuNum": 2 }, { "_defaultOrder": 5, "_isFastLaunch": false, "category": "General purpose", "gpuNum": 0, "hideHardwareSpecs": false, "memoryGiB": 16, "name": "ml.m5.xlarge", "vcpuNum": 4 }, { "_defaultOrder": 6, "_isFastLaunch": false, "category": "General purpose", "gpuNum": 0, "hideHardwareSpecs": false, "memoryGiB": 32, "name": "ml.m5.2xlarge", "vcpuNum": 8 }, { "_defaultOrder": 7, "_isFastLaunch": false, "category": "General purpose", "gpuNum": 0, "hideHardwareSpecs": false, "memoryGiB": 64, "name": "ml.m5.4xlarge", "vcpuNum": 16 }, { "_defaultOrder": 8, "_isFastLaunch": false, "category": "General purpose", "gpuNum": 0, "hideHardwareSpecs": false, "memoryGiB": 128, "name": "ml.m5.8xlarge", "vcpuNum": 32 }, { "_defaultOrder": 9, "_isFastLaunch": false, "category": "General purpose", "gpuNum": 0, "hideHardwareSpecs": false, "memoryGiB": 192, "name": "ml.m5.12xlarge", "vcpuNum": 48 }, { "_defaultOrder": 10, "_isFastLaunch": false, "category": "General purpose", "gpuNum": 0, "hideHardwareSpecs": false, "memoryGiB": 256, "name": "ml.m5.16xlarge", "vcpuNum": 64 }, { "_defaultOrder": 11, "_isFastLaunch": false, "category": "General purpose", "gpuNum": 0, "hideHardwareSpecs": false, "memoryGiB": 384, "name": "ml.m5.24xlarge", "vcpuNum": 96 }, { "_defaultOrder": 12, "_isFastLaunch": false, "category": "General purpose", "gpuNum": 0, "hideHardwareSpecs": false, "memoryGiB": 8, "name": "ml.m5d.large", "vcpuNum": 2 }, { "_defaultOrder": 13, "_isFastLaunch": false, "category": "General purpose", "gpuNum": 0, "hideHardwareSpecs": false, "memoryGiB": 16, "name": "ml.m5d.xlarge", "vcpuNum": 4 }, { "_defaultOrder": 14, "_isFastLaunch": false, "category": "General purpose", "gpuNum": 0, "hideHardwareSpecs": false, "memoryGiB": 32, "name": "ml.m5d.2xlarge", "vcpuNum": 8 }, { "_defaultOrder": 15, "_isFastLaunch": false, "category": "General purpose", "gpuNum": 0, "hideHardwareSpecs": false, "memoryGiB": 64, "name": "ml.m5d.4xlarge", "vcpuNum": 16 }, { "_defaultOrder": 16, "_isFastLaunch": false, "category": "General purpose", "gpuNum": 0, "hideHardwareSpecs": false, "memoryGiB": 128, "name": "ml.m5d.8xlarge", "vcpuNum": 32 }, { "_defaultOrder": 17, "_isFastLaunch": false, "category": "General purpose", "gpuNum": 0, "hideHardwareSpecs": false, "memoryGiB": 192, "name": "ml.m5d.12xlarge", "vcpuNum": 48 }, { "_defaultOrder": 18, "_isFastLaunch": false, "category": "General purpose", "gpuNum": 0, "hideHardwareSpecs": false, "memoryGiB": 256, "name": "ml.m5d.16xlarge", "vcpuNum": 64 }, { "_defaultOrder": 19, "_isFastLaunch": false, "category": "General purpose", "gpuNum": 0, "hideHardwareSpecs": false, "memoryGiB": 384, "name": "ml.m5d.24xlarge", "vcpuNum": 96 }, { "_defaultOrder": 20, "_isFastLaunch": false, "category": "General purpose", "gpuNum": 0, "hideHardwareSpecs": true, "memoryGiB": 0, "name": "ml.geospatial.interactive", "supportedImageNames": [ "sagemaker-geospatial-v1-0" ], "vcpuNum": 0 }, { "_defaultOrder": 21, "_isFastLaunch": true, "category": "Compute optimized", "gpuNum": 0, "hideHardwareSpecs": false, "memoryGiB": 4, "name": "ml.c5.large", "vcpuNum": 2 }, { "_defaultOrder": 22, "_isFastLaunch": false, "category": "Compute optimized", "gpuNum": 0, "hideHardwareSpecs": false, "memoryGiB": 8, "name": "ml.c5.xlarge", "vcpuNum": 4 }, { "_defaultOrder": 23, "_isFastLaunch": false, "category": "Compute optimized", "gpuNum": 0, "hideHardwareSpecs": false, "memoryGiB": 16, "name": "ml.c5.2xlarge", "vcpuNum": 8 }, { "_defaultOrder": 24, "_isFastLaunch": false, "category": "Compute optimized", "gpuNum": 0, "hideHardwareSpecs": false, "memoryGiB": 32, "name": "ml.c5.4xlarge", "vcpuNum": 16 }, { "_defaultOrder": 25, "_isFastLaunch": false, "category": "Compute optimized", "gpuNum": 0, "hideHardwareSpecs": false, "memoryGiB": 72, "name": "ml.c5.9xlarge", "vcpuNum": 36 }, { "_defaultOrder": 26, "_isFastLaunch": false, "category": "Compute optimized", "gpuNum": 0, "hideHardwareSpecs": false, "memoryGiB": 96, "name": "ml.c5.12xlarge", "vcpuNum": 48 }, { "_defaultOrder": 27, "_isFastLaunch": false, "category": "Compute optimized", "gpuNum": 0, "hideHardwareSpecs": false, "memoryGiB": 144, "name": "ml.c5.18xlarge", "vcpuNum": 72 }, { "_defaultOrder": 28, "_isFastLaunch": false, "category": "Compute optimized", "gpuNum": 0, "hideHardwareSpecs": false, "memoryGiB": 192, "name": "ml.c5.24xlarge", "vcpuNum": 96 }, { "_defaultOrder": 29, "_isFastLaunch": true, "category": "Accelerated computing", "gpuNum": 1, "hideHardwareSpecs": false, "memoryGiB": 16, "name": "ml.g4dn.xlarge", "vcpuNum": 4 }, { "_defaultOrder": 30, "_isFastLaunch": false, "category": "Accelerated computing", "gpuNum": 1, "hideHardwareSpecs": false, "memoryGiB": 32, "name": "ml.g4dn.2xlarge", "vcpuNum": 8 }, { "_defaultOrder": 31, "_isFastLaunch": false, "category": "Accelerated computing", "gpuNum": 1, "hideHardwareSpecs": false, "memoryGiB": 64, "name": "ml.g4dn.4xlarge", "vcpuNum": 16 }, { "_defaultOrder": 32, "_isFastLaunch": false, "category": "Accelerated computing", "gpuNum": 1, "hideHardwareSpecs": false, "memoryGiB": 128, "name": "ml.g4dn.8xlarge", "vcpuNum": 32 }, { "_defaultOrder": 33, "_isFastLaunch": false, "category": "Accelerated computing", "gpuNum": 4, "hideHardwareSpecs": false, "memoryGiB": 192, "name": "ml.g4dn.12xlarge", "vcpuNum": 48 }, { "_defaultOrder": 34, "_isFastLaunch": false, "category": "Accelerated computing", "gpuNum": 1, "hideHardwareSpecs": false, "memoryGiB": 256, "name": "ml.g4dn.16xlarge", "vcpuNum": 64 }, { "_defaultOrder": 35, "_isFastLaunch": false, "category": "Accelerated computing", "gpuNum": 1, "hideHardwareSpecs": false, "memoryGiB": 61, "name": "ml.p3.2xlarge", "vcpuNum": 8 }, { "_defaultOrder": 36, "_isFastLaunch": false, "category": "Accelerated computing", "gpuNum": 4, "hideHardwareSpecs": false, "memoryGiB": 244, "name": "ml.p3.8xlarge", "vcpuNum": 32 }, { "_defaultOrder": 37, "_isFastLaunch": false, "category": "Accelerated computing", "gpuNum": 8, "hideHardwareSpecs": false, "memoryGiB": 488, "name": "ml.p3.16xlarge", "vcpuNum": 64 }, { "_defaultOrder": 38, "_isFastLaunch": false, "category": "Accelerated computing", "gpuNum": 8, "hideHardwareSpecs": false, "memoryGiB": 768, "name": "ml.p3dn.24xlarge", "vcpuNum": 96 }, { "_defaultOrder": 39, "_isFastLaunch": false, "category": "Memory Optimized", "gpuNum": 0, "hideHardwareSpecs": false, "memoryGiB": 16, "name": "ml.r5.large", "vcpuNum": 2 }, { "_defaultOrder": 40, "_isFastLaunch": false, "category": "Memory Optimized", "gpuNum": 0, "hideHardwareSpecs": false, "memoryGiB": 32, "name": "ml.r5.xlarge", "vcpuNum": 4 }, { "_defaultOrder": 41, "_isFastLaunch": false, "category": "Memory Optimized", "gpuNum": 0, "hideHardwareSpecs": false, "memoryGiB": 64, "name": "ml.r5.2xlarge", "vcpuNum": 8 }, { "_defaultOrder": 42, "_isFastLaunch": false, "category": "Memory Optimized", "gpuNum": 0, "hideHardwareSpecs": false, "memoryGiB": 128, "name": "ml.r5.4xlarge", "vcpuNum": 16 }, { "_defaultOrder": 43, "_isFastLaunch": false, "category": "Memory Optimized", "gpuNum": 0, "hideHardwareSpecs": false, "memoryGiB": 256, "name": "ml.r5.8xlarge", "vcpuNum": 32 }, { "_defaultOrder": 44, "_isFastLaunch": false, "category": "Memory Optimized", "gpuNum": 0, "hideHardwareSpecs": false, "memoryGiB": 384, "name": "ml.r5.12xlarge", "vcpuNum": 48 }, { "_defaultOrder": 45, "_isFastLaunch": false, "category": "Memory Optimized", "gpuNum": 0, "hideHardwareSpecs": false, "memoryGiB": 512, "name": "ml.r5.16xlarge", "vcpuNum": 64 }, { "_defaultOrder": 46, "_isFastLaunch": false, "category": "Memory Optimized", "gpuNum": 0, "hideHardwareSpecs": false, "memoryGiB": 768, "name": "ml.r5.24xlarge", "vcpuNum": 96 }, { "_defaultOrder": 47, "_isFastLaunch": false, "category": "Accelerated computing", "gpuNum": 1, "hideHardwareSpecs": false, "memoryGiB": 16, "name": "ml.g5.xlarge", "vcpuNum": 4 }, { "_defaultOrder": 48, "_isFastLaunch": false, "category": "Accelerated computing", "gpuNum": 1, "hideHardwareSpecs": false, "memoryGiB": 32, "name": "ml.g5.2xlarge", "vcpuNum": 8 }, { "_defaultOrder": 49, "_isFastLaunch": false, "category": "Accelerated computing", "gpuNum": 1, "hideHardwareSpecs": false, "memoryGiB": 64, "name": "ml.g5.4xlarge", "vcpuNum": 16 }, { "_defaultOrder": 50, "_isFastLaunch": false, "category": "Accelerated computing", "gpuNum": 1, "hideHardwareSpecs": false, "memoryGiB": 128, "name": "ml.g5.8xlarge", "vcpuNum": 32 }, { "_defaultOrder": 51, "_isFastLaunch": false, "category": "Accelerated computing", "gpuNum": 1, "hideHardwareSpecs": false, "memoryGiB": 256, "name": "ml.g5.16xlarge", "vcpuNum": 64 }, { "_defaultOrder": 52, "_isFastLaunch": false, "category": "Accelerated computing", "gpuNum": 4, "hideHardwareSpecs": false, "memoryGiB": 192, "name": "ml.g5.12xlarge", "vcpuNum": 48 }, { "_defaultOrder": 53, "_isFastLaunch": false, "category": "Accelerated computing", "gpuNum": 4, "hideHardwareSpecs": false, "memoryGiB": 384, "name": "ml.g5.24xlarge", "vcpuNum": 96 }, { "_defaultOrder": 54, "_isFastLaunch": false, "category": "Accelerated computing", "gpuNum": 8, "hideHardwareSpecs": false, "memoryGiB": 768, "name": "ml.g5.48xlarge", "vcpuNum": 192 } ], "instance_type": "ml.t3.medium", "kernelspec": { "display_name": "Python 3 (Data Science)", "language": "python", "name": "python3__SAGEMAKER_INTERNAL__arn:aws:sagemaker:us-west-2:236514542706:image/datascience-1.0" }, "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.10" } }, "nbformat": 4, "nbformat_minor": 5 }