add save strategy

Author: chenyushuo

docs/sphinx_doc/source/tutorial/trinity_configs.md

@@ -160,7 +160,7 @@ model:
 
 - `model_path`: Path to the model being trained.
 - `critic_model_path`: Optional path to a separate critic model. If empty, defaults to `model_path`.
-- `max_model_len`: Maximum number of tokens in a sequence. It is recommended to set this value manually. If not set, it will be inferred from the model configuration.
+- `max_model_len`: Maximum number of tokens in a sequence. It is recommended to set this value manually. If not set, it will default to `max_prompt_tokens` + `max_response_tokens`. However, if either `max_prompt_tokens` or `max_response_tokens` is not set, we will raise an error.
 - `max_response_tokens`: Maximum number of tokens allowed in generated responses. Only for `chat` and `generate` methods in `InferenceModel`.
 - `max_prompt_tokens`: Maximum number of tokens allowed in prompts. Only for `chat` and `generate` methods in `InferenceModel`.
 - `min_response_tokens`: Minimum number of tokens allowed in generated responses. Only for `chat` and `generate` methods in `InferenceModel`. Default is `1`. It must be less than `max_response_tokens`.

docs/sphinx_doc/source_zh/tutorial/trinity_configs.md

@@ -160,7 +160,7 @@ model:
 
 - `model_path`: 被训练模型的路径。
 - `critic_model_path`: 可选的独立 critic 模型路径。若为空，则默认为 `model_path`。
-- `max_model_len`: 该模型所支持的单个序列最大 token 数。
+- `max_model_len`: 表示模型所支持的单个序列最大 token 数。如果未设置该值，则会尝试将其默认设为 `max_prompt_tokens + max_response_tokens`。但如果 `max_prompt_tokens` 或 `max_response_tokens` 中有任何一个未设置，代码将会报错。
 - `max_prompt_tokens`: 输入 prompt 中允许的最大 token 数。仅对 `InferenceModel` 中的 `chat` 和 `generate` 方法生效。
 - `max_response_tokens`: 模型生成的回复中允许的最大 token 数。仅对 `InferenceModel` 中的 `chat` 和 `generate` 方法生效。
 - `min_response_tokens`: 模型生成的回复中允许的最小 token 数。仅对 `InferenceModel` 中的 `chat` 和 `generate` 方法生效。

trinity/common/config.py

@@ -19,6 +19,7 @@
     PLUGIN_DIRS_ENV_VAR,
     TRAINER_NAME,
     PromptType,
+    SaveStrategy,
     StorageType,
     SyncMethod,
     SyncStyle,
@@ -470,6 +471,8 @@ class TrainerConfig:
     actor_grad_clip: Optional[float] = None
     # TODO: extract more train-related params from underlying trainer engine
 
+    save_strategy: SaveStrategy = SaveStrategy.UNRESTRICTED
+
     # Only one needs to be set for `trainer_config` and `trainer_config_path`
     trainer_config: Any = field(default_factory=dict)
     trainer_config_path: str = ""

trinity/common/constants.py

@@ -104,3 +104,10 @@ class SyncStyle(CaseInsensitiveEnum):
     FIXED = "fixed"
     DYNAMIC_BY_TRAINER = "dynamic_by_trainer"
     DYNAMIC_BY_EXPLORER = "dynamic_by_explorer"
+
+
+class SaveStrategy(CaseInsensitiveEnum):
+    SINGLE_THREAD = "single_thread"
+    SINGLE_PROCESS = "single_process"
+    SINGLE_NODE = "single_node"
+    UNRESTRICTED = "unrestricted"

trinity/trainer/verl/fsdp_checkpoint_manager.py

