Method comparison: Support more options for the optimizer (#2479)

Allow setting a different optimizer, including PEFT specific ones like
LoRA+.

Add experiment for LoRA-FA

Update param name, rm obsolete directories

Author: BenjaminBossan

method_comparison/MetaMathQA/default_training_params.json

@@ -9,6 +9,7 @@
   "compile": false,
   "seed": 0,
   "grad_norm_clip": 1.0,
+  "optimizer_type": "AdamW",
   "optimizer_kwargs": {
     "lr": 1e-4,
     "weight_decay": 0.1

method_comparison/MetaMathQA/experiments/lora/llama-3.2-3B-rank32-lorafa/adapter_config.json

@@ -0,0 +1,30 @@
+{
+  "alpha_pattern": {},
+  "auto_mapping": null,
+  "base_model_name_or_path": null,
+  "bias": "none",
+  "corda_config": null,
+  "eva_config": null,
+  "exclude_modules": null,
+  "fan_in_fan_out": false,
+  "inference_mode": false,
+  "init_lora_weights": true,
+  "layer_replication": null,
+  "layers_pattern": null,
+  "layers_to_transform": null,
+  "loftq_config": {},
+  "lora_alpha": 64,
+  "lora_bias": false,
+  "lora_dropout": 0.0,
+  "megatron_config": null,
+  "megatron_core": "megatron.core",
+  "modules_to_save": null,
+  "peft_type": "LORA",
+  "r": 32,
+  "rank_pattern": {},
+  "revision": null,
+  "target_modules": null,
+  "task_type": "CAUSAL_LM",
+  "use_dora": false,
+  "use_rslora": false
+}

method_comparison/MetaMathQA/experiments/lora/llama-3.2-3B-rank32-lorafa/training_params.json

@@ -0,0 +1,9 @@
+{
+  "optimizer_type": "lora-fa",
+  "optimizer_kwargs": {
+    "r": 32,
+    "lora_alpha": 64,
+    "lr": 1e-4,
+    "weight_decay": 0.1
+  }
+}

method_comparison/MetaMathQA/run.py

@@ -34,7 +34,7 @@
 from torch import nn
 from torch.amp import GradScaler, autocast
 from tqdm import tqdm
-from transformers import GenerationConfig, get_cosine_schedule_with_warmup, set_seed
+from transformers import GenerationConfig, set_seed
 from utils import (
     FILE_NAME_TRAIN_PARAMS,
     BucketIterator,
@@ -44,6 +44,7 @@
     get_base_model_info,
     get_dataset_info,
     get_model,
+    get_optimizer_and_scheduler,
     get_tokenizer,
     get_train_config,
     init_cuda,
@@ -63,7 +64,6 @@
 
 dtype_to_bytes_linear = {"float32": 4, "float16": 2, "bfloat16": 2, "int8": 1, "int4": 0.5}
 # if lr scheduler with warmup is used, the ratio of warmup steps to total steps
-WARMUP_STEP_RATIO = 0.1
 BUCKET_FACTOR = 20  # number of batches per bucket, increasing this further has diminishing returns
 
 
@@ -98,18 +98,6 @@ def evaluate(model, tokenizer, ds, batch_size, generate_kwargs, use_tqdm: bool =
     return predictions, responses
 
 
-class DummyScheduler:
-    # if no lr scheduler is being used
-    def __init__(self, lr):
-        self.lr = lr
-
-    def get_last_lr(self):
-        return [self.lr]
-
-    def step(self):
-        pass
-
-
 class DummyGradScaler:
     # if no mixed precision is being used
     def scale(self, loss):
@@ -136,6 +124,7 @@ def train(
     eval_steps: int,
     generation_kwargs: dict[str, Any],
     grad_norm_clip: float,
+    optimizer_type: str,
     optimizer_kwargs: dict[str, Any],
     query_template: str,
     lr_scheduler_arg: Optional[Literal["cosine"]],
@@ -156,16 +145,20 @@ def train(
     else:
         grad_scaler = DummyGradScaler()
         autocast_ctx = nullcontext
-    optimizer = torch.optim.AdamW(model.parameters(), **optimizer_kwargs)
-    if lr_scheduler_arg == "cosine":
-        warmup_steps = int(WARMUP_STEP_RATIO * max_steps)
-        lr_scheduler = get_cosine_schedule_with_warmup(
-            optimizer, num_warmup_steps=warmup_steps, num_training_steps=max_steps
-        )
-    elif lr_scheduler_arg is None:
-        lr_scheduler = DummyScheduler(optimizer_kwargs["lr"])
-    else:
-        raise ValueError(f"Invalid lr_scheduler argument: {lr_scheduler_arg}")
+
+    optimizer, lr_scheduler = get_optimizer_and_scheduler(
+        model,
+        optimizer_type=optimizer_type,
+        max_steps=max_steps,
+        lr_scheduler_arg=lr_scheduler_arg,
+        **optimizer_kwargs,
+    )
+    # print this after getting the optimizer, in case it modifies requires_gard
+    num_trainable_params, num_params = model.get_nb_trainable_parameters()
+    print_verbose(
+        f"trainable params: {num_trainable_params:,d} || all params: {num_params:,d} || "
+        f"trainable: {100 * num_trainable_params / num_params:.4f}%"
+    )
 
     status = TrainStatus.FAILED
     tic_train = time.perf_counter()
@@ -371,11 +364,6 @@ def main(*, path_experiment: str, experiment_name: str, clean: bool) -> None:
         autocast_adapter_dtype=train_config.autocast_adapter_dtype,
     )
     print_verbose(model)
-    num_trainable_params, num_params = model.get_nb_trainable_parameters()
-    print_verbose(
-        f"trainable params: {num_trainable_params:,d} || all params: {num_params:,d} || "
-        f"trainable: {100 * num_trainable_params / num_params:.4f}%"
-    )
 
     # train model
     try:
@@ -389,6 +377,7 @@ def main(*, path_experiment: str, experiment_name: str, clean: bool) -> None:
             eval_steps=train_config.eval_steps,
             generation_kwargs=train_config.generation_kwargs,
             grad_norm_clip=train_config.grad_norm_clip,
+            optimizer_type=train_config.optimizer_type,
             optimizer_kwargs=train_config.optimizer_kwargs,
             query_template=train_config.query_template,
             lr_scheduler_arg=train_config.lr_scheduler,

method_comparison/MetaMathQA/utils.py

@@ -37,10 +37,12 @@
     AutoModelForCausalLM,
     AutoTokenizer,
     BitsAndBytesConfig,
+    get_cosine_schedule_with_warmup,
 )
 
 import peft
 from peft import PeftConfig, get_peft_model, prepare_model_for_kbit_training
+from peft.optimizers import create_lorafa_optimizer, create_loraplus_optimizer
 from peft.utils import CONFIG_NAME
 
 
@@ -58,6 +60,7 @@
 # cancelled results
 RESULT_PATH_CANCELLED = os.path.join(os.path.dirname(__file__), "cancelled_results")
 hf_api = huggingface_hub.HfApi()
+WARMUP_STEP_RATIO = 0.1
 
 
 @dataclass
@@ -76,6 +79,7 @@ class TrainConfig:
         query_template: The template for the query
         seed: The random seed
         grad_norm_clip: The gradient norm clipping value (set to 0 to skip)
+        optimizer_type: The name of a torch optimizer (e.g. AdamW) or a PEFT method ("lora+", "lora-fa")
         optimizer_kwargs: The optimizer keyword arguments (lr etc.)
         lr_scheduler: The learning rate scheduler (currently only None or 'cosine' are supported)
         use_amp: Whether to use automatic mixed precision
@@ -95,6 +99,7 @@ class TrainConfig:
     query_template: str
     seed: int
     grad_norm_clip: float  # set to 0 to skip
+    optimizer_type: str
     optimizer_kwargs: dict[str, Any]
     lr_scheduler: Optional[Literal["cosine"]]
     use_amp: bool
@@ -121,6 +126,8 @@ def __post_init__(self) -> None:
             raise ValueError(f"Invalid eval_steps: {self.eval_steps} > max_steps: {self.max_steps}")
         if self.grad_norm_clip < 0:
             raise ValueError(f"Invalid grad_norm_clip: {self.grad_norm_clip}")
+        if self.optimizer_type not in ["lora+", "lora-fa"] and not hasattr(torch.optim, self.optimizer_type):
+            raise ValueError(f"Invalid optimizer_type: {self.optimizer_type}")
         if self.lr_scheduler not in [None, "cosine"]:
             raise ValueError(f"Invalid lr_scheduler: {self.lr_scheduler}, must be None or 'cosine'")
         if "{query}" not in self.query_template:
@@ -246,6 +253,42 @@ def get_model(
     return model
 
 
+class DummyScheduler:
+    # if no lr scheduler is being used
+    def __init__(self, lr):
+        self.lr = lr
+
+    def get_last_lr(self):
+        return [self.lr]
+
+    def step(self):
+        pass
+
+
+def get_optimizer_and_scheduler(
+    model, *, optimizer_type: str, max_steps: int, lr_scheduler_arg: Optional[Literal["cosine"]], **optimizer_kwargs
+) -> tuple[torch.optim.Optimizer, Any]:
+    if optimizer_type == "lora+":
+        optimizer = create_loraplus_optimizer(model, optimizer_cls=torch.optim.AdamW, **optimizer_kwargs)
+    elif optimizer_type == "lora-fa":
+        optimizer = create_lorafa_optimizer(model, **optimizer_kwargs)
+    else:
+        cls = getattr(torch.optim, optimizer_type)
+        optimizer = cls(model.parameters(), **optimizer_kwargs)
+
+    if lr_scheduler_arg == "cosine":
+        warmup_steps = int(WARMUP_STEP_RATIO * max_steps)
+        lr_scheduler = get_cosine_schedule_with_warmup(
+            optimizer, num_warmup_steps=warmup_steps, num_training_steps=max_steps
+        )
+    elif lr_scheduler_arg is None:
+        lr_scheduler = DummyScheduler(optimizer_kwargs["lr"])
+    else:
+        raise ValueError(f"Invalid lr_scheduler argument: {lr_scheduler_arg}")
+
+    return optimizer, lr_scheduler
+
+
 class BucketIterator:
     """
     Iterator that yields batches of data from a torch Dataset, grouped in buckets by sequence length