fix optimizer config

docs/examples/config.rst

@@ -377,6 +377,7 @@ ____________________________________________________
 .. code:: yaml
 
     optim:
+      optimizer: adam
       lr: 1e-6
       clip_grad: 1.0
       total_training_steps: -1  # must be override by program

verl/trainer/config/ppo_megatron_trainer.yaml

@@ -59,6 +59,7 @@ actor_rollout_ref:
     data_loader_seed: null
     shuffle: False
     optim:
+      optimizer: adam
       lr: 1e-6
       clip_grad: 1.0
       total_training_steps: -1  # must be override by program

verl/utils/megatron/optimizer.py

@@ -37,7 +37,7 @@ def get_megatron_optimizer(
 
 def get_megatron_optimizer_param_scheduler(
     optimizer,
-    config: OptimizerConfig,
+    config,
 ):
     """
     Get the optimizer parameter scheduler for Megatron.

verl/utils/megatron_utils.py

@@ -200,10 +200,11 @@ def convert_config(hf_config: PretrainedConfig, megatron_config) -> TransformerC
 
 def init_megatron_optim_config(optim_config: Dict) -> OptimizerConfig:
     config = OptimizerConfig(
-        optimizer="adam",
+        optimizer=optim_config.get("optimizer", "adam"),
         lr=optim_config.get("lr"),
-        clip_grad=optim_config.get("clip_grad"),
-        weight_decay=optim_config.get("weight_decay"),
+        min_lr=optim_config.get("min_lr", None),
+        clip_grad=optim_config.get("clip_grad", 1.0),
+        weight_decay=optim_config.get("weight_decay", 0.01),
         bf16=True,
         params_dtype=torch.bfloat16,
         use_distributed_optimizer=True,

verl/workers/megatron_workers.py

@@ -192,8 +192,8 @@ def megatron_actor_model_provider(pre_process, post_process):
 
         # TODO: add more optimizer args into config
         if self._is_actor:
-            optim_config = init_megatron_optim_config(optim_config)
-            actor_optimizer = get_megatron_optimizer(model=actor_module, config=optim_config)
+            optim_config_megatron = init_megatron_optim_config(optim_config)
+            actor_optimizer = get_megatron_optimizer(model=actor_module, config=optim_config_megatron)
             actor_optimizer_scheduler = get_megatron_optimizer_param_scheduler(optimizer=actor_optimizer, config=optim_config)
         else:
             optim_config = None
@@ -666,8 +666,8 @@ def megatron_critic_model_provider(pre_process, post_process):
             print_model_size(critic_module[0])
 
         # TODO: add more optimizer args into config
-        optim_config = init_megatron_optim_config(optim_config)
-        critic_optimizer = get_megatron_optimizer(model=critic_module, config=optim_config)
+        optim_config_megatron = init_megatron_optim_config(optim_config)
+        critic_optimizer = get_megatron_optimizer(model=critic_module, config=optim_config_megatron)
         critic_optimizer_scheduler = get_megatron_optimizer_param_scheduler(optimizer=critic_optimizer, config=optim_config)
         torch.cuda.empty_cache()
         return critic_module, critic_optimizer, critic_optimizer_scheduler, self.hf_config, optim_config