From 1665b068841b287013b58b9a2dbb3819bebf829f Mon Sep 17 00:00:00 2001 From: Daniel Han-Chen Date: Sun, 31 Dec 2023 17:57:45 +1100 Subject: [PATCH] DPO --- README.md | 20 ++++++------ unsloth/models/dpo.py | 68 ++++++++++++++++++++++++++++++++++++++++ unsloth/models/loader.py | 2 ++ 3 files changed, 80 insertions(+), 10 deletions(-) create mode 100644 unsloth/models/dpo.py diff --git a/README.md b/README.md index 47ca0a3187..b2bcab865e 100644 --- a/README.md +++ b/README.md @@ -8,25 +8,25 @@ | Llama 7b | Mistral 7b | CodeLlama 34b | Llama 7b Kaggle 2x T4 | |-----------------------------|-----------------------------|-------------------------|------------------------| | **2.2x faster, -43% VRAM** | **2.2x faster, -62% VRAM** | **1.9x faster, -27% VRAM** | **5.5x faster, -44% VRAM** | -| [Free Colab Alpaca dataset example](https://colab.research.google.com/drive/1lBzz5KeZJKXjvivbYvmGarix9Ao6Wxe5?usp=sharing) | [Free Colab Alpaca dataset example](https://colab.research.google.com/drive/1Dyauq4kTZoLewQ1cApceUQVNcnnNTzg_?usp=sharing) | [Colab A100 example](https://colab.research.google.com/drive/1y7A0AxE3y8gdj4AVkl2aZX47Xu3P1wJT?usp=sharing) | [Kaggle Alpaca example](https://www.kaggle.com/danielhanchen/unsloth-alpaca-t4-ddp) | +| [Free Colab Llama + Alpaca example](https://colab.research.google.com/drive/1lBzz5KeZJKXjvivbYvmGarix9Ao6Wxe5?usp=sharing) | [Free Colab Mistral + Alpaca example](https://colab.research.google.com/drive/1Dyauq4kTZoLewQ1cApceUQVNcnnNTzg_?usp=sharing) | [Colab A100 example](https://colab.research.google.com/drive/1y7A0AxE3y8gdj4AVkl2aZX47Xu3P1wJT?usp=sharing) | [Kaggle Alpaca example](https://www.kaggle.com/danielhanchen/unsloth-alpaca-t4-ddp) | | [Colab A100 example](https://colab.research.google.com/drive/1YIPY_18xm-K0iJDgvNkRoJsgkPMPAO3G?usp=sharing) | [Colab A100 example](https://colab.research.google.com/drive/1SKrKGV-BZoU4kv5q3g0jtE_OhRgPtrrQ?usp=sharing) | (59 more examples if you scroll down) | [Kaggle Slim Orca example](https://www.kaggle.com/danielhanchen/unsloth-slimorca-t4-ddp) | * Supports Llama, Yi, Mistral, CodeLlama, Qwen (llamafied), Deepseek and their derived models (Open Hermes etc). -* All kernels written in [OpenAI's Triton](https://openai.com/research/triton) language. **Manual backpropagation engine**. +* All kernels written in [OpenAI's Triton](https://openai.com/research/triton) language. **Manual backprop engine**. * **0% loss in accuracy** - no approximation methods - all exact. * No change of hardware necessary. Supports NVIDIA GPUs since 2018+. Minimum CUDA Compute Capability 7.0 (V100, T4, Titan V, RTX 20, 30, 40x, A100, H100, L40 etc) [Check your GPU](https://developer.nvidia.com/cuda-gpus) * **NEW!** Works on **Linux** and **Windows** via WSL. -* **NEW!** Support for [DPO (Direct Preference Optimization)](https://arxiv.org/abs/2305.18290), PPO and Reward Modelling via [TRL](https://huggingface.co/docs/trl/dpo_trainer). +* **NEW!** [DPO](https://arxiv.org/abs/2305.18290), PPO and Reward Modelling support via [TRL](https://huggingface.co/docs/trl/dpo_trainer). Example [DPO Colab notebook](https://colab.research.google.com/drive/15vttTpzzVXv_tJwEk-hIcQ0S9FcEWvwP?usp=sharing). * **NEW!** Download 4 bit models 4x faster from Huggingface! Eg: `unsloth/mistral-7b-bnb-4bit`. * Supports 4bit and 16bit QLoRA / LoRA finetuning via [bitsandbytes](https://github.com/TimDettmers/bitsandbytes). -* Open source version trains 5x faster - check out [Unsloth Max](https://unsloth.ai/) for **30x faster training**! +* Open source trains 5x faster - see [Unsloth Pro](https://unsloth.ai/) for **30x faster training**! -| 1 A100 40GB | Hugging Face | Flash Attention | Unsloth Open | Unsloth Equal | Unsloth Pro | Unsloth Max | -|--------------|-------------|-------------|-----------------|--------------|---------------|-------------| -| Alpaca | 1x | 1.04x | 1.98x | 2.48x | 5.32x | **15.64x** | -| LAION Chip2 | 1x | 0.92x | 1.61x | 1.84x | 7.05x | **20.73x** | -| OASST | 1x | 1.19x | 2.17x | 2.66x | 5.04x | **14.83x** | -| Slim Orca | 1x | 1.18x | 2.22x | 2.64x | 5.04x | **14.82x** | +| 1 A100 40GB | Hugging Face | Flash Attention | Unsloth Open Source | [Unsloth Pro](https://unsloth.ai/pricing) | +|--------------|--------------|-----------------|---------------------|-----------------| +| Alpaca | 1x | 1.04x | 1.98x | **15.64x** | +| LAION Chip2 | 1x | 0.92x | 1.61x | **20.73x** | +| OASST | 1x | 1.19x | 2.17x | **14.83x** | +| Slim Orca | 1x | 1.18x | 2.22x | **14.82x** | Join our [Discord](https://discord.gg/nsS4V5Z6ge)! If you trained a model with Unsloth, we made a cool sticker if you want to use it! diff --git a/unsloth/models/dpo.py b/unsloth/models/dpo.py new file mode 100644 index 0000000000..59e614983f --- /dev/null +++ b/unsloth/models/dpo.py @@ -0,0 +1,68 @@ +# Copyright 2023-present Daniel Han-Chen & the Unsloth team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from transformers.utils.notebook import ( + IntervalStrategy, + NotebookTrainingTracker, + NotebookProgressCallback, +) + +DPOTrainer_metrics = [ + "rewards/chosen", + "rewards/rejected", + "rewards/accuracies", + "rewards/margins", + "logps/rejected", + "logps/chosen", + "logits/rejected", + "logits/chosen", +] + +def NotebookProgressCallback_on_train_begin(self, args, state, control, **kwargs): + self.first_column = "Epoch" if args.evaluation_strategy == IntervalStrategy.EPOCH else "Step" + self.training_loss = 0 + self.last_log = 0 + column_names = [self.first_column] + ["Training Loss"] + if args.evaluation_strategy != IntervalStrategy.NO: + column_names.append("Validation Loss") + column_names += [x.replace("/", " / ") for x in DPOTrainer_metrics] + self.training_tracker = NotebookTrainingTracker(state.max_steps, column_names) +pass + + +def NotebookProgressCallback_on_log(self, args, state, control, logs=None, **kwargs): + # Only for when there is no evaluation + if args.evaluation_strategy == IntervalStrategy.NO and "loss" in logs: + values = {"Training Loss": logs["loss"]} + for metric in DPOTrainer_metrics: + values[metric.replace("/", " / ")] = logs[metric] + # First column is necessarily Step since we're not in epoch eval strategy + values["Step"] = state.global_step + self.training_tracker.write_line(values) +pass + + +def patch_dpo_trainer(): + # We patch Jupyter Notebook's printing to include all columns for DPO. + NotebookProgressCallback.on_train_begin = NotebookProgressCallback_on_train_begin + NotebookProgressCallback.on_log = NotebookProgressCallback_on_log +pass +# Patch DPO notebook printing +patch_dpo_trainer() + + +from trl import DPOTrainer +class FastDPOTrainer(DPOTrainer): + pass +pass diff --git a/unsloth/models/loader.py b/unsloth/models/loader.py index 421b743299..baa55f9f00 100644 --- a/unsloth/models/loader.py +++ b/unsloth/models/loader.py @@ -14,6 +14,7 @@ from .llama import FastLlamaModel, logger from .mistral import FastMistralModel +from .dpo import FastDPOTrainer from transformers import AutoConfig from transformers import __version__ as transformers_version @@ -23,6 +24,7 @@ FOURBIT_MAPPER = \ "unsloth/llama-2-7b-bnb-4bit" : "unsloth/llama-2-7b", "unsloth/llama-2-13b-bnb-4bit" : "unsloth/llama-13-7b", "unsloth/codellama-34b-bnb-4bit" : "codellama/CodeLlama-34b-hf", + "unsloth/zephyr-sft-bnb-4bit" : "unsloth/zephyr-sft", } # https://github.com/huggingface/transformers/pull/26037 allows 4 bit loading!