deepseek/deepseek-r1-distilled-llama-8b-fsdp-qlora-remote-decorator_qa.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "f50a1de6-c757-4392-a48d-e5840348537c",
   "metadata": {},
   "source": [
    "# Fine-tune DeepSeek-R1-Distill-Llama-8B with PyTorch FSDP, QLoRA and SageMaker remote decorator\n",
    "\n",
    "## Instruction-based fine-tuning"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a54122cf-17ff-4d98-bf4a-ac2330e8e63f",
   "metadata": {},
   "source": [
    "---\n",
    "\n",
    "In this demo notebook, we demonstrate how to fine-tune the DeepSeek-R1-Distill-Llama-8B model using PyTorch FSDP, QLoRA, Hugging Face PEFT, and bitsandbytes.\n",
    "\n",
    "We are using SageMaker remote decorator for runinng the fine-tuning job on Amazon SageMaker Training job\n",
    "---\n",
    "\n",
    "JupyterLab Instance Type: ml.t3.medium\n",
    "\n",
    "Python version: 3.11\n",
    "\n",
    "Fine-Tuning:\n",
    "* Instance Type: ml.g5.12xlarge"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ac5833a9-a34a-4c5c-a25f-7bcec69317f6",
   "metadata": {},
   "source": [
    "Install the required libriaries, including the Hugging Face libraries, and restart the kernel."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "505031fd-34d1-4261-a3f6-88ede3fd16a7",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-01-26T09:07:37.160628Z",
     "start_time": "2024-01-26T09:07:24.214358Z"
    },
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    },
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "%pip install -r requirements.txt --upgrade\n",
    "%pip install -q -U python-dotenv"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "95c59da5-3c66-48a1-8436-e59878e218eb",
   "metadata": {},
   "source": [
    "## Setup Configuration file path\n",
    "\n",
    "We are setting the directory in which the config.yaml file resides so that remote decorator can make use of the settings through [SageMaker Defaults](https://sagemaker.readthedocs.io/en/stable/overview.html#configuring-and-using-defaults-with-the-sagemaker-python-sdk).\n",
    "\n",
    "This notebook is using the Hugging Face container for the `us-east-1` region. Make sure you are using the right image for your AWS region, otherwise edit [config.yaml](./config.yaml). Container Images are available [here](https://github.com/aws/deep-learning-containers/blob/master/available_images.md)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "95f6fd4d-c4e6-4439-bff5-0ffe05cc5260",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-11-15T09:24:51.291677Z",
     "start_time": "2023-11-15T09:24:51.282905Z"
    },
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "from dotenv import load_dotenv\n",
    "import os\n",
    "\n",
    "# Use .env in case of hidden variables\n",
    "load_dotenv()\n",
    "\n",
    "# Set path to config file\n",
    "os.environ[\"SAGEMAKER_USER_CONFIG_OVERRIDE\"] = os.getcwd()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "82089d28-b97a-4956-83fb-d8c46d44fdb5",
   "metadata": {},
   "source": [
    "## Visualize and upload the dataset\n",
    "\n",
    "We are going to load [FreedomIntelligence/medical-o1-reasoning-SFT](https://huggingface.co/datasets/FreedomIntelligence/medical-o1-reasoning-SFT) dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "63c1c1f7-eec0-441b-88be-ad6ecd19352e",
   "metadata": {},
   "outputs": [],
   "source": [
    "import sagemaker"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1ff9c879-adec-4f5d-a5f5-cce7f7597407",
   "metadata": {},
   "outputs": [],
   "source": [
    "sagemaker_session = sagemaker.Session()\n",
    "bucket_name = sagemaker_session.default_bucket()\n",
    "default_prefix = sagemaker_session.default_bucket_prefix"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1f6da4af-4041-4e63-b2e8-a085e4e62652",
   "metadata": {},
   "outputs": [],
   "source": [
    "from datasets import load_dataset\n",
    "import pandas as pd\n",
    "\n",
    "dataset = load_dataset(\"FreedomIntelligence/medical-o1-reasoning-SFT\", \"en\")\n",
    "\n",
    "df = pd.DataFrame(dataset['train'])\n",
    "df = df.iloc[0:2000]\n",
    "\n",
    "df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "766e8efe-4dbc-4f32-a439-67f1f45b3fa9",
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.model_selection import train_test_split\n",
    "\n",
    "train, test = train_test_split(df, test_size=0.1, random_state=42)\n",
    "\n",
    "print(\"Number of train elements: \", len(train))\n",
    "print(\"Number of test elements: \", len(test))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f368c020-e9a3-48b3-a53b-45404bba9482",
   "metadata": {},
   "source": [
    "Create a prompt template and load the dataset with a random sample to try summarization."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "86eb8edd-35c0-4cf1-82d3-54417bdabd6a",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-09-03T00:02:01.435195Z",
     "start_time": "2023-09-03T00:02:01.429794Z"
    },
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "# custom instruct prompt start\n",
    "prompt_template = f\"\"\"\n",
    "<｜begin▁of▁sentence｜>\n",
    "You are a medical expert with advanced knowledge in clinical reasoning, diagnostics, and treatment planning. \n",
    "Below is an instruction that describes a task, paired with an input that provides further context. \n",
    "Write a response that appropriately completes the request.\n",
    "Before answering, think carefully about the question and create a step-by-step chain of thoughts to ensure a logical and accurate response.\n",
    "<｜User｜>\n",
    "{{question}}\n",
    "<｜Assistant｜>\n",
    "<think>\n",
    "{{complex_cot}}\n",
    "</think>\n",
    "\n",
    "{{answer}}\n",
    "<｜end▁of▁sentence｜>\n",
    "\"\"\"\n",
    "\n",
    "# template dataset to add prompt to each sample\n",
    "def template_dataset(sample):\n",
    "    sample[\"text\"] = prompt_template.format(question=sample[\"Question\"],\n",
    "                                            complex_cot=sample[\"Complex_CoT\"],\n",
    "                                            answer=sample[\"Response\"])\n",
    "    return sample"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1e7f0a32-96c2-4285-b593-3de211f79933",
   "metadata": {},
   "source": [
    "Use the Hugging Face Trainer class to fine-tune the model. Define the hyperparameters we want to use. We also create a DataCollator that will take care of padding our inputs and labels."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2e9cbedd-7403-467e-8cc6-1d2550d8b8e8",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-09-03T00:02:10.364459Z",
     "start_time": "2023-09-03T00:02:09.672705Z"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "from datasets import Dataset, DatasetDict\n",
    "from random import randint\n",
    "\n",
    "train_dataset = Dataset.from_pandas(train)\n",
    "test_dataset = Dataset.from_pandas(test)\n",
    "\n",
    "dataset = DatasetDict({\"train\": train_dataset, \"test\": test_dataset})\n",
    "\n",
    "train_dataset = dataset[\"train\"].map(template_dataset, remove_columns=list(dataset[\"train\"].features))\n",
    "\n",
    "print(train_dataset[randint(0, len(dataset))][\"text\"])\n",
    "\n",
    "test_dataset = dataset[\"test\"].map(template_dataset, remove_columns=list(dataset[\"test\"].features))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a22e9a05-0b47-4344-9ada-8d982a2ae178",
   "metadata": {},
   "source": [
    "Utility function for initializing the distribution across multiple GPUs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9ddbe850-93dd-4169-af53-40ca634231ea",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "\n",
    "def init_distributed():\n",
    "    # Initialize the process group\n",
    "    torch.distributed.init_process_group(\n",
    "        backend=\"nccl\", # Use \"gloo\" backend for CPU\n",
    "        timeout=datetime.timedelta(seconds=5400)\n",
    "    )\n",
    "    local_rank = int(os.environ[\"LOCAL_RANK\"])\n",
    "    torch.cuda.set_device(local_rank)\n",
    "\n",
    "    return local_rank"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2d356590-c57a-493c-b87f-23f9b948b2ad",
   "metadata": {},
   "source": [
    "Utility function for model download"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cefacb97-b8ff-401e-886b-06e0e72d2486",
   "metadata": {},
   "outputs": [],
   "source": [
    "from huggingface_hub import snapshot_download\n",
    "import os\n",
    "\n",
    "def download_model(model_name):\n",
    "    print(\"Downloading model \", model_name)\n",
    "\n",
    "    os.makedirs(\"/tmp/tmp_folder\", exist_ok=True)\n",
    "\n",
    "    snapshot_download(repo_id=model_name, local_dir=\"/tmp/tmp_folder\")\n",
    "\n",
    "    print(f\"Model {model_name} downloaded under /tmp/tmp_folder\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2f4b538e-5469-4058-b4c3-9b1699a0091b",
   "metadata": {},
   "source": [
    "To train our model, we need to convert our inputs (text) to token IDs. This is done by a Hugging Face Transformers Tokenizer. In addition to Lora, we will use bitsanbytes 4-bit precision to quantize out frozen LLM to 4-bit and attach LoRA adapters on it.\n",
    "\n",
    "Define the train function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b6ee05aa-df07-4824-b13a-0d4395de460c",
   "metadata": {},
   "outputs": [],
   "source": [
    "model_id = \"deepseek-ai/DeepSeek-R1-Distill-Llama-8B\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "10cabb4d-b0b2-498c-95cb-41ed7d05ee65",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-09-03T00:02:21.382486Z",
     "start_time": "2023-09-03T00:02:20.962208Z"
    },
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    },
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "from accelerate import Accelerator\n",
    "import datetime\n",
    "import os\n",
    "from peft import AutoPeftModelForCausalLM, LoraConfig, get_peft_model, prepare_model_for_kbit_training\n",
    "from sagemaker.remote_function import remote\n",
    "import torch\n",
    "from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig, set_seed\n",
    "import transformers\n",
    "import traceback\n",
    "\n",
    "# Start training\n",
    "@remote(\n",
    "    keep_alive_period_in_seconds=0, #Warm-pool instance. Put 0 for avoiding additional costs\n",
    "    volume_size=100,\n",
    "    job_name_prefix=f\"train-{model_id.split('/')[-1].replace('.', '-')}\",\n",
    "    use_torchrun=True,\n",
    ")\n",
    "def train_fn(\n",
    "    model_name,             # Name or path of the base model to fine-tune\n",
    "    train_ds,               # Training dataset\n",
    "    test_ds=None,           # Optional test/validation dataset\n",
    "    torch_dtype=torch.bfloat16,  # Precision type for training\n",
    "    lora_r=8,               # LoRA rank - controls capacity of adaptations\n",
    "    lora_alpha=16,          # LoRA alpha - scales the adaptations\n",
    "    lora_dropout=0.1,       # Dropout probability for LoRA layers\n",
    "    per_device_train_batch_size=8,  # Batch size for training\n",
    "    per_device_eval_batch_size=8,   # Batch size for evaluation\n",
    "    gradient_accumulation_steps=1,  # Number of steps to accumulate gradients\n",
    "    learning_rate=2e-4,     # Learning rate for training\n",
    "    num_train_epochs=1,     # Number of training epochs\n",
    "    fsdp=\"\",                # Fully Sharded Data Parallel configuration\n",
    "    fsdp_config=None,       # Additional FSDP configurations\n",
    "    gradient_checkpointing=False,  # Whether to use gradient checkpointing\n",
    "    merge_weights=False,    # Whether to merge LoRA weights with base model\n",
    "    seed=42,                # Random seed for reproducibility\n",
    "    token=None              # HuggingFace token for model access\n",
    "):\n",
    "    # Initialize distributed training if multiple GPUs are available\n",
    "    if torch.cuda.is_available() and (torch.cuda.device_count() > 1 or int(os.environ.get(\"SM_HOST_COUNT\", 1)) > 1):\n",
    "        # Call this function at the beginning of your script\n",
    "        local_rank = init_distributed()\n",
    "\n",
    "        # Now you can use distributed functionalities\n",
    "        torch.distributed.barrier(device_ids=[local_rank])\n",
    "\n",
    "    # Enable HuggingFace transfer for model downloading\n",
    "    os.environ.update({\"HF_HUB_ENABLE_HF_TRANSFER\": \"1\"})\n",
    "\n",
    "    set_seed(seed)\n",
    "\n",
    "    accelerator = Accelerator()\n",
    "\n",
    "    # Set up HuggingFace token if provided\n",
    "    if token is not None:\n",
    "        os.environ.update({\"HF_TOKEN\": token})\n",
    "        accelerator.wait_for_everyone()\n",
    "\n",
    "    # Download model based on training setup (single or multi-node)\n",
    "    if int(os.environ.get(\"SM_HOST_COUNT\", 1)) == 1:\n",
    "        if accelerator.is_main_process:\n",
    "            download_model(model_name)\n",
    "    else:\n",
    "        download_model(model_name)\n",
    "\n",
    "    accelerator.wait_for_everyone()\n",
    "\n",
    "    model_name = \"/tmp/tmp_folder\"\n",
    "\n",
    "    tokenizer = AutoTokenizer.from_pretrained(model_name)\n",
    "\n",
    "    # Set Tokenizer pad Token\n",
    "    tokenizer.pad_token = tokenizer.eos_token\n",
    "\n",
    "    with accelerator.main_process_first():\n",
    "        # tokenize and chunk dataset\n",
    "        lm_train_dataset = train_ds.map(\n",
    "            lambda sample: tokenizer(sample[\"text\"]), remove_columns=list(train_ds.features)\n",
    "        )\n",
    "\n",
    "        print(f\"Total number of train samples: {len(lm_train_dataset)}\")\n",
    "\n",
    "        if test_ds is not None:\n",
    "            lm_test_dataset = test_ds.map(\n",
    "                lambda sample: tokenizer(sample[\"text\"]), remove_columns=list(test_ds.features)\n",
    "            )\n",
    "\n",
    "            print(f\"Total number of test samples: {len(lm_test_dataset)}\")\n",
    "        else:\n",
    "            lm_test_dataset = None\n",
    "\n",
    "    # Configure model settings for bfloat16 precision\n",
    "    # Setup flash_attention_2 for memory-efficient attention computation\n",
    "    if torch_dtype == torch.bfloat16:\n",
    "        print(\"flash_attention_2 init\")\n",
    "\n",
    "        model_configs = {\n",
    "            \"attn_implementation\": \"flash_attention_2\",\n",
    "            \"torch_dtype\": torch_dtype,\n",
    "        }\n",
    "    else:\n",
    "        model_configs = dict()\n",
    "\n",
    "    # Configure training settings based on FSDP usage\n",
    "    # Set up trainer configurations for FSDP or standard training\n",
    "    if fsdp != \"\" and fsdp_config is not None:\n",
    "        print(\"Configurations for FSDP\")\n",
    "\n",
    "        bnb_config_params = {\n",
    "            \"bnb_4bit_quant_storage\": torch_dtype\n",
    "        }\n",
    "\n",
    "        trainer_configs = {\n",
    "            \"fsdp\": fsdp,\n",
    "            \"fsdp_config\": fsdp_config,\n",
    "            \"gradient_checkpointing_kwargs\": {\n",
    "                \"use_reentrant\": False\n",
    "            }\n",
    "        }\n",
    "    else:\n",
    "        bnb_config_params = dict()\n",
    "        trainer_configs = {\n",
    "            \"gradient_checkpointing\": gradient_checkpointing, # Enable in case of DDP\n",
    "        }\n",
    "\n",
    "    # Enable Quantization\n",
    "    bnb_config = BitsAndBytesConfig(\n",
    "        load_in_4bit=True,\n",
    "        bnb_4bit_use_double_quant=True,\n",
    "        bnb_4bit_quant_type=\"nf4\",\n",
    "        bnb_4bit_compute_dtype=torch_dtype,\n",
    "        **bnb_config_params\n",
    "    )\n",
    "\n",
    "    model = AutoModelForCausalLM.from_pretrained(\n",
    "        model_name,\n",
    "        trust_remote_code=True,\n",
    "        quantization_config=bnb_config,\n",
    "        use_cache=not gradient_checkpointing,\n",
    "        cache_dir=\"/tmp/.cache\",\n",
    "        **model_configs\n",
    "    )\n",
    "\n",
    "    # Configure gradient checkpointing based on FSDP usage\n",
    "    if fsdp == \"\" and fsdp_config is None:\n",
    "        print(\"Prepare model for quantization\")\n",
    "        model = prepare_model_for_kbit_training(model, use_gradient_checkpointing=gradient_checkpointing)\n",
    "\n",
    "        if gradient_checkpointing:\n",
    "            print(\"gradient_checkpointing enabled\")\n",
    "            model.gradient_checkpointing_enable()\n",
    "    else:\n",
    "        if gradient_checkpointing:\n",
    "            print(\"gradient_checkpointing enabled\")\n",
    "            model.gradient_checkpointing_enable(gradient_checkpointing_kwargs={\"use_reentrant\": False})\n",
    "\n",
    "    config = LoraConfig(\n",
    "        r=lora_r,\n",
    "        lora_alpha=lora_alpha,\n",
    "        target_modules=\"all-linear\",\n",
    "        lora_dropout=lora_dropout,\n",
    "        bias=\"none\",\n",
    "        task_type=\"CAUSAL_LM\"\n",
    "    )\n",
    "\n",
    "    model = get_peft_model(model, config)\n",
    "\n",
    "    trainer = transformers.Trainer(\n",
    "        model=model,\n",
    "        train_dataset=lm_train_dataset,\n",
    "        eval_dataset=lm_test_dataset if lm_test_dataset is not None else None,\n",
    "        args=transformers.TrainingArguments(\n",
    "            per_device_train_batch_size=per_device_train_batch_size,\n",
    "            per_device_eval_batch_size=per_device_eval_batch_size,\n",
    "            gradient_accumulation_steps=gradient_accumulation_steps,\n",
    "            logging_strategy=\"steps\",\n",
    "            logging_steps=1,\n",
    "            log_on_each_node=False,\n",
    "            num_train_epochs=num_train_epochs,\n",
    "            learning_rate=learning_rate,\n",
    "            bf16=(\n",
    "                True if torch_dtype == torch.bfloat16 else False\n",
    "            ),  # Enable mixed-precision training\n",
    "            tf32=False,\n",
    "            ddp_find_unused_parameters=False,\n",
    "            save_strategy=\"no\",\n",
    "            output_dir=\"outputs\",\n",
    "            **trainer_configs\n",
    "        ),\n",
    "        data_collator=transformers.DataCollatorForLanguageModeling(tokenizer, mlm=False),\n",
    "    )\n",
    "\n",
    "    if trainer.accelerator.is_main_process:\n",
    "        trainer.model.print_trainable_parameters()\n",
    "\n",
    "    trainer.train()\n",
    "\n",
    "    if trainer.is_fsdp_enabled:\n",
    "        trainer.accelerator.state.fsdp_plugin.set_state_dict_type(\"FULL_STATE_DICT\")\n",
    "\n",
    "    if merge_weights:\n",
    "        output_dir = \"/tmp/model\"\n",
    "\n",
    "        # merge adapter weights with base model and save\n",
    "        # save int 4 model\n",
    "        trainer.model.save_pretrained(output_dir, safe_serialization=False)\n",
    "\n",
    "        if accelerator.is_main_process:\n",
    "            # clear memory\n",
    "            del model\n",
    "            del trainer\n",
    "\n",
    "            torch.cuda.empty_cache()\n",
    "\n",
    "            # load PEFT model\n",
    "            model = AutoPeftModelForCausalLM.from_pretrained(\n",
    "                output_dir,\n",
    "                torch_dtype=torch.float16,\n",
    "                low_cpu_mem_usage=True,\n",
    "                trust_remote_code=True,\n",
    "                use_cache=True,\n",
    "                cache_dir=\"/tmp/.cache\",\n",
    "            )\n",
    "\n",
    "            # Merge LoRA and base model and save\n",
    "            model = model.merge_and_unload()\n",
    "            model.save_pretrained(\n",
    "                os.environ.get(\"SM_MODEL_DIR\", \"/opt/ml/model\"),\n",
    "                safe_serialization=True,\n",
    "                max_shard_size=\"2GB\"\n",
    "            )\n",
    "    else:\n",
    "        trainer.model.save_pretrained(\n",
    "            os.environ.get(\"SM_MODEL_DIR\", \"/opt/ml/model\"),\n",
    "            safe_serialization=True\n",
    "        )\n",
    "\n",
    "    if accelerator.is_main_process:\n",
    "        tokenizer.save_pretrained(os.environ.get(\"SM_MODEL_DIR\", \"/opt/ml/model\"))\n",
    "\n",
    "    accelerator.wait_for_everyone()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2a386bd9-172c-485c-af45-ebc1d126470b",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    },
    "scrolled": true,
    "tags": []
   },
   "outputs": [],
   "source": [
    "train_fn(\n",
    "    model_id,\n",
    "    train_ds=train_dataset,\n",
    "    test_ds=test_dataset,\n",
    "    per_device_train_batch_size=2,\n",
    "    per_device_eval_batch_size=1,\n",
    "    gradient_accumulation_steps=2,\n",
    "    gradient_checkpointing=True,\n",
    "    num_train_epochs=2,\n",
    "    fsdp=\"full_shard auto_wrap offload\",\n",
    "    fsdp_config={\n",
    "        'backward_prefetch': 'backward_pre',\n",
    "        'cpu_ram_efficient_loading': True,\n",
    "        'offload_params': True,\n",
    "        'forward_prefetch': False,\n",
    "        'use_orig_params': False,\n",
    "    },\n",
    "    merge_weights=True\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "da81686e-d27f-4c7b-bec6-f596e7dbaa32",
   "metadata": {},
   "source": [
    "***"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cb892b35-ab33-4964-9947-f9487a1e50cb",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Load Fine-Tuned model\n",
    "\n",
    "Note: Run `train_fn` with `merge_weights=True` for merging the trained adapter"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "288e2807-fe0e-46d8-8e79-8dd65f45757a",
   "metadata": {},
   "source": [
    "### Download model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6b50ff8a-842c-45bc-aa55-4a5e87f2b190",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-11-20T18:41:08.515277Z",
     "start_time": "2023-11-20T18:41:08.503555Z"
    },
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "import boto3\n",
    "import sagemaker"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "57abfc30-b9f0-4cd2-8c2e-99bb77c80767",
   "metadata": {},
   "outputs": [],
   "source": [
    "sagemaker_session = sagemaker.Session()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "35481eff-1142-46f3-8e38-50a1bdadba7b",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-09-03T13:45:11.757861Z",
     "start_time": "2023-09-03T13:45:11.747993Z"
    },
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "model_id = \"deepseek-ai/DeepSeek-R1-Distill-Llama-8B\"\n",
    "\n",
    "bucket_name = sagemaker_session.default_bucket()\n",
    "default_prefix = sagemaker_session.default_bucket_prefix\n",
    "job_prefix = f\"train-{model_id.split('/')[-1].replace('.', '-')}\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7ed118e7-1c80-4392-8ea5-147b63fc2f03",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_last_job_name(job_name_prefix):\n",
    "    sagemaker_client = boto3.client('sagemaker')\n",
    "\n",
    "    matching_jobs = []\n",
    "    next_token = None\n",
    "\n",
    "    while True:\n",
    "        # Prepare the search parameters\n",
    "        search_params = {\n",
    "            'Resource': 'TrainingJob',\n",
    "            'SearchExpression': {\n",
    "                'Filters': [\n",
    "                    {\n",
    "                        'Name': 'TrainingJobName',\n",
    "                        'Operator': 'Contains',\n",
    "                        'Value': job_name_prefix\n",
    "                    },\n",
    "                    {\n",
    "                        'Name': 'TrainingJobStatus',\n",
    "                        'Operator': 'Equals',\n",
    "                        'Value': \"Completed\"\n",
    "                    }\n",
    "                ]\n",
    "            },\n",
    "            'SortBy': 'CreationTime',\n",
    "            'SortOrder': 'Descending',\n",
    "            'MaxResults': 100\n",
    "        }\n",
    "\n",
    "        # Add NextToken if we have one\n",
    "        if next_token:\n",
    "            search_params['NextToken'] = next_token\n",
    "\n",
    "        # Make the search request\n",
    "        search_response = sagemaker_client.search(**search_params)\n",
    "\n",
    "        # Filter and add matching jobs\n",
    "        matching_jobs.extend([\n",
    "            job['TrainingJob']['TrainingJobName'] \n",
    "            for job in search_response['Results']\n",
    "            if job['TrainingJob']['TrainingJobName'].startswith(job_name_prefix)\n",
    "        ])\n",
    "\n",
    "        # Check if we have more results to fetch\n",
    "        next_token = search_response.get('NextToken')\n",
    "        if not next_token or matching_jobs:  # Stop if we found at least one match or no more results\n",
    "            break\n",
    "\n",
    "    if not matching_jobs:\n",
    "        raise ValueError(f\"No completed training jobs found starting with prefix '{job_name_prefix}'\")\n",
    "\n",
    "    return matching_jobs[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "18f4e9bd-de61-4806-b314-6bcf988a2c86",
   "metadata": {},
   "outputs": [],
   "source": [
    "job_name = get_last_job_name(job_prefix)\n",
    "\n",
    "job_name"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "36eeba4b-cfc9-4aae-be86-8d2f3c4e5cb0",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "source": [
    "#### Inference configurations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b89ef485-b13b-43a7-a523-eacbdddbb81d",
   "metadata": {},
   "outputs": [],
   "source": [
    "import sagemaker\n",
    "from sagemaker import get_execution_role\n",
    "from sagemaker import Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7f29f2d8-b3b3-4b27-a0d2-e49812f56a21",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-11-20T18:41:10.492877Z",
     "start_time": "2023-11-20T18:41:10.488495Z"
    },
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "instance_count = 1\n",
    "instance_type = \"ml.g5.12xlarge\"\n",
    "health_check_timeout = 700"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "67f062e4-e829-40f4-a613-bf16ad503829",
   "metadata": {},
   "outputs": [],
   "source": [
    "image_uri = sagemaker.image_uris.retrieve(\n",
    "    framework=\"djl-lmi\",\n",
    "    region=sagemaker_session.boto_session.region_name,\n",
    "    version=\"latest\"\n",
    ")\n",
    "\n",
    "image_uri"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8bdd6ef9-484a-4234-ba92-3c0ef95ccad0",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-11-20T18:41:12.433399Z",
     "start_time": "2023-11-20T18:41:11.963091Z"
    },
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "if default_prefix:\n",
    "    model_data = f\"s3://{bucket_name}/{default_prefix}/{job_name}/{job_name}/output/model.tar.gz\"\n",
    "else:\n",
    "    model_data = f\"s3://{bucket_name}/{job_name}/{job_name}/output/model.tar.gz\"\n",
    "\n",
    "model = Model(\n",
    "    image_uri=image_uri,\n",
    "    model_data=model_data,\n",
    "    role=get_execution_role(),\n",
    "    env={\n",
    "        'HF_MODEL_ID': \"/opt/ml/model\", # path to where sagemaker stores the model\n",
    "        'OPTION_TRUST_REMOTE_CODE': 'true',\n",
    "        'OPTION_ROLLING_BATCH': \"vllm\",\n",
    "        'OPTION_DTYPE': 'bf16',\n",
    "        'OPTION_TENSOR_PARALLEL_DEGREE': 'max',\n",
    "        'OPTION_MAX_ROLLING_BATCH_SIZE': '1',\n",
    "        'OPTION_MODEL_LOADING_TIMEOUT': '3600',\n",
    "        'OPTION_MAX_MODEL_LEN': '4096'\n",
    "    }\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a85d8484-69c7-4a91-934b-7cc12613d563",
   "metadata": {},
   "outputs": [],
   "source": [
    "model_id = \"deepseek-ai/DeepSeek-R1-Distill-Llama-8B\"\n",
    "\n",
    "endpoint_name = f\"{model_id.split('/')[-1].replace('.', '-')}-djl\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a9035076-ac69-4859-9824-dcbf07c0f2b2",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-11-20T18:45:49.265298Z",
     "start_time": "2023-11-20T18:41:14.621743Z"
    },
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "predictor = model.deploy(\n",
    "    endpoint_name=endpoint_name,\n",
    "    initial_instance_count=instance_count,\n",
    "    instance_type=instance_type,\n",
    "    container_startup_health_check_timeout=health_check_timeout,\n",
    "    model_data_download_timeout=3600\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "545d0cac-4af1-4034-9f88-35861396c228",
   "metadata": {},
   "source": [
    "#### Predict"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a2d6e8c1-5a3b-459f-8763-febab2b8f094",
   "metadata": {},
   "outputs": [],
   "source": [
    "import sagemaker"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "625af54c-198f-4a1a-8bf7-8e6d8dcfa460",
   "metadata": {},
   "outputs": [],
   "source": [
    "sagemaker_session = sagemaker.Session()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a61094a1-1fa9-495f-8343-aa27d9c4ba2f",
   "metadata": {},
   "outputs": [],
   "source": [
    "model_id = \"deepseek-ai/DeepSeek-R1-Distill-Llama-8B\"\n",
    "\n",
    "endpoint_name = f\"{model_id.split('/')[-1].replace('.', '-')}-djl\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8e7b040d-b8f4-4be0-9687-b28eff520236",
   "metadata": {},
   "outputs": [],
   "source": [
    "predictor = sagemaker.Predictor(\n",
    "    endpoint_name=endpoint_name,\n",
    "    sagemaker_session=sagemaker_session,\n",
    "    serializer=sagemaker.serializers.JSONSerializer(),\n",
    "    deserializer=sagemaker.deserializers.JSONDeserializer(),\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1d5b03ad-a4aa-4107-be89-5d1160cd2d01",
   "metadata": {},
   "outputs": [],
   "source": [
    "base_prompt = f\"\"\"\n",
    "<｜begin▁of▁sentence｜>\n",
    "You are a medical expert with advanced knowledge in clinical reasoning, diagnostics, and treatment planning. \n",
    "Below is an instruction that describes a task, paired with an input that provides further context. \n",
    "Write a response that appropriately completes the request.\n",
    "Before answering, think carefully about the question and create a step-by-step chain of thoughts to ensure a logical and accurate response.\n",
    "<｜User｜>\n",
    "{{question}}\n",
    "<｜Assistant｜>\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1586b889-6554-4d75-b295-e7dc99673cf0",
   "metadata": {},
   "outputs": [],
   "source": [
    "prompt = base_prompt.format(\n",
    "    question=\"A 3-week-old child has been diagnosed with late onset perinatal meningitis, and the CSF culture shows gram-positive bacilli. What characteristic of this bacterium can specifically differentiate it from other bacterial agents?\"\n",
    ")\n",
    "\n",
    "response = predictor.predict({\n",
    "\t\"inputs\": prompt,\n",
    "    \"parameters\": {\n",
    "        \"temperature\": 0.2,\n",
    "        \"top_p\": 0.9,\n",
    "        \"return_full_text\": False,\n",
    "        \"max_new_tokens\": 1024,\n",
    "        \"stop\": ['<｜end▁of▁sentence｜>']\n",
    "    }\n",
    "})\n",
    "\n",
    "response = response[\"generated_text\"].split('<｜end▁of▁sentence｜>')[0]\n",
    "\n",
    "response"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c6ff5d47-62d5-4751-828c-c6f998c33b16",
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "source": [
    "#### Delete Endpoint"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "06aebbc6-5761-44a7-949d-45b2fffd2522",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-11-20T18:48:55.153276Z",
     "start_time": "2023-11-20T18:48:54.165351Z"
    },
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "predictor.delete_model()\n",
    "predictor.delete_endpoint(delete_endpoint_config=True)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "mlops-demo",
   "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.11.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}