       prepare_data/sm-ground_truth_rlhf_llm_finetuning.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "b85eaf69-9088-4cdb-b208-6f8fc5a249a8",
   "metadata": {
    "tags": []
   },
   "source": [
    "# Improving Your LLMs with RLHF on SageMaker"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "03ccf43d",
   "metadata": {},
   "source": [
    "---\n",
    "\n",
    "This notebook's CI test result for us-west-2 is as follows. CI test results in other regions can be found at the end of the notebook. \n",
    "\n",
    "![This us-west-2 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://prod.us-west-2.tcx-beacon.docs.aws.dev/sagemaker-nb/us-west-2/prepare_data|sm-ground_truth_rlhf_llm_finetuning.ipynb)\n",
    "\n",
    "---"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "f671bd8a-ce99-4ead-bc05-e255c9a4cbff",
   "metadata": {},
   "source": [
    "Reinforcement learning from human feedback (RLHF) has proven to be essential to recent large language models (LLMs), e.g. ChatGPT, Claude’s impressive capability and fast adoption. Gone are the days that you need unnatural prompt engineering to get base models, e.g. GPT3 to solve your tasks. Thanks to RLHF, Large Language Models (LLMs) are now much more aligned with human value. \n",
    "\n",
    "However, well-known to the reinforcement learning community, RLHF is notoriously hard to get right. Until very recently only a small number of ML scientists have mastered the skill. In this notebook, we demystify and bring the technique at the disposal of any ML scientist. We describe how to train a base model with RLHF on Amazon SageMaker step by step. \n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "61ef7ec2-61d2-4e03-81c1-3658d48a0156",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Install Prerequisites"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cc9adffd-dcc7-48d0-88f9-25d47ed1d25b",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "!git clone https://github.com/CarperAI/trlx.git\n",
    "!pip install torch==2.0.0 --extra-index-url https://download.pytorch.org/whl/cu116 # for cuda\n",
    "%cd trlx\n",
    "!git checkout 355c9741f2e606de796f5c6f9b682f7dd00f97c5\n",
    "!pip install -e .\n",
    "!pip install transformers==4.27.1 accelerate==0.19.0\n",
    "%cd .."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "45ba4ecc-1c03-42d1-9d2e-c8415879352f",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Supervised Fine-tuning a Base LLM"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "baae70fe-811b-4343-8b96-d3418a19217d",
   "metadata": {},
   "source": [
    "We first make some changes to the code in `trlx/examples/hh/sft_hh.py`, so that the trained model weights are stored in the `checkpoints/sft_hh_gptj_6b` folder\n",
    "\n",
    "```\n",
    "        from itertools import islice\n",
    "\n",
    "        ...\n",
    "        \n",
    "        train=TrainConfig(\n",
    "            seq_length=1024,\n",
    "            epochs=100,\n",
    "            total_steps=10000,\n",
    "            batch_size=1,\n",
    "            checkpoint_interval=10000,\n",
    "            eval_interval=1000,\n",
    "            pipeline=\"PromptPipeline\",\n",
    "            trainer=\"AccelerateSFTTrainer\",\n",
    "            checkpoint_dir=\"checkpoints/sft_hh_gptj_6b\",           # <-- changes\n",
    "            tracker=\"tensorboard\",                                 # <-- changes\n",
    "            logging_dir=\"checkpoints/sft_hh_gptj_6b\"               # <-- changes\n",
    "        ),\n",
    "        \n",
    "        ...\n",
    "\n",
    "        trlx.train(\n",
    "            config=config,\n",
    "            samples=dataset[\"train\"][\"chosen_sample\"],\n",
    "            eval_prompts = [{\"prompt\": x[\"prompt\"], \"original_output\": x[\"chosen\"]} for x in islice(dataset[\"test\"], 280)]  # <-- changes\n",
    "            metric_fn=lambda **kwargs: {\"reward\": reward_fn(**kwargs)},\n",
    "            stop_sequences=[\"Human:\", \"human:\", \"Assistant:\", \"assistant:\"],\n",
    "        )\n",
    "```\n",
    "\n",
    "Prepend the following line to line 600 of `trlx/trainer/accelerate_base_trainer.py`:\n",
    "\n",
    "```\n",
    "        self.save_pretrained(directory+‘/pretrained’)             # <-- changes\n",
    "```\n",
    "\n",
    "We then perform training using the following command:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "316830cd-45ea-412a-89e8-f322135854f4",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "!cd trlx/examples/hh ; accelerate launch --num_processes 7 --config_file ../../configs/accelerate/zero2-bf16.yaml sft_hh.py"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "7199ddfa-395b-48da-aa49-ccca8a6211da",
   "metadata": {},
   "source": [
    "## RLHF Training"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "e0c0a0ea-5c27-4c33-9118-eef0789016bb",
   "metadata": {},
   "source": [
    "Now we have acquired all the required components for RLHF training, i.e. a Supervised Fine-Tuned base model (SFT), a reward model (RM), we are ready to begin optimizing the policy using RLHF. To do this, we will modify the path to SFT model in `examples/hh/ppo_hh.py` to the model weights trained in the previous section (i.e. `checkpoints/sft_hh_gptj_6b/best_checkpoint`), and the final weights would be stored in `checkpoints/ppo_hh_6B`:\n",
    "\n",
    "```\n",
    "    elif config_name == \"6B\":\n",
    "        ...\n",
    "        default_config.model.model_path = \"checkpoints/sft_hh_gptj_6b/best_checkpoint\"     # <-- changes\n",
    "        default_config.train.checkpoint_dir = \"checkpoints/ppo_hh_6B\"      # <-- changes\n",
    "        ...\n",
    "```\n",
    "\n",
    "We then run the training command to start RLHF training:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f9d6e2ee-098f-434a-832e-1eb96f7d6e76",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "!cd trlx/examples/hh ; CONFIG_NAME=6B accelerate launch --num_processes 7 --config_file ../../configs/accelerate/zero2-bf16.yaml ppo_hh.py"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "07341dc9-9463-4fd9-998f-65bc550592c7",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Running Inference"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "506b6b7f-c11f-4fcc-b8bd-a2462820a6fe",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import json\n",
    "import math\n",
    "import torch\n",
    "import random\n",
    "import numpy as np\n",
    "from itertools import islice\n",
    "from datasets import load_dataset\n",
    "from transformers import AutoModelForCausalLM, AutoTokenizer\n",
    "from transformers import AutoConfig\n",
    "\n",
    "\n",
    "def set_seed(seed_val=42):\n",
    "    random.seed(seed_val)\n",
    "    np.random.seed(seed_val)\n",
    "    torch.manual_seed(seed_val)\n",
    "    torch.cuda.manual_seed_all(seed_val)\n",
    "\n",
    "\n",
    "def create_hf_model(model_class, model_name_or_path, tokenizer, disable_dropout=False):\n",
    "    model_config = AutoConfig.from_pretrained(model_name_or_path)\n",
    "    if disable_dropout:\n",
    "        model_config.dropout = 0.0\n",
    "\n",
    "    model = model_class.from_pretrained(\n",
    "        model_name_or_path,\n",
    "        from_tf=bool(\".ckpt\" in model_name_or_path),\n",
    "        config=model_config,\n",
    "        torch_dtype=torch.float16,\n",
    "    )\n",
    "\n",
    "    model.config.end_token_id = tokenizer.eos_token_id\n",
    "    model.config.pad_token_id = model.config.eos_token_id\n",
    "    model.resize_token_embeddings(\n",
    "        int(8 * math.ceil(len(tokenizer) / 8.0))\n",
    "    )  # make the vocab size multiple of 8\n",
    "\n",
    "    return model\n",
    "\n",
    "\n",
    "def load_hf_tokenizer(model_name_or_path, truncation_side=\"left\", padding_side=\"left\"):\n",
    "    if os.path.exists(model_name_or_path):\n",
    "        # Locally tokenizer loading has some issue, so we need to force download\n",
    "        model_json = os.path.join(model_name_or_path, \"config.json\")\n",
    "        if os.path.exists(model_json):\n",
    "            model_json_file = json.load(open(model_json))\n",
    "            model_name = (\n",
    "                \"/home/ec2-user/SageMaker/trlx/examples/hh/\" + model_json_file[\"_name_or_path\"]\n",
    "            )\n",
    "            print(model_name)\n",
    "            tokenizer = AutoTokenizer.from_pretrained(\n",
    "                model_name,\n",
    "                fast_tokenizer=True,\n",
    "                truncation_side=truncation_side,\n",
    "                padding_side=padding_side,\n",
    "            )\n",
    "    else:\n",
    "        tokenizer = AutoTokenizer.from_pretrained(\n",
    "            model_name_or_path,\n",
    "            fast_tokenizer=True,\n",
    "            truncation_side=truncation_side,\n",
    "            padding_side=padding_side,\n",
    "        )\n",
    "\n",
    "    tokenizer.pad_token = tokenizer.eos_token\n",
    "    tokenizer.pad_token_id = tokenizer.eos_token_id\n",
    "    return tokenizer\n",
    "\n",
    "\n",
    "def generate(\n",
    "    model,\n",
    "    tokenizer,\n",
    "    inputs,\n",
    "    num_beams=1,\n",
    "    num_beam_groups=1,\n",
    "    num_return_sequences=1,\n",
    "    max_new_tokens=128,\n",
    "):\n",
    "    prompt_length = inputs.input_ids.shape[1]\n",
    "    generate_ids = model.generate(\n",
    "        inputs.input_ids,\n",
    "        attention_mask=inputs.attention_mask,\n",
    "        num_beams=num_beams,\n",
    "        num_beam_groups=num_beam_groups,\n",
    "        num_return_sequences=num_return_sequences,\n",
    "        max_new_tokens=max_new_tokens,\n",
    "        top_k=0,\n",
    "        top_p=1.0,\n",
    "        do_sample=True,\n",
    "    )\n",
    "\n",
    "    result = tokenizer.batch_decode(\n",
    "        generate_ids[:, prompt_length:],\n",
    "        skip_special_tokens=True,\n",
    "        clean_up_tokenization_spaces=False,\n",
    "    )\n",
    "    return result\n",
    "\n",
    "\n",
    "def run_inference(model, tokenizer, device, prompts, max_prompt_length):\n",
    "    for prompt in prompts:\n",
    "        inputs = tokenizer(\n",
    "            prompt[\"prompt\"],\n",
    "            return_tensors=\"pt\",\n",
    "            truncation=True,\n",
    "            padding=False,\n",
    "            max_length=max_prompt_length,\n",
    "            add_special_tokens=False,\n",
    "        ).to(device)\n",
    "\n",
    "        print(f\"\\n\\n\\n\\tPrompt------------------------\\n  {prompt['prompt']}\")\n",
    "        r_base = generate(\n",
    "            model, tokenizer, inputs, num_beams=1, num_return_sequences=1, max_new_tokens=128\n",
    "        )[0]\n",
    "        for stop in [\"Human:\", \"human:\", \"Assistant:\", \"assistant:\"]:\n",
    "            stop_ix = r_base.find(stop)\n",
    "            if stop_ix >= 0:\n",
    "                r_base = r_base[:stop_ix].rstrip()\n",
    "        print(f\"\\tResponse------------------------\\n {r_base}\")\n",
    "\n",
    "        prompt[\"machine_output\"] = r_base\n",
    "\n",
    "\n",
    "def main(model_name_or_path_baseline):\n",
    "    set_seed(seed_val=42)\n",
    "\n",
    "    device = torch.device(\"cuda:0\")\n",
    "\n",
    "    max_new_tokens = 128\n",
    "    seq_length = 512\n",
    "    max_prompt_length = seq_length - max_new_tokens\n",
    "    tokenizer_pth = model_name_or_path_baseline\n",
    "\n",
    "    tokenizer = load_hf_tokenizer(tokenizer_pth, truncation_side=\"left\", padding_side=\"left\")\n",
    "\n",
    "    model = create_hf_model(AutoModelForCausalLM, model_name_or_path_baseline, tokenizer=tokenizer)\n",
    "    model.to(device)\n",
    "    dataset = load_dataset(\"Dahoas/rm-static\")\n",
    "    eval_prompts = [\n",
    "        {\"prompt\": x[\"prompt\"], \"original_output\": x[\"chosen\"]}\n",
    "        for x in islice(dataset[\"test\"], 100)\n",
    "    ]\n",
    "\n",
    "    run_inference(model, tokenizer, device, eval_prompts, max_prompt_length=max_prompt_length)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "28a4c8f2-b7b4-4180-b1df-ae04c30c303a",
   "metadata": {},
   "source": [
    "The following code generates outputs for the HH test set:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c0cba00e-31ea-45b7-baf6-f5086552efe5",
   "metadata": {},
   "outputs": [],
   "source": [
    "model_name_or_path_baseline = (\n",
    "    \"/home/ec2-user/SageMaker/trlx/examples/hh/checkpoints/ppo_hh_6B/best_checkpoint/pretrained/\"\n",
    ")\n",
    "main(model_name_or_path_baseline)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e9c4dfcb",
   "metadata": {},
   "source": [
    "## Notebook CI Test Results\n",
    "\n",
    "This notebook was tested in multiple regions. The test results are as follows, except for us-west-2 which is shown at the top of the notebook.\n",
    "\n",
    "![This us-east-1 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://prod.us-west-2.tcx-beacon.docs.aws.dev/sagemaker-nb/us-east-1/prepare_data|sm-ground_truth_rlhf_llm_finetuning.ipynb)\n",
    "\n",
    "![This us-east-2 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://prod.us-west-2.tcx-beacon.docs.aws.dev/sagemaker-nb/us-east-2/prepare_data|sm-ground_truth_rlhf_llm_finetuning.ipynb)\n",
    "\n",
    "![This us-west-1 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://prod.us-west-2.tcx-beacon.docs.aws.dev/sagemaker-nb/us-west-1/prepare_data|sm-ground_truth_rlhf_llm_finetuning.ipynb)\n",
    "\n",
    "![This ca-central-1 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://prod.us-west-2.tcx-beacon.docs.aws.dev/sagemaker-nb/ca-central-1/prepare_data|sm-ground_truth_rlhf_llm_finetuning.ipynb)\n",
    "\n",
    "![This sa-east-1 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://prod.us-west-2.tcx-beacon.docs.aws.dev/sagemaker-nb/sa-east-1/prepare_data|sm-ground_truth_rlhf_llm_finetuning.ipynb)\n",
    "\n",
    "![This eu-west-1 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://prod.us-west-2.tcx-beacon.docs.aws.dev/sagemaker-nb/eu-west-1/prepare_data|sm-ground_truth_rlhf_llm_finetuning.ipynb)\n",
    "\n",
    "![This eu-west-2 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://prod.us-west-2.tcx-beacon.docs.aws.dev/sagemaker-nb/eu-west-2/prepare_data|sm-ground_truth_rlhf_llm_finetuning.ipynb)\n",
    "\n",
    "![This eu-west-3 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://prod.us-west-2.tcx-beacon.docs.aws.dev/sagemaker-nb/eu-west-3/prepare_data|sm-ground_truth_rlhf_llm_finetuning.ipynb)\n",
    "\n",
    "![This eu-central-1 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://prod.us-west-2.tcx-beacon.docs.aws.dev/sagemaker-nb/eu-central-1/prepare_data|sm-ground_truth_rlhf_llm_finetuning.ipynb)\n",
    "\n",
    "![This eu-north-1 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://prod.us-west-2.tcx-beacon.docs.aws.dev/sagemaker-nb/eu-north-1/prepare_data|sm-ground_truth_rlhf_llm_finetuning.ipynb)\n",
    "\n",
    "![This ap-southeast-1 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://prod.us-west-2.tcx-beacon.docs.aws.dev/sagemaker-nb/ap-southeast-1/prepare_data|sm-ground_truth_rlhf_llm_finetuning.ipynb)\n",
    "\n",
    "![This ap-southeast-2 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://prod.us-west-2.tcx-beacon.docs.aws.dev/sagemaker-nb/ap-southeast-2/prepare_data|sm-ground_truth_rlhf_llm_finetuning.ipynb)\n",
    "\n",
    "![This ap-northeast-1 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://prod.us-west-2.tcx-beacon.docs.aws.dev/sagemaker-nb/ap-northeast-1/prepare_data|sm-ground_truth_rlhf_llm_finetuning.ipynb)\n",
    "\n",
    "![This ap-northeast-2 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://prod.us-west-2.tcx-beacon.docs.aws.dev/sagemaker-nb/ap-northeast-2/prepare_data|sm-ground_truth_rlhf_llm_finetuning.ipynb)\n",
    "\n",
    "![This ap-south-1 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://prod.us-west-2.tcx-beacon.docs.aws.dev/sagemaker-nb/ap-south-1/prepare_data|sm-ground_truth_rlhf_llm_finetuning.ipynb)\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "temp",
   "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.10.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}