Flexrlhf

README_cn.md

@@ -58,6 +58,8 @@ OpenSPG核心能力模型包括：
 
 # Cite
 
+None
+
 # License
 
 [Apache License 2.0](LICENSE)

python/nn4k/nn4k/executor/base.py

@@ -152,10 +152,34 @@ def execute_sft(self, args=None, callbacks=None, **kwargs):
         """
         raise NotImplementedError(f"{self.__class__.__name__} does not support SFT.")
 
-    def execute_rl_tuning(self, args=None, callbacks=None, **kwargs):
+
+class AlignExecutor(LLMExecutor):
+    """An executor for X-Alignment"""
+
+    @classmethod
+    def from_config(cls, nn_config: Union[str, dict]) -> "NNExecutor":
+        return super().from_config(nn_config)
+
+    def execute_rm(self, args=None, callbacks=None, **kwargs):
         """
-        The entry point of SFT execution in a certain pod.
+        Execute reward model(rm) training.
+        """
+        raise NotImplementedError(
+            f"{self.__class__.__name__} does not support execute_rm yet."
+        )
+
+    def execute_dpo(self, args=None, callbacks=None, **kwargs):
+        """
+        Execute DPO training.
+        """
+        raise NotImplementedError(
+            f"{self.__class__.__name__} does not support execute_dpo yet."
+        )
+
+    def execute_ppo(self, args=None, callbacks=None, **kwargs):
+        """
+        Execute PPO training.
         """
         raise NotImplementedError(
-            f"{self.__class__.__name__} does not support RL-Tuning."
+            f"{self.__class__.__name__} does not support execute_ppo yet."
         )

python/nn4k/nn4k/invoker/base.py

@@ -17,6 +17,7 @@
 
 
 class SubmitMode(Enum):
+    Local = "local"
     K8s = "k8s"
     Docker = "docker"
 
@@ -138,14 +139,6 @@ def submit_sft(self, submit_mode: SubmitMode = SubmitMode.K8s):
         """
         raise NotImplementedError(f"{self.__class__.__name__} does not support SFT.")
 
-    def submit_rl_tuning(self, submit_mode: SubmitMode = SubmitMode.K8s):
-        """
-        Submit remote RL-Tuning execution.
-        """
-        raise NotImplementedError(
-            f"{self.__class__.__name__} does not support RL-Tuning."
-        )
-
     def local_inference(self, data, **kwargs):
         """
         Implement local inference for local invoker.
@@ -184,3 +177,38 @@ def from_config(cls, nn_config: dict) -> "LLMInvoker":
         """
         invoker = cls(nn_config)
         return invoker
+
+
+class AlignInvoker(NNInvoker):
+    """A invoker for X-Alignment tuning"""
+    @classmethod
+    def from_config(cls, nn_config: Union[str, dict]) -> "NNInvoker":
+        """
+        Create an instance from `nn_config`.
+        """
+        invoker = cls(nn_config)
+        return invoker
+
+    def submit_rm(self, submit_mode: SubmitMode = SubmitMode.K8s):
+        """
+        Submit remote tuning execution for reward model(rm).
+        """
+        raise NotImplementedError(
+            f"{self.__class__.__name__} does not support submit_rm yet."
+        )
+
+    def submit_dpo(self, submit_mode: SubmitMode = SubmitMode.K8s):
+        """
+        Submit remote execution for DPO.
+        """
+        raise NotImplementedError(
+            f"{self.__class__.__name__} does not support submit_dpo yet."
+        )
+
+    def submit_ppo(self, submit_mode: SubmitMode = SubmitMode.K8s):
+        """
+        Submit remote execution for PPO.
+        """
+        raise NotImplementedError(
+            f"{self.__class__.__name__} does not support submit_ppo yet."
+        )