notebooks/train-deploy-llama3.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Fine-tune Llama 3 with PyTorch FSDP and Q-Lora on Amazon SageMaker\n",
    "\n",
    "This blog post walks you thorugh how to fine-tune a Llama 3 using PyTorch FSDP and Q-Lora with the help of Hugging Face [TRL](https://huggingface.co/docs/trl/index), [Transformers](https://huggingface.co/docs/transformers/index), [peft](https://huggingface.co/docs/peft/index) & [datasets](https://huggingface.co/docs/datasets/index) on Amazon SageMAker. In addition to FSDP we will use [Flash Attention v2](https://github.com/Dao-AILab/flash-attention) implementation. \n",
    "\n",
    "This blog is an extension and dedicated version to my [Efficiently fine-tune Llama 3 with PyTorch FSDP and Q-Lora](https://www.philschmid.de/fsdp-qlora-llama3) version, specifically tailored to run on Amazon SageMaker.\n",
    "\n",
    "1. [Setup development environment](#2-setup-development-environment)\n",
    "2. [Create and prepare the dataset](#3-create-and-prepare-the-dataset)\n",
    "3. [Fine-tune Llama 3 on Amazon SageMaker](#4-fine-tune-llm-using-trl-and-the-sfttrainer)\n",
    "4. [Deploy & Test fine-tuned Llama 3 on Amazon SageMaker](#5-test-and-evaluate-the-llm)\n",
    "\n",
    "_Note: This blog was created and validated on `ml.p4d.24xlarge` and `ml.g5.48.xlarge` instances. The configurations and code are optimized for `ml.p4d.24xlarge` with 8xA100 GPUs each with 40GB of Memory. We tried `ml.g5.12xlarge` but Amazon SageMaker reserves more memory than EC2. We plan to add support for `trn1` in the coming weeks._\n",
    "\n",
    "**FSDP + Q-Lora Background**\n",
    "\n",
    "In a collaboration between [Answer.AI](https://www.answer.ai/posts/2024-03-06-fsdp-qlora.html), Tim Dettmers [Q-Lora creator](https://github.com/TimDettmers/bitsandbytes) and [Hugging Face](https://huggingface.co/), we are proud to announce to share the support of Q-Lora and PyTorch FSDP (Fully Sharded Data Parallel). FSDP and Q-Lora allows you now to fine-tune Llama 2 70b or Mixtral 8x7B on 2x consumer GPUs (24GB). If you want to learn more about the background of this collaboration take a look at [You can now train a 70b language model at home](https://www.answer.ai/posts/2024-03-06-fsdp-qlora.html). Hugging Face PEFT is were the magic happens for this happens, read more about it in the [PEFT documentation](https://huggingface.co/docs/peft/v0.10.0/en/accelerate/fsdp).\n",
    "\n",
    "* [PyTorch FSDP](https://pytorch.org/blog/introducing-pytorch-fully-sharded-data-parallel-api/) is a data/model parallelism technique that shards model across GPUs, reducing memory requirements and enabling the training of larger models more efficiently​​​​​​.\n",
    "* Q-LoRA is a fine-tuning method that leverages quantization and Low-Rank Adapters to efficiently reduced computational requirements and memory footprint. \n",
    "\n",
    "This blog post walks you thorugh how to fine-tune open LLMs from Hugging Face using Amazon SageMaker. This blog is an extension and dedicated version to my [How to Fine-Tune LLMs in 2024 with Hugging Face](https://www.philschmid.de/fine-tune-llms-in-2024-with-trl) version, specifically tailored to run on Amazon SageMaker.\n",
    "\n",
    "\n",
    "\n",
    "## 1. Setup Development Environment\n",
    "\n",
    "Our first step is to install Hugging Face Libraries we need on the client to correctly prepare our dataset and start our training/evaluations jobs. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install transformers \"datasets[s3]==2.18.0\" \"sagemaker>=2.190.0\" \"huggingface_hub[cli]\" --upgrade --quiet"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next we need to login into Hugging Face to access the Llama 3 70b model and store our trained model on Hugging Face. If you don't have an account yet and accepted the terms, you can create one [here](https://huggingface.co/join). "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!huggingface-cli login --token YOUR_TOKEN"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If you are going to use Sagemaker in a local environment. You need access to an IAM Role with the required permissions for Sagemaker. You can find [here](https://docs.aws.amazon.com/sagemaker/latest/dg/sagemaker-roles.html) more about it.\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sagemaker\n",
    "import boto3\n",
    "sess = sagemaker.Session()\n",
    "# sagemaker session bucket -> used for uploading data, models and logs\n",
    "# sagemaker will automatically create this bucket if it not exists\n",
    "sagemaker_session_bucket=None\n",
    "if sagemaker_session_bucket is None and sess is not None:\n",
    "    # set to default bucket if a bucket name is not given\n",
    "    sagemaker_session_bucket = sess.default_bucket()\n",
    "\n",
    "try:\n",
    "    role = sagemaker.get_execution_role()\n",
    "except ValueError:\n",
    "    iam = boto3.client('iam')\n",
    "    role = iam.get_role(RoleName='sagemaker_execution_role')['Role']['Arn']\n",
    "\n",
    "sess = sagemaker.Session(default_bucket=sagemaker_session_bucket)\n",
    "\n",
    "print(f\"sagemaker role arn: {role}\")\n",
    "print(f\"sagemaker bucket: {sess.default_bucket()}\")\n",
    "print(f\"sagemaker session region: {sess.boto_region_name}\")\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. Create and prepare the dataset\n",
    "\n",
    "\n",
    "After our environment is set up, we can start creating and preparing our dataset. A fine-tuning dataset should have a diverse set of demonstrations of the task you want to solve. If you want to learn more about how to create a dataset, take a look at the [How to Fine-Tune LLMs in 2024 with Hugging Face](https://www.philschmid.de/fine-tune-llms-in-2024-with-trl#3-create-and-prepare-the-dataset).\n",
    "\n",
    "We will use the [HuggingFaceH4/no_robots](https://huggingface.co/datasets/HuggingFaceH4/no_robots) dataset a high-quality dataset of 10,000 instructions and demonstrations created by skilled human annotators. This data can be used for supervised fine-tuning (SFT) to make language models follow instructions better. No Robots was modelled after the instruction dataset described in OpenAI's [InstructGPT paper](https://huggingface.co/papers/2203.02155), and is comprised mostly of single-turn instructions.\n",
    "\n",
    "```json\n",
    "{\"messages\": [{\"role\": \"system\", \"content\": \"You are...\"}, {\"role\": \"user\", \"content\": \"...\"}, {\"role\": \"assistant\", \"content\": \"...\"}]}\n",
    "{\"messages\": [{\"role\": \"system\", \"content\": \"You are...\"}, {\"role\": \"user\", \"content\": \"...\"}, {\"role\": \"assistant\", \"content\": \"...\"}]}\n",
    "{\"messages\": [{\"role\": \"system\", \"content\": \"You are...\"}, {\"role\": \"user\", \"content\": \"...\"}, {\"role\": \"assistant\", \"content\": \"...\"}]}\n",
    "```\n",
    "\n",
    "The [no_robots](https://huggingface.co/datasets/HuggingFaceH4/no_robots) dataset has 10,000 split into 9,500 training and  500 test examples. Some samples are not including a `system` message. We will load the dataset with the `datasets` library, add a missing `system` message and save them to separate json files."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from datasets import load_dataset\n",
    "\n",
    "# Convert dataset to OAI messages\n",
    "system_message = \"\"\"You are Llama, an AI assistant created by Philipp to be helpful and honest. Your knowledge spans a wide range of topics, allowing you to engage in substantive conversations and provide analysis on complex subjects.\"\"\"\n",
    "\n",
    "def create_conversation(sample):\n",
    "    if sample[\"messages\"][0][\"role\"] == \"system\":\n",
    "        return sample\n",
    "    else:\n",
    "      sample[\"messages\"] = [{\"role\": \"system\", \"content\": system_message}] + sample[\"messages\"]\n",
    "      return sample\n",
    "\n",
    "# Load dataset from the hub\n",
    "dataset = load_dataset(\"HuggingFaceH4/no_robots\")\n",
    "\n",
    "# Add system message to each conversation\n",
    "columns_to_remove = list(dataset[\"train\"].features)\n",
    "columns_to_remove.remove(\"messages\")\n",
    "dataset = dataset.map(create_conversation, remove_columns=columns_to_remove,batched=False)\n",
    "\n",
    "# Filter out conversations which are corrupted with wrong turns, keep which have even number of turns after adding system message\n",
    "dataset[\"train\"] = dataset[\"train\"].filter(lambda x: len(x[\"messages\"][1:]) % 2 == 0)\n",
    "dataset[\"test\"] = dataset[\"test\"].filter(lambda x: len(x[\"messages\"][1:]) % 2 == 0)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "After we processed the datasets we are going to use the [FileSystem integration](https://huggingface.co/docs/datasets/filesystems) to upload our dataset to S3. We are using the `sess.default_bucket()`, adjust this if you want to store the dataset in a different S3 bucket. We will use the S3 path later in our training script."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# save train_dataset to s3 using our SageMaker session\n",
    "input_path = f's3://{sess.default_bucket()}/datasets/llama3'\n",
    "\n",
    "# save datasets to s3\n",
    "dataset[\"train\"].to_json(f\"{input_path}/train/dataset.json\", orient=\"records\")\n",
    "train_dataset_s3_path = f\"{input_path}/train/dataset.json\"\n",
    "dataset[\"test\"].to_json(f\"{input_path}/test/dataset.json\", orient=\"records\")\n",
    "test_dataset_s3_path = f\"{input_path}/test/dataset.json\"\n",
    "\n",
    "print(f\"Training data uploaded to:\")\n",
    "print(train_dataset_s3_path)\n",
    "print(test_dataset_s3_path)\n",
    "print(f\"https://s3.console.aws.amazon.com/s3/buckets/{sess.default_bucket()}/?region={sess.boto_region_name}&prefix={input_path.split('/', 3)[-1]}/\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. Fine-tune Llama 3 on Amazon SageMaker\n",
    "\n",
    "We are now ready to fine-tune our model. We will use the [SFTTrainer](https://huggingface.co/docs/trl/sft_trainer) from `trl` to fine-tune our model. The `SFTTrainer` makes it straightfoward to supervise fine-tune open LLMs. The `SFTTrainer` is a subclass of the `Trainer` from the `transformers`. We prepared a script [run_fsdp_qlora.py](../scripts/fsdp/run_fsdp_qlora.py) which will loads the dataset from disk, prepare the model, tokenizer and start the training. It usees the [SFTTrainer](https://huggingface.co/docs/trl/sft_trainer) from `trl` to fine-tune our model. The `SFTTrainer` makes it straightfoward to supervise fine-tune open LLMs supporting:\n",
    "* Dataset formatting, including conversational and instruction format (✅ used)\n",
    "* Training on completions only, ignoring prompts (❌ not used)\n",
    "* Packing datasets for more efficient training (✅ used)\n",
    "* PEFT (parameter-efficient fine-tuning) support including Q-LoRA (✅ used)\n",
    "* Preparing the model and tokenizer for conversational fine-tuning (❌ not used, see below)\n",
    "\n",
    "For configuration we use the new `TrlParser`, that allows us to provide hyperparameters in a yaml file. This `yaml` will be uploaded and provided to Amazon SageMaker similar to our datasets. Below is the config file for fine-tuning Llama 3 70B on 8x A100 GPUs or 4x24GB GPUs. We are saving the config file as `fsdp_qlora_llama3_70b.yaml` and upload it to S3.\n",
    "\n",
    "For the chat template we use the Anthropic/Vicuna template, not the official one. Since we then would need to train and save the embedding layer as well leading to more memory requirements. If you wnat to use the official Llama 3 template comment in the `LLAMA_3_CHAT_TEMPLATE` in the `run_fsdp_qlora.py` script and make sure to add modules_to_save. The template used will look like this.\n",
    "\n",
    "```\n",
    "You are a helpful Assistant. \n",
    "\n",
    "Human: What is 2+2? \n",
    "\n",
    "Assistant: 2+2 equals 4.\n",
    "```\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%writefile llama_3_70b_fsdp_qlora.yaml\n",
    "# script parameters\n",
    "model_id: \"meta-llama/Meta-Llama-3-70b\"# Hugging Face model id\n",
    "max_seq_length:  512 # 2048              # max sequence length for model and packing of the dataset\n",
    "# sagemaker specific parameters\n",
    "train_dataset_path: \"/opt/ml/input/data/train/\" # path to where SageMaker saves train dataset\n",
    "test_dataset_path: \"/opt/ml/input/data/test/\"   # path to where SageMaker saves test dataset\n",
    "# output_dir: \"/opt/ml/model\"            # path to where SageMaker will upload the model \n",
    "output_dir: \"/tmp/llama3\"            # path to where SageMaker will upload the model \n",
    "# training parameters\n",
    "report_to: \"tensorboard\"               # report metrics to tensorboard\n",
    "learning_rate: 0.0002                  # learning rate 2e-4\n",
    "lr_scheduler_type: \"constant\"          # learning rate scheduler\n",
    "num_train_epochs: 2                    # number of training epochs\n",
    "per_device_train_batch_size: 8         # batch size per device during training\n",
    "per_device_eval_batch_size: 1          # batch size for evaluation\n",
    "gradient_accumulation_steps: 2         # number of steps before performing a backward/update pass\n",
    "optim: adamw_torch                     # use torch adamw optimizer\n",
    "logging_steps: 10                      # log every 10 steps\n",
    "save_strategy: epoch                   # save checkpoint every epoch\n",
    "evaluation_strategy: epoch             # evaluate every epoch\n",
    "max_grad_norm: 0.3                     # max gradient norm\n",
    "warmup_ratio: 0.03                     # warmup ratio\n",
    "bf16: true                             # use bfloat16 precision\n",
    "tf32: true                             # use tf32 precision\n",
    "gradient_checkpointing: true           # use gradient checkpointing to save memory\n",
    "# FSDP parameters: https://huggingface.co/docs/transformers/main/en/fsdp\n",
    "fsdp: \"full_shard auto_wrap offload\" # remove offload if enough GPU memory\n",
    "fsdp_config:\n",
    "  backward_prefetch: \"backward_pre\"\n",
    "  forward_prefetch: \"false\"\n",
    "  use_orig_params: \"false\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Lets upload the config file to S3. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sagemaker.s3 import S3Uploader\n",
    "\n",
    "# upload the model yaml file to s3\n",
    "model_yaml = \"llama_3_70b_fsdp_qlora.yaml\"\n",
    "train_config_s3_path = S3Uploader.upload(local_path=model_yaml, desired_s3_uri=f\"{input_path}/config\")\n",
    "\n",
    "print(f\"Training config uploaded to:\")\n",
    "print(train_config_s3_path)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In order to create a sagemaker training job we need an `HuggingFace` Estimator. The Estimator handles end-to-end Amazon SageMaker training and deployment tasks. The Estimator manages the infrastructure use. Amazon SagMaker takes care of starting and managing all the required ec2 instances for us, provides the correct huggingface container, uploads the provided scripts and downloads the data from our S3 bucket into the container at `/opt/ml/input/data`. Then, it starts the training job by running.\n",
    "\n",
    "> Note: Make sure that you include the `requirements.txt` in the `source_dir` if you are using a custom training script. We recommend to just clone the whole repository.\n",
    "\n",
    "To use `torchrun` to execute our scripts, we only have to define the `distribution` parameter in our Estimator and set it to `{\"torch_distributed\": {\"enabled\": True}}`. This tells SageMaker to launch our training job with.\n",
    "\n",
    "```python\n",
    "torchrun --nnodes 2 --nproc_per_node 8 --master_addr algo-1 --master_port 7777 --node_rank 1 run_fsdp_qlora.py --config /opt/ml/input/data/config/config.yaml\n",
    "```\n",
    "\n",
    "The `HuggingFace` configuration below will start a training job on 1x `p4d.24xlarge` using 8x A100 GPUs. The amazing part about SageMaker is that you can easily scale up to 2x `p4d.24xlarge` by modifying the `instance_count`. SageMaker will take care of the rest for you. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sagemaker.huggingface import HuggingFace\n",
    "from huggingface_hub import HfFolder\n",
    "\n",
    "# define Training Job Name \n",
    "job_name = f'llama3-70b-exp1'\n",
    "\n",
    "# create the Estimator\n",
    "huggingface_estimator = HuggingFace(\n",
    "    entry_point          = 'run_fsdp_qlora.py',      # train script\n",
    "    source_dir           = '../scripts/fsdp',  # directory which includes all the files needed for training\n",
    "    instance_type        = 'ml.p4d.24xlarge',  # instances type used for the training job\n",
    "    instance_count       = 1,                 # the number of instances used for training\n",
    "    max_run              = 2*24*60*60,        # maximum runtime in seconds (days * hours * minutes * seconds)\n",
    "    base_job_name        = job_name,          # the name of the training job\n",
    "    role                 = role,              # Iam role used in training job to access AWS ressources, e.g. S3\n",
    "    volume_size          = 500,               # the size of the EBS volume in GB\n",
    "    transformers_version = '4.36.0',          # the transformers version used in the training job\n",
    "    pytorch_version      = '2.1.0',           # the pytorch_version version used in the training job\n",
    "    py_version           = 'py310',           # the python version used in the training job\n",
    "    hyperparameters      =  {\n",
    "        \"config\": \"/opt/ml/input/data/config/llama_3_70b_fsdp_qlora.yaml\" # path to TRL config which was uploaded to s3\n",
    "    },\n",
    "    disable_output_compression = True,        # not compress output to save training time and cost\n",
    "    distribution={\"torch_distributed\": {\"enabled\": True}},   # enables torchrun\n",
    "    environment  = {\n",
    "        \"HUGGINGFACE_HUB_CACHE\": \"/tmp/.cache\", # set env variable to cache models in /tmp\n",
    "        \"HF_TOKEN\": HfFolder.get_token(),       # huggingface token to access gated models, e.g. llama 3\n",
    "        \"ACCELERATE_USE_FSDP\": \"1\",             # enable FSDP\n",
    "        \"FSDP_CPU_RAM_EFFICIENT_LOADING\": \"1\"   # enable CPU RAM efficient loading\n",
    "    }, \n",
    ")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "_Note: When using QLoRA, we only train adapters and not the full model. The [run_fsdp_qlora.py](../scripts/fsdp/run_fsdp_qlora.py) merges the `base_model`with the `adapter` at the end of the training to directly be able to deploy to Amazon SageMaker._\n",
    "\n",
    "We can now start our training job, with the `.fit()` method passing our S3 path to the training script."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# define a data input dictonary with our uploaded s3 uris\n",
    "data = {\n",
    "  'train': train_dataset_s3_path,\n",
    "  'test': test_dataset_s3_path,\n",
    "  'config': train_config_s3_path\n",
    "  }\n",
    "\n",
    "# starting the train job with our uploaded datasets as input\n",
    "huggingface_estimator.fit(data, wait=True)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In our example the training Llama 3 70B with Flash Attention for 2 epochs with a dataset of 10k samples takes 5052 seconds (~84minutes) on a `ml.p4d.24xlarge` or ~$50."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4. Deploy & Test fine-tuned Llama 3 on Amazon SageMaker\n",
    "\n",
    "Evaluating LLMs is crucial for understanding their capabilities and limitations, yet it poses significant challenges due to their complex and opaque nature. There are multiple ways to evaluate a fine-tuned model. You could either use an additional Training job to evaluate the model as we demonstrated in [Evaluate LLMs with Hugging Face Lighteval on Amazon SageMaker](https://www.philschmid.de/sagemaker-evaluate-llm-lighteval) or you can deploy the model to an endpoint and interactively test the model. We are going to use the latter approach in this example. We will load our eval dataset and evaluate the model on those samples, using a simple loop and accuracy as our metric.\n",
    "\n",
    "_Note: Evaluating Generative AI models is not a trivial task since 1 input can have multiple correct outputs. If you want to learn more about evaluating generative models, check out [Evaluate LLMs and RAG a practical example using Langchain and Hugging Face](https://www.philschmid.de/evaluate-llm) blog post._\n",
    "\n",
    "We are going to use the [Hugging Face LLM Inference DLC](https://huggingface.co/blog/sagemaker-huggingface-llm#what-is-hugging-face-llm-inference-dlc) a purpose-built Inference Container to easily deploy LLMs in a secure and managed environment. The DLC is powered by [Text Generation Inference (TGI)](https://huggingface.co/docs/text-generation-inference/index) solution for deploying and serving Large Language Models (LLMs)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sagemaker.huggingface import get_huggingface_llm_image_uri\n",
    "\n",
    "# retrieve the llm image uri\n",
    "llm_image = get_huggingface_llm_image_uri(\n",
    "  \"huggingface\",\n",
    "  session=sess,\n",
    "  version=\"2.0.2\",\n",
    ")\n",
    "\n",
    "# print ecr image uri\n",
    "print(f\"llm image uri: {llm_image}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can now create a `HuggingFaceModel` using the container uri and the S3 path to our model. We also need to set our TGI configuration including the number of GPUs, max input tokens. You can find a full list of configuration options [here](https://huggingface.co/docs/text-generation-inference/basic_tutorials/launcher)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from huggingface_hub import HfFolder\n",
    "from sagemaker.huggingface import HuggingFaceModel\n",
    "\n",
    "# sagemaker config\n",
    "instance_type = \"ml.p4d.24xlarge\"\n",
    "health_check_timeout = 1200 # 20 minutes\n",
    "\n",
    "# Define Model and Endpoint configuration parameter\n",
    "config = {\n",
    "  'HF_MODEL_ID': \"/opt/ml/model\",       # Path to the model in the container\n",
    "  'SM_NUM_GPUS': \"8\",                   # Number of GPU used per replica\n",
    "  'MAX_INPUT_LENGTH': \"8000\",           # Max length of input text\n",
    "  'MAX_TOTAL_TOKENS': \"8096\",           # Max length of the generation (including input text)\n",
    "  'MAX_BATCH_PREFILL_TOKENS': \"16182\",  # Limits the number of tokens that can be processed in parallel during the generation\n",
    "  'MESSAGES_API_ENABLED': \"true\",       # Enable the OpenAI Messages API\n",
    "}\n",
    "\n",
    "# create HuggingFaceModel with the image uri\n",
    "llm_model = HuggingFaceModel(\n",
    "  role=role,\n",
    "  # path to s3 bucket with model, we are not using a compressed model\n",
    "  # {'S3DataSource':{'S3Uri': \"s3://...\",'S3DataType': 'S3Prefix','CompressionType': 'None'}},\n",
    "  model_data=huggingface_estimator.model_data,\n",
    "  image_uri=llm_image,\n",
    "  env=config\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "After we have created the HuggingFaceModel we can deploy it to Amazon SageMaker using the deploy method."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Deploy model to an endpoint\n",
    "llm = llm_model.deploy(\n",
    "  initial_instance_count=1,\n",
    "  instance_type=instance_type,\n",
    "  container_startup_health_check_timeout=health_check_timeout, # 20 minutes to give SageMaker the time to download and merge model\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "SageMaker will now create our endpoint and deploy the model to it. This can takes a 15-20 minutes. Afterwards, we can test our model by sending some example inputs to the endpoint. We will use the `predict` method of the predictor to send the input to the model and get the output. \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Prompt to generate\n",
    "messages=[\n",
    "    { \"role\": \"system\", \"content\": \"You are a helpful assistant.\" },\n",
    "    { \"role\": \"user\", \"content\": \"Tell me something about Amazon SageMaker?\" }\n",
    "  ]\n",
    "\n",
    "# Generation arguments\n",
    "parameters = {\n",
    "    \"model\": \"meta-llama-3-fine-tuned\", # placeholder, needed\n",
    "    \"top_p\": 0.6,\n",
    "    \"temperature\": 0.9,\n",
    "    \"max_tokens\": 512,\n",
    "    \"stop\": [\"<|eot_id|>\"],\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "chat = llm.predict({\"messages\" :messages, **parameters})\n",
    "\n",
    "print(chat[\"choices\"][0][\"message\"][\"content\"].strip())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "_Note: If you want to learn more about streaming responses or benchmarking your endpoint checkout [Deploy Llama 3 on Amazon SageMaker](https://www.philschmid.de/sagemaker-llama3) and [Deploy Mixtral 8x7B on Amazon SageMaker\n",
    "](https://www.philschmid.de/sagemaker-deploy-mixtral)._\n",
    "\n",
    "To clean up, we can delete the model and endpoint."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "llm.delete_model()\n",
    "llm.delete_endpoint()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "pytorch",
   "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.9.13"
  },
  "vscode": {
   "interpreter": {
    "hash": "2d58e898dde0263bc564c6968b04150abacfd33eed9b19aaa8e45c040360e146"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}