sagemaker/24_train_bloom_peft_lora/sagemaker-notebook.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Efficient Large Language Model training with LoRA and Hugging Face\n",
    "\n",
    "In this sagemaker example, we are going to learn how to apply [Low-Rank Adaptation of Large Language Models (LoRA)](https://arxiv.org/abs/2106.09685) to fine-tune BLOOMZ (7 billion parameter version instruction tuned version of BLOOM) on a single GPU. We are going to leverage Hugging Face [Transformers](https://huggingface.co/docs/transformers/index), [Accelerate](https://huggingface.co/docs/accelerate/index), and [PEFT](https://github.com/huggingface/peft). \n",
    "\n",
    "You will learn how to:\n",
    "\n",
    "1. Setup Development Environment\n",
    "2. Load and prepare the dataset\n",
    "3. Fine-Tune BLOOM with LoRA and bnb int-8 on Amazon SageMaker\n",
    "4. Deploy the model to Amazon SageMaker Endpoint\n",
    "\n",
    "### Quick intro: PEFT or Parameter Efficient Fine-tunin\n",
    "\n",
    "[PEFT](https://github.com/huggingface/peft), or Parameter Efficient Fine-tuning, is a new open-source library from Hugging Face to enable efficient adaptation of pre-trained language models (PLMs) to various downstream applications without fine-tuning all the model's parameters. PEFT currently includes techniques for:\n",
    "\n",
    "- LoRA: [LORA: LOW-RANK ADAPTATION OF LARGE LANGUAGE MODELS](https://arxiv.org/pdf/2106.09685.pdf)\n",
    "- Prefix Tuning: [P-Tuning v2: Prompt Tuning Can Be Comparable to Fine-tuning Universally Across Scales and Tasks](https://arxiv.org/pdf/2110.07602.pdf)\n",
    "- P-Tuning: [GPT Understands, Too](https://arxiv.org/pdf/2103.10385.pdf)\n",
    "- Prompt Tuning: [The Power of Scale for Parameter-Efficient Prompt Tuning](https://arxiv.org/pdf/2104.08691.pdf)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install \"transformers==4.26.0\" \"datasets[s3]==2.9.0\" sagemaker py7zr --upgrade --quiet"
   ]
  },
  {
   "attachments": {},
   "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. Load and prepare the dataset\n",
    "\n",
    "we will use the [samsum](https://huggingface.co/datasets/samsum) dataset, a collection of about 16k messenger-like conversations with summaries. Conversations were created and written down by linguists fluent in English.\n",
    "\n",
    "```python\n",
    "{\n",
    "  \"id\": \"13818513\",\n",
    "  \"summary\": \"Amanda baked cookies and will bring Jerry some tomorrow.\",\n",
    "  \"dialogue\": \"Amanda: I baked cookies. Do you want some?\\r\\nJerry: Sure!\\r\\nAmanda: I'll bring you tomorrow :-)\"\n",
    "}\n",
    "```\n",
    "\n",
    "To load the `samsum` dataset, we use the `load_dataset()` method from the 🤗 Datasets library."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from datasets import load_dataset\n",
    "\n",
    "# Load dataset from the hub\n",
    "dataset = load_dataset(\"samsum\", split=\"train\")\n",
    "\n",
    "print(f\"Train dataset size: {len(dataset)}\")\n",
    "# Train dataset size: 14732"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To train our model, we need to convert our inputs (text) to token IDs. This is done by a 🤗 Transformers Tokenizer. If you are not sure what this means, check out **[chapter 6](https://huggingface.co/course/chapter6/1?fw=tf)** of the Hugging Face Course."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import AutoTokenizer\n",
    "\n",
    "model_id=\"bigscience/bloomz-7b1\"\n",
    "\n",
    "# Load tokenizer of BLOOMZ\n",
    "tokenizer = AutoTokenizer.from_pretrained(model_id)\n",
    "tokenizer.model_max_length = 2048 # overwrite wrong value"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Before we can start training, we need to preprocess our data. Abstractive Summarization is a text-generation task. Our model will take a text as input and generate a summary as output. We want to understand how long our input and output will take to batch our data efficiently.\n",
    "\n",
    "We defined a `prompt_template` which we will use to construct an instruct prompt for better performance of our model. Our `prompt_template` has a “fixed” start and end, and our document is in the middle. This means we need to ensure that the “fixed” template parts + document are not exceeding the max length of the model. \n",
    "We preprocess our dataset before training and save it to disk to then upload it to S3. You could run this step on your local machine or a CPU and upload it to the [Hugging Face Hub](https://huggingface.co/docs/hub/datasets-overview)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from random import randint\n",
    "from itertools import chain\n",
    "from functools import partial\n",
    "\n",
    "# custom instruct prompt start\n",
    "prompt_template = f\"Summarize the chat dialogue:\\n{{dialogue}}\\n---\\nSummary:\\n{{summary}}{{eos_token}}\"\n",
    "\n",
    "# template dataset to add prompt to each sample\n",
    "def template_dataset(sample):\n",
    "    sample[\"text\"] = prompt_template.format(dialogue=sample[\"dialogue\"],\n",
    "                                            summary=sample[\"summary\"],\n",
    "                                            eos_token=tokenizer.eos_token)\n",
    "    return sample\n",
    "\n",
    "\n",
    "# apply prompt template per sample\n",
    "dataset = dataset.map(template_dataset, remove_columns=list(dataset.features))\n",
    "\n",
    "print(dataset[randint(0, len(dataset))][\"text\"])\n",
    "\n",
    "# empty list to save remainder from batches to use in next batch\n",
    "remainder = {\"input_ids\": [], \"attention_mask\": []}\n",
    "\n",
    "\n",
    "def chunk(sample, chunk_length=2048):\n",
    "    # define global remainder variable to save remainder from batches to use in next batch\n",
    "    global remainder\n",
    "    # Concatenate all texts and add remainder from previous batch\n",
    "    concatenated_examples = {k: list(chain(*sample[k])) for k in sample.keys()}\n",
    "    concatenated_examples = {k: remainder[k] + concatenated_examples[k] for k in concatenated_examples.keys()}\n",
    "    # get total number of tokens for batch\n",
    "    batch_total_length = len(concatenated_examples[list(sample.keys())[0]])\n",
    "\n",
    "    # get max number of chunks for batch\n",
    "    if batch_total_length >= chunk_length:\n",
    "        batch_chunk_length = (batch_total_length // chunk_length) * chunk_length\n",
    "\n",
    "    # Split by chunks of max_len.\n",
    "    result = {\n",
    "        k: [t[i : i + chunk_length] for i in range(0, batch_chunk_length, chunk_length)]\n",
    "        for k, t in concatenated_examples.items()\n",
    "    }\n",
    "    # add remainder to global variable for next batch\n",
    "    remainder = {k: concatenated_examples[k][batch_chunk_length:] for k in concatenated_examples.keys()}\n",
    "    # prepare labels\n",
    "    result[\"labels\"] = result[\"input_ids\"].copy()\n",
    "    return result\n",
    "\n",
    "\n",
    "# tokenize and chunk dataset\n",
    "lm_dataset = dataset.map(\n",
    "    lambda sample: tokenizer(sample[\"text\"]), batched=True, remove_columns=list(dataset.features)\n",
    ").map(\n",
    "    partial(chunk, chunk_length=2048),\n",
    "    batched=True,\n",
    ")\n",
    "\n",
    "# Print total number of samples\n",
    "print(f\"Total number of samples: {len(lm_dataset)}\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "After we processed the datasets we are going to use the new [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\n",
    "training_input_path = f's3://{sess.default_bucket()}/processed/samsum-sagemaker/train'\n",
    "lm_dataset.save_to_disk(training_input_path)\n",
    "\n",
    "print(\"uploaded data to:\")\n",
    "print(f\"training dataset to: {training_input_path}\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. Fine-Tune BLOOM with LoRA and bnb int-8 on Amazon SageMaker\n",
    "\n",
    "In addition to the LoRA technique, we will use [bitsanbytes LLM.int8()](https://huggingface.co/blog/hf-bitsandbytes-integration) to quantize out frozen LLM to int8. This allows us to reduce the needed memory for BLOOMZ ~4x.  \n",
    "\n",
    "We prepared a [run_clm.py](./scripts/run_clm.py), which implements uses PEFT to train our model. If you are interested in how this works check-out [Efficient Large Language Model training with LoRA and Hugging Face](https://www.philschmid.de/fine-tune-flan-t5-peft) blog, where we explain the training script in detail. T\n",
    "\n",
    "\n",
    "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. \n",
    "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"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import time\n",
    "# define Training Job Name \n",
    "job_name = f'huggingface-peft-{time.strftime(\"%Y-%m-%d-%H-%M-%S\", time.localtime())}'\n",
    "\n",
    "from sagemaker.huggingface import HuggingFace\n",
    "\n",
    "# hyperparameters, which are passed into the training job\n",
    "hyperparameters ={\n",
    "  'model_id': model_id,                                # pre-trained model\n",
    "  'dataset_path': '/opt/ml/input/data/training', # path where sagemaker will save training dataset\n",
    "  'epochs': 3,                                         # number of training epochs\n",
    "  'per_device_train_batch_size': 1,                    # batch size for training\n",
    "  'lr': 2e-4,                                          # learning rate used during training\n",
    "}\n",
    "\n",
    "# create the Estimator\n",
    "huggingface_estimator = HuggingFace(\n",
    "    entry_point          = 'run_clm.py',      # train script\n",
    "    source_dir           = 'scripts',         # directory which includes all the files needed for training\n",
    "    instance_type        = 'ml.g5.2xlarge', # instances type used for the training job\n",
    "    instance_count       = 1,                 # the number of instances used for training\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          = 300,               # the size of the EBS volume in GB\n",
    "    transformers_version = '4.26',            # the transformers version used in the training job\n",
    "    pytorch_version      = '1.13',            # the pytorch_version version used in the training job\n",
    "    py_version           = 'py39',            # the python version used in the training job\n",
    "    hyperparameters      =  hyperparameters\n",
    ")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "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 = {'training': training_input_path}\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": [
    "The trainign took `20632` seconds, which is about `5.7` hours. The `ml.g5.2xlarge` instance we used costs `$1.515` per hour. So the total cost for training BLOOMZ 7B was is `$8.63`. We could reduce the cost by using a spot instance, but the training time could increase, by waiting or restarts."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4. Deploy the model to Amazon SageMaker Endpoint\n",
    "\n",
    "When using `peft` for training, you normally end up with adapter weights. We added the `merge_and_unload()` method to merge the base model with the adatper to make it easier to deploy the model. Since we can now use the `pipelines` feature of the `transformers` library. \n",
    "\n",
    "We can now deploy our model using the `deploy()` on our HuggingFace estimator object, passing in our desired number of instances and instance type.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sagemaker.huggingface import HuggingFaceModel\n",
    "\n",
    "# create Hugging Face Model Class\n",
    "huggingface_model = HuggingFaceModel(\n",
    "   model_data=huggingface_estimator.model_data,\n",
    "   #model_data=\"s3://hf-sagemaker-inference/model.tar.gz\",  # Change to your model path\n",
    "   role=role, \n",
    "   transformers_version=\"4.26\", \n",
    "   pytorch_version=\"1.13\", \n",
    "   py_version=\"py39\",\n",
    "   model_server_workers=1\n",
    ")\n",
    "# deploy model to SageMaker Inference\n",
    "predictor = huggingface_model.deploy(\n",
    "   initial_instance_count=1,\n",
    "   instance_type= \"ml.g5.4xlarge\"\n",
    ")\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "SageMaker starts the deployment process by creating a SageMaker Endpoint Configuration and a SageMaker Endpoint. The Endpoint Configuration defines the model and the instance type.\n",
    "\n",
    "Lets test by using a example from the `test` split."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from random import randint\n",
    "from datasets import load_dataset\n",
    "\n",
    "# Load dataset from the hub\n",
    "test_dataset = load_dataset(\"samsum\", split=\"test\")\n",
    "\n",
    "# select a random test sample\n",
    "sample = test_dataset[randint(0,len(test_dataset))]\n",
    "\n",
    "# format sample\n",
    "prompt_template = f\"Summarize the chat dialogue:\\n{{dialogue}}\\n---\\nSummary:\\n\"\n",
    "\n",
    "fomatted_sample = {\n",
    "  \"inputs\": prompt_template.format(dialogue=sample[\"dialogue\"]),\n",
    "  \"parameters\": {\n",
    "    \"do_sample\": True,\n",
    "    \"top_p\": 0.9,\n",
    "    \"temperature\": 0.1,\n",
    "    \"max_new_tokens\": 100,\n",
    "  }\n",
    "}\n",
    "\n",
    "# predict\n",
    "res = predictor.predict(fomatted_sample)\n",
    "\n",
    "print(res[0][\"generated_text\"].split(\"Summary:\")[-1])\n",
    "# Kirsten and Alex are going bowling this Friday at 7 pm. They will meet up and then go together.\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Lets compare it to the test result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(sample[\"summary\"])\n",
    "# Kirsten reminds Alex that the youth group meets this Friday at 7 pm and go bowling."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, we delete the endpoint again."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "predictor.delete_model()\n",
    "predictor.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"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "2d58e898dde0263bc564c6968b04150abacfd33eed9b19aaa8e45c040360e146"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}