Update README.md

README.md

@@ -3,13 +3,235 @@ library_name: transformers
 tags: []
 ---
 
-# Model Card for Model ID
+# Fine-tune Llama 3 with ORPO 
+
+ORPO is a new exciting fine-tuning technique that combines the traditional supervised fine-tuning and preference alignment stages into a single process. This reduces the computational resources and time required for training. Moreover, empirical results demonstrate that ORPO outperforms other alignment methods on various model sizes and benchmarks.
+
+In this article, we will fine-tune the new Llama 3 8B model using ORPO with the TRL library. 
 
 <!-- Provide a quick summary of what the model is/does. -->
 
+## ORPO
+Instruction tuning and preference alignment are essential techniques for adapting Large Language Models (LLMs) to specific tasks. Traditionally, this involves a multi-stage process: 1/ Supervised Fine-Tuning (SFT) on instructions to adapt the model to the target domain, followed by 2/ preference alignment methods like Reinforcement Learning with Human Feedback (RLHF) or Direct Preference Optimization (DPO) to increase the likelihood of generating preferred responses over rejected ones.
+
+However, researchers have identified a limitation in this approach. While SFT effectively adapts the model to the desired domain, it inadvertently increases the probability of generating undesirable answers alongside preferred ones. This is why the preference alignment stage is necessary to widen the gap between the likelihoods of preferred and rejected outputs.
+
+see more on ORPO [link](https://arxiv.org/abs/2403.07691)
+
+## Fine-tuning Llama 3 with ORPO 
+
+[Llama 3](https://github.com/meta-llama/llama3/tree/main) is the latest family of LLMs developed by Meta. The models were trained on an extensive dataset of 15 trillion tokens (compared to 2T tokens for Llama 2). Two model sizes have been released: a 70 billion parameter model and a smaller 8 billion parameter model. The 70B model has already demonstrated impressive performance, scoring 82 on the MMLU benchmark and 81.7 on the HumanEval benchmark.
+
+Llama 3 models also increased the context length up to 8,192 tokens (4,096 tokens for Llama 2), and potentially scale up to 32k with RoPE. Additionally, the models use a new tokenizer with a 128K-token vocabulary, reducing the number of tokens required to encode text by 15%. This vocabulary also explains the bump from 7B to 8B parameters.
+
+# Required packages
+```bash
+pip install -U transformers datasets accelerate peft trl bitsandbytes wandb
+pip install -qqq flash-attn
+pip install -qU transformers accelerate
+```
+
+Once it's installed, we can import the necessary libraries and log in to W&B (optional):
+
+```python
+
+"""
+wandb
+https://wandb.ai/aliforootani-UFZ
+you need wb_token
+"""
+
+import gc
+import os
+
+import torch
+import wandb
+from datasets import load_dataset
+
+
+
+# Directly insert your Weights & Biases API key here
+wb_token = 'your_wb_token'
+wandb.login(key=wb_token)
+
+
+from peft import LoraConfig, PeftModel, prepare_model_for_kbit_training
+
+from transformers import (
+    AutoModelForCausalLM,
+    AutoTokenizer,
+    BitsAndBytesConfig,
+    TrainingArguments,
+    pipeline,)
+
+from trl import ORPOConfig, ORPOTrainer, setup_chat_format
+```
+
+If you have a recent GPU, you should also be able to use the Flash Attention library to replace the default eager attention implementation with a more efficient one.
+
+
+```python
+
+if torch.cuda.get_device_capability()[0] >= 128:
+    
+    attn_implementation = "flash_attention_2"
+    torch_dtype = torch.bfloat16
+else:
+    attn_implementation = "eager"
+    torch_dtype = torch.float16
+
+
+##################################
+
+import sys
+import os
+
+cwd = os.getcwd()
+# sys.path.append(cwd + '/my_directory')
+sys.path.append(cwd)
+
+
+def setting_directory(depth):
+    current_dir = os.path.abspath(os.getcwd())
+    root_dir = current_dir
+    for i in range(depth):
+        root_dir = os.path.abspath(os.path.join(root_dir, os.pardir))
+        sys.path.append(os.path.dirname(root_dir))
+    return root_dir
+
+# I load the model from local directory!
+model_path = "/data/bio-eng-llm/llm_repo/mlabonne/OrpoLlama-3-8B"
+```
+
+In the following, we will load the OrpoLlama-3-8B in 4-bit precision thanks to bitsandbytes. We then set the LoRA configuration using PEFT for QLoRA. I'm also using the convenient setup_chat_format() function to modify the model and tokenizer for ChatML support. It automatically applies this chat template, adds special tokens, and resizes the model's embedding layer to match the new vocabulary size.
+
+
+
+```python
+
+# QLoRA config
+bnb_config = BitsAndBytesConfig(
+    load_in_4bit=True,
+    bnb_4bit_quant_type="nf4",
+    bnb_4bit_compute_dtype= torch_dtype,
+    bnb_4bit_use_double_quant=True,
+)
+
+# LoRA config
+peft_config = LoraConfig(
+    r=16,
+    lora_alpha=32,
+    lora_dropout=0.05,
+    bias="none",
+    task_type="CAUSAL_LM",
+    target_modules=['up_proj', 'down_proj', 'gate_proj', 'k_proj', 'q_proj', 'v_proj', 'o_proj']
+)
+
+# Load tokenizer
+tokenizer = AutoTokenizer.from_pretrained(model_path)
+
+# Load model
+model = AutoModelForCausalLM.from_pretrained(
+    model_path,
+    quantization_config=bnb_config,
+    device_map="auto",
+    attn_implementation= attn_implementation
+)
+
+
+model, tokenizer = setup_chat_format(model, tokenizer)
+model = prepare_model_for_kbit_training(model)
+```
+
+Now that the model is ready for training, we can take care of the dataset. We load mlabonne/orpo-dpo-mix-40k and use the apply_chat_template() function to convert the "chosen" and "rejected" columns into the ChatML format. Note that I'm only using 1,000 samples and not the entire dataset, as it would take too long to run.
+
+First, we need to set a few hyperparameters:
+
+    learning_rate: ORPO uses very low learning rates compared to traditional SFT or even DPO. This value of 8e-6 comes from the original paper, and roughly corresponds to an SFT learning rate of 1e-5 and a DPO learning rate of 5e-6. I would recommend increasing it around 1e-6 for a real fine-tune.
+    beta: It is the $\lambda$ parameter in the paper, with a default value of 0.1. An appendix from the original paper shows how it's been selected with an ablation study.
+    Other parameters, like max_length and batch size are set to use as much VRAM as available (~20 GB in this configuration). Ideally, we would train the model for 3-5 epochs, but we'll stick to 1 here.
+
+Finally, we can train the model using the ORPOTrainer, which acts as a wrapper.
+
+
+```python
+
+dataset_name = "/data/bio-eng-llm/llm_repo/mlabonne/OrpoLlama-3-8B"
+
+dataset = load_dataset(dataset_name, split="all")
+dataset = dataset.shuffle(seed=42).select(range(1000))
+
+
+def format_chat_template(row):
+    row["chosen"] = tokenizer.apply_chat_template(row["chosen"], tokenize=False)
+    row["rejected"] = tokenizer.apply_chat_template(row["rejected"], tokenize=False)
+    return row
+
+dataset = dataset.map(
+    format_chat_template,
+    num_proc= os.cpu_count(),
+)
+dataset = dataset.train_test_split(test_size=0.01)
+
+epochs=20
+
+orpo_args = ORPOConfig(
+    learning_rate=8e-6,
+    beta=0.1,
+    lr_scheduler_type="linear",
+    max_length=1024,
+    max_prompt_length=512,
+    per_device_train_batch_size=2,
+    per_device_eval_batch_size=2,
+    gradient_accumulation_steps=4,
+    optim="paged_adamw_8bit",
+    num_train_epochs=epochs,
+    evaluation_strategy="steps",
+    eval_steps=0.2,
+    logging_steps=1,
+    warmup_steps=10,
+    report_to="wandb",
+    output_dir="./results/",
+)
+
+trainer = ORPOTrainer(
+    model=model,
+    args=orpo_args,
+    train_dataset=dataset["train"],
+    eval_dataset=dataset["test"],
+    peft_config=peft_config,
+    tokenizer=tokenizer,
+)
+trainer.train()
+
+import os
+
+# Define the directory where you want to save the model
+# 
+
+root_dir = setting_directory(0)
+
+save_dir = root_dir + f"models/fine_tuned_models/OrpoLlama-3-8B_{epochs}e_qa_qa"
+#trainer.save_model(save_dir)
+
+
+# Create the directory if it doesn't exist
+os.makedirs(save_dir, exist_ok=True)
+
+# Combine the directory path with the model name
+#new_model_path = os.path.join(save_dir, "OrpoLlama-3-8B")
+
+# Save the model to the specified directory
+trainer.save_model(save_dir)
+
+
+#new_model = "OrpoLlama-3-8B"
+#trainer.save_model(new_model)
+```
+
+Training the model on these 1,000 samples and 20 epochs took about 22 hours on an Nvidia-A100 80GB GPU, but based on the Wnadb graphs only 34GB has been used. Let's check the W&B plots:
 
 
-## Model Details
 
 ### Model Description