2_finetune.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# FineTuning LLM with Model-As-Service Serverless\n",
    "\n",
    "This sample shows how to create a serverless FineTuning job to fine tune a model using a dataset generated synthetically using RAFT.\n",
    "\n",
    "#### Training data\n",
    "We use the dataset generated in the previous [gen.ipynb](./0_gen.ipynb) notebook using the RAFT method. The dataset has three splits, suitable for:\n",
    "* Fine-tuning\n",
    "* Validation\n",
    "* Evaluation\n",
    "\n",
    "We will use the fine-tuning and validation splits in this notebook\n",
    "\n",
    "#### Model\n",
    "We will use the smaller model of the Llama family to show how a user can finetune a model for chat-completion task.\n",
    "\n",
    "#### Outline\n",
    "1. Setup pre-requisites\n",
    "2. Pick a model to fine-tune.\n",
    "3. Create training and validation datasets.\n",
    "4. Configure the fine tuning job.\n",
    "5. Submit the fine tuning job."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Running time and cost\n",
    "\n",
    "The fine-tuning job usually takes roughly 1.5 hours. Serverless fine-tuning is billed on a time basis and is in the $35/hour range so the cost at this date of running this notebook will be roughly in the $50 range."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Pre-requisites\n",
    "\n",
    "#### Quickstart\n",
    "\n",
    "* Authenticate to Azure using the `az login --use-device-code` command and select an account and subscription configured with **MaaS Serverless**\n",
    "* Copy file `config.json.sample` to `config.json` and replace  `<WORKSPACE_NAME>`, `<RESOURCE_GROUP>` and `<SUBSCRIPTION_ID>` in `config.json`\n",
    "\n",
    "#### In case of issue\n",
    "\n",
    "Checkout the following documentation walking you through setting up the subscription and the Azure ML workspace:\n",
    "* Follow the prerequisites section of article [MS Learn article \"Fine-tune Meta Llama models in Azure AI Studio\"](https://aka.ms/c/learn-ft-prereq)\n",
    "* Connect to AzureML Workspace. Learn more at [set up SDK authentication](https://learn.microsoft.com/en-us/azure/machine-learning/how-to-setup-authentication?tabs=sdk)."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Install dependencies by running below cell. This is not an optional step if running in a new environment.**\n",
    "\n",
    "The requirements should have been automatically installed if you opened the project in Dev Container or Codespaces, but if not, uncomment the following cell to install the requirements"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "#%pip install azure-ai-ml\n",
    "#%pip install azure-identity\n",
    "\n",
    "#%pip install mlflow\n",
    "#%pip install azureml-mlflow"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create AzureML Workspace connections"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from azure.ai.ml import MLClient\n",
    "from azure.identity import (\n",
    "    DefaultAzureCredential,\n",
    "    InteractiveBrowserCredential,\n",
    ")\n",
    "\n",
    "try:\n",
    "    credential = DefaultAzureCredential()\n",
    "    credential.get_token(\"https://management.azure.com/.default\")\n",
    "except Exception as ex:\n",
    "    credential = InteractiveBrowserCredential()\n",
    "\n",
    "# Expects config.json in the same directory as this notebook.\n",
    "workspace_ml_client = MLClient.from_config(credential=credential)\n",
    "\n",
    "# the models, fine tuning pipelines and environments are available in the AzureML system registry, \"azureml\"\n",
    "registry_ml_client = MLClient(credential, registry_name=\"azureml\")\n",
    "registry_ml_client_meta = MLClient(credential, registry_name=\"azureml-meta\")\n",
    "\n",
    "# Get AzureML workspace object.\n",
    "workspace = workspace_ml_client._workspaces.get(workspace_ml_client.workspace_name)\n",
    "workspace_id = workspace._workspace_id"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2. Pick a foundation model to fine tune\n",
    "\n",
    "We will be fine-tuning a `Llama` model for this recipe.\n",
    "\n",
    "At the date of writing this notebook, the following student models have been tested successfully"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pathlib import Path\n",
    "from strictyaml import load\n",
    "\n",
    "for model in [load(x.read_text()).data for x in Path(\"parameters\").glob(\"*.yaml\")]:\n",
    "    print(f\"- {model['model_name']} ({model['format']} model)\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You may want to try other Llama models available in the model registry. Please note that your experience may vary, some of them might not yet be fine-tunable through the Python SDK.\n",
    "\n",
    "If you find another model, feel free to submit a [PR on this repository](https://github.com/Azure-Samples/raft-distillation-recipe/pulls)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for m in filter(lambda x: \"llama\" in x.name.lower() and \"-hf\" not in x.name.lower(), registry_ml_client_meta.models.list()):\n",
    "    print(f\"- {m.name}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Notebook parameters\n",
    "\n",
    "Those parameters are introspected by Papermill in the [`./run_all.sh`](./run_all.sh) script and can be used to parametize the headless execution of this notebook through the command line with parameter files from the [`./parameters`](./parameters/) folder."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "tags": [
     "parameters"
    ]
   },
   "outputs": [],
   "source": [
    "model_name: str = \"Meta-Llama-3.1-8B-Instruct\"\n",
    "format: str = \"chat\"\n",
    "learning_rate: float = 0.0002"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "foundation_model = registry_ml_client_meta.models.get(model_name, label=\"latest\")\n",
    "print(\"\\n\\nUsing model name: {0}, version: {1}, id: {2} for fine tuning\".format(foundation_model.name, foundation_model.version, foundation_model.id))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "from azure.ai.ml.constants._common import AssetTypes\n",
    "from azure.ai.ml.entities._inputs_outputs import Input\n",
    "\n",
    "model_to_finetune = Input(type=AssetTypes.MLFLOW_MODEL, path=foundation_model.id)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 3. Prepare data\n",
    "\n",
    "We are using the data generated previously using RAFT"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Create data inputs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "from dotenv import load_dotenv\n",
    "\n",
    "# Variables passed by previous notebooks\n",
    "load_dotenv(\".env.state\")\n",
    "\n",
    "ds_name = os.getenv(\"DATASET_NAME\")\n",
    "ds_path = f\"dataset/{ds_name}\"\n",
    "dataset_path_ft_train = f\"{ds_path}-files/{ds_name}-ft.train.jsonl\"\n",
    "dataset_path_ft_valid = f\"{ds_path}-files/{ds_name}-ft.valid.jsonl\"\n",
    "\n",
    "print(f\"Using dataset {ds_name} for fine tuning\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Preview of training split"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "\n",
    "pd.read_json(dataset_path_ft_train, lines=True).head(2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Preview of validation split"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "\n",
    "pd.read_json(dataset_path_ft_valid, lines=True).head(2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's calculate the hashes of dataset files, we'll use those hashes combined with dataset names to distinguish between different generated datasets and track which ones have already been uploaded"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "from utils import file_sha256\n",
    "\n",
    "train_hash = file_sha256(dataset_path_ft_train)[:4]\n",
    "valid_hash = file_sha256(dataset_path_ft_valid)[:4]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Upload training dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from azure.ai.ml.entities import Data\n",
    "\n",
    "dataset_version = \"1\"\n",
    "train_dataset_name = f\"{ds_name}_train_{train_hash}\"\n",
    "try:\n",
    "    train_data_created = workspace_ml_client.data.get(train_dataset_name, version=dataset_version)\n",
    "    print(f\"Dataset {train_dataset_name} already exists\")\n",
    "except:\n",
    "    print(f\"Creating dataset {train_dataset_name}\")\n",
    "    train_data = Data(\n",
    "        path=dataset_path_ft_train,\n",
    "        type=AssetTypes.URI_FILE,\n",
    "        description=f\"{ds_name} training dataset\",\n",
    "        name=train_dataset_name,\n",
    "        version=dataset_version,\n",
    "    )\n",
    "    train_data_created = workspace_ml_client.data.create_or_update(train_data)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Upload validation dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from azure.ai.ml.entities import Data\n",
    "\n",
    "dataset_version = \"1\"\n",
    "validation_dataset_name = f\"{ds_name}_validation_{valid_hash}\"\n",
    "try:\n",
    "    validation_data_created = workspace_ml_client.data.get(validation_dataset_name, version=dataset_version)\n",
    "    print(f\"Dataset {validation_dataset_name} already exists\")\n",
    "except:\n",
    "    print(f\"Creating dataset {validation_dataset_name}\")\n",
    "    validation_data = Data(\n",
    "        path=dataset_path_ft_valid,\n",
    "        type=AssetTypes.URI_FILE,\n",
    "        description=f\"{ds_name} validation dataset\",\n",
    "        name=validation_dataset_name,\n",
    "        version=dataset_version,\n",
    "    )\n",
    "    validation_data_created = workspace_ml_client.data.create_or_update(validation_data)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Create training and validation inputs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "from azure.ai.ml.entities._inputs_outputs import Input\n",
    "\n",
    "training_data = Input(\n",
    "    type=train_data_created.type, path=f\"azureml://locations/{workspace.location}/workspaces/{workspace._workspace_id}/data/{train_data_created.name}/versions/{train_data_created.version}\"\n",
    ")\n",
    "validation_data = Input(\n",
    "    type=validation_data_created.type,\n",
    "    path=f\"azureml://locations/{workspace.location}/workspaces/{workspace._workspace_id}/data/{validation_data_created.name}/versions/{validation_data_created.version}\",\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Subscribe to Marketplace offer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from azure.ai.ml.entities import MarketplaceSubscription\n",
    "\n",
    "model_id = \"/\".join(foundation_model.id.split(\"/\")[:-2])\n",
    "subscription_name = model_id.split(\"/\")[-1].replace(\".\", \"-\").replace(\"_\", \"-\")\n",
    "\n",
    "print(f\"Subscribing to Marketplace model: {model_id}\")\n",
    "\n",
    "from azure.core.exceptions import ResourceExistsError\n",
    "marketplace_subscription = MarketplaceSubscription(\n",
    "    model_id=model_id,\n",
    "    name=subscription_name,\n",
    ")\n",
    "\n",
    "try:\n",
    "    marketplace_subscription = workspace_ml_client.marketplace_subscriptions.begin_create_or_update(marketplace_subscription).result()\n",
    "except ResourceExistsError as ex:\n",
    "    print(f\"Marketplace subscription {subscription_name} already exists for model {model_id}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 3. Submit the fine tuning job using the the model and data as inputs\n",
    " \n",
    "Create FineTuning job using all the data that we have so far."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Define finetune parameters\n",
    "\n",
    "##### There are following set of parameters that are required.\n",
    "\n",
    "1. `model` - Base model to finetune.\n",
    "2. `training_data` - Training data for finetuning the base model.\n",
    "3. `task` - FineTuning task to perform. eg. TEXT_COMPLETION for text-generation/text-generation finetuning jobs.\n",
    "4. `outputs`- Output registered model name.\n",
    "\n",
    "##### Following parameters are optional:\n",
    "\n",
    "1. `hyperparameters` - Parameters that control the FineTuning behavior at runtime.\n",
    "2. `name`- FineTuning job name\n",
    "3. `experiment_name` - Experiment name for FineTuning job.\n",
    "4. `display_name` - FineTuning job display name."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from azure.ai.ml.entities._job.finetuning.custom_model_finetuning_job import CustomModelFineTuningJob\n",
    "import uuid\n",
    "from azure.ai.ml._restclient.v2024_01_01_preview.models import (\n",
    "    FineTuningTaskType,\n",
    ")\n",
    "from azure.ai.ml.entities._inputs_outputs import Output\n",
    "\n",
    "guid = uuid.uuid4()\n",
    "short_guid = str(guid)[:4]\n",
    "experiment_name = f\"ft-raft-{ds_name}\"\n",
    "registered_model_name = f\"ft-raft-{model_name}-{ds_name}-{train_hash}-v{dataset_version}\".replace(\".\", \"_\")\n",
    "job_name = f\"{registered_model_name}-{short_guid}\"\n",
    "\n",
    "from utils import update_state\n",
    "\n",
    "task = FineTuningTaskType.CHAT_COMPLETION if format == \"chat\" else FineTuningTaskType.TEXT_COMPLETION\n",
    "print(f\"Model format: {format}\")\n",
    "print(f\"Fine tuning task: {task}\")\n",
    "\n",
    "finetuning_job = CustomModelFineTuningJob(\n",
    "    task=task,\n",
    "    training_data=training_data,\n",
    "    validation_data=validation_data,\n",
    "    hyperparameters={\n",
    "        \"per_device_train_batch_size\": \"1\",\n",
    "        \"learning_rate\": str(learning_rate),\n",
    "        \"num_train_epochs\": \"1\",\n",
    "        \"registered_model_name\": registered_model_name,\n",
    "    },\n",
    "    model=model_to_finetune,\n",
    "    display_name=job_name,\n",
    "    name=job_name,\n",
    "    experiment_name=experiment_name,\n",
    "    outputs={\"registered_model\": Output(type=\"mlflow_model\", name=f\"ft-job-finetune-registered-{short_guid}\")},\n",
    ")\n",
    "\n",
    "update_state(\"FINETUNED_MODEL_NAME\", finetuning_job.outputs['registered_model'].name)\n",
    "update_state(\"FINETUNED_MODEL_FORMAT\", format)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Submit the fine-tuning job"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "try:\n",
    "    print(f\"Submitting job {finetuning_job.name}\")\n",
    "    created_job = workspace_ml_client.jobs.create_or_update(finetuning_job)\n",
    "    print(f\"Successfully created job {finetuning_job.name}\")\n",
    "    print(f\"Studio URL is {created_job.studio_url}\")\n",
    "    print(f\"Registered model name will be {registered_model_name}\")\n",
    "except Exception as e:\n",
    "    print(\"Error creating job\", e)\n",
    "    raise e"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.12.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}