fix unittest and optimizer config

Author: chenyushuo

tests/common/config_test.py

@@ -6,6 +6,8 @@
 import shutil
 import unittest
 
+import torch
+
 from tests.tools import get_template_config, get_unittest_dataset_config
 from trinity.common.config import InferenceModelConfig, load_config
 
@@ -143,10 +145,9 @@ def test_optimizer_config_propagation(self):
         config.algorithm.optimizer.lr = 1e-4
         config.algorithm.optimizer.weight_decay = 0.05
         config.algorithm.optimizer.clip_grad = 2.0
-        config.algorithm.optimizer.lr_decay_steps = 1000
-        config.algorithm.optimizer.lr_decay_style = "cosine"
-        config.algorithm.optimizer.lr_warmup_init = 1e-7
-        config.algorithm.optimizer.min_lr = 1e-6
+        config.trainer.total_steps = 1000
+        config.algorithm.optimizer.lr_scheduler_type = "cosine"
+        config.algorithm.optimizer.min_lr_ratio = 1e-2
         config.check_and_update()
         self.assertEqual(config.trainer.trainer_config.actor_rollout_ref.actor.optim.lr, 1e-4)
         self.assertEqual(
@@ -159,10 +160,20 @@ def test_optimizer_config_propagation(self):
         self.assertEqual(
             config.trainer.trainer_config.actor_rollout_ref.actor.optim.lr_decay_style, "cosine"
         )
-        self.assertEqual(
-            config.trainer.trainer_config.actor_rollout_ref.actor.optim.lr_warmup_init, 1e-7
+        self.assertTrue(
+            torch.allclose(
+                torch.tensor(
+                    config.trainer.trainer_config.actor_rollout_ref.actor.optim.lr_warmup_init
+                ),
+                torch.tensor(1e-6),
+            )
+        )
+        self.assertTrue(
+            torch.allclose(
+                torch.tensor(config.trainer.trainer_config.actor_rollout_ref.actor.optim.min_lr),
+                torch.tensor(1e-6),
+            )
         )
-        self.assertEqual(config.trainer.trainer_config.actor_rollout_ref.actor.optim.min_lr, 1e-6)
         # critic optimizer should not be affected
         self.assertEqual(config.trainer.trainer_config.critic.optim.lr, 1e-5)
         self.assertEqual(config.trainer.trainer_config.critic.optim.weight_decay, 0.01)

tests/trainer/trainer_test.py

@@ -1437,8 +1437,8 @@ def tearDown(self):
         shutil.rmtree(self.config.checkpoint_job_dir, ignore_errors=True)
 
 
+@unittest.skipIf("TINKER_API_KEY" not in os.environ, "TINKER_API_KEY is not set")
 class TestTinkerTrainer(BaseTrainerCase):
-    @unittest.skipIf("TINKER_API_KEY" not in os.environ, "TINKER_API_KEY is not set")
     def test_trainer(self):
         """Test GSM8K on tinker."""
         # test both mode

trinity/common/config.py

@@ -93,17 +93,13 @@ class OptimizerConfig:
     lr: float = 1e-6
     lr_warmup_steps: int = -1
     lr_warmup_steps_ratio: float = 0.0
-    min_lr_ratio: Optional[float] = 0.0
+    min_lr_ratio: float = 0.0
     warmup_style: Optional[str] = None  # deprecated !
     lr_scheduler_type: str = "constant"
     optimizer_type: str = "adam"
     betas: List[float] = field(default_factory=lambda: [0.9, 0.999])
     weight_decay: float = 0.01
     clip_grad: float = 1.0
-    lr_warmup_init: float = 0.0  # used in megatron
-    lr_decay_steps: Optional[int] = None  # used in megatron
-    lr_decay_style: str = "constant"  # used in megatron, duplicated with lr_scheduler_type in veRL
-    min_lr: float = 0.0
 
 
 @dataclass

trinity/common/models/tinker_model.py

@@ -95,7 +95,8 @@ async def generate(self, prompt: str, **kwargs) -> Sequence[Experience]:
         if with_chat_completion:
             create_time = int(time.time())
         output = await self._generate_internal(prompt={"prompt_token_ids": token_ids}, **kwargs)
-        return_logprobs = kwargs.get("logprobs", self.config.logprobs is not None)
+        logprobs = kwargs.get("logprobs", self.config.logprobs)
+        return_logprobs = logprobs is not None and logprobs is not False
         experiences = [
             Experience(
                 tokens=torch.tensor(token_ids + sequence.tokens, dtype=torch.int32),

trinity/common/verl_config.py

@@ -66,10 +66,10 @@ class Optim:
     total_training_steps: int = -1  # ! DO NOT SET, use trainer.total_steps
     betas: List[float] = field(default_factory=lambda: [0.9, 0.999])
     clip_grad: float = 1.0
-    lr_warmup_init: float = 0.0
+    lr_warmup_init: Optional[float] = None  # 0.0
     lr_decay_steps: Optional[int] = None
-    lr_decay_style: str = "constant"
-    min_lr: float = 0.0
+    lr_decay_style: Optional[str] = None  # "constant"
+    min_lr: Optional[float] = None  # 0.0
     weight_decay: float = 0.01
     weight_decay_incr_style: str = "constant"
     lr_wsd_decay_style: str = "exponential"
@@ -607,22 +607,32 @@ def synchronize_config(self, config: Config) -> None:  # noqa: C901
                 self.critic.strategy = "fsdp"
 
         # Algorithm related config
-        for field_name in config.algorithm.optimizer.__dataclass_fields__:
-            field_value = getattr(config.algorithm.optimizer, field_name)
+        actor_optim = self.actor_rollout_ref.actor.optim
+        critic_optim = self.critic.optim
+        optim_config = config.algorithm.optimizer
+        for field_name in optim_config.__dataclass_fields__:
+            field_value = getattr(optim_config, field_name)
             if field_name == "optimizer_type":
-                setattr(self.actor_rollout_ref.actor.optim, "optimizer", field_value)
-            elif hasattr(self.actor_rollout_ref.actor.optim, field_name):
-                setattr(self.actor_rollout_ref.actor.optim, field_name, field_value)
+                setattr(actor_optim, "optimizer", field_value)
+            elif hasattr(actor_optim, field_name):
+                setattr(actor_optim, field_name, field_value)
+        # ensure megatron optimizer config compatibility
+        set_if_none(actor_optim, "lr_warmup_init", optim_config.min_lr_ratio * optim_config.lr)
+        set_if_none(actor_optim, "lr_decay_steps", self.trainer.total_training_steps)
+        set_if_none(actor_optim, "lr_decay_style", optim_config.lr_scheduler_type)
+        set_if_none(actor_optim, "min_lr", optim_config.min_lr_ratio * optim_config.lr)
+        set_if_none(critic_optim, "lr_warmup_init", 0.0)
+        set_if_none(critic_optim, "lr_decay_steps", self.trainer.total_training_steps)
+        set_if_none(critic_optim, "lr_decay_style", "constant")
+        set_if_none(critic_optim, "min_lr", 0.0)
         # fix optimizer type for fsdp
         if config.trainer.trainer_strategy.startswith("fsdp"):
             optim_map = {
                 "adam": "AdamW",
                 "adamw": "AdamW",
                 "sgd": "SGD",
             }
-            actor_optim = self.actor_rollout_ref.actor.optim
             actor_optim.optimizer = optim_map.get(actor_optim.optimizer, actor_optim.optimizer)
-            critic_optim = self.critic.optim
             critic_optim.optimizer = optim_map.get(critic_optim.optimizer, critic_optim.optimizer)
         self.actor_rollout_ref.actor.use_kl_loss = config.algorithm.kl_loss_fn != "none"
         self.algorithm.use_kl_in_reward = config.algorithm.kl_penalty_fn != "none"