@@ -115,7 +115,10 @@ def _save_with_thread(
             thread.join()
 
         def _save():
-            ray.get(self.checkpoint_monitor.notify_started.remote())
+            runtime_context = ray.get_runtime_context()
+            node_id = runtime_context.get_node_id()
+            job_id = runtime_context.get_job_id()
+            ray.get(self.checkpoint_monitor.notify_started.remote(node_id=node_id, job_id=job_id))
             torch.save(obj, path)
             log_with_rank(
                 f"Saved {prefix} to {os.path.abspath(path)}",
@@ -358,7 +361,14 @@ def save_checkpoint(  # noqa: C901
                     self._save_model_thread.join()
 
                 def _save_model():
-                    ray.get(self.checkpoint_monitor.notify_started.remote())
+                    runtime_context = ray.get_runtime_context()
+                    node_id = runtime_context.get_node_id()
+                    job_id = runtime_context.get_job_id()
+                    ray.get(
+                        self.checkpoint_monitor.notify_started.remote(
+                            node_id=node_id, job_id=job_id
+                        )
+                    )
                     save_model.save_pretrained(hf_local_path, state_dict=state_dict)
                     log_with_rank(
                         f"Saved hf_model to {os.path.abspath(hf_local_path)}",

trinity/trainer/verl/megatron_checkpoint_manager.py

@@ -125,7 +125,10 @@ def _save_state_dict(self, local_path, global_step):
 
         def finalize_save_fn():
             # Rank 0 uploads checkpoint to HDFS if hdfs_path is provided
-            ray.get(self.checkpoint_monitor.notify_started.remote())
+            runtime_context = ray.get_runtime_context()
+            node_id = runtime_context.get_node_id()
+            job_id = runtime_context.get_job_id()
+            ray.get(self.checkpoint_monitor.notify_started.remote(node_id=node_id, job_id=job_id))
             log_with_rank(
                 f"Dist checkpointing save completed for {dist_checkpoint_path}",
                 rank=self.rank,

trinity/trainer/verl_trainer.py

@@ -34,14 +34,17 @@
 from trinity.algorithm.algorithm import ALGORITHM_TYPE
 from trinity.algorithm.utils import prefix_metrics
 from trinity.common.config import Config
+from trinity.common.constants import SaveStrategy
 from trinity.common.experience import Experiences
 from trinity.trainer.trainer import TrainEngineWrapper
 from trinity.trainer.verl.utils import compute_data_metrics, to_data_proto
 from trinity.utils.log import get_logger
 
 
 class CheckpointMonitor:
-    def __init__(self, default_local_dir: str, default_hdfs_dir: str = None):
+    def __init__(
+        self, save_strategy: SaveStrategy, default_local_dir: str, default_hdfs_dir: str = None
+    ):
         self.logger = get_logger("checkpoint_monitor", in_ray_actor=True)
         self.default_local_dir = default_local_dir
         self.default_hdfs_dir = default_hdfs_dir
@@ -58,7 +61,9 @@ def __init__(self, default_local_dir: str, default_hdfs_dir: str = None):
         self.latest_checkpoint_step = 0
         self.latest_state_dict_step = 0
 
+        self.save_strategy = save_strategy
         self.condition = asyncio.Condition()
+        self.current_identifier = 0
         self.saving_count = 0
 
     def update_latest_checkpoint_step(self, step: int):
@@ -113,16 +118,28 @@ async def monitor_step(self, step: int, is_state_dict: bool = False):
             if self.checkpoint_counter[step] == 0 and self.state_dict_counter[step] == 0:
                 self.update_latest_checkpoint_step(step)
 
-    async def notify_started(self):
+    async def notify_started(self, node_id: str, job_id: str):
+        if self.save_strategy == SaveStrategy.SINGLE_THREAD:
+            identifier = self.current_identifier + 1
+        elif self.save_strategy == SaveStrategy.SINGLE_PROCESS:
+            identifier = f"{node_id}_{job_id}"
+        elif self.save_strategy == SaveStrategy.SINGLE_NODE:
+            identifier = node_id
+        elif self.save_strategy == SaveStrategy.UNRESTRICTED:
+            return
+        else:
+            raise ValueError(f"Invalid save strategy: {self.save_strategy}")
+
         async with self.condition:
-            while self.saving_count > 0:
+            if identifier != self.current_identifier and self.saving_count > 0:
                 await self.condition.wait_for(lambda: self.saving_count == 0)
+            self.current_identifier = identifier
             self.saving_count += 1
 
     async def notify_finished(self, step: int, is_state_dict: bool = False):
         async with self.condition:
             self.saving_count -= 1
-            self.condition.notify()
+            self.condition.notify_all()
         if is_state_dict:
             self.state_dict_counter[step] -= 1
             if (
@@ -144,6 +161,7 @@ async def notify_finished(self, step: int, is_state_dict: bool = False):
     def get_actor(
         cls,
         namespace: str,
+        save_strategy: Optional[SaveStrategy] = None,
         default_local_dir: Optional[str] = None,
         default_hdfs_dir: Optional[str] = None,
     ):
@@ -154,7 +172,11 @@ def get_actor(
                 namespace=namespace,
                 get_if_exists=True,
             )
-            .remote(default_local_dir=default_local_dir, default_hdfs_dir=default_hdfs_dir)
+            .remote(
+                save_strategy=save_strategy,
+                default_local_dir=default_local_dir,
+                default_hdfs_dir=default_hdfs_dir,
+            )
         )
 
 
@@ -204,6 +226,7 @@ def __init__(
 
         self.checkpoint_monitor = CheckpointMonitor.get_actor(
             namespace=global_config.synchronizer.ray_namespace,
+            save_strategy=global_config.trainer.save_strategy,
             default_local_dir=config.trainer.default_local_dir,
             default_hdfs_dir=config.trainer.default_hdfs_dir,
         )