Check Explorer GPU Number

trinity/common/config.py

@@ -1314,6 +1314,27 @@ def check_and_update(self) -> Config:  # noqa: C901
                 for args in model_args:
                     set_if_none(aux_model, args, getattr(self.model, args))
 
+            # check gpu number
+            rollout_gpu_num = (
+                self.explorer.rollout_model.tensor_parallel_size
+                * self.explorer.rollout_model.engine_num
+                + sum(
+                    [
+                        model.tensor_parallel_size * model.engine_num
+                        for model in self.explorer.auxiliary_models
+                    ]
+                )
+            )
+            assert self.cluster.node_num is not None
+            assert self.cluster.gpu_per_node is not None
+            total_gpu_num = self.cluster.node_num * self.cluster.gpu_per_node
+            if (self.mode in ["explore", "bench"] and rollout_gpu_num > total_gpu_num) or (
+                self.mode == "both" and rollout_gpu_num >= total_gpu_num
+            ):
+                raise ValueError(
+                    f"Total GPU number ({total_gpu_num}) is less than the number of GPUs required for rollout ({rollout_gpu_num})."
+                )
+
             if self.explorer.over_rollout.ratio > 0.0:
                 if not (0.0 <= self.explorer.over_rollout.ratio < 1.0):
                     raise ValueError("over_rollout_ratio should be in [0.0, 1.0)")