oumi-ai · optas · Jan 24, 2025 · Jan 24, 2025
diff --git a/docs/user_guides/train/configuration.md b/docs/user_guides/train/configuration.md
@@ -228,6 +228,7 @@ peft:
   lora_modules_to_save: null         # Modules to unfreeze and train
   lora_bias: "none"                  # Bias training type
   lora_task_type: "CAUSAL_LM"        # Task type for adaptation
+  init_lora_weights: "DEFAULT"       # Initialization of LoRA weights
 
   # Q-LoRA settings
   q_lora: false                      # Enable quantization

diff --git a/docs/user_guides/train/train.md b/docs/user_guides/train/train.md
@@ -112,7 +112,7 @@ In the following sections, we'll cover some common workflows for training.
 
 ### Fine-tuning a Pre-trained Model
 
-The simplest workflow is to fine-tune a pre-trained model on a dataset. This will perform fully finetune the model using SFT (supervised fine-tuning).
+The simplest workflow is to fine-tune a pre-trained model on a dataset. The following will fully finetune the model using SFT (supervised fine-tuning).
 
 ```yaml
 model:
@@ -133,6 +133,42 @@ training:
   max_steps: 10  # Number of training steps
 ```
 
+### Using Parameter-Efficient Fine-tuning (PEFT)
+
+Excellent results can be achieved at a fraction of the computational cost by fine-tuning your network with [Low Rank (LoRA) adapters](https://arxiv.org/abs/2106.09685) instead of updating all original parameters. The following adaptation enables _parameter efficient fine-tuning_ with very few additions:
+
+
+```yaml
+model:
+  model_name: "meta-llama/Meta-Llama-3.2-3B-Instruct"  # Replace with your model
+  trust_remote_code: true
+  dtype: "bfloat16"
+
+data:
+  train:  # Training dataset mixture, can be a single dataset or a list of datasets
+    datasets:
+      - dataset_name: "yahma/alpaca-cleaned" # Replace with your dataset, or add more datasets
+        split: "train"
+
+training:
+  output_dir: "output/llama-finetuned"  # Where to save outputs
+  optimizer: "adamw_torch_fused"
+  learning_rate: 2e-5
+  max_steps: 10  # Number of training steps
+  use_peft: True  # Activate Parameter Efficient Fine-Tuning
+
+peft: # Control key hyper-parameters of the PEFT training process
+  lora_r: 64
+  lora_alpha: 128
+  lora_target_modules: # Select the modules for which adapters will be added
+    - "q_proj"
+    - "v_proj"
+    - "o_proj"
+    - "gate_proj"
+    - "up_proj"
+    - "down_proj"
+```
+
 ### Fine-tuning a Vision-Language Model
 
 Multimodal support in Oumi is similar to support for text-only models with few config changes e.g., data collation.

diff --git a/docs/user_guides/train/training_methods.md b/docs/user_guides/train/training_methods.md
@@ -25,7 +25,7 @@ Supervised Fine-Tuning (SFT) is the most common approach for adapting a pre-trai
 - **Agent development:** Training language models to act as agents that can interact with their environment and perform tasks autonomously. This involves fine-tuning the model on data that demonstrates how to complete tasks, communicate effectively, and make decisions.
 - **Tool use:** Fine-tuning models to effectively use external tools (e.g., calculators, APIs, databases) to augment their capabilities. This involves training on data that shows how to call tools, interpret their outputs, and integrate them into problem-solving.
 - **Structured data extraction:** Training models to extract structured information from unstructured text. This can be used to extract entities, relationships, or key events from documents, enabling automated data analysis and knowledge base construction.
-- **Text generation:** Generating coherent text, code, scripts, email, etc based on a prompt.
+- **Text generation:** Generating coherent text, code, scripts, email, etc. based on a prompt.
 
 ### Data Format
 SFT uses the {class}`~oumi.core.types.conversation.Conversation` format, which represents a conversation between a user and an assistant. Each turn in the conversation is represented by a message with a role ("user" or "assistant") and content.