Update LoRA fine-tune example - more target_modules, lower LR, bf16

README.md

@@ -96,31 +96,40 @@ model = AutoModelForCausalLM.from_pretrained("ai21labs/Jamba-v0.1",
 </details>
 
 ### Fine-tuning example
-Jamba is a base model that can be fine-tuned for custom solutions (including for chat/instruct versions). You can fine-tune it using any technique of your choice. Here is an example of fine-tuning with the [PEFT](https://huggingface.co/docs/peft/index) library:
+Jamba is a base model that can be fine-tuned for custom solutions (including for chat/instruct versions). You can fine-tune it using any technique of your choice. Here is an example of fine-tuning with the [PEFT](https://huggingface.co/docs/peft/index) library (requires ~120GB GPU RAM, in example 2xA100 80GB):
 
 ```python
+import torch
 from datasets import load_dataset
-from trl import SFTTrainer
+from trl import SFTTrainer, SFTConfig
 from peft import LoraConfig
 from transformers import AutoTokenizer, AutoModelForCausalLM, TrainingArguments
 
 tokenizer = AutoTokenizer.from_pretrained("ai21labs/Jamba-v0.1")
-model = AutoModelForCausalLM.from_pretrained("ai21labs/Jamba-v0.1", device_map='auto')
+model = AutoModelForCausalLM.from_pretrained(
+    "ai21labs/Jamba-v0.1", device_map='auto', torch_dtype=torch.bfloat16)
+
+lora_config = LoraConfig(
+    r=8,
+    target_modules=[
+        "embed_tokens", 
+        "x_proj", "in_proj", "out_proj", # mamba
+        "gate_proj", "up_proj", "down_proj", # mlp
+        "q_proj", "k_proj", "v_proj" # attention
+    ],
+    task_type="CAUSAL_LM",
+    bias="none"
+)
 
 dataset = load_dataset("Abirate/english_quotes", split="train")
-training_args = TrainingArguments(
+training_args = SFTConfig(
     output_dir="./results",
-    num_train_epochs=3,
+    num_train_epochs=2,
     per_device_train_batch_size=4,
     logging_dir='./logs',
     logging_steps=10,
-    learning_rate=2e-3
-)
-lora_config = LoraConfig(
-    r=8,
-    target_modules=["embed_tokens", "x_proj", "in_proj", "out_proj"],
-    task_type="CAUSAL_LM",
-    bias="none"
+    learning_rate=1e-5,
+    dataset_text_field="quote",
 )
 trainer = SFTTrainer(
     model=model,
@@ -128,9 +137,7 @@ trainer = SFTTrainer(
     args=training_args,
     peft_config=lora_config,
     train_dataset=dataset,
-    dataset_text_field="quote",
 )
-
 trainer.train()
 ```