[trainer] support auto resume (#425)

Author: hiyouga

examples/config.yaml

@@ -69,7 +69,8 @@ worker:
     tensor_parallel_size: 2
     disable_tqdm: false
     val_override_config:
-      temperature: 1.0
+      temperature: 0.6
+      top_p: 0.95
       n: 1
 
   ref:
@@ -102,3 +103,4 @@ trainer:
   save_model_only: false
   save_checkpoint_path: null
   load_checkpoint_path: null
+  find_last_checkpoint: true

tests/test_checkpoint.py

@@ -0,0 +1,50 @@
+# Copyright 2024 Bytedance Ltd. and/or its affiliates
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+import json
+import os
+import shutil
+import uuid
+
+import pytest
+
+from verl.utils.checkpoint import CHECKPOINT_TRACKER, find_latest_ckpt, remove_obsolete_ckpt
+
+
+@pytest.fixture
+def save_checkpoint_path():
+    ckpt_dir = os.path.join("checkpoints", str(uuid.uuid4()))
+    os.makedirs(ckpt_dir, exist_ok=True)
+    yield ckpt_dir
+    shutil.rmtree(ckpt_dir, ignore_errors=True)
+
+
+def test_find_latest_ckpt(save_checkpoint_path):
+    with open(os.path.join(save_checkpoint_path, CHECKPOINT_TRACKER), "w") as f:
+        json.dump({"last_global_step": 10}, f, ensure_ascii=False, indent=2)
+
+    assert find_latest_ckpt(save_checkpoint_path) is None
+    os.makedirs(os.path.join(save_checkpoint_path, "global_step_10"), exist_ok=True)
+    assert find_latest_ckpt(save_checkpoint_path) == os.path.join(save_checkpoint_path, "global_step_10")
+
+
+def test_remove_obsolete_ckpt(save_checkpoint_path):
+    for step in range(5, 30, 5):
+        os.makedirs(os.path.join(save_checkpoint_path, f"global_step_{step}"), exist_ok=True)
+
+    remove_obsolete_ckpt(save_checkpoint_path, global_step=30, best_global_step=10, save_limit=3)
+    for step in range(5, 30, 5):
+        is_exist = step in [10, 25]
+        assert os.path.exists(os.path.join(save_checkpoint_path, f"global_step_{step}")) == is_exist

verl/trainer/config.py

@@ -141,6 +141,8 @@ class TrainerConfig:
     """save checkpoint path, if not specified, use `checkpoints/project_name/experiment_name`"""
     load_checkpoint_path: Optional[str] = None
     """load checkpoint path"""
+    find_last_checkpoint: bool = True
+    """automatically find the last checkpoint in the save checkpoint path to resume training"""
 
     def post_init(self):
         if self.save_checkpoint_path is None:

verl/trainer/metrics.py

@@ -26,27 +26,22 @@ def reduce_metrics(metrics: Dict[str, List[Any]]) -> Dict[str, Any]:
 
 def compute_length_metrics(batch: DataProto) -> Dict[str, Any]:
     max_response_length = batch.batch["responses"].size(-1)
+    max_prompt_length = batch.batch["attention_mask"].size(-1) - max_response_length
 
-    prompt_mask = batch.batch["attention_mask"][:, :-max_response_length].bool()
-    response_mask = batch.batch["attention_mask"][:, -max_response_length:].bool()
-
-    max_prompt_length = prompt_mask.size(-1)
-    prompt_length = prompt_mask.sum(-1).float()
-    response_length = response_mask.sum(-1).float()
+    prompt_length = batch.batch["attention_mask"][:, :-max_response_length].sum(-1).float()
+    response_length = batch.batch["attention_mask"][:, -max_response_length:].sum(-1).float()
 
     return {
         # response length
         "response_length/mean": torch.mean(response_length).detach().item(),
         "response_length/max": torch.max(response_length).detach().item(),
         "response_length/min": torch.min(response_length).detach().item(),
-        "response_length/clip_ratio": torch.mean(torch.eq(response_length, max_response_length).float())
-        .detach()
-        .item(),
+        "response_length/clip_ratio": torch.eq(response_length, max_response_length).float().mean().detach().item(),
         # prompt length
         "prompt_length/mean": torch.mean(prompt_length).detach().item(),
         "prompt_length/max": torch.max(prompt_length).detach().item(),
         "prompt_length/min": torch.min(prompt_length).detach().item(),
-        "prompt_length/clip_ratio": torch.mean(torch.eq(prompt_length, max_prompt_length).float()).detach().item(),
+        "prompt_length/clip_ratio": torch.eq(prompt_length, max_prompt_length).float().mean().detach().item(),
     }
 
 

verl/trainer/ray_trainer.py

@@ -37,7 +37,7 @@
 from ..single_controller.ray import RayClassWithInitArgs, RayResourcePool, RayWorkerGroup
 from ..single_controller.ray.base import create_colocated_worker_cls
 from ..utils import torch_functional as VF
-from ..utils.checkpoint import CHECKPOINT_TRACKER, remove_obsolete_ckpt
+from ..utils.checkpoint import CHECKPOINT_TRACKER, find_latest_ckpt, remove_obsolete_ckpt
 from ..utils.logger import Tracker
 from ..utils.py_functional import convert_dict_to_str, timer
 from ..utils.seqlen_balancing import get_seqlen_balanced_partitions, log_seqlen_unbalance
@@ -342,21 +342,28 @@ def _save_checkpoint(self) -> None:
             json.dump(checkpointer_tracker_info, f, ensure_ascii=False, indent=2)
 
     def _load_checkpoint(self) -> None:
-        if self.config.trainer.load_checkpoint_path is None:
+        if self.config.trainer.load_checkpoint_path is not None:
+            load_checkpoint_path = self.config.trainer.load_checkpoint_path
+        elif self.config.trainer.find_last_checkpoint:
+            load_checkpoint_path = find_latest_ckpt(self.config.trainer.save_checkpoint_path)
+        else:
+            load_checkpoint_path = None
+
+        if load_checkpoint_path is None:
             return
 
-        if "global_step_" not in self.config.trainer.load_checkpoint_path.strip(os.path.sep).split(os.path.sep)[-1]:
+        if "global_step_" not in load_checkpoint_path.strip(os.path.sep).split(os.path.sep)[-1]:
             raise ValueError("`load_checkpoint_path` should end with `global_step_*`.")
 
-        print(f"Load from checkpoint: {self.config.trainer.load_checkpoint_path}.")
-        self.global_step = int(self.config.trainer.load_checkpoint_path.strip(os.path.sep).split("global_step_")[-1])
-        actor_path = os.path.join(self.config.trainer.load_checkpoint_path, "actor")
+        print(f"Load from checkpoint: {load_checkpoint_path}.")
+        self.global_step = int(load_checkpoint_path.strip(os.path.sep).split("global_step_")[-1])
+        actor_path = os.path.join(load_checkpoint_path, "actor")
         self.actor_rollout_ref_wg.load_checkpoint(actor_path)
         if self.use_critic:
-            critic_path = os.path.join(self.config.trainer.load_checkpoint_path, "critic")
+            critic_path = os.path.join(load_checkpoint_path, "critic")
             self.critic_wg.load_checkpoint(critic_path)
 
-        dataloader_path = os.path.join(self.config.trainer.load_checkpoint_path, "dataloader.pt")
+        dataloader_path = os.path.join(load_checkpoint_path, "dataloader.pt")
         if os.path.exists(dataloader_path):
             dataloader_state_dict = torch.load(dataloader_path, weights_only=False)
             self.train_dataloader.load_state_dict(dataloader_state_dict)

verl/utils/checkpoint/__init__.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from .checkpoint_manager import CHECKPOINT_TRACKER, remove_obsolete_ckpt
+from .checkpoint_manager import CHECKPOINT_TRACKER, find_latest_ckpt, remove_obsolete_ckpt
 
 
-__all__ = ["CHECKPOINT_TRACKER", "remove_obsolete_ckpt"]
+__all__ = ["CHECKPOINT_TRACKER", "find_latest_ckpt", "remove_obsolete_ckpt"]

verl/utils/checkpoint/checkpoint_manager.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import json
 import os
 import random
 import re
@@ -107,39 +108,38 @@ def load_rng_state(rng_state: Dict[str, Any]):
         random.setstate(rng_state["random"])
 
 
-def find_latest_ckpt_path(path: Optional[str] = None, directory_format: str = "global_step_{}") -> Optional[str]:
-    if path is None:
-        return None
+def get_checkpoint_tracker_filename(root_path: str) -> str:
+    """
+    Tracker file rescords the latest chckpoint during training to restart from.
+    """
+    return os.path.join(root_path, CHECKPOINT_TRACKER)
 
+
+def find_latest_ckpt(path: str, directory_format: str = "global_step_{}") -> Optional[str]:
+    """
+    Find the latest checkpoint in the save path.
+    """
     tracker_file = get_checkpoint_tracker_filename(path)
     if not os.path.exists(tracker_file):
-        print("Checkpoint tracker file does not exist: %s", tracker_file)
         return None
 
     with open(tracker_file, "rb") as f:
-        iteration = int(f.read().decode())
+        checkpointer_tracker_info = json.load(f)
 
-    ckpt_path = os.path.join(path, directory_format.format(iteration))
+    ckpt_path = os.path.join(path, directory_format.format(checkpointer_tracker_info["last_global_step"]))
     if not os.path.exists(ckpt_path):
-        print("Checkpoint does not exist: %s", ckpt_path)
+        print(f"Checkpoint does not exist: {ckpt_path}")
         return None
 
-    print("Found checkpoint: %s", ckpt_path)
+    print(f"Found latest checkpoint: {ckpt_path}, will resume from it. Turn off `find_last_checkpoint` to disable it.")
     return ckpt_path
 
 
-def get_checkpoint_tracker_filename(root_path: str) -> str:
-    """
-    Tracker file rescords the latest chckpoint during training to restart from.
-    """
-    return os.path.join(root_path, CHECKPOINT_TRACKER)
-
-
 def remove_obsolete_ckpt(
     path: str, global_step: int, best_global_step: int, save_limit: int = -1, directory_format: str = "global_step_{}"
 ):
     """
-    Remove the obsolete checkpoints that exceed the save_limit.
+    Remove the obsolete checkpoints that exceed the save limit.
     """
     if save_limit <= 0 or not os.path.exists(path):
         return

verl/workers/reward/function.py

@@ -84,20 +84,21 @@ def compute_reward(self, data: DataProto) -> Tuple[torch.Tensor, Dict[str, List[
         reward_tensor = torch.zeros_like(data.batch["responses"], dtype=torch.float32)
         reward_metrics = defaultdict(list)
         response_ids = data.batch["responses"]
-        response_length = data.batch["response_mask"].sum(dim=-1)
+        response_length = torch.sum(data.batch["response_mask"], dim=-1)
         for i in range(len(data)):
-            valid_response_ids = response_ids[i][: response_length[i]]
+            cur_response_length = int(response_length[i].item())  # avoid tensor indexing error
+            valid_response_ids = response_ids[i][:cur_response_length]
             response_str = self.tokenizer.decode(
                 valid_response_ids, skip_special_tokens=self.config.skip_special_tokens
             )
             score = self.reward_fn(
                 {
                     "response": response_str,
-                    "response_length": response_length[i],
+                    "response_length": cur_response_length,
                     "ground_truth": data.non_tensor_batch["ground_truth"][i],
                 }
             )
-            reward_tensor[i, response_length[i] - 1] = score["overall"]
+            reward_tensor[i, cur_response_length - 1] = score["overall"]
             for key, value in score.items():
                 reward_metrics[key].append(value)
 
@@ -110,16 +111,17 @@ class BatchFunctionRewardManager(FunctionRewardManager):
     def compute_reward(self, data: DataProto) -> Tuple[torch.Tensor, Dict[str, List[float]]]:
         reward_inputs = []
         response_ids = data.batch["responses"]
-        response_length = data.batch["response_mask"].sum(dim=-1)
+        response_length = torch.sum(data.batch["response_mask"], dim=-1)
         for i in range(len(data)):
-            valid_response_ids = response_ids[i][: response_length[i]]
+            cur_response_length = int(response_length[i].item())  # avoid tensor indexing error
+            valid_response_ids = response_ids[i][:cur_response_length]
             response_str = self.tokenizer.decode(
                 valid_response_ids, skip_special_tokens=self.config.skip_special_tokens
             )
             reward_inputs.append(
                 {
                     "response": response_str,
-                    "response_length": response_length[i],
+                    "response_length": cur_response_length,
                     "ground_truth": data.non_tensor_batch["ground_truth"][i],
                 }
             )
@@ -128,7 +130,8 @@ def compute_reward(self, data: DataProto) -> Tuple[torch.Tensor, Dict[str, List[
         reward_tensor = torch.zeros_like(data.batch["responses"], dtype=torch.float32)
         reward_metrics = defaultdict(list)
         for i, score in enumerate(scores):
-            reward_tensor[i, response_length[i] - 1] = score["overall"]
+            cur_response_length = int(response_length[i].item())  # avoid tensor indexing error
+            reward_tensor[i, cur_response_length - 1] = score["overall"]
             for key, value in score.items():
                 reward_metrics[key].append(value)