fix unittest

Author: chenyushuo

trinity/common/config_validator.py

@@ -777,6 +777,14 @@ def validate(self, config: Config) -> None:
             config.buffer.train_batch_size = (
                 config.buffer.batch_size * config.algorithm.repeat_times
             )
+        if (
+            config.mode in {"train", "both"}
+            and config.buffer.train_batch_size % config.cluster.trainer_gpu_num != 0
+        ):
+            raise ValueError(
+                f"batch_size ({config.buffer.train_batch_size}) must be "
+                f"divisible by ({config.cluster.trainer_gpu_num})."
+            )
 
         # create buffer.cache_dir at <checkpoint_root_dir>/<project>/<name>/buffer
         config.buffer.cache_dir = os.path.abspath(os.path.join(config.checkpoint_job_dir, "buffer"))

trinity/common/verl_config.py

@@ -421,13 +421,9 @@ def _adjust_token_len_if_needed(
 
     def synchronize_config(self, config: Config) -> None:  # noqa: C901
         """Synchronize config."""
+        # Trainer Config
         self.trainer.nnodes = config.cluster.trainer_node_num
         self.trainer.n_gpus_per_node = config.cluster.trainer_gpu_num_per_node
-        world_size = config.cluster.trainer_gpu_num
-        if config.buffer.train_batch_size % world_size != 0:
-            raise ValueError(
-                f"batch_size ({config.buffer.train_batch_size}) must be divisible by ({world_size})"
-            )
         self.trainer.total_training_steps = config.trainer.total_steps or sys.maxsize
         self.trainer.sync_freq = config.synchronizer.sync_interval
         self.trainer.save_freq = config.trainer.save_interval
@@ -443,9 +439,8 @@ def synchronize_config(self, config: Config) -> None:  # noqa: C901
         else:
             self.trainer.resume_mode = "auto"
 
-        self.data.train_batch_size = (
-            config.buffer.train_batch_size
-        )  # kept to pass RayPPOTrainer._validate_config
+        # kept to pass RayPPOTrainer._validate_config
+        self.data.train_batch_size = config.buffer.train_batch_size
 
         self.synchronizer = config.synchronizer
         self.actor_rollout_ref.nccl_timeout = config.synchronizer.sync_timeout