QOL improvements and doc updates

docs/source/developer_guides/lora.md

@@ -135,4 +135,4 @@ model.unload()
 
 # delete adapter
 model.delete_adapter("dpo")
-```
+```

src/peft/__init__.py

@@ -83,5 +83,6 @@
     set_peft_model_state_dict,
     shift_tokens_right,
     load_peft_weights,
+    cast_non_trainable_to_dtype,
 )
 from .config import PeftConfig, PromptLearningConfig

src/peft/utils/__init__.py

@@ -47,5 +47,6 @@
     get_auto_gptq_quant_linear,
     get_quantization_config,
     id_tensor_storage,
+    cast_non_trainable_to_dtype,
 )
 from .save_and_load import get_peft_model_state_dict, set_peft_model_state_dict, load_peft_weights

src/peft/utils/other.py

@@ -498,3 +498,23 @@ def id_tensor_storage(tensor: torch.Tensor) -> Tuple[torch.device, int, int]:
         unique_id = storage_ptr(tensor)
 
     return tensor.device, unique_id, storage_size(tensor)
+
+
+def cast_non_trainable_to_dtype(model, dtype):
+    """
+    Cast all non-trainable parameters of the model to the given `dtype`. The trainable parameters are casted to full
+    precision. This is meant to reduce the GPU memory usage when using PEFT methods by using half-precision dtype for
+    non-trainable parameters. Having the trainable parameters in full-precision preserves training stability when using
+    automatic mixed precision training.
+
+    Args:
+        model (`torch.nn.Module`):
+            The model to cast the non-trainable parameters of.
+        dtype (`torch.dtype`):
+            The dtype to cast the non-trainable parameters to.
+    """
+    for p in model.parameters():
+        if not p.requires_grad:
+            p.data = p.to(dtype)
+        else:
+            p.data = p.to(torch.float32)

src/peft/utils/peft_types.py

@@ -20,7 +20,22 @@
 
 
 class PeftType(str, enum.Enum):
-    """Enum class for the different types of adapters in PEFT."""
+    """
+    Enum class for the different types of adapters in PEFT.
+
+    Supported PEFT types:
+    - PROMPT_TUNING
+    - MULTITASK_PROMPT_TUNING
+    - P_TUNING
+    - PREFIX_TUNING
+    - LORA
+    - ADALORA
+    - ADAPTION_PROMPT
+    - IA3
+    - LOHA
+    - LOKR
+    - OFT
+    """
 
     PROMPT_TUNING = "PROMPT_TUNING"
     MULTITASK_PROMPT_TUNING = "MULTITASK_PROMPT_TUNING"
@@ -36,7 +51,18 @@ class PeftType(str, enum.Enum):
 
 
 class TaskType(str, enum.Enum):
-    """Enum class for the different types of tasks supported by PEFT."""
+    """
+    Enum class for the different types of tasks supported by PEFT.
+
+    Overview of the supported task types:
+    - SEQ_CLS: Text classification.
+    - SEQ_2_SEQ_LM: Sequence-to-sequence language modeling.
+    - Causal LM: Causal language modeling.
+    - TOKEN_CLS: Token classification.
+    - QUESTION_ANS: Question answering.
+    - FEATURE_EXTRACTION: Feature extraction. Provides the hidden states which can be used as embeddings or features
+      for downstream tasks.
+    """
 
     SEQ_CLS = "SEQ_CLS"
     SEQ_2_SEQ_LM = "SEQ_2_SEQ_LM"