FEAT Add OLoRA initialization strategy to LoRA (#1828)

docs/source/developer_guides/lora.md

@@ -54,6 +54,15 @@ lora_config = LoraConfig(init_lora_weights="pissa_niter_[number of iters]", ...)
 ```
 For detailed instruction on using PiSSA, please follow [these instructions](https://github.com/fxmeng/peft/tree/main/examples/pissa_finetuning).
 
+### OLoRA
+[OLoRA](https://arxiv.org/abs/2406.01775) utilizes QR decomposition to initialize the LoRA adapters. OLoRA translates the base weights of the model by a factor of their QR decompositions, i.e., it mutates the weights before performing any training on them. This approach significantly improves stability, accelerates convergence speed, and ultimately achieves superior performance.
+
+You just need to pass a single additional option to use OLoRA:
+```python
+from peft import LoraConfig
+config = LoraConfig(init_lora_weights="olora", ...)
+```
+For more advanced usage, please refer to our [documentation](https://github.com/huggingface/peft/tree/main/examples/olora_finetuning).
 ### LoftQ
 
 #### Standard approach

examples/olora_finetuning/README.md

@@ -0,0 +1,84 @@
+# OLoRA: Orthonormal Low Rank Adaptation of Large Language Models
+
+## Introduction
+[OLoRA](https://arxiv.org/abs/2406.01775) is a novel approach that leverages orthonormal low rank adaptation through QR decomposition. Unlike the default LoRA implementation, OLoRA decomposes original weights into their $\mathbf{Q}$ and $\mathbf{R}$ parts, and then uses the first `rank` rows of $\mathbf{R}$ and the first `rank` columns of $\mathbf{Q}$ to initialize $\mathbf{A}$ and $\mathbf{B}$, respectively. This results in significantly faster convergence, more stable training, and superior performance.
+
+## Quick start
+```python
+import torch
+from peft import LoraConfig, get_peft_model
+from transformers import AutoTokenizer, AutoModelForCausalLM
+from trl import SFTTrainer
+from datasets import load_dataset
+
+model = AutoModelForCausalLM.from_pretrained("facebook/opt-350m", torch_dtype=torch.bfloat16, device_map="auto")
+tokenizer = AutoTokenizer.from_pretrained("facebook/opt-350m")
+dataset = load_dataset("imdb", split="train[:1%]")
+lora_config = LoraConfig(
+    init_lora_weights="olora"
+)
+peft_model = get_peft_model(model, lora_config)
+trainer = SFTTrainer(
+    model=peft_model,
+    train_dataset=dataset,
+    dataset_text_field="text",
+    max_seq_length=512,
+    tokenizer=tokenizer,
+)
+trainer.train()
+peft_model.save_pretrained("olora-opt-350m")
+```
+
+There is no additional change needed to your standard LoRA procedure, except for specifying `init_lora_weights = "olora"` option in your lora configuration.
+
+Additionally you can refer to olora finetuning script.
+Run the script simply by running:
+```bash
+python3 examples/olora_finetuning/olora_finetuning.py --base_model facebook/opt-350m
+```
+OLoRA also supports quantization. To use 4-bit quantization try:
+```bash
+python3 examples/olora_finetuning/olora_finetuning.py --base_model facebook/opt-350m --quantize
+```
+
+
+## Use the model
+You can load and use the model as any other 🤗 PEFT model
+```python
+from peft import PeftModel
+model = AutoModelForCausalLM.from_pretrained("facebook/opt-350m")
+tokenizer = AutoTokenizer.from_pretrained("facebook/opt-350m")
+olora_model = PeftModel.from_pretrained(model, "olora-opt-350m")
+```
+
+## OLoRA and LoRA
+OLoRA differs from LoRA in that it mutates the original weights. To utilize multiple adapters simultaneously, you can leverage the `path_initial_model_for_weight_conversion` option. Below is a simple template illustrating how to convert OLoRA to conventional LoRA:
+```python
+base_model = AutoModel.from_pretrained("facebook/opt-350m")
+olora_config = LoraConfig(
+    ...
+    init_lora_weights = "olora" # Initialize the model with OLoRA
+)
+olora_model = get_peft_model(base_model, olora_config)
+init_path = <path-to-untrained-olora-model>
+olora_model.save_pretrained(init_path) # Save the model *before* performing any training
+
+# Train the model
+train(olora_model) # Your training loop
+
+#Save the model after training
+olora_model.save_pretrained(output_dir, path_initial_model_for_weight_conversion=init_path) 
+```
+After completing training, you can save and convert your OLoRA model to a conventional LoRA model by setting `path_initial_model_for_weight_conversion` to `init_path`, that is the path of your untrained OLoRA model. This conversion enables you to use multiple adapters with your LoRA model.
+
+## Citation
+```
+@misc{büyükakyüz2024olora,
+      title={OLoRA: Orthonormal Low-Rank Adaptation of Large Language Models}, 
+      author={Kerim Büyükakyüz},
+      year={2024},
+      eprint={2406.01775},
+      archivePrefix={arXiv},
+      primaryClass={cs.CL}
+}
+```

examples/olora_finetuning/olora_finetuning.py

@@ -0,0 +1,184 @@
+# Copyright 2024-present the HuggingFace Inc. team.
+#
+# 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 typing import List
+
+import torch
+import transformers
+from datasets import load_dataset
+from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
+
+from peft import (
+    LoraConfig,
+    get_peft_model,
+)
+
+
+def train(
+    base_model: str = "path/to/model",
+    data_path: str = "yahma/alpaca-cleaned",
+    output_dir: str = "olora",
+    batch_size: int = 16,
+    num_epochs: int = 1,
+    learning_rate: float = 3e-4,
+    cutoff_len: int = 256,
+    val_set_size: int = 16,
+    quantize: bool = False,
+    eval_step: int = 100,
+    save_step: int = 100,
+    device_map: str = "auto",
+    lora_r: int = 32,
+    lora_alpha: int = 16,
+    lora_dropout: float = 0.05,
+    lora_target_modules: List[str] = None,
+    init_lora_weights="olora",
+):
+    model = AutoModelForCausalLM.from_pretrained(
+        base_model,
+        device_map=device_map,
+        quantization_config=BitsAndBytesConfig(
+            load_in_4bit=True,
+            bnb_4bit_compute_dtype=torch.bfloat16,
+            bnb_4bit_use_double_quant=True,
+            bnb_4bit_quant_type="nf4",
+        )
+        if quantize
+        else None,
+        torch_dtype=torch.float16,
+    )
+
+    tokenizer = AutoTokenizer.from_pretrained(base_model, trust_remote_code=True)
+
+    def tokenize(prompt, add_eos_token=True):
+        result = tokenizer(
+            prompt,
+            truncation=True,
+            max_length=cutoff_len,
+            padding=False,
+            return_tensors=None,
+        )
+        if (
+            result["input_ids"][-1] != tokenizer.eos_token_id
+            and len(result["input_ids"]) < cutoff_len
+            and add_eos_token
+        ):
+            result["input_ids"].append(tokenizer.eos_token_id)
+            result["attention_mask"].append(1)
+
+        result["labels"] = result["input_ids"].copy()
+
+        return result
+
+    def generate_and_tokenize_prompt(example):
+        full_prompt = generate_prompt(example)
+        tokenized_full_prompt = tokenize(full_prompt)
+        return tokenized_full_prompt
+
+    config = LoraConfig(
+        r=lora_r,
+        lora_alpha=lora_alpha,
+        target_modules=lora_target_modules,
+        lora_dropout=lora_dropout,
+        bias="none",
+        task_type="CAUSAL_LM",
+        init_lora_weights=init_lora_weights,
+    )
+    model = get_peft_model(model, config)
+
+    data = load_dataset(data_path)
+
+    train_val = data["train"].train_test_split(test_size=val_set_size, shuffle=True, seed=42)
+    train_data = train_val["train"].shuffle().map(generate_and_tokenize_prompt)
+    val_data = train_val["test"].shuffle().map(generate_and_tokenize_prompt)
+
+    trainer = transformers.Trainer(
+        model=model,
+        train_dataset=train_data,
+        eval_dataset=val_data,
+        args=transformers.TrainingArguments(
+            per_device_train_batch_size=batch_size,
+            warmup_steps=100,
+            num_train_epochs=num_epochs,
+            learning_rate=learning_rate,
+            fp16=True,
+            logging_steps=100,
+            optim="adamw_torch",
+            evaluation_strategy="steps",
+            save_strategy="steps",
+            eval_steps=eval_step,
+            save_steps=save_step,
+            output_dir=output_dir,
+            save_total_limit=3,
+            load_best_model_at_end=True,
+        ),
+        data_collator=transformers.DataCollatorForSeq2Seq(
+            tokenizer, pad_to_multiple_of=8, return_tensors="pt", padding=True
+        ),
+    )
+    trainer.train()
+    model.save_pretrained(output_dir)
+
+
+def generate_prompt(example):
+    return f"""Below is an instruction that describes a task. Write a response that appropriately completes the request.
+            ### Instruction:
+            {example["instruction"]}
+            ### Response:
+            {example["output"]}"""
+
+
+if __name__ == "__main__":
+    import argparse
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--base_model", type=str, default="path/to/model")
+    parser.add_argument("--data_path", type=str, default="yahma/alpaca-cleaned")
+    parser.add_argument("--output_dir", type=str, default="olora")
+    parser.add_argument("--batch_size", type=int, default=16)
+    parser.add_argument("--num_epochs", type=int, default=1)
+    parser.add_argument("--learning_rate", type=float, default=3e-4)
+    parser.add_argument("--cutoff_len", type=int, default=256)
+    parser.add_argument("--val_set_size", type=int, default=16)
+    parser.add_argument("--quantize", action="store_true")
+    parser.add_argument("--eval_step", type=int, default=100)
+    parser.add_argument("--save_step", type=int, default=100)
+    parser.add_argument("--device_map", type=str, default="auto")
+    parser.add_argument("--lora_r", type=int, default=32)
+    parser.add_argument("--lora_alpha", type=int, default=16)
+    parser.add_argument("--lora_dropout", type=float, default=0.05)
+    parser.add_argument("--lora_target_modules", type=str, default=None)
+    parser.add_argument("--init_lora_weights", type=str, default="olora")
+
+    args = parser.parse_args()
+
+    train(
+        base_model=args.base_model,
+        data_path=args.data_path,
+        output_dir=args.output_dir,
+        batch_size=args.batch_size,
+        num_epochs=args.num_epochs,
+        learning_rate=args.learning_rate,
+        cutoff_len=args.cutoff_len,
+        val_set_size=args.val_set_size,
+        quantize=args.quantize,
+        eval_step=args.eval_step,
+        save_step=args.save_step,
+        device_map=args.device_map,
+        lora_r=args.lora_r,
+        lora_alpha=args.lora_alpha,
+        lora_dropout=args.lora_dropout,
+        lora_target_modules=args.lora_target_modules,
+        init_lora_weights=args.init_lora_weights,
+    )