Add LoRA LR for DS Chat steps 1-3 (deepspeedai#685)

This PR adds an explicit LoRA learning rate argument for DS Chat steps 1 through 3.

Step 1:
- lora_learning_rate

Step 2:
- lora_learning_rate

Step 3:
- actor_lora_learning_rate
- critic_lora_learning_rate

Author: lekurile

applications/DeepSpeed-Chat/training/step1_supervised_finetuning/main.py

@@ -161,6 +161,13 @@ def parse_args():
     parser.add_argument('--only_optimize_lora',
                         action='store_true',
                         help='Only optimize the LoRA parameters.')
+    parser.add_argument(
+        "--lora_learning_rate",
+        type=float,
+        default=5e-4,
+        help=
+        "Initial LoRA learning rate (after the potential warmup period) to use."
+    )
     ## Tensorboard logging
     parser.add_argument('--enable_tensorboard',
                         action='store_true',
@@ -274,7 +281,7 @@ def evaluation(model, eval_dataloader):
 
     # Split weights in two groups, one with weight decay and the other not.
     optimizer_grouped_parameters = get_optimizer_grouped_parameters(
-        model, args.weight_decay)
+        model, args.weight_decay, args.lora_learning_rate)
 
     AdamOptimizer = DeepSpeedCPUAdam if args.offload else FusedAdam
     optimizer = AdamOptimizer(optimizer_grouped_parameters,

applications/DeepSpeed-Chat/training/step2_reward_model_finetuning/main.py

@@ -161,6 +161,13 @@ def parse_args():
     parser.add_argument('--only_optimize_lora',
                         action='store_true',
                         help='Only optimize the LoRA parameters.')
+    parser.add_argument(
+        "--lora_learning_rate",
+        type=float,
+        default=5e-4,
+        help=
+        "Initial LoRA learning rate (after the potential warmup period) to use."
+    )
     ## Tensorboard logging
     parser.add_argument('--enable_tensorboard',
                         action='store_true',
@@ -271,7 +278,7 @@ def evaluation_reward(model, eval_dataloader):
 
     # Split weights in two groups, one with weight decay and the other not.
     optimizer_grouped_parameters = get_optimizer_grouped_parameters(
-        rm_model, args.weight_decay)
+        rm_model, args.weight_decay, args.lora_learning_rate)
 
     AdamOptimizer = DeepSpeedCPUAdam if args.offload else FusedAdam
     optimizer = AdamOptimizer(optimizer_grouped_parameters,

applications/DeepSpeed-Chat/training/step3_rlhf_finetuning/main.py

@@ -286,6 +286,20 @@ def parse_args():
     parser.add_argument('--only_optimize_lora',
                         action='store_true',
                         help='Only optimize the LoRA parameters.')
+    parser.add_argument(
+        "--actor_lora_learning_rate",
+        type=float,
+        default=5e-4,
+        help=
+        "Initial actor LoRA learning rate (after the potential warmup period) to use."
+    )
+    parser.add_argument(
+        "--critic_lora_learning_rate",
+        type=float,
+        default=5e-4,
+        help=
+        "Initial critic LoRA learning rate (after the potential warmup period) to use."
+    )
     ## Make EMA as an optional feature
     parser.add_argument('--enable_ema',
                         action='store_true',

applications/DeepSpeed-Chat/training/step3_rlhf_finetuning/rlhf_engine.py

@@ -105,7 +105,8 @@ def _init_actor(self, actor_model_name_or_path):
         # Optimizer
         AdamOptimizer = DeepSpeedCPUAdam if self.args.offload else FusedAdam
         optim_params = get_optimizer_grouped_parameters(
-            actor_model, self.args.actor_weight_decay)
+            actor_model, self.args.actor_weight_decay,
+            self.args.actor_lora_learning_rate)
         optim = AdamOptimizer(optim_params,
                               lr=self.args.actor_learning_rate,
                               betas=(0.9, 0.95))
@@ -231,9 +232,10 @@ def _init_critic(self, critic_model_name_or_path):
 
         # Optimizer
         AdamOptimizer = DeepSpeedCPUAdam if self.args.offload else FusedAdam
-        optim_pararms = get_optimizer_grouped_parameters(
-            critic_model, self.args.critic_weight_decay)
-        optim = AdamOptimizer(optim_pararms,
+        optim_params = get_optimizer_grouped_parameters(
+            critic_model, self.args.critic_weight_decay,
+            self.args.critic_lora_learning_rate)
+        optim = AdamOptimizer(optim_params,
                               lr=self.args.critic_learning_rate,
                               betas=(0.9, 0.95))