LLM-FineTuning-Large-Language-Models
/
Mistral_7b_FineTuning_with_DPO_Direct_Preference_Optimization.ipynb
762 строки · 28.0 Кб
1{2"cells": [3{4"cell_type": "markdown",5"metadata": {6"id": "Pa8905-YsHAn"7},8"source": [9"# Fine-tune a Mistral-7b model with DPO (Direct Preference Optimization)\n",10"\n",11"### Checkout my [Twitter(@rohanpaul_ai)](https://twitter.com/rohanpaul_ai) for daily LLM bits"12]13},14{15"cell_type": "code",16"execution_count": null,17"metadata": {18"id": "_zIBL8IssExG"19},20"outputs": [],21"source": [22"!pip install --upgrade trl peft accelerate bitsandbytes datasets auto-gptq optimum huggingface-hub sentencepiece wandb autoawq -q"23]24},25{26"cell_type": "code",27"execution_count": null,28"metadata": {29"colab": {30"base_uri": "https://localhost:8080/"31},32"id": "YpdkZsMNylvp",33"outputId": "6c2df234-1ce7-4cd2-a7e3-567e7536319f"34},35"outputs": [],36"source": [37"import os\n",38"import gc\n",39"import torch\n",40"from dataclasses import dataclass, field\n",41"from typing import Any, Dict, List, NewType, Optional, Tuple\n",42"import transformers\n",43"from transformers import AutoModelForCausalLM, AutoTokenizer, TrainingArguments, BitsAndBytesConfig\n",44"from datasets import load_dataset\n",45"from peft import LoraConfig, PeftModel, get_peft_model, prepare_model_for_kbit_training\n",46"from trl import DPOTrainer\n",47"import bitsandbytes as bnb\n",48"from google.colab import\n",49"\n",50"\n",51"model_name = \"teknium/OpenHermes-2.5-Mistral-7B\"\n",52"\n",53"new_model = \"OpenHermes-2.5-Mistral-7B-DPO-Math\""54]55},56{57"cell_type": "markdown",58"metadata": {59"id": "d8CvUgROUDw-"60},61"source": [62"## Format DPO dataset\n",63"\n",64"First take a note of our Dataset here\n",65"\n",66"\n",67"\n",68"📌 DPO (Direct Preference Optimization) datasets for LLM training, typically consist of a collection of answers that are ranked by humans. This ranking is essential, as the RLHF process fine-tunes LLMs to output the preferred answer. \n",69"\n",70"📌 The structure of the dataset is straightforward: for each row, there is one chosen (preferred) answer, and one rejected answer. The goal of RLHF is to guide the model to output the preferred answer.\n",71"\n",72"📌 And Huggingface's `DPOTrainer` expects a very specific format for the dataset. \n",73"\n",74"📌 Since the model will be trained to directly optimize the preference of which sentence is the most relevant, given two sentences. We provide an example from the Anthropic/hh-rlhf dataset below.\n",75"\n",76"📌 To synthetically create DPO datasets for a set of prompts, you can create the answers with GPT-4/3.5 which will be your preferred answers, and with Llama-2-13b or similar class of models, create the rejected responses. \n",77"\n",78"It’s a smart way to bypass human feedback and only rely on models with different levels of size/performance."79]80},81{82"cell_type": "code",83"execution_count": null,84"metadata": {85"colab": {86"base_uri": "https://localhost:8080/"87},88"id": "MCD77GZ60DOT",89"outputId": "c7c6773c-5545-4fee-bfa3-6fa6d69c0f3f"90},91"outputs": [],92"source": [93"def format_message(role, content, tokenizer, add_generation_prompt=False):\n",94" if content:\n",95" message = {\"role\": role, \"content\": content}\n",96" return tokenizer.apply_chat_template([message], tokenize=False, add_generation_prompt=add_generation_prompt)\n",97" return \"\"\n",98"\n",99"def format_prompt_into_chatml(row_sample):\n",100" # Error handling for missing keys in row_sample\n",101" required_keys = ['system', 'question', 'chosen', 'rejected']\n",102" for key in required_keys:\n",103" if key not in row_sample:\n",104" # Handle missing key appropriately, e.g., raise an error or return a default value\n",105" raise ValueError(f\"Key '{key}' missing in row sample\")\n",106"\n",107" system = format_message(\"system\", row_sample['system'], tokenizer)\n",108" prompt = format_message(\"user\", row_sample['question'], tokenizer, add_generation_prompt=True)\n",109" chosen = row_sample['chosen'] + \"\\n\"\n",110" rejected = row_sample['rejected'] + \"\\n\"\n",111"\n",112" return {\n",113" \"prompt\": system + prompt,\n",114" \"chosen\": chosen,\n",115" \"rejected\": rejected,\n",116" }\n",117"\n",118"\n",119"dataset = load_dataset(\"Intel/orca_dpo_pairs\")['train']\n",120"\n",121"# Save columns\n",122"original_columns = dataset.column_names\n",123"\n",124"# Tokenizer\n",125"tokenizer = AutoTokenizer.from_pretrained(model_name)\n",126"tokenizer.pad_token = tokenizer.eos_token\n",127"tokenizer.padding_side = \"left\"\n",128"\n",129"# Format dataset\n",130"dataset = dataset.map(\n",131" format_prompt_into_chatml,\n",132" remove_columns=original_columns\n",133")\n",134"\n",135"# Print sample\n",136"dataset[1]"137]138},139{140"cell_type": "markdown",141"metadata": {142"id": "DeT5eUK_UJgK"143},144"source": [145"## Train model with DPO\n"146]147},148{149"cell_type": "code",150"execution_count": null,151"metadata": {152"id": "rKPILNOLR-aK"153},154"outputs": [],155"source": [156"\n",157"@dataclass\n",158"class DPOConfig(transformers.TrainingArguments):\n",159" \"\"\"\n",160" Arguments related to the DPO training process itself.\n",161" For all parameters, see:\n",162" https://huggingface.co/docs/transformers/v4.26.1/en/main_classes/trainer#transformers.TrainingArguments\n",163" \"\"\"\n",164"\n",165" beta: Optional[float] = field(\n",166" default=0.1,\n",167" metadata={\"help\": \"The beta factor in DPO loss. Higher beta means less divergence from the initial policy.\"},\n",168" )\n",169" hub_model_revision: Optional[str] = field(\n",170" default=\"main\",\n",171" metadata={\"help\": (\"The Hub model branch to push the model to.\")},\n",172" )\n",173" logging_first_step: bool = field(\n",174" default=True,\n",175" metadata={\"help\": (\"Whether to log and evaluate the first global_step or not.\")},\n",176" )\n",177" max_prompt_length: Optional[int] = field(\n",178" default=None,\n",179" metadata={\"help\": (\"For DPO, the maximum length of the prompt to use for conditioning the model.\")},\n",180" )\n",181" max_length: Optional[int] = field(\n",182" default=None,\n",183" metadata={\"help\": (\"Used by TRL for reward model training, which tries to read this parameter in init.\")},\n",184" )\n",185" optim: Optional[str] = field(default=\"rmsprop\")\n",186" remove_unused_columns: bool = field(default=False)\n"187]188},189{190"cell_type": "markdown",191"metadata": {},192"source": [193"## Note on `DPOConfig` class - \n",194"\n",195"📌 When you instantiate `DPOConfig` with parameters such as `save_strategy=\"steps\"`, which are not explicitly defined in `DPOConfig`, Python's inheritance mechanism comes into play. Since `DPOConfig` is a subclass of `transformers.TrainingArguments`, it inherits all attributes and methods of the base class. Therefore, even if certain parameters like `save_strategy` are not explicitly defined in `DPOConfig`, they are valid as long as they are part of the `transformers.TrainingArguments`.\n",196"\n",197"📌 The `DPOConfig` class, derived from `transformers.TrainingArguments`, allows customization of training parameters specific to your fine-tuning task. \n",198"\n",199"📌 In the `DPOConfig` class, you've defined certain parameters like `beta`, `hub_model_revision`, `logging_first_step`, etc. These are additional or overridden parameters on top of the standard ones provided by `transformers.TrainingArguments`.\n",200"\n",201"📌 This mechanism allows your instance `training_args` to use parameters from both `DPOConfig` and `transformers.TrainingArguments`. \n",202"\n",203"📌 However, it's important to ensure that the parameters you are using in `training_args` are indeed valid and recognized by `transformers.TrainingArguments`."204]205},206{207"cell_type": "code",208"execution_count": null,209"metadata": {},210"outputs": [],211"source": [212"def train(model_name,\n",213" dataset,\n",214" tokenizer,\n",215" new_model,\n",216" #wandb_project: str = \"\",\n",217" #wandb_run_name: str = \"\",\n",218" #wandb_watch: str = \"\", # options: false | gradients | all\n",219" #wandb_log_model: str = \"\", # options: false | true\n",220" ):\n",221" peft_config = LoraConfig(\n",222" r=16,\n",223" lora_alpha=16,\n",224" lora_dropout=0.05,\n",225" bias=\"none\",\n",226" task_type=\"CAUSAL_LM\",\n",227" target_modules=List[str] =['k_proj', 'gate_proj', 'v_proj', 'up_proj', 'q_proj', 'o_proj', 'down_proj']\n",228" )\n",229" assert (\n",230" model_name\n",231" ), \"Please specify a --base_model, e.g. --base_model='huggyllama/llama-7b'\"\n",232"\n",233" # Check if parameter passed or if set within environ\n",234" '''\n",235" use_wandb = len(wandb_project) > 0 or (\n",236" \"WANDB_PROJECT\" in os.environ and len(os.environ[\"WANDB_PROJECT\"]) > 0\n",237" )\n",238" # Only overwrite environ if wandb param passed\n",239" if len(wandb_project) > 0:\n",240" os.environ[\"WANDB_PROJECT\"] = wandb_project\n",241" if len(wandb_watch) > 0:\n",242" os.environ[\"WANDB_WATCH\"] = wandb_watch\n",243" if len(wandb_log_model) > 0:\n",244" os.environ[\"WANDB_LOG_MODEL\"] = wandb_log_model\n",245" '''\n",246"\n",247" # Base Model\n",248" model = AutoModelForCausalLM.from_pretrained(\n",249" model_name,\n",250" torch_dtype=torch.float16,\n",251" load_in_4bit=True\n",252" )\n",253" model.config.use_cache = False\n",254"\n",255" # Reference model\n",256" ref_model = AutoModelForCausalLM.from_pretrained(\n",257" model_name,\n",258" torch_dtype=torch.float16,\n",259" load_in_4bit=True\n",260" )\n",261"\n",262" # Training arguments\n",263" training_args = DPOConfig(\n",264" num_train_epochs=3,\n",265" per_device_train_batch_size=1,\n",266" gradient_accumulation_steps=4,\n",267" gradient_checkpointing=True,\n",268" learning_rate=5e-5,\n",269" lr_scheduler_type=\"linear\",\n",270" max_steps=200,\n",271" save_strategy=\"no\",\n",272" logging_steps=1,\n",273" output_dir=new_model,\n",274" optim=\"paged_adamw_32bit\",\n",275" warmup_steps=100,\n",276" fp16=True,\n",277" # report_to=\"wandb\",\n",278" )\n",279"\n",280" dpo_trainer = DPOTrainer(\n",281" model,\n",282" ref_model,\n",283" args=training_args,\n",284" train_dataset=dataset,\n",285" tokenizer=tokenizer,\n",286" peft_config=peft_config,\n",287" beta=0.1,\n",288" max_prompt_length=1024,\n",289" max_length=1536,\n",290" )"291]292},293{294"cell_type": "code",295"execution_count": null,296"metadata": {},297"outputs": [],298"source": [299"dpo_trainer.train(model_name, dataset, tokenizer, new_model)"300]301},302{303"cell_type": "markdown",304"metadata": {305"id": "3LdhPpcrUM3H"306},307"source": [308"## Save the finetuned model"309]310},311{312"cell_type": "code",313"execution_count": null,314"metadata": {315"id": "h7cIvxcTfBC4"316},317"outputs": [],318"source": [319"# Save artifacts\n",320"dpo_trainer.model.save_pretrained(\"final_checkpoint\")\n",321"tokenizer.save_pretrained(\"final_checkpoint\")\n",322"\n",323"# Clean up memory\n",324"del dpo_trainer, model, ref_model\n",325"gc.collect()\n",326"torch.cuda.empty_cache()\n",327"\n",328"# Reload model in FP16 (instead of NF4)\n",329"base_model = AutoModelForCausalLM.from_pretrained(\n",330" model_name,\n",331" return_dict=True,\n",332" torch_dtype=torch.float16,\n",333")\n",334"tokenizer = AutoTokenizer.from_pretrained(model_name)\n",335"\n",336"# Merge base model with the adapter\n",337"model = PeftModel.from_pretrained(base_model, \"final_checkpoint\")\n",338"model = model.merge_and_unload()\n",339"\n",340"# Save model and tokenizer\n",341"model.save_pretrained(new_model)\n",342"tokenizer.save_pretrained(new_model)\n",343"\n",344"# Push them to the HF Hub\n",345"# model.push_to_hub(new_model, use_temp_dir=False, token=hf_token)\n",346"# tokenizer.push_to_hub(new_model, use_temp_dir=False, token=hf_token)"347]348},349{350"cell_type": "markdown",351"metadata": {352"id": "G6EFsmS4UOgV"353},354"source": [355"## Inference"356]357},358{359"cell_type": "code",360"execution_count": null,361"metadata": {},362"outputs": [],363"source": [364"def generate_chat_response(message, new_model):\n",365" tokenizer = AutoTokenizer.from_pretrained(new_model)\n",366" prompt = tokenizer.apply_chat_template(message, add_generation_prompt=True, tokenize=False)\n",367"\n",368" chat_pipeline = pipeline(\n",369" \"text-generation\",\n",370" model=new_model,\n",371" tokenizer=tokenizer\n",372" )\n",373"\n",374" sequences = chat_pipeline(\n",375" prompt,\n",376" do_sample=True,\n",377" temperature=0.8,\n",378" top_p=0.8,\n",379" num_return_sequences=1,\n",380" max_length=250,\n",381" )\n",382"\n",383" return sequences[0]['generated_text']\n",384"\n",385"# Usage\n",386"message = [\n",387" {\"role\": \"system\", \"content\": \"You are a friendly AI chatbot.\"},\n",388" {\"role\": \"user\", \"content\": \"Plan a holiday for the summer in Europe?\"}\n",389"]\n",390"\n",391"generated_text = generate_chat_response(message, new_model)\n",392"\n",393"print(generated_text)\n"394]395}396],397"metadata": {398"accelerator": "GPU",399"colab": {400"authorship_tag": "ABX9TyOJJCuqxZQnS1q+Fvz5+URG",401"gpuType": "A100",402"include_colab_link": true,403"machine_shape": "hm",404"provenance": []405},406"kernelspec": {407"display_name": "Python 3",408"name": "python3"409},410"language_info": {411"name": "python"412},413"widgets": {414"application/vnd.jupyter.widget-state+json": {415"173769f6f465485f8848a11bf269850b": {416"model_module": "@jupyter-widgets/controls",417"model_module_version": "1.5.0",418"model_name": "HTMLModel",419"state": {420"_dom_classes": [],421"_model_module": "@jupyter-widgets/controls",422"_model_module_version": "1.5.0",423"_model_name": "HTMLModel",424"_view_count": null,425"_view_module": "@jupyter-widgets/controls",426"_view_module_version": "1.5.0",427"_view_name": "HTMLView",428"description": "",429"description_tooltip": null,430"layout": "IPY_MODEL_9083029642744c43b7705532cbe0cf79",431"placeholder": "",432"style": "IPY_MODEL_d028a98caa13425b907ceb513119006e",433"value": " 3/3 [00:11<00:00, 2.89s/it]"434}435},436"22773c721a7c4221a9c14cd388461d4c": {437"model_module": "@jupyter-widgets/controls",438"model_module_version": "1.5.0",439"model_name": "HBoxModel",440"state": {441"_dom_classes": [],442"_model_module": "@jupyter-widgets/controls",443"_model_module_version": "1.5.0",444"_model_name": "HBoxModel",445"_view_count": null,446"_view_module": "@jupyter-widgets/controls",447"_view_module_version": "1.5.0",448"_view_name": "HBoxView",449"box_style": "",450"children": [451"IPY_MODEL_6b54841f5de1482694c360095dae3039",452"IPY_MODEL_448ccbc85e624ec3b3e71931a7ee4ff6",453"IPY_MODEL_173769f6f465485f8848a11bf269850b"454],455"layout": "IPY_MODEL_60978b9b4e8348f0a71ce3e35c73bcff"456}457},458"448ccbc85e624ec3b3e71931a7ee4ff6": {459"model_module": "@jupyter-widgets/controls",460"model_module_version": "1.5.0",461"model_name": "FloatProgressModel",462"state": {463"_dom_classes": [],464"_model_module": "@jupyter-widgets/controls",465"_model_module_version": "1.5.0",466"_model_name": "FloatProgressModel",467"_view_count": null,468"_view_module": "@jupyter-widgets/controls",469"_view_module_version": "1.5.0",470"_view_name": "ProgressView",471"bar_style": "success",472"description": "",473"description_tooltip": null,474"layout": "IPY_MODEL_6e32854952b340008edca0139d3471d6",475"max": 3,476"min": 0,477"orientation": "horizontal",478"style": "IPY_MODEL_db6d7cfcdade4b4baa213a5d0abc07d7",479"value": 3480}481},482"60978b9b4e8348f0a71ce3e35c73bcff": {483"model_module": "@jupyter-widgets/base",484"model_module_version": "1.2.0",485"model_name": "LayoutModel",486"state": {487"_model_module": "@jupyter-widgets/base",488"_model_module_version": "1.2.0",489"_model_name": "LayoutModel",490"_view_count": null,491"_view_module": "@jupyter-widgets/base",492"_view_module_version": "1.2.0",493"_view_name": "LayoutView",494"align_content": null,495"align_items": null,496"align_self": null,497"border": null,498"bottom": null,499"display": null,500"flex": null,501"flex_flow": null,502"grid_area": null,503"grid_auto_columns": null,504"grid_auto_flow": null,505"grid_auto_rows": null,506"grid_column": null,507"grid_gap": null,508"grid_row": null,509"grid_template_areas": null,510"grid_template_columns": null,511"grid_template_rows": null,512"height": null,513"justify_content": null,514"justify_items": null,515"left": null,516"margin": null,517"max_height": null,518"max_width": null,519"min_height": null,520"min_width": null,521"object_fit": null,522"object_position": null,523"order": null,524"overflow": null,525"overflow_x": null,526"overflow_y": null,527"padding": null,528"right": null,529"top": null,530"visibility": null,531"width": null532}533},534"6a38dcbaf4674b448329ac0a16587d2a": {535"model_module": "@jupyter-widgets/base",536"model_module_version": "1.2.0",537"model_name": "LayoutModel",538"state": {539"_model_module": "@jupyter-widgets/base",540"_model_module_version": "1.2.0",541"_model_name": "LayoutModel",542"_view_count": null,543"_view_module": "@jupyter-widgets/base",544"_view_module_version": "1.2.0",545"_view_name": "LayoutView",546"align_content": null,547"align_items": null,548"align_self": null,549"border": null,550"bottom": null,551"display": null,552"flex": null,553"flex_flow": null,554"grid_area": null,555"grid_auto_columns": null,556"grid_auto_flow": null,557"grid_auto_rows": null,558"grid_column": null,559"grid_gap": null,560"grid_row": null,561"grid_template_areas": null,562"grid_template_columns": null,563"grid_template_rows": null,564"height": null,565"justify_content": null,566"justify_items": null,567"left": null,568"margin": null,569"max_height": null,570"max_width": null,571"min_height": null,572"min_width": null,573"object_fit": null,574"object_position": null,575"order": null,576"overflow": null,577"overflow_x": null,578"overflow_y": null,579"padding": null,580"right": null,581"top": null,582"visibility": null,583"width": null584}585},586"6b54841f5de1482694c360095dae3039": {587"model_module": "@jupyter-widgets/controls",588"model_module_version": "1.5.0",589"model_name": "HTMLModel",590"state": {591"_dom_classes": [],592"_model_module": "@jupyter-widgets/controls",593"_model_module_version": "1.5.0",594"_model_name": "HTMLModel",595"_view_count": null,596"_view_module": "@jupyter-widgets/controls",597"_view_module_version": "1.5.0",598"_view_name": "HTMLView",599"description": "",600"description_tooltip": null,601"layout": "IPY_MODEL_6a38dcbaf4674b448329ac0a16587d2a",602"placeholder": "",603"style": "IPY_MODEL_7eaeada2158e493189449af91f643553",604"value": "Loading checkpoint shards: 100%"605}606},607"6e32854952b340008edca0139d3471d6": {608"model_module": "@jupyter-widgets/base",609"model_module_version": "1.2.0",610"model_name": "LayoutModel",611"state": {612"_model_module": "@jupyter-widgets/base",613"_model_module_version": "1.2.0",614"_model_name": "LayoutModel",615"_view_count": null,616"_view_module": "@jupyter-widgets/base",617"_view_module_version": "1.2.0",618"_view_name": "LayoutView",619"align_content": null,620"align_items": null,621"align_self": null,622"border": null,623"bottom": null,624"display": null,625"flex": null,626"flex_flow": null,627"grid_area": null,628"grid_auto_columns": null,629"grid_auto_flow": null,630"grid_auto_rows": null,631"grid_column": null,632"grid_gap": null,633"grid_row": null,634"grid_template_areas": null,635"grid_template_columns": null,636"grid_template_rows": null,637"height": null,638"justify_content": null,639"justify_items": null,640"left": null,641"margin": null,642"max_height": null,643"max_width": null,644"min_height": null,645"min_width": null,646"object_fit": null,647"object_position": null,648"order": null,649"overflow": null,650"overflow_x": null,651"overflow_y": null,652"padding": null,653"right": null,654"top": null,655"visibility": null,656"width": null657}658},659"7eaeada2158e493189449af91f643553": {660"model_module": "@jupyter-widgets/controls",661"model_module_version": "1.5.0",662"model_name": "DescriptionStyleModel",663"state": {664"_model_module": "@jupyter-widgets/controls",665"_model_module_version": "1.5.0",666"_model_name": "DescriptionStyleModel",667"_view_count": null,668"_view_module": "@jupyter-widgets/base",669"_view_module_version": "1.2.0",670"_view_name": "StyleView",671"description_width": ""672}673},674"9083029642744c43b7705532cbe0cf79": {675"model_module": "@jupyter-widgets/base",676"model_module_version": "1.2.0",677"model_name": "LayoutModel",678"state": {679"_model_module": "@jupyter-widgets/base",680"_model_module_version": "1.2.0",681"_model_name": "LayoutModel",682"_view_count": null,683"_view_module": "@jupyter-widgets/base",684"_view_module_version": "1.2.0",685"_view_name": "LayoutView",686"align_content": null,687"align_items": null,688"align_self": null,689"border": null,690"bottom": null,691"display": null,692"flex": null,693"flex_flow": null,694"grid_area": null,695"grid_auto_columns": null,696"grid_auto_flow": null,697"grid_auto_rows": null,698"grid_column": null,699"grid_gap": null,700"grid_row": null,701"grid_template_areas": null,702"grid_template_columns": null,703"grid_template_rows": null,704"height": null,705"justify_content": null,706"justify_items": null,707"left": null,708"margin": null,709"max_height": null,710"max_width": null,711"min_height": null,712"min_width": null,713"object_fit": null,714"object_position": null,715"order": null,716"overflow": null,717"overflow_x": null,718"overflow_y": null,719"padding": null,720"right": null,721"top": null,722"visibility": null,723"width": null724}725},726"d028a98caa13425b907ceb513119006e": {727"model_module": "@jupyter-widgets/controls",728"model_module_version": "1.5.0",729"model_name": "DescriptionStyleModel",730"state": {731"_model_module": "@jupyter-widgets/controls",732"_model_module_version": "1.5.0",733"_model_name": "DescriptionStyleModel",734"_view_count": null,735"_view_module": "@jupyter-widgets/base",736"_view_module_version": "1.2.0",737"_view_name": "StyleView",738"description_width": ""739}740},741"db6d7cfcdade4b4baa213a5d0abc07d7": {742"model_module": "@jupyter-widgets/controls",743"model_module_version": "1.5.0",744"model_name": "ProgressStyleModel",745"state": {746"_model_module": "@jupyter-widgets/controls",747"_model_module_version": "1.5.0",748"_model_name": "ProgressStyleModel",749"_view_count": null,750"_view_module": "@jupyter-widgets/base",751"_view_module_version": "1.2.0",752"_view_name": "StyleView",753"bar_color": null,754"description_width": ""755}756}757}758}759},760"nbformat": 4,761"nbformat_minor": 0762}763