{ "cells": [ { "cell_type": "markdown", "id": "ef297e12", "metadata": {}, "source": [ "# gemma-2b-orpo Training notebook\n", "\n", "gemma-2b-orpo is ORPO fine-tune of [google/gemma-2b](https://huggingface.co/google/gemma-2b) with\n", "[`alvarobartt/dpo-mix-7k-simplified`](https://huggingface.co/datasets/alvarobartt/dpo-mix-7k-simplified).\n", "\n", "Some good resources:\n", "- [HF Transformers Trainer docs](https://huggingface.co/docs/transformers/main_classes/trainer)\n", "- [Docs on training with ORPO using HF TRL](https://huggingface.co/docs/trl/main/en/orpo_trainer)\n", "- [TRL example script for ORPO](https://github.com/huggingface/trl/blob/main/examples/scripts/orpo.py)\n", "- [How to fine-tune Google Gemma with ChatML and Hugging Face TRL](https://www.philschmid.de/fine-tune-google-gemma)" ] }, { "cell_type": "code", "execution_count": null, "id": "a34c07e3-396b-4a9f-83c1-ba9ae128c83a", "metadata": {}, "outputs": [], "source": [ "! pip install git+https://github.com/huggingface/trl.git # install TRL from the main branch to use the ORPOTrainer\n", "! pip install bitsandbytes accelerate\n", "! pip install ninja packaging\n", "! MAX_JOBS=6 pip install flash-attn --no-build-isolation --upgrade # flash-attn speeds up the training on compatible GPUs\n", "! pip install wandb" ] }, { "cell_type": "code", "execution_count": null, "id": "594f15fa-38a3-4898-88f8-65eb4cf22531", "metadata": {}, "outputs": [], "source": [ "# Login to the Hugging Face Hub to save the model\n", "from huggingface_hub import login\n", "\n", "login(token=\"YOUR_TOKEN\")" ] }, { "cell_type": "code", "execution_count": 4, "id": "d155473a-b979-46fd-b858-04907f046a5e", "metadata": {}, "outputs": [], "source": [ "# https://huggingface.co/docs/trl/main/en/orpo_trainer#trl.ORPOConfig\n", "# https://www.philschmid.de/fine-tune-google-gemma\n", "\n", "from trl import ORPOConfig, ORPOTrainer\n", "\n", "# in the following config, we combine the usual HF Trainer args with the ORPOConfig args (beta)\n", "\n", "cfg = ORPOConfig(\n", " output_dir='content/gemma-2b-orpo', # usual HF Trainer args: https://huggingface.co/docs/transformers/main_classes/trainer#transformers.Trainer.args\n", " num_train_epochs=3, # number of training epochs\n", " per_device_train_batch_size=2, # batch size per device during training\n", " gradient_accumulation_steps=2, # number of steps before performing a backward/update pass\n", " gradient_checkpointing=True, # use gradient checkpointing to save memory\n", " optim=\"adamw_torch_fused\", # use fused adamw optimizer\n", " logging_steps=20, # log every 20 steps\n", " bf16=True, # use bfloat16 precision\n", " tf32=True, # use tf32 \n", " learning_rate=5e-5, # learning rate\n", " warmup_ratio=0.1,\n", " warmup_steps=100,\n", " lr_scheduler_type=\"cosine\",\n", " max_prompt_length=512,\n", " remove_unused_columns=False,\n", " max_length=1024,\n", " beta=0.1, # ORPO beta\n", " save_total_limit=3, # args related to saving the model...\n", " save_strategy=\"epoch\",\n", " push_to_hub=True, \n", " report_to=['wandb'], # report metrics to Weights & Biases\n", " hub_model_id='anakin87/gemma-2b-orpo',\n", ")" ] }, { "cell_type": "code", "execution_count": 5, "id": "7f3fcb58-3cb5-4898-b4df-8023f85e9b1e", "metadata": {}, "outputs": [], "source": [ "import torch\n", "from transformers import AutoTokenizer, AutoModelForCausalLM" ] }, { "cell_type": "code", "execution_count": null, "id": "a8c20d77-38ee-4a29-9baf-ba75fd6c4c72", "metadata": {}, "outputs": [], "source": [ "model_id = \"google/gemma-2b\"\n", "tokenizer_id = \"philschmid/gemma-tokenizer-chatml\"\n", "\n", "\n", "# Load model and tokenizer\n", "model = AutoModelForCausalLM.from_pretrained(\n", " model_id,\n", " device_map=\"auto\",\n", " torch_dtype=torch.bfloat16,\n", " attn_implementation=\"flash_attention_2\",\n", ")\n", "tokenizer = AutoTokenizer.from_pretrained(tokenizer_id)\n", "tokenizer.padding_side = 'right' # to prevent warnings" ] }, { "cell_type": "code", "execution_count": 7, "id": "eabf18eb-8159-4712-8a4a-d69a29de794f", "metadata": {}, "outputs": [], "source": [ "from datasets import load_dataset\n", "import multiprocessing" ] }, { "cell_type": "code", "execution_count": null, "id": "5915fc40-96cd-4f86-8b7c-d394e801201a", "metadata": {}, "outputs": [], "source": [ "# https://github.com/huggingface/trl/blob/main/examples/scripts/orpo.py\n", "\n", "ds = load_dataset(\"alvarobartt/dpo-mix-7k-simplified\")\n", "\n", "def process(row):\n", " row[\"prompt\"] = tokenizer.apply_chat_template(row[\"prompt\"], tokenize=False)\n", " row[\"chosen\"] = tokenizer.apply_chat_template(row[\"chosen\"], tokenize=False)\n", " row[\"rejected\"] = tokenizer.apply_chat_template(row[\"rejected\"], tokenize=False)\n", " return row\n", "\n", "ds = ds.map(\n", " process,\n", " num_proc=multiprocessing.cpu_count(),\n", " load_from_cache_file=False,\n", ")\n", "train_dataset = ds[\"train\"]\n", "eval_dataset = ds[\"test\"]" ] }, { "cell_type": "code", "execution_count": 9, "id": "fff5ba61-7f23-4971-8472-cc339a2f803c", "metadata": {}, "outputs": [ { "data": { "text/plain": [ "{'dataset': 'argilla/distilabel-intel-orca-dpo-pairs',\n", " 'prompt': '<|im_start|>user\\nQ:Question: how old julio cesar chavez when he fought de la hoya I found the following answer on Google: He holds records for most successful consecutive defenses of world titles (27), most title fights (37), most title-fight victories (31) and he is after Joe Louis with (23) for most title defenses won by knockout (21). Is that a correct answer? Yes or no.\\nA:<|im_end|>\\n',\n", " 'chosen': \"<|im_start|>assistant\\n Yes, the information you found on Google is correct. Julio César Chávez holds several records related to world title defenses and victories, and he is considered one of the greatest boxers in history. Here is a detailed answer to your question:\\n\\nJulio César Chávez was born on July 12, 1962, in Ciudad Obregón, Sonora, Mexico. He began boxing at a young age and quickly made a name for himself in the sport, winning his first world title in 1984 when he defeated Mario Miranda for the WBC super featherweight championship.\\n\\nOver the course of his career, Chávez would go on to hold titles in three different weight classes (super featherweight, lightweight, and junior welterweight) and defend his titles a record 27 times, including 21 times by knockout. This record for most successful consecutive defenses of world titles has never been broken, and it is a testament to Chávez's skill, determination, and durability as a fighter.\\n\\nIn addition to his record for most successful title defenses, Chávez also holds the records for most title fights (37) and most title-fight victories (31). These records are also unbroken and demonstrate Chávez's consistent success and dominance in the ring.\\n\\nChávez's impressive record and achievements have earned him a place among the greatest boxers of all time, and he is often compared to other boxing legends such as Joe Louis, who holds the record for most title defenses won by knockout (23). While Chávez did not quite match Louis's record for most title defenses won by knockout, he is still widely regarded as one of the greatest boxers in history and a true legend of the sport.\\n\\nIn conclusion, the information you found on Google is correct: Julio César Chávez holds several records related to world title defenses and victories, and he is considered one of the greatest boxers in history. His impressive record and achievements have earned him a place among the greatest boxers of all time, and he will always be remembered as a true legend of the sport.<|im_end|>\\n\",\n", " 'rejected': '<|im_start|>assistant\\nNo, the answer provided does not directly answer the question regarding the age of Julio Cesar Chavez when he fought Oscar De La Hoya. The provided information describes some general records held by Julio Cesar Chavez throughout his career. To answer your original question, let me provide the relevant information:\\n\\nJulio Cesar Chavez fought Oscar De La Hoya on June 7, 1996, in a match called \"Ultimate Glory.\" Chavez was born on July 12, 1962. To calculate his age at the time of the fight, we need to find the difference between the fight date and his birthdate.\\n\\nFrom July 12, 1962, to June 7, 1996, there are:\\n- 33 years (from 1962 to 1995)\\n- An additional year from his birthday in 1995 (July 12, 1995) to the fight date in 1996 (June 7, 1996), which is approximately 10 months and 26 days.\\n\\nTherefore, Julio Cesar Chavez was about 33 years and 10 months old when he fought Oscar De La Hoya.<|im_end|>\\n'}" ] }, "execution_count": 9, "metadata": {}, "output_type": "execute_result" } ], "source": [ "train_dataset[0]" ] }, { "cell_type": "code", "execution_count": null, "id": "e3d29888-f3db-46eb-8302-6dff2ebf27f3", "metadata": {}, "outputs": [], "source": [ "import wandb\n", "run = wandb.init(project=\"YOUR_PROJECT_NAME\")" ] }, { "cell_type": "code", "execution_count": null, "id": "7a92cb46-4d17-4910-8cf0-40970d5f7193", "metadata": {}, "outputs": [], "source": [ "orpo_trainer = ORPOTrainer(\n", " model=model,\n", " args=cfg,\n", " train_dataset=train_dataset,\n", " eval_dataset=eval_dataset,\n", " tokenizer=tokenizer\n", ")" ] }, { "cell_type": "code", "execution_count": null, "id": "15d24895-204c-4313-943a-d6c31fcde6e5", "metadata": {}, "outputs": [], "source": [ "orpo_trainer.train()" ] }, { "cell_type": "code", "execution_count": null, "id": "ade7f4ed-339b-4609-ab41-8aa736001474", "metadata": {}, "outputs": [], "source": [ "orpo_trainer.push_to_hub()" ] } ], "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.10.12" } }, "nbformat": 4, "nbformat_minor": 5 }