0402 update

requirements.txt

 accelerate
 codetiming
 datasets
-flash-attn>=2.4.3
 liger-kernel
 mathruler
 numpy
@@ -12,8 +11,10 @@ pillow
 pyarrow>=15.0.0
 pylatexenc
 qwen-vl-utils
-ray
+ray[default]
 tensordict
+torchdata
 transformers>=4.49.0
-vllm>=0.7.3
 wandb
+orjson
+tensorboard
\ No newline at end of file

setup.py

@@ -12,11 +12,21 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import os
+import re
 
 from setuptools import find_packages, setup
 
 
-def get_requires():
+def get_version() -> str:
+    with open(os.path.join("verl", "__init__.py"), encoding="utf-8") as f:
+        file_content = f.read()
+        pattern = r"__version__\W*=\W*\"([^\"]+)\""
+        (version,) = re.findall(pattern, file_content)
+        return version
+
+
+def get_requires() -> list[str]:
     with open("requirements.txt", encoding="utf-8") as f:
         file_content = f.read()
         lines = [line.strip() for line in file_content.strip().split("\n") if not line.startswith("#")]
@@ -31,18 +41,19 @@ extra_require = {
 def main():
     setup(
         name="verl",
-        version="0.2.0.dev0",
-        package_dir={"": "."},
-        packages=find_packages(where="."),
-        url="https://github.com/volcengine/verl",
-        license="Apache 2.0",
+        version=get_version(),
+        description="An Efficient, Scalable, Multi-Modality RL Training Framework based on veRL",
+        long_description=open("README.md", encoding="utf-8").read(),
+        long_description_content_type="text/markdown",
         author="verl",
         author_email="zhangchi.usc1992@bytedance.com, gmsheng@connect.hku.hk, hiyouga@buaa.edu.cn",
-        description="",
+        license="Apache 2.0 License",
+        url="https://github.com/volcengine/verl",
+        package_dir={"": "."},
+        packages=find_packages(where="."),
+        python_requires=">=3.9.0",
         install_requires=get_requires(),
         extras_require=extra_require,
-        long_description=open("README.md", encoding="utf-8").read(),
-        long_description_content_type="text/markdown",
     )
 
 

verl/__init__.py

@@ -12,8 +12,4 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from .protocol import DataProto
-
-
-__all__ = ["DataProto"]
 __version__ = "0.2.0.dev"

verl/models/monkey_patch.py

+# 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.
+
+
+from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS
+
+from .transformers.flash_attention_utils import flash_attention_forward
+from .transformers.qwen2_vl import qwen2_vl_attn_forward
+
+
+def apply_ulysses_patch(model_type: str) -> None:
+    if model_type in ("llama", "gemma", "gemma2", "mistral", "qwen2"):
+        ALL_ATTENTION_FUNCTIONS["flash_attention_2"] = flash_attention_forward
+    elif model_type in ("qwen2_vl", "qwen2_5_vl"):
+        from transformers.models.qwen2_5_vl.modeling_qwen2_5_vl import Qwen2_5_VLFlashAttention2
+        from transformers.models.qwen2_vl.modeling_qwen2_vl import Qwen2VLFlashAttention2
+
+        Qwen2VLFlashAttention2.forward = qwen2_vl_attn_forward
+        Qwen2_5_VLFlashAttention2.forward = qwen2_vl_attn_forward
+    else:
+        raise NotImplementedError(f"Model architecture {model_type} is not supported yet.")

verl/models/transformers/flash_attention_utils.py

+# Copyright 2024 The Fairseq Authors and the HuggingFace Inc. team
+# Copyright 2024 Bytedance Ltd. and/or its affiliates
+# Based on https://github.com/huggingface/transformers/blob/v4.49.0/src/transformers/modeling_flash_attention_utils.py
+#
+# 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 inspect
+import os
+from typing import Optional, Tuple
+
+import torch
+import torch.distributed as dist
+from transformers.modeling_flash_attention_utils import _flash_attention_forward, fa_peft_integration_check
+from transformers.utils import is_flash_attn_2_available, is_flash_attn_greater_or_equal_2_10
+
+from ...utils.ulysses import (
+    gather_heads_scatter_seq,
+    gather_seq_scatter_heads,
+    get_ulysses_sequence_parallel_group,
+    get_ulysses_sequence_parallel_world_size,
+)
+
+
+if is_flash_attn_2_available():
+    from flash_attn import flash_attn_func, flash_attn_varlen_func
+
+    _flash_supports_window_size = "window_size" in inspect.signature(flash_attn_func).parameters
+    _flash_supports_deterministic = "deterministic" in inspect.signature(flash_attn_func).parameters
+    _flash_deterministic_enabled = os.environ.get("FLASH_ATTENTION_DETERMINISTIC", "0") == "1"
+    _flash_use_top_left_mask = not is_flash_attn_greater_or_equal_2_10()
+
+
+def prepare_fa2_from_position_ids(
+    query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, position_ids: torch.Tensor
+):
+    query = query.view(-1, query.size(-2), query.size(-1))
+    key = key.contiguous().view(-1, key.size(-2), key.size(-1))
+    value = value.contiguous().view(-1, value.size(-2), value.size(-1))
+    position_ids = position_ids.flatten()
+    indices_q = torch.arange(position_ids.size(0), device=position_ids.device, dtype=torch.int32)
+    cu_seqlens = torch.cat(
+        (
+            indices_q[position_ids == 0],
+            torch.tensor(position_ids.size(), device=position_ids.device, dtype=torch.int32),
+        )
+    )
+    max_length = cu_seqlens.diff().max()  # use cu_seqlens to infer max_length for qwen2vl mrope
+    return (query, key, value, indices_q, (cu_seqlens, cu_seqlens), (max_length, max_length))
+
+
+def _custom_flash_attention_forward(
+    query_states: torch.Tensor,
+    key_states: torch.Tensor,
+    value_states: torch.Tensor,
+    attention_mask: Optional[torch.Tensor],
+    query_length: int,
+    is_causal: bool = True,
+    position_ids: Optional[torch.Tensor] = None,
+    sliding_window: Optional[int] = None,
+    use_top_left_mask: bool = False,
+    deterministic: Optional[bool] = None,
+    **kwargs,
+):
+    """
+    Patches flash attention forward to handle 3D position ids in mrope. (3, batch_size, seq_length)
+    """
+    if not use_top_left_mask:
+        causal = is_causal
+    else:
+        causal = is_causal and query_length != 1
+
+    # Assuming 4D tensors, key_states.shape[1] is the key/value sequence length (source length).
+    use_sliding_windows = (
+        _flash_supports_window_size and sliding_window is not None and key_states.shape[1] > sliding_window
+    )
+    flash_kwargs = {"window_size": (sliding_window, sliding_window)} if use_sliding_windows else {}
+
+    if _flash_supports_deterministic:
+        flash_kwargs["deterministic"] = deterministic if deterministic is not None else _flash_deterministic_enabled
+
+    if kwargs.get("softcap") is not None:
+        flash_kwargs["softcap"] = kwargs.pop("softcap")
+
+    query_states, key_states, value_states = fa_peft_integration_check(
+        query_states, key_states, value_states, target_dtype=torch.bfloat16
+    )
+
+    sp_size = get_ulysses_sequence_parallel_world_size()
+    if sp_size > 1:
+        # (batch_size, seq_length, num_head, head_size)
+        query_states = gather_seq_scatter_heads(query_states, seq_dim=1, head_dim=2)
+        key_states = gather_seq_scatter_heads(key_states, seq_dim=1, head_dim=2)
+        value_states = gather_seq_scatter_heads(value_states, seq_dim=1, head_dim=2)
+        position_ids_lst = [torch.empty_like(position_ids) for _ in range(sp_size)]
+        position_ids = dist.all_gather(position_ids_lst, position_ids, group=get_ulysses_sequence_parallel_group())
+        position_ids = torch.cat(position_ids_lst, dim=-1)  # (..., batch_size, seq_length)
+
+    if position_ids is not None and position_ids.dim() == 3:  # qwen2vl mrope
+        position_ids = position_ids[0]
+
+    if position_ids is not None and query_length != 1 and not (torch.diff(position_ids, dim=-1) >= 0).all():
+        batch_size = query_states.size(0)
+        query_states, key_states, value_states, _, cu_seq_lens, max_seq_lens = prepare_fa2_from_position_ids(
+            query_states, key_states, value_states, position_ids
+        )
+        cu_seqlens_q, cu_seqlens_k = cu_seq_lens
+        max_seqlen_in_batch_q, max_seqlen_in_batch_k = max_seq_lens
+        attn_output = flash_attn_varlen_func(
+            query_states,
+            key_states,
+            value_states,
+            cu_seqlens_q=cu_seqlens_q,
+            cu_seqlens_k=cu_seqlens_k,
+            max_seqlen_q=max_seqlen_in_batch_q,
+            max_seqlen_k=max_seqlen_in_batch_k,
+            dropout_p=kwargs.pop("dropout", 0.0),
+            softmax_scale=kwargs.pop("softmax_scale", None),
+            causal=causal,
+            **flash_kwargs,
+        )
+        attn_output = attn_output.view(batch_size, -1, attn_output.size(-2), attn_output.size(-1))
+    else:
+        attn_output = _flash_attention_forward(
+            query_states,
+            key_states,
+            value_states,
+            attention_mask,
+            query_length,
+            is_causal=is_causal,
+            sliding_window=sliding_window,
+            use_top_left_mask=use_top_left_mask,
+            deterministic=deterministic,
+            **kwargs,
+        )  # do not pass position_ids to old flash_attention_forward
+
+    if sp_size > 1:
+        # (batch_size, seq_length, num_head, head_size)
+        attn_output = gather_heads_scatter_seq(attn_output, head_dim=2, seq_dim=1)
+
+    return attn_output
+
+
+def flash_attention_forward(
+    module: torch.nn.Module,
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    attention_mask: Optional[torch.Tensor],
+    dropout: float = 0.0,
+    scaling: Optional[float] = None,
+    sliding_window: Optional[int] = None,
+    softcap: Optional[float] = None,
+    **kwargs,
+) -> Tuple[torch.Tensor, None]:
+    # This is before the transpose
+    q_len = query.shape[2]
+
+    # FA2 uses non-transposed inputs
+    query = query.transpose(1, 2)
+    key = key.transpose(1, 2)
+    value = value.transpose(1, 2)
+
+    # FA2 always relies on the value set in the module, so remove it if present in kwargs to avoid passing it twice
+    kwargs.pop("is_causal", None)
+
+    attn_output = _custom_flash_attention_forward(
+        query,
+        key,
+        value,
+        attention_mask,
+        query_length=q_len,
+        is_causal=True,
+        dropout=dropout,
+        softmax_scale=scaling,
+        sliding_window=sliding_window,
+        softcap=softcap,
+        use_top_left_mask=_flash_use_top_left_mask,
+        **kwargs,
+    )
+
+    return attn_output, None

verl/models/transformers/qwen2_vl.py

+# Copyright 2024 The Qwen team, Alibaba Group and the HuggingFace Inc. team
+# Copyright 2024 Bytedance Ltd. and/or its affiliates
+# Based on:
+# https://github.com/huggingface/transformers/blob/v4.49.0/src/transformers/models/qwen2_vl/modeling_qwen2_vl.py
+#
+# 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.
+
+from typing import Optional, Tuple
+
+import torch
+
+from .flash_attention_utils import flash_attention_forward
+
+
+try:
+    from transformers.models.qwen2_vl.modeling_qwen2_vl import (
+        Qwen2VLAttention,
+        apply_multimodal_rotary_pos_emb,
+        repeat_kv,
+    )
+    from transformers.models.qwen2_vl.processing_qwen2_vl import Qwen2VLProcessor
+except ImportError:
+    pass
+
+
+def get_rope_index(
+    processor: "Qwen2VLProcessor",
+    input_ids: torch.Tensor,
+    image_grid_thw: Optional[torch.Tensor] = None,
+    video_grid_thw: Optional[torch.Tensor] = None,
+    second_per_grid_ts: Optional[torch.Tensor] = None,
+    attention_mask: Optional[torch.Tensor] = None,
+) -> torch.Tensor:
+    """
+    Gets the position ids for Qwen2-VL, it should be generated before sharding the sequence.
+    The batch dim has been removed and the input_ids should be a 1D tensor representing a single example.
+    https://github.com/huggingface/transformers/blob/v4.49.0/src/transformers/models/qwen2_5_vl/modeling_qwen2_5_vl.py#L1546
+    """
+    spatial_merge_size = processor.image_processor.merge_size
+    tokens_per_second = 2
+    image_token_id = processor.tokenizer.convert_tokens_to_ids("<|image_pad|>")
+    video_token_id = processor.tokenizer.convert_tokens_to_ids("<|video_pad|>")
+    vision_start_token_id = processor.tokenizer.convert_tokens_to_ids("<|vision_start|>")
+    if input_ids is not None and (image_grid_thw is not None or video_grid_thw is not None):
+        if attention_mask is None:
+            attention_mask = torch.ones_like(input_ids)
+
+        position_ids = torch.ones(3, input_ids.size(0), dtype=input_ids.dtype, device=input_ids.device)  # (3, seqlen)
+        image_index, video_index = 0, 0
+        input_ids = input_ids[attention_mask == 1]
+        image_nums, video_nums = 0, 0
+        vision_start_indices = torch.argwhere(input_ids == vision_start_token_id)
+        vision_tokens = input_ids[vision_start_indices + 1]
+        image_nums = (vision_tokens == image_token_id).sum()
+        video_nums = (vision_tokens == video_token_id).sum()
+        input_tokens = input_ids.tolist()
+        llm_pos_ids_list: list = []
+        st = 0
+        remain_images, remain_videos = image_nums, video_nums
+        for _ in range(image_nums + video_nums):
+            if image_token_id in input_tokens and remain_images > 0:
+                ed_image = input_tokens.index(image_token_id, st)
+            else:
+                ed_image = len(input_tokens) + 1
+            if video_token_id in input_tokens and remain_videos > 0:
+                ed_video = input_tokens.index(video_token_id, st)
+            else:
+                ed_video = len(input_tokens) + 1
+            if ed_image < ed_video:
+                t, h, w = (
+                    image_grid_thw[image_index][0],
+                    image_grid_thw[image_index][1],
+                    image_grid_thw[image_index][2],
+                )
+                second_per_grid_t = 0
+                image_index += 1
+                remain_images -= 1
+                ed = ed_image
+            else:
+                t, h, w = (
+                    video_grid_thw[video_index][0],
+                    video_grid_thw[video_index][1],
+                    video_grid_thw[video_index][2],
+                )
+                if second_per_grid_ts is not None:
+                    second_per_grid_t = second_per_grid_ts[video_index]
+                else:
+                    second_per_grid_t = 1.0
+
+                video_index += 1
+                remain_videos -= 1
+                ed = ed_video
+
+            llm_grid_t, llm_grid_h, llm_grid_w = (
+                t.item(),
+                h.item() // spatial_merge_size,
+                w.item() // spatial_merge_size,
+            )
+            text_len = ed - st
+
+            st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0
+            llm_pos_ids_list.append(torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx)
+
+            t_index = torch.arange(llm_grid_t).view(-1, 1).expand(-1, llm_grid_h * llm_grid_w)
+            t_index = (t_index * second_per_grid_t * tokens_per_second).long().flatten()
+            h_index = torch.arange(llm_grid_h).view(1, -1, 1).expand(llm_grid_t, -1, llm_grid_w).flatten()
+            w_index = torch.arange(llm_grid_w).view(1, 1, -1).expand(llm_grid_t, llm_grid_h, -1).flatten()
+            llm_pos_ids_list.append(torch.stack([t_index, h_index, w_index]) + text_len + st_idx)
+            st = ed + llm_grid_t * llm_grid_h * llm_grid_w
+
+        if st < len(input_tokens):
+            st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0
+            text_len = len(input_tokens) - st
+            llm_pos_ids_list.append(torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx)
+
+        llm_positions = torch.cat(llm_pos_ids_list, dim=1).reshape(3, -1)
+        position_ids[..., attention_mask == 1] = llm_positions.to(position_ids.device)
+    else:
+        if attention_mask is not None:
+            position_ids = attention_mask.long().cumsum(-1) - 1
+            position_ids.masked_fill_(attention_mask == 0, 1)
+            position_ids = position_ids.unsqueeze(0).expand(3, -1).to(input_ids.device)
+        else:
+            position_ids = torch.arange(input_ids.shape[1], device=input_ids.device).view(1, -1).expand(3, -1)
+
+    return position_ids
+
+
+def qwen2_vl_attn_forward(
+    self: "Qwen2VLAttention",
+    hidden_states: torch.Tensor,
+    attention_mask: Optional[torch.Tensor] = None,
+    position_ids: Optional[torch.LongTensor] = None,
+    position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,  # will become mandatory in v4.46
+    **kwargs,
+) -> Tuple[torch.Tensor, None, None]:
+    bsz, q_len, _ = hidden_states.size()  # q_len = seq_length / sp_size
+    query_states = self.q_proj(hidden_states)  # (batch_size, seq_length / sp_size, num_heads * head_size)
+    key_states = self.k_proj(hidden_states)
+    value_states = self.v_proj(hidden_states)
+
+    query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
+    key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
+    value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
+
+    # Because the input can be padded, the absolute sequence length depends on the max position id.
+    if position_embeddings is None:
+        cos, sin = self.rotary_emb(value_states, position_ids)
+    else:
+        cos, sin = position_embeddings
+
+    query_states, key_states = apply_multimodal_rotary_pos_emb(
+        query_states, key_states, cos, sin, self.rope_scaling["mrope_section"]
+    )
+    key_states = repeat_kv(key_states, self.num_key_value_groups)
+    value_states = repeat_kv(value_states, self.num_key_value_groups)
+    dropout_rate = 0.0 if not self.training else self.attention_dropout
+
+    sliding_window = None
+    if (
+        self.config.use_sliding_window
+        and getattr(self.config, "sliding_window", None) is not None
+        and self.layer_idx >= self.config.max_window_layers
+    ):
+        sliding_window = self.config.sliding_window
+
+    attn_output, _ = flash_attention_forward(
+        self,
+        query_states,
+        key_states,
+        value_states,
+        attention_mask,
+        dropout=dropout_rate,
+        sliding_window=sliding_window,
+        position_ids=position_ids,  # important: pass position ids
+    )  # (batch_size, seq_length, num_head / sp_size, head_size)
+    attn_output = attn_output.reshape(bsz, q_len, self.hidden_size).contiguous()
+    attn_output = self.o_proj(attn_output)
+    return attn_output, None, None

verl/protocol.py

@@ -26,13 +26,12 @@ from typing import Any, Callable, Dict, List, Optional, Tuple, Union
 import numpy as np
 import ray
 import torch
-import torch.distributed as dist
 from numpy.typing import NDArray
 from tensordict import TensorDict
 from torch.distributed import ProcessGroup
 from torch.utils.data import DataLoader
 
-from verl.utils.py_functional import union_two_dict
+from .utils.py_functional import union_two_dict
 
 
 try:
@@ -89,21 +88,22 @@ def union_tensor_dict(tensor_dict1: TensorDict, tensor_dict2: TensorDict) -> Ten
             f"Two tensor dict must have identical batch size. Got {tensor_dict1.batch_size} and {tensor_dict2.batch_size}"
         )
 
-    for key, value in tensor_dict2.items():
-        if key in tensor_dict1 and not torch.equal(tensor_dict1[key], value):
+    for key in tensor_dict2.keys():
+        if key in tensor_dict1 and not torch.equal(tensor_dict1[key], tensor_dict2[key]):
             raise ValueError(f"Key already exists: {key}.")
 
-        tensor_dict1[key] = value
+        tensor_dict1[key] = tensor_dict2[key]
 
     return tensor_dict1
 
 
-def union_numpy_dict(
-    tensor_dict1: Dict[str, Union[List, NDArray]], tensor_dict2: Dict[str, Union[List, NDArray]]
-) -> Dict[str, Union[List, NDArray]]:
-    for key, value in tensor_dict2.items():
-        if key in tensor_dict1 and isinstance(value, np.ndarray) and not np.all(tensor_dict1[key] == value):
-            raise ValueError(f"Key already exists: {key}.")
+def union_numpy_dict(tensor_dict1: Dict[str, NDArray], tensor_dict2: Dict[str, NDArray]) -> Dict[str, NDArray]:
+    for key in tensor_dict2.keys():
+        if key in tensor_dict1:
+            assert isinstance(tensor_dict2[key], np.ndarray)
+            assert isinstance(tensor_dict1[key], np.ndarray)
+            if not np.all(tensor_dict1[key] == tensor_dict2[key]):
+                raise ValueError(f"Key already exists: {key}.")
 
         tensor_dict1[key] = tensor_dict2[key]
 
@@ -151,6 +151,7 @@ def collate_fn(data_items: list["DataProtoItem"]):
 
     batch = torch.stack(batch).contiguous()
     non_tensor_batch = batch_collate(non_tensor_batch)
+    non_tensor_batch = {key: np.array(value, dtype=object) for key, value in non_tensor_batch.items()}
     return DataProto(batch=batch, non_tensor_batch=non_tensor_batch)
 
 
@@ -189,7 +190,8 @@ class DataProto:
     def __getitem__(self, item):
         tensor_data = self.batch[item]
         non_tensor_data = {key: val[item] for key, val in self.non_tensor_batch.items()}
-        return DataProtoItem(batch=tensor_data, non_tensor_batch=non_tensor_data, meta_info=self.meta_info)
+        return_type = DataProto if isinstance(item, slice) else DataProtoItem
+        return return_type(batch=tensor_data, non_tensor_batch=non_tensor_data, meta_info=self.meta_info)
 
     def __getstate__(self):
         buffer = io.BytesIO()
@@ -229,8 +231,7 @@ class DataProto:
 
         size_of_numpy_array = 0
         for value in self.non_tensor_batch.values():
-            if isinstance(value, np.ndarray):
-                size_of_numpy_array += value.nbytes
+            size_of_numpy_array += value.nbytes
 
         size_of_numpy_array /= 1024**3
         size_of_tensordict /= 1024**3
@@ -254,13 +255,13 @@ class DataProto:
                 assert len(val) == batch_size, f"key {key} length {len(val)} is not equal to batch size {batch_size}."
 
     @classmethod
-    def from_single_dict(cls, data: Dict[str, Union[torch.Tensor, list, NDArray]], meta_info=None):
+    def from_single_dict(cls, data: Dict[str, Union[torch.Tensor, NDArray]], meta_info=None):
         tensors = {}
         non_tensors = {}
         for key, value in data.items():
             if isinstance(value, torch.Tensor):
                 tensors[key] = value
-            elif isinstance(value, (list, np.ndarray)):
+            elif isinstance(value, np.ndarray):
                 non_tensors[key] = value
             else:
                 raise ValueError(f"Unsupported type in data {type(value)}")
@@ -472,8 +473,6 @@ class DataProto:
         assert len(self) % chunks == 0, (
             f"only support equal chunk. Got size of DataProto {len(self)} and chunk {chunks}."
         )
-        chunk_size = len(self) // chunks
-
         if self.batch is not None:
             batch_lst = self.batch.chunk(chunks=chunks, dim=0)
         else:
@@ -481,12 +480,8 @@ class DataProto:
 
         non_tensor_batch_lst = [{} for _ in range(chunks)]
         for key, value in self.non_tensor_batch.items():
-            assert isinstance(value, (list, np.ndarray))
-            if isinstance(value, np.ndarray):
-                non_tensor_lst = np.array_split(value, chunks)
-            else:
-                non_tensor_lst = [value[i : i + chunk_size] for i in range(0, len(self), chunk_size)]
-
+            assert isinstance(value, np.ndarray)
+            non_tensor_lst = np.array_split(value, chunks)
             assert len(non_tensor_lst) == chunks
             for i in range(chunks):
                 non_tensor_batch_lst[i][key] = non_tensor_lst[i]
@@ -524,12 +519,7 @@ class DataProto:
 
         non_tensor_batch = batch_collate([d.non_tensor_batch for d in data])
         for key, value in non_tensor_batch.items():
-            if isinstance(value[0], np.ndarray):
-                non_tensor_batch[key] = np.concatenate(value, axis=0)
-            else:
-                non_tensor_batch[key] = []
-                for item in value:
-                    non_tensor_batch[key].extend(item)
+            non_tensor_batch[key] = np.concatenate(value, axis=0)
 
         return DataProto(batch=new_batch, non_tensor_batch=non_tensor_batch, meta_info=data[0].meta_info)
 
@@ -574,16 +564,10 @@ class DataProto:
 
         repeated_non_tensor_batch = {}
         for key, value in self.non_tensor_batch.items():
-            if isinstance(value, np.ndarray):
-                if interleave:
-                    repeated_non_tensor_batch[key] = np.repeat(value, repeat_times, axis=0)
-                else:
-                    repeated_non_tensor_batch[key] = np.tile(value, (repeat_times,) + (1,) * (value.ndim - 1))
+            if interleave:
+                repeated_non_tensor_batch[key] = np.repeat(value, repeat_times, axis=0)
             else:
-                if interleave:
-                    repeated_non_tensor_batch[key] = [item for item in value for _ in range(repeat_times)]
-                else:
-                    repeated_non_tensor_batch[key] = [item for _ in range(repeat_times) for item in value]
+                repeated_non_tensor_batch[key] = np.tile(value, (repeat_times,) + (1,) * (value.ndim - 1))
 
         return DataProto(
             batch=repeated_batch,
@@ -591,39 +575,6 @@ class DataProto:
             meta_info=self.meta_info,
         )
 
-    def broadcast(self, src: int, group: Optional[ProcessGroup] = None):
-        for key in self.batch.sorted_keys:
-            dist.broadcast(self.batch[key], src=src, group=group, async_op=False)
-
-        object_list = [self.non_tensor_batch]
-        dist.broadcast_object_list(object_list, src=src, group=group)
-        self.non_tensor_batch = object_list[0]
-
-    def all_gather(self, group: Optional[ProcessGroup] = None):
-        world_size = dist.get_world_size(group)
-        output = {}
-        for key in self.batch.sorted_keys:
-            value = self.batch[key].contiguous()
-            output[key] = [torch.empty_like(value) for _ in range(world_size)]
-            dist.all_gather(output[key], value, group=group, async_op=False)
-            output[key] = torch.cat(output[key], dim=0)
-
-        self.batch = TensorDict(output, batch_size=self.batch.batch_size[0] * world_size)
-
-        # all gather non_tensor_batch
-        all_non_tensor_batch = [None for _ in range(world_size)]
-        dist.all_gather_object(all_non_tensor_batch, self.non_tensor_batch, group=group)
-        non_tensor_batch = defaultdict(list)
-        for key, value in self.non_tensor_batch.items():
-            if isinstance(value, np.ndarray):
-                non_tensor_batch[key] = np.concatenate([batch[key] for batch in all_non_tensor_batch])
-            else:
-                for batch in all_non_tensor_batch:
-                    non_tensor_batch[key].extend(batch[key])
-
-        self.non_tensor_batch = non_tensor_batch
-        self.check_consistency()
-
 
 @dataclass
 class DataProtoFuture:
@@ -664,10 +615,53 @@ class DataProtoFuture:
         return arg_future_lst
 
     def get(self):
-        output = ray.get(self.futures)  # dp_size.
-        for o in output:
-            assert isinstance(o, DataProto)
-        output = self.collect_fn(output)  # select dp, concat
+        outputs = ray.get(self.futures)  # dp_size.
+        for output in outputs:
+            assert isinstance(output, DataProto)
+
+        outputs = self.collect_fn(outputs)  # select dp, concat
         if self.dispatch_fn is not None:
-            output = self.dispatch_fn(output)  # split in batch dim, select using dp
-        return output
+            outputs = self.dispatch_fn(outputs)  # split in batch dim, select using dp
+
+        return outputs
+
+
+def allgather_dict_tensors(
+    tensors: Union[Dict[str, torch.Tensor], TensorDict], size: int, group: ProcessGroup, dim: int = 0
+) -> Union[Dict[str, torch.Tensor], TensorDict]:
+    """
+    TODO: optimize this.
+    - We can use async ops
+    - We can use only one allgather
+    """
+    if isinstance(tensors, TensorDict):
+        is_tensor_dict = True
+        tensors_as_dict = tensors.to_dict()
+    else:
+        tensors_as_dict = tensors
+        is_tensor_dict = False
+
+    output = {}
+    sorted_keys = sorted(tensors_as_dict.keys())
+    for key in sorted_keys:
+        val = tensors_as_dict[key]
+        output[key] = [torch.empty_like(val) for _ in range(size)]
+        torch.distributed.all_gather(output[key], val, group=group, async_op=False)
+        output[key] = torch.cat(output[key], dim=dim)
+
+    if is_tensor_dict:
+        output = TensorDict(source=output, batch_size=tensors.batch_size[0] * size)
+
+    return output
+
+
+def all_gather_data_proto(data: DataProto, size: int, group: ProcessGroup) -> None:
+    # Note that this is an inplace operator just like torch.distributed.all_gather
+    prev_device = data.batch.device
+    data.batch = data.batch.cuda(device=torch.cuda.current_device())
+    data.batch = allgather_dict_tensors(data.batch.contiguous(), size=size, group=group, dim=0)
+    data.batch = data.batch.to(prev_device)
+    # all gather non_tensor_batch
+    all_non_tensor_batch = [None for _ in range(size)]
+    torch.distributed.all_gather_object(all_non_tensor_batch, data.non_tensor_batch, group=group)
+    data.non_tensor_batch = {k: np.concatenate([d[k] for d in all_non_tensor_batch]) for k in data.non_tensor_batch}

verl/single_controller/base/__init__.py

@@ -14,3 +14,6 @@
 
 from .worker import Worker
 from .worker_group import ClassWithInitArgs, ResourcePool, WorkerGroup
+
+
+__all__ = ["ClassWithInitArgs", "ResourcePool", "Worker", "WorkerGroup"]

verl/single_controller/base/decorator.py

@@ -12,14 +12,18 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from enum import Enum
+from enum import Enum, auto
 from functools import wraps
 from types import FunctionType
-from typing import Dict, List
+from typing import TYPE_CHECKING, Dict, List, Literal, Union
 
 import ray
 
-from verl.protocol import DataProto, DataProtoFuture
+from ...protocol import DataProto, DataProtoFuture
+
+
+if TYPE_CHECKING:
+    from .worker_group import WorkerGroup
 
 
 # here we add a magic number of avoid user-defined function already have this attribute
@@ -27,13 +31,13 @@ MAGIC_ATTR = "attrs_3141562937"
 
 
 class Dispatch(Enum):
-    RANK_ZERO = 0
-    ONE_TO_ALL = 1
-    ALL_TO_ALL = 2
-    DP_COMPUTE = 3
-    DP_COMPUTE_PROTO = 4
-    DP_COMPUTE_PROTO_WITH_FUNC = 5
-    DP_COMPUTE_METRIC = 6
+    RANK_ZERO = auto()
+    ONE_TO_ALL = auto()
+    ALL_TO_ALL = auto()
+    DP_COMPUTE = auto()
+    DP_COMPUTE_PROTO = auto()
+    DP_COMPUTE_PROTO_WITH_FUNC = auto()
+    DP_COMPUTE_METRIC = auto()
 
 
 class Execute(Enum):
@@ -41,100 +45,85 @@ class Execute(Enum):
     RANK_ZERO = 1
 
 
-def _split_args_kwargs_data_proto(chunks, *args, **kwargs):
+def _split_args_kwargs_data_proto(chunks: int, *args, **kwargs):
     splitted_args = []
     for arg in args:
         assert isinstance(arg, (DataProto, DataProtoFuture))
         splitted_args.append(arg.chunk(chunks=chunks))
 
     splitted_kwargs = {}
-    for key, val in kwargs.items():
-        assert isinstance(val, (DataProto, DataProtoFuture))
-        splitted_kwargs[key] = val.chunk(chunks=chunks)
+    for key, value in kwargs.items():
+        assert isinstance(value, (DataProto, DataProtoFuture))
+        splitted_kwargs[key] = value.chunk(chunks=chunks)
 
     return splitted_args, splitted_kwargs
 
 
-def dispatch_one_to_all(worker_group, *args, **kwargs):
-    from verl.single_controller.base.worker_group import WorkerGroup
-
-    assert isinstance(worker_group, WorkerGroup)
+def dispatch_one_to_all(worker_group: "WorkerGroup", *args, **kwargs):
     args = tuple([arg] * worker_group.world_size for arg in args)
     kwargs = {k: [v] * worker_group.world_size for k, v in kwargs.items()}
     return args, kwargs
 
 
-def dispatch_all_to_all(worker_group, *args, **kwargs):
-    from verl.single_controller.base.worker_group import WorkerGroup
-
-    assert isinstance(worker_group, WorkerGroup)
+def dispatch_all_to_all(worker_group: "WorkerGroup", *args, **kwargs):
     return args, kwargs
 
 
-def collect_all_to_all(worker_group, output):
-    from verl.single_controller.base.worker_group import WorkerGroup
-
-    assert isinstance(worker_group, WorkerGroup)
+def collect_all_to_all(worker_group: "WorkerGroup", output):
     return output
 
 
-def _concat_data_proto_or_future(output: List):
+def _concat_data_proto_or_future(outputs: List[DataProto]) -> DataProto:
     # make sure all the elements in output has the same type
-    for o in output:
-        assert type(o) is type(output[0])
+    for output in outputs:
+        assert type(output) is type(outputs[0])
 
-    o = output[0]
+    output = outputs[0]
 
-    if isinstance(o, DataProto):
-        return DataProto.concat(output)
-    elif isinstance(o, ray.ObjectRef):
-        return DataProtoFuture.concat(output)
+    if isinstance(output, DataProto):
+        return DataProto.concat(outputs)
+    elif isinstance(output, ray.ObjectRef):
+        return DataProtoFuture.concat(outputs)
     else:
         raise NotImplementedError
 
 
-def dispatch_dp_compute(worker_group, *args, **kwargs):
-    from verl.single_controller.base.worker_group import WorkerGroup
-
-    assert isinstance(worker_group, WorkerGroup)
-    return args, kwargs
+def dispatch_dp_compute(worker_group: "WorkerGroup", *args, **kwargs):
+    for arg in args:
+        assert isinstance(arg, (tuple, list)) and len(arg) == worker_group.world_size
 
+    for value in kwargs.values():
+        assert isinstance(value, (tuple, list)) and len(value) == worker_group.world_size
 
-def collect_dp_compute(worker_group, output):
-    from verl.single_controller.base.worker_group import WorkerGroup
+    return args, kwargs
 
-    assert isinstance(worker_group, WorkerGroup)
-    assert len(output) == worker_group.world_size
-    return output
 
+def collect_dp_compute(worker_group: "WorkerGroup", outputs: List[DataProto]) -> List[DataProto]:
+    assert len(outputs) == worker_group.world_size
+    return outputs
 
-def dispatch_dp_compute_data_proto(worker_group, *args, **kwargs):
-    from verl.single_controller.base.worker_group import WorkerGroup
 
-    assert isinstance(worker_group, WorkerGroup)
+def dispatch_dp_compute_data_proto(worker_group: "WorkerGroup", *args, **kwargs):
     splitted_args, splitted_kwargs = _split_args_kwargs_data_proto(worker_group.world_size, *args, **kwargs)
     return splitted_args, splitted_kwargs
 
 
-def dispatch_dp_compute_data_proto_with_func(worker_group, *args, **kwargs):
-    from verl.single_controller.base.worker_group import WorkerGroup
-
-    assert isinstance(worker_group, WorkerGroup)
+def dispatch_dp_compute_data_proto_with_func(worker_group: "WorkerGroup", *args, **kwargs):
     assert type(args[0]) is FunctionType  # NOTE: The first one args is a function!
     splitted_args, splitted_kwargs = _split_args_kwargs_data_proto(worker_group.world_size, *args[1:], **kwargs)
     splitted_args_with_func = [[args[0]] * worker_group.world_size] + splitted_args
     return splitted_args_with_func, splitted_kwargs
 
 
-def collect_dp_compute_data_proto(worker_group, output):
-    for o in output:
-        assert isinstance(o, (DataProto, ray.ObjectRef)), f"expecting {o} to be DataProto, but got {type(o)}"
+def collect_dp_compute_data_proto(worker_group: "WorkerGroup", outputs: List[DataProto]) -> DataProto:
+    for output in outputs:
+        assert isinstance(output, (DataProto, ray.ObjectRef)), f"Expect a DataProto, but got {type(output)}"
 
-    output = collect_dp_compute(worker_group, output)
-    return _concat_data_proto_or_future(output)
+    outputs = collect_dp_compute(worker_group, outputs)
+    return _concat_data_proto_or_future(outputs)
 
 
-def get_predefined_dispatch_fn(dispatch_mode):
+def get_predefined_dispatch_fn(dispatch_mode: Dispatch):
     predefined_dispatch_mode_fn = {
         Dispatch.ONE_TO_ALL: {
             "dispatch_fn": dispatch_one_to_all,
@@ -144,6 +133,10 @@ def get_predefined_dispatch_fn(dispatch_mode):
             "dispatch_fn": dispatch_all_to_all,
             "collect_fn": collect_all_to_all,
         },
+        Dispatch.DP_COMPUTE: {
+            "dispatch_fn": dispatch_dp_compute,
+            "collect_fn": collect_dp_compute,
+        },
         Dispatch.DP_COMPUTE_PROTO: {
             "dispatch_fn": dispatch_dp_compute_data_proto,
             "collect_fn": collect_dp_compute_data_proto,
@@ -152,11 +145,15 @@ def get_predefined_dispatch_fn(dispatch_mode):
             "dispatch_fn": dispatch_dp_compute_data_proto_with_func,
             "collect_fn": collect_dp_compute_data_proto,
         },
+        Dispatch.DP_COMPUTE_METRIC: {
+            "dispatch_fn": dispatch_dp_compute_data_proto,
+            "collect_fn": collect_dp_compute,
+        },
     }
     return predefined_dispatch_mode_fn[dispatch_mode]
 
 
-def get_predefined_execute_fn(execute_mode):
+def get_predefined_execute_fn(execute_mode: Execute):
     """
     Note that here we only asks execute_all and execute_rank_zero to be implemented
     Leave the choice of how these two functions handle argument 'blocking' to users
@@ -168,17 +165,17 @@ def get_predefined_execute_fn(execute_mode):
     return predefined_execute_mode_fn[execute_mode]
 
 
-def _check_dispatch_mode(dispatch_mode):
-    assert isinstance(dispatch_mode, (Dispatch, Dict)), (
+def _check_dispatch_mode(dispatch_mode: Union[Dispatch, Dict[Literal["dispatch_fn", "collect_fn"], FunctionType]]):
+    assert isinstance(dispatch_mode, (Dispatch, dict)), (
         f"dispatch_mode must be a Dispatch or a Dict. Got {dispatch_mode}"
     )
-    if isinstance(dispatch_mode, Dict):
+    if isinstance(dispatch_mode, dict):
         necessary_keys = ["dispatch_fn", "collect_fn"]
         for key in necessary_keys:
             assert key in dispatch_mode, f"key {key} should be in dispatch_mode if it is a dictionary"
 
 
-def _check_execute_mode(execute_mode):
+def _check_execute_mode(execute_mode: Execute):
     assert isinstance(execute_mode, Execute), f"execute_mode must be a Execute. Got {execute_mode}"
 
 
@@ -189,9 +186,10 @@ def _materialize_futures(*args, **kwargs):
             arg = arg.get()
         # add more type to materialize
         new_args.append(arg)
-    for k, v in kwargs.items():
-        if isinstance(v, DataProtoFuture):
-            kwargs[k] = v.get()
+
+    for key, value in kwargs.items():
+        if isinstance(value, DataProtoFuture):
+            kwargs[key] = value.get()
 
     new_args = tuple(new_args)
     return new_args, kwargs

verl/single_controller/base/worker.py

@@ -18,11 +18,13 @@ the class for Worker
 import os
 import socket
 from dataclasses import dataclass
+from typing import Tuple
 
 import ray
+import torch
 
-from verl.single_controller.base.decorator import Dispatch, Execute, register
-from verl.single_controller.base.register_center.ray import create_worker_group_register_center
+from .decorator import Dispatch, Execute, register
+from .register_center.ray import create_worker_group_register_center
 
 
 @dataclass
@@ -40,7 +42,7 @@ class DistGlobalInfo:
 
 
 class WorkerHelper:
-    def _get_node_ip(self):
+    def _get_node_ip(self) -> str:
         host_ipv4 = os.getenv("MY_HOST_IP", None)
         host_ipv6 = os.getenv("MY_HOST_IPV6", None)
         host_ip_by_env = host_ipv4 or host_ipv6
@@ -49,12 +51,12 @@ class WorkerHelper:
         host_ip = host_ip_by_env or host_ip_by_sdk
         return host_ip
 
-    def _get_free_port(self):
+    def _get_free_port(self) -> int:
         with socket.socket() as sock:
             sock.bind(("", 0))
             return sock.getsockname()[1]
 
-    def get_availale_master_addr_port(self):
+    def get_availale_master_addr_port(self) -> Tuple[str, str]:
         return self._get_node_ip(), str(self._get_free_port())
 
     def _get_pid(self):
@@ -81,16 +83,26 @@ class WorkerMeta:
 
 # we assume that in each WorkerGroup, there is a Master Worker
 class Worker(WorkerHelper):
+    """A (distributed) worker."""
+
+    _world_size: int
+    _rank: int
+    _local_world_size: int
+    _local_rank: int
+    _master_addr: str
+    _master_port: str
+    _cuda_visible_devices: str
+
     def __new__(cls, *args, **kwargs):
         instance = super().__new__(cls)
 
         # note that here we use int to distinguish
-        disable_worker_init = int(os.environ.get("DISABLE_WORKER_INIT", 0))
+        disable_worker_init = int(os.getenv("DISABLE_WORKER_INIT", 0))
         if disable_worker_init:
             return instance
 
-        rank = os.environ.get("RANK", None)
-        worker_group_prefix = os.environ.get("WG_PREFIX", None)
+        rank = os.getenv("RANK", None)
+        worker_group_prefix = os.getenv("WG_PREFIX", None)
 
         # when decorator @ray.remote applies, __new__ will be called while we don't want to apply _configure_before_init
         if None not in [rank, worker_group_prefix] and "ActorClass(" not in cls.__name__:
@@ -112,13 +124,19 @@ class Worker(WorkerHelper):
 
     def __init__(self, cuda_visible_devices=None) -> None:
         # construct a meta from envrionment variable. Note that the import must be inside the class because it is executed remotely
-        world_size = int(os.environ["WORLD_SIZE"])
-        rank = int(os.environ["RANK"])
+        world_size = int(os.getenv("WORLD_SIZE"))
+        rank = int(os.getenv("RANK"))
         self._rank = rank
         self._world_size = world_size
 
-        master_addr = os.environ["MASTER_ADDR"]
-        master_port = os.environ["MASTER_PORT"]
+        if "AMD" in torch.cuda.get_device_name():
+            os.environ["CUDA_VISIBLE_DEVICES"] = os.getenv("ROCR_VISIBLE_DEVICES")
+            os.environ["LOCAL_RANK"] = os.getenv("RAY_LOCAL_RANK")
+            cuda_visible_devices = os.getenv("LOCAL_RANK", "0")
+            torch.cuda.set_device(int(cuda_visible_devices))
+
+        master_addr = os.getenv("MASTER_ADDR")
+        master_port = os.getenv("MASTER_PORT")
 
         local_world_size = int(os.getenv("LOCAL_WORLD_SIZE", "1"))
         local_rank = int(os.getenv("LOCAL_RANK", "0"))
@@ -149,6 +167,7 @@ class Worker(WorkerHelper):
             if val is not None:
                 # print(f"set {key} to {val}")
                 os.environ[key] = str(val)
+
         os.environ["REDIS_STORE_SERVER_HOST"] = (
             str(self._master_addr).replace("[", "").replace("]", "") if self._master_addr else ""
         )
@@ -157,7 +176,7 @@ class Worker(WorkerHelper):
         return self._master_addr, self._master_port
 
     def get_cuda_visible_devices(self):
-        cuda_visible_devices = os.environ.get("CUDA_VISIBLE_DEVICES", "not set")
+        cuda_visible_devices = os.getenv("CUDA_VISIBLE_DEVICES", "not set")
         return cuda_visible_devices
 
     def print_rank0(self, *args, **kwargs):

verl/single_controller/base/worker_group.py

@@ -19,20 +19,20 @@ import logging
 import signal
 import threading
 import time
-from typing import Any, Callable, Dict, List
+from typing import Any, Callable, Dict, List, Optional
 
-from verl.single_controller.base.decorator import (
-    MAGIC_ATTR,
-    Dispatch,
-    get_predefined_dispatch_fn,
-    get_predefined_execute_fn,
-)
+from .decorator import MAGIC_ATTR, Dispatch, get_predefined_dispatch_fn, get_predefined_execute_fn
 
 
 class ResourcePool:
-    def __init__(self, process_on_nodes=None, max_collocate_count: int = 10, n_gpus_per_node=8) -> None:
+    """The resource pool with meta info such as world size."""
+
+    def __init__(
+        self, process_on_nodes: Optional[Any] = None, max_collocate_count: int = 10, n_gpus_per_node: int = 8
+    ) -> None:
         if process_on_nodes is None:
             process_on_nodes = []
+
         self._store = process_on_nodes
         self.max_collocate_count = max_collocate_count
         self.n_gpus_per_node = n_gpus_per_node  # this is left for future huawei GPU that contains 16 GPUs per node
@@ -73,28 +73,23 @@ class ClassWithInitArgs:
         self.args = args
         self.kwargs = kwargs
 
-    # def add_arg(self, arg):
-    #     self.args += (arg,)
-
-    # def add_kwarg(self, key, value):
-    #     self.kwargs[key] = value
-
     def __call__(self) -> Any:
         return self.cls(*self.args, **self.kwargs)
 
 
 def check_workers_alive(workers: List, is_alive: Callable, gap_time: float = 1) -> None:
-    import time
-
     while True:
         for worker in workers:
             if not is_alive(worker):
                 logging.warning(f"Worker {worker} is not alive, sending signal to main thread")
                 signal.raise_signal(signal.SIGABRT)
+
         time.sleep(gap_time)
 
 
 class WorkerGroup:
+    """A group of workers"""
+
     def __init__(self, resource_pool: ResourcePool, **kwargs) -> None:
         self._is_init_with_detached_workers = True if resource_pool is None else False
 
@@ -136,14 +131,10 @@ class WorkerGroup:
     def world_size(self):
         return len(self._workers)
 
-    # execute_all_async and execute_rank_zero_async should be implemented by RayWorkerGroup, TorchRPCWorkerGroup,
-    # MegatronWorkerGroup, XperfWorkerGroup should skip
-
     def _bind_worker_method(self, user_defined_cls, func_generator):
         """
         Bind the worker method to the WorkerGroup
         """
-
         for method_name in dir(user_defined_cls):
             try:
                 method = getattr(user_defined_cls, method_name)

verl/single_controller/ray/base.py

@@ -13,27 +13,28 @@
 # limitations under the License.
 
 import os
+import random
+import re
+import string
 import time
-from typing import Any, Dict, List, Tuple
+from typing import Any, Dict, List, Optional, Tuple
 from unittest.mock import patch
 
 import ray
+from ray.actor import ActorHandle
 from ray.experimental.state.api import get_actor
 from ray.util import list_named_actors
 from ray.util.placement_group import PlacementGroup, placement_group
 from ray.util.scheduling_strategies import NodeAffinitySchedulingStrategy, PlacementGroupSchedulingStrategy
 
-from verl.single_controller.base import ClassWithInitArgs, ResourcePool, Worker, WorkerGroup
-from verl.single_controller.base.decorator import MAGIC_ATTR
+from ..base import ClassWithInitArgs, ResourcePool, Worker, WorkerGroup
+from ..base.decorator import MAGIC_ATTR
 
 
 __all__ = ["Worker"]
 
 
 def get_random_string(length: int) -> str:
-    import random
-    import string
-
     letters_digits = string.ascii_letters + string.digits
     return "".join(random.choice(letters_digits) for _ in range(length))
 
@@ -50,6 +51,27 @@ def func_generator(self, method_name, dispatch_fn, collect_fn, execute_fn, block
     return func
 
 
+def sort_placement_group_by_node_ip(pgs: List[PlacementGroup]) -> List[PlacementGroup]:
+    """
+    Sort the placement groups by node ip, all bundles in a single placement group should be on the same node.
+
+    FSDPCheckpointManager saves sharded model states and optimizer states in local storage, which requires RANK
+    to be consistent across nodes when resume from checkpoint.
+
+    With this function, if there's only one resource pool and there's no node change, RANK should be consistent
+    across nodes in multiple ray jobs, even if the whole ray cluster is restarted.
+    """
+    node_ip = {node["NodeID"]: node["NodeManagerAddress"] for node in ray.nodes()}
+    pg_ip = {}
+    for pg in pgs:
+        specs = ray._private.state.state.placement_group_table(pg.id)
+        # all bunles should be on the same node
+        node_id = specs["bundles_to_node_id"][0]
+        pg_ip[pg.id] = node_ip[node_id]
+
+    return sorted(pgs, key=lambda pg: pg_ip[pg.id])
+
+
 class RayResourcePool(ResourcePool):
     def __init__(
         self,
@@ -57,7 +79,7 @@ class RayResourcePool(ResourcePool):
         use_gpu: bool = True,
         name_prefix: str = "",
         max_colocate_count: int = 5,
-        detached=False,
+        detached: bool = False,
     ) -> None:
         super().__init__(process_on_nodes, max_colocate_count)
         self.use_gpu = use_gpu
@@ -66,7 +88,7 @@ class RayResourcePool(ResourcePool):
         self.pgs = None
         self.detached = detached
 
-    def get_placement_groups(self, strategy="STRICT_PACK", name=None):
+    def get_placement_groups(self, strategy: str = "STRICT_PACK", name: Optional[str] = None) -> List[PlacementGroup]:
         if self.pgs is not None:
             return self.pgs
 
@@ -97,7 +119,7 @@ class RayResourcePool(ResourcePool):
 
 def extract_pg_from_exist(
     resource_pools: Dict[str, RayResourcePool], src_role_names: List[str], resource_pool: RayResourcePool
-) -> List:
+) -> List[PlacementGroup]:
     src_pgs = [
         pg
         for role_name, resource_pool in resource_pools.items()
@@ -151,7 +173,12 @@ class RayClassWithInitArgs(ClassWithInitArgs):
         self._options.update(options)
 
     def __call__(
-        self, placement_group, placement_group_bundle_idx, use_gpu: bool = True, num_gpus=1, sharing_with=None
+        self,
+        placement_group: PlacementGroup,
+        placement_group_bundle_idx: int,
+        use_gpu: bool = True,
+        num_gpus: int = 1,
+        sharing_with: Worker = None,
     ) -> Any:
         if sharing_with is not None:
             target_node_id = ray.get(sharing_with.get_node_id.remote())
@@ -188,8 +215,8 @@ class RayWorkerGroup(WorkerGroup):
         ray_cls_with_init: RayClassWithInitArgs = None,
         bin_pack: bool = True,
         name_prefix: str = None,
-        detached=False,
-        worker_names=None,
+        detached: bool = False,
+        worker_names: List[str] = None,
         **kwargs,
     ) -> None:
         super().__init__(resource_pool=resource_pool, **kwargs)
@@ -210,21 +237,24 @@ class RayWorkerGroup(WorkerGroup):
         if ray_cls_with_init is not None:
             self._bind_worker_method(self.ray_cls_with_init.cls, func_generator)
 
-    def _is_worker_alive(self, worker: ray.actor.ActorHandle):
+    def _is_worker_alive(self, worker: ActorHandle) -> bool:
         worker_state_dict = get_actor(worker._actor_id.hex())
         return worker_state_dict.get("state", "undefined") == "ALIVE" if worker_state_dict is not None else False
 
-    def _init_with_detached_workers(self, worker_names):
+    def _init_with_detached_workers(self, worker_names: List[str]) -> None:
         workers = [ray.get_actor(name=name) for name in worker_names]
         self._workers = workers
         self._world_size = len(worker_names)
 
-    def _init_with_resource_pool(self, resource_pool, ray_cls_with_init, bin_pack, detached):
+    def _init_with_resource_pool(
+        self, resource_pool: RayResourcePool, ray_cls_with_init: RayClassWithInitArgs, bin_pack: bool, detached: bool
+    ):
         use_gpu = resource_pool.use_gpu
 
         strategy = "PACK"
         if bin_pack:
             strategy = "STRICT_PACK"
+
         pgs = resource_pool.get_placement_groups(strategy=strategy)
         world_size = resource_pool.world_size
         self._world_size = world_size
@@ -232,8 +262,8 @@ class RayWorkerGroup(WorkerGroup):
         num_gpus = 1 / resource_pool.max_collocate_count
 
         rank = -1
-        for pg_idx, local_world_size in enumerate(resource_pool.store):
-            pg = pgs[pg_idx]
+        local_world_size = resource_pool.store[0]
+        for pg_idx, pg in enumerate(sort_placement_group_by_node_ip(pgs)):
             assert local_world_size <= pg.bundle_count, f"when generating for {self.name_prefix}, for the "
             for local_rank in range(local_world_size):
                 rank += 1
@@ -251,8 +281,6 @@ class RayWorkerGroup(WorkerGroup):
                     env_vars["MASTER_ADDR"] = self._master_addr
                     env_vars["MASTER_PORT"] = self._master_port
 
-                import re
-
                 cia_name = type(ray_cls_with_init.cls).__name__
                 match = re.search(r"ActorClass\(([^)]+)\)", cia_name)  # ray.remote(Obj) -> "ActorClass(Obj)"
                 cia_name = match.group(1) if match else cia_name  # "ActorClass(Obj)" -> "Obj"
@@ -272,7 +300,7 @@ class RayWorkerGroup(WorkerGroup):
 
                 if rank == 0:
                     register_center_actor = None
-                    for _ in range(120):
+                    for _ in range(360):
                         if f"{self.name_prefix}_register_center" not in list_named_actors():
                             time.sleep(1)
                         else:

verl/trainer/config.py

@@ -19,7 +19,7 @@ import os
 from dataclasses import asdict, dataclass, field, fields, is_dataclass
 from typing import Optional, Tuple
 
-from verl.workers.config import WorkerConfig
+from ..workers.config import WorkerConfig
 
 
 def recursive_post_init(dataclass_obj):
@@ -36,12 +36,13 @@ class DataConfig:
     train_files: str = ""
     val_files: str = ""
     prompt_key: str = "prompt"
+    answer_key: str = "answer"
+    image_key: str = "images"
     max_prompt_length: int = 512
     max_response_length: int = 512
     rollout_batch_size: int = 512
-    return_raw_input_ids: bool = False
-    return_raw_prompt: bool = False
-    system_prompt: str = r"Please reason step by step, and put your final answer within \boxed{}."
+    val_batch_size: int = -1
+    system_prompt: Optional[str] = None
     shuffle: bool = True
     seed: int = 1
     max_pixels: int = 4194304
@@ -52,10 +53,12 @@ class DataConfig:
 class AlgorithmConfig:
     gamma: float = 1.0
     lam: float = 1.0
-    adv_estimator: str = "gae"
+    adv_estimator: str = "grpo"
+    disable_kl: bool = False
+    use_kl_loss: bool = False
     kl_penalty: str = "kl"
-    kl_type: str = "fixed"
     kl_coef: float = 1e-3
+    kl_type: str = "fixed"
     kl_horizon: float = 0.0
     kl_target: float = 0.0
 
@@ -67,18 +70,17 @@ class TrainerConfig:
     project_name: str = "easy_r1"
     experiment_name: str = "demo"
     logger: Tuple[str] = ("console", "wandb")
-    val_generations_to_log_to_wandb: int = 0
     nnodes: int = 1
     n_gpus_per_node: int = 8
-    save_freq: int = -1
-    load_checkpoint_path: Optional[str] = None
+    critic_warmup: int = 0
+    val_freq: int = -1
     val_before_train: bool = True
     val_only: bool = False
-    test_freq: int = -1
-    critic_warmup: int = 0
-    remove_previous_ckpt: bool = False
-    del_local_ckpt_after_load: bool = False
+    val_generations_to_log: int = 0
+    save_freq: int = -1
+    save_limit: int = -1
     save_checkpoint_path: Optional[str] = None
+    load_checkpoint_path: Optional[str] = None
 
     def post_init(self):
         if self.save_checkpoint_path is None:
@@ -95,6 +97,10 @@ class PPOConfig:
     def post_init(self):
         self.worker.rollout.prompt_length = self.data.max_prompt_length
         self.worker.rollout.response_length = self.data.max_response_length
+        self.worker.actor.disable_kl = self.algorithm.disable_kl
+        self.worker.actor.use_kl_loss = self.algorithm.use_kl_loss
+        self.worker.actor.kl_penalty = self.algorithm.kl_penalty
+        self.worker.actor.kl_coef = self.algorithm.kl_coef
 
     def deep_post_init(self):
         recursive_post_init(self)

verl/trainer/core_algos.py

+# Copyright 2022 The HuggingFace Team
 # Copyright 2024 Bytedance Ltd. and/or its affiliates
-# Copyright 2022 The HuggingFace Team. All rights reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -18,50 +18,55 @@ The function implemented in this file should be used by trainer with different d
 implement PPO
 """
 
+from abc import ABC, abstractmethod
 from collections import defaultdict
 from typing import TYPE_CHECKING, Tuple
 
 import numpy as np
 import torch
+import torch.nn.functional as F
 
-import verl.utils.torch_functional as verl_F
+from ..utils import torch_functional as VF
 
 
 if TYPE_CHECKING:
-    from verl.trainer.config import AlgorithmConfig
+    from .config import AlgorithmConfig
 
 
-class AdaptiveKLController:
-    """
-    Adaptive KL controller described in the paper:
-    https://arxiv.org/pdf/1909.08593.pdf
-    """
+class KLController(ABC):
+    @abstractmethod
+    def update(self, current_kl: float, n_steps: int) -> None: ...
+
+
+class AdaptiveKLController(KLController):
+    """Adaptive KL controller described in: https://arxiv.org/pdf/1909.08593.pdf"""
 
     def __init__(self, init_kl_coef: float, target_kl: float, horizon: float):
         self.value = init_kl_coef
         self.target = target_kl
         self.horizon = horizon
 
-    def update(self, current_kl, n_steps):
+    def update(self, current_kl: float, n_steps: int) -> None:
         target = self.target
         proportional_error = np.clip(current_kl / target - 1, -0.2, 0.2)
         mult = 1 + proportional_error * n_steps / self.horizon
         self.value *= mult
 
 
-class FixedKLController:
+class FixedKLController(KLController):
     """Fixed KL controller."""
 
-    def __init__(self, kl_coef: float):
-        self.value = kl_coef
+    def __init__(self, init_kl_coef: float):
+        self.value = init_kl_coef
 
-    def update(self, current_kl, n_steps):
+    def update(self, current_kl: float, n_steps: int) -> None:
         pass
 
 
-def get_kl_controller(algorithm_config: "AlgorithmConfig"):
+def get_kl_controller(algorithm_config: "AlgorithmConfig") -> KLController:
+    """Adapted from https://github.com/huggingface/trl/blob/v0.11.0/trl/trainer/ppo_trainer.py#L319"""
     if algorithm_config.kl_type == "fixed":
-        kl_ctrl = FixedKLController(kl_coef=algorithm_config.kl_coef)
+        kl_ctrl = FixedKLController(init_kl_coef=algorithm_config.kl_coef)
     elif algorithm_config.kl_type == "adaptive":
         assert algorithm_config.kl_horizon > 0, f"horizon must be larger than 0. Got {algorithm_config.kl_horizon}."
         kl_ctrl = AdaptiveKLController(
@@ -70,19 +75,20 @@ def get_kl_controller(algorithm_config: "AlgorithmConfig"):
             horizon=algorithm_config.kl_horizon,
         )
     else:
-        raise ValueError("Unknown kl_ctrl type")
+        raise ValueError(f"Unknown kl type: {algorithm_config.kl_type}.")
 
     return kl_ctrl
 
 
+@torch.no_grad()
 def compute_gae_advantage_return(
     token_level_rewards: torch.Tensor,
     values: torch.Tensor,
     eos_mask: torch.Tensor,
     gamma: torch.Tensor,
     lam: torch.Tensor,
-):
-    """Adapted from https://github.com/huggingface/trl/blob/main/trl/trainer/ppo_trainer.py
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """Adapted from https://github.com/huggingface/trl/blob/v0.16.0/trl/trainer/ppo_trainer.py#L513
 
     Args:
         token_level_rewards: `(torch.Tensor)`
@@ -103,27 +109,26 @@ def compute_gae_advantage_return(
             shape: (bs, response_length)
 
     """
-    with torch.no_grad():
-        lastgaelam = 0
-        advantages_reversed = []
-        gen_len = token_level_rewards.shape[-1]
-
-        for t in reversed(range(gen_len)):
-            nextvalues = values[:, t + 1] if t < gen_len - 1 else 0.0
-            delta = token_level_rewards[:, t] + gamma * nextvalues - values[:, t]
-            lastgaelam = delta + gamma * lam * lastgaelam
-            advantages_reversed.append(lastgaelam)
-        advantages = torch.stack(advantages_reversed[::-1], dim=1)
-
-        returns = advantages + values
-        advantages = verl_F.masked_whiten(advantages, eos_mask)
+    lastgaelam = 0
+    advantages_reversed = []
+    gen_len = token_level_rewards.shape[-1]
+    for t in reversed(range(gen_len)):
+        nextvalues = values[:, t + 1] if t < gen_len - 1 else 0.0
+        delta = token_level_rewards[:, t] + gamma * nextvalues - values[:, t]
+        lastgaelam = delta + gamma * lam * lastgaelam
+        advantages_reversed.append(lastgaelam)
+
+    advantages = torch.stack(advantages_reversed[::-1], dim=1)
+    returns = (advantages + values) * eos_mask
+    advantages = VF.masked_whiten(advantages, eos_mask) * eos_mask
     return advantages, returns
 
 
 # NOTE(sgm): this implementation only consider outcome supervision, where the reward is a scalar.
+@torch.no_grad()
 def compute_grpo_outcome_advantage(
-    token_level_rewards: torch.Tensor, eos_mask: torch.Tensor, index: torch.Tensor, epsilon: float = 1e-6
-):
+    token_level_rewards: torch.Tensor, eos_mask: torch.Tensor, index: torch.Tensor, eps: float = 1e-6
+) -> Tuple[torch.Tensor, torch.Tensor]:
     """
     Compute advantage for GRPO, operating only on Outcome reward
     (with only one scalar reward for each response).
@@ -133,6 +138,50 @@ def compute_grpo_outcome_advantage(
         eos_mask: `(torch.Tensor)`
             shape: (bs, response_length)
 
+    Returns:
+        advantages: `(torch.Tensor)`
+            shape: (bs, response_length)
+        Returns: `(torch.Tensor)`
+            shape: (bs, response_length)
+    """
+    response_length = token_level_rewards.shape[-1]
+    scores = token_level_rewards.sum(dim=-1)
+    id2score = defaultdict(list)
+    id2mean, id2std = {}, {}
+
+    bsz = scores.shape[0]
+    for i in range(bsz):
+        id2score[index[i]].append(scores[i])
+
+    for idx in id2score:
+        if len(id2score[idx]) == 1:
+            id2mean[idx] = torch.tensor(0.0)
+            id2std[idx] = torch.tensor(1.0)
+        elif len(id2score[idx]) > 1:
+            id2mean[idx] = torch.mean(torch.tensor(id2score[idx]))
+            id2std[idx] = torch.std(torch.tensor([id2score[idx]]))
+        else:
+            raise ValueError(f"no score in prompt index: {idx}")
+
+    for i in range(bsz):
+        scores[i] = (scores[i] - id2mean[index[i]]) / (id2std[index[i]] + eps)
+
+    scores = scores.unsqueeze(-1).tile([1, response_length]) * eos_mask
+    return scores, scores
+
+
+@torch.no_grad()
+def compute_rloo_outcome_advantage(
+    token_level_rewards: torch.Tensor, eos_mask: torch.Tensor, index: torch.Tensor
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """
+    Compute advantage for RLOO based on https://arxiv.org/abs/2402.14740
+    Args:
+        token_level_rewards: `(torch.Tensor)`
+            shape: (bs, response_length)
+        eos_mask: `(torch.Tensor)`
+            shape: (bs, response_length)
+
     Returns:
         advantages: `(torch.Tensor)`
             shape: (bs, response_length)
@@ -144,31 +193,33 @@ def compute_grpo_outcome_advantage(
 
     id2score = defaultdict(list)
     id2mean = {}
-    id2std = {}
-
-    with torch.no_grad():
-        bsz = scores.shape[0]
-        for i in range(bsz):
-            id2score[index[i]].append(scores[i])
-        for idx in id2score:
-            if len(id2score[idx]) == 1:
-                id2mean[idx] = torch.tensor(0.0)
-                id2std[idx] = torch.tensor(1.0)
-            elif len(id2score[idx]) > 1:
-                id2mean[idx] = torch.mean(torch.tensor(id2score[idx]))
-                id2std[idx] = torch.std(torch.tensor([id2score[idx]]))
-            else:
-                raise ValueError(f"no score in prompt index: {idx}")
-        for i in range(bsz):
-            scores[i] = (scores[i] - id2mean[index[i]]) / (id2std[index[i]] + epsilon)
-        scores = scores.unsqueeze(-1).tile([1, response_length]) * eos_mask
-
+    bsz = scores.shape[0]
+    for i in range(bsz):
+        id2score[index[i]].append(scores[i])
+
+    for idx in id2score:
+        if len(id2score[idx]) == 1:
+            id2mean[idx] = torch.tensor(0.0)
+        elif len(id2score[idx]) > 1:
+            id2mean[idx] = torch.mean(torch.tensor(id2score[idx]))
+        else:
+            raise ValueError(f"no score in prompt index: {idx}.")
+
+    for i in range(bsz):
+        response_num = len(id2score[index[i]])
+        if response_num > 1:
+            scores[i] = scores[i] * response_num / (response_num - 1) - id2mean[index[i]] * response_num / (
+                response_num - 1
+            )
+
+    scores = scores.unsqueeze(-1).tile([1, response_length]) * eos_mask
     return scores, scores
 
 
+@torch.no_grad()
 def compute_reinforce_plus_plus_outcome_advantage(
     token_level_rewards: torch.Tensor, eos_mask: torch.Tensor, gamma: torch.Tensor
-):
+) -> Tuple[torch.Tensor, torch.Tensor]:
     """
     Compute advantage for REINFORCE++.
     This implementation is based on the paper: https://arxiv.org/abs/2501.03262
@@ -184,26 +235,24 @@ def compute_reinforce_plus_plus_outcome_advantage(
         Returns: `(torch.Tensor)`
             shape: (bs, response_length)
     """
-
-    with torch.no_grad():
-        returns = torch.zeros_like(token_level_rewards)
-        running_return = 0
-
-        for t in reversed(range(token_level_rewards.shape[1])):
-            running_return = token_level_rewards[:, t] + gamma * running_return
-            returns[:, t] = running_return
-            # Reset after EOS
-            running_return = running_return * eos_mask[:, t]
-
-        advantages = verl_F.masked_whiten(returns, eos_mask)
-        advantages = advantages * eos_mask
-
+    returns = torch.zeros_like(token_level_rewards)
+    running_return = 0
+    for t in reversed(range(token_level_rewards.shape[1])):
+        running_return = token_level_rewards[:, t] + gamma * running_return
+        returns[:, t] = running_return
+        # Reset after EOS
+        running_return = running_return * eos_mask[:, t]
+
+    advantages = VF.masked_whiten(returns, eos_mask)
+    advantages *= eos_mask
+    returns *= eos_mask
     return advantages, returns
 
 
+@torch.no_grad()
 def compute_remax_outcome_advantage(
     token_level_rewards: torch.Tensor, reward_baselines: torch.Tensor, eos_mask: torch.Tensor
-):
+) -> Tuple[torch.Tensor, torch.Tensor]:
     """
     Compute advantage for ReMax, operating only on Outcome reward
     This implementation is based on the paper: https://arxiv.org/abs/2310.10505
@@ -225,23 +274,31 @@ def compute_remax_outcome_advantage(
     """
     response_length = token_level_rewards.shape[-1]
     # scores = token_level_rewards.sum(dim=-1)
-
-    with torch.no_grad():
-        returns = (token_level_rewards * eos_mask).flip(dims=[-1]).cumsum(dim=-1).flip(dims=[-1])
-        advantages = returns - reward_baselines.unsqueeze(-1).tile([1, response_length]) * eos_mask
-
+    returns = (token_level_rewards * eos_mask).flip(dims=[-1]).cumsum(dim=-1).flip(dims=[-1]) * eos_mask
+    advantages = returns - reward_baselines.unsqueeze(-1).tile([1, response_length]) * eos_mask
     return advantages, returns
 
 
-def compute_rewards(token_level_scores, old_log_prob, ref_log_prob, kl_ratio):
-    kl = old_log_prob - ref_log_prob
+def compute_rewards(
+    token_level_scores: torch.Tensor,
+    log_probs: torch.Tensor,
+    ref_log_probs: torch.Tensor,
+    kl_ratio: float,
+) -> torch.Tensor:
+    kl = log_probs - ref_log_probs
     return token_level_scores - kl * kl_ratio
 
 
 def compute_policy_loss(
-    old_log_prob, log_prob, advantages, eos_mask, cliprange
+    old_log_probs: torch.Tensor,
+    log_probs: torch.Tensor,
+    advantages: torch.Tensor,
+    eos_mask: torch.Tensor,
+    cliprange: float,
 ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
-    """Adapted from https://github.com/huggingface/trl/blob/main/trl/trainer/ppo_trainer.py#L1122
+    """Compute the policy loss.
+
+    Adapted from https://github.com/huggingface/trl/blob/v0.15.0/trl/trainer/ppo_trainer.py#L568
 
     Args:
         old_log_prob: `(torch.Tensor)`
@@ -260,95 +317,88 @@ def compute_policy_loss(
             policy gradient loss computed via PPO
         pg_clipfrac: (float)
             a float number indicating the fraction of policy gradient loss being clipped
-
     """
-    negative_approx_kl = log_prob - old_log_prob
+    negative_approx_kl = log_probs - old_log_probs
+    # clamp the ratio before exp to avoid nan
+    # see: https://github.com/pytorch/pytorch/issues/10729
     ratio = torch.exp(negative_approx_kl)
-    ppo_kl = verl_F.masked_mean(-negative_approx_kl, eos_mask)
+    clipped_ratio = torch.exp(torch.clamp(negative_approx_kl, np.log(1.0 - cliprange), np.log(1.0 + cliprange)))
+    ppo_kl = VF.masked_mean(-negative_approx_kl, eos_mask)
 
     pg_losses = -advantages * ratio
-    pg_losses2 = -advantages * torch.clamp(ratio, 1.0 - cliprange, 1.0 + cliprange)
+    pg_losses2 = -advantages * clipped_ratio
 
-    pg_loss = verl_F.masked_mean(torch.max(pg_losses, pg_losses2), eos_mask)
-    pg_clipfrac = verl_F.masked_mean(torch.gt(pg_losses2, pg_losses).float(), eos_mask)
+    pg_loss = VF.masked_mean(torch.max(pg_losses, pg_losses2), eos_mask)
+    pg_clipfrac = VF.masked_mean(torch.gt(pg_losses2, pg_losses).float(), eos_mask)
     return pg_loss, pg_clipfrac, ppo_kl
 
 
-def compute_entropy_loss(logits, eos_mask):
-    """Compute Categorical entropy loss
-
-    Args:
-        logits: `(torch.Tensor)`
-            shape: (bs, response_length, vocab_size)
-        eos_mask: `(torch.Tensor)`
-            shape: (bs, response_length)
-
-    Returns:
-        entropy: a scalar torch.Tensor
-
-    """
-    # compute entropy
-    entropy = verl_F.entropy_from_logits(logits)  # (bs, response_len)
-    entropy_loss = verl_F.masked_mean(entropy, mask=eos_mask)
-    return entropy_loss
-
+def compute_value_loss(
+    vpreds: torch.Tensor,
+    returns: torch.Tensor,
+    values: torch.Tensor,
+    eos_mask: torch.Tensor,
+    cliprange_value: float,
+) -> Tuple[torch.Tensor, float]:
+    """Compute the value loss.
 
-def compute_value_loss(vpreds, returns, values, eos_mask, cliprange_value):
-    """Compute the value loss. Copied from https://github.com/huggingface/trl/blob/main/trl/trainer/ppo_trainer.py#L1151
+    Copied from https://github.com/huggingface/trl/blob/v0.15.0/trl/trainer/ppo_trainer.py#L556
 
     Args:
         vpreds (`torch.FloatTensor`):
             Predicted values of the value head, shape (`batch_size`, `response_length`)
-        values (`torch.FloatTensor`):
-            Old values of value head, shape (`batch_size`, `response_length`)
         returns: (`torch.FloatTensor`):
             Ground truth returns, shape (`batch_size`, `response_length`)
+        values (`torch.FloatTensor`):
+            Old values of value head, shape (`batch_size`, `response_length`)
+        eos_mask: `(torch.Tensor)`
+            shape: (bs, response_length)
+        cliprange_value: (float)
+            The clip range for value net used in PPO. See https://arxiv.org/abs/1707.06347
 
     Returns:
         vf_loss: a scalar (`torch.FloatTensor`):
             value function loss
         vf_clipfrac: a float
             The ratio of vf being clipped
-
     """
-    vpredclipped = verl_F.clip_by_value(vpreds, values - cliprange_value, values + cliprange_value)
-    vf_losses1 = (vpreds - returns) ** 2
-    vf_losses2 = (vpredclipped - returns) ** 2
-    vf_loss = 0.5 * verl_F.masked_mean(torch.max(vf_losses1, vf_losses2), eos_mask)
-    vf_clipfrac = verl_F.masked_mean(torch.gt(vf_losses2, vf_losses1).float(), eos_mask)
+    vpredclipped = torch.clamp(vpreds, values - cliprange_value, values + cliprange_value)
+    vf_losses1 = torch.square(vpreds - returns)
+    vf_losses2 = torch.square(vpredclipped - returns)
+    vf_loss = 0.5 * VF.masked_mean(torch.max(vf_losses1, vf_losses2), eos_mask)
+    vf_clipfrac = VF.masked_mean(torch.gt(vf_losses2, vf_losses1).float(), eos_mask)
     return vf_loss, vf_clipfrac
 
 
-def kl_penalty(logprob: torch.FloatTensor, ref_logprob: torch.FloatTensor, kl_penalty) -> torch.Tensor:
-    """Compute KL divergence given logprob and ref_logprob.
-    Copied from https://github.com/huggingface/trl/blob/main/trl/trainer/ppo_trainer.py#L1104
+def kl_penalty(log_probs: torch.FloatTensor, ref_log_probs: torch.FloatTensor, kl_penalty: str) -> torch.Tensor:
+    """Compute KL divergence given log_probs and ref_log_probs.
+    Copied from https://github.com/huggingface/trl/blob/v0.11.0/trl/trainer/ppo_trainer.py#L1150
 
     Args:
-        logprob:
-        ref_logprob:
+        log_probs: torch.Tensor
+        ref_log_probs: torch.Tensor
 
     Returns:
-
+        kl_div: torch.Tensor
     """
+    log_probs, ref_log_probs = log_probs.float(), ref_log_probs.float()
     if kl_penalty == "kl":
-        return logprob - ref_logprob
+        return log_probs - ref_log_probs
 
     if kl_penalty == "abs":
-        return (logprob - ref_logprob).abs()
+        return (log_probs - ref_log_probs).abs()
 
     if kl_penalty == "mse":
-        return 0.5 * (logprob - ref_logprob).square()
+        return 0.5 * (log_probs - ref_log_probs).square()
 
     # J. Schulman. Approximating kl divergence, 2020.
-    # # URL http://joschu.net/blog/kl-approx.html.
+    # URL http://joschu.net/blog/kl-approx.html
     if kl_penalty == "low_var_kl":
-        kl = ref_logprob - logprob
-        ratio = torch.exp(kl)
-        kld = (ratio - kl - 1).contiguous()
+        kl = ref_log_probs - log_probs
+        kld = (kl.exp() - kl - 1).contiguous()
         return torch.clamp(kld, min=-10, max=10)
 
     if kl_penalty == "full":
-        # so, here logprob and ref_logprob should contain the logits for every token in vocabulary
-        raise NotImplementedError
+        return F.kl_div(ref_log_probs, log_probs, log_target=True, reduction="none").sum(-1)
 
-    raise NotImplementedError
+    raise NotImplementedError(f"Unknown KL penalty: {kl_penalty}.")

verl/trainer/main.py

@@ -16,81 +16,100 @@ Note that we don't combine the main with ray_trainer as ray_trainer is used by o
 """
 
 import json
-
+import torch
 import ray
+
 from omegaconf import OmegaConf
 
-from verl.single_controller.ray import RayWorkerGroup
-from verl.trainer.config import PPOConfig
-from verl.trainer.ray_trainer import RayPPOTrainer, ResourcePoolManager, Role
-from verl.utils import get_processor, get_tokenizer
-from verl.workers.fsdp_workers import FSDPWorker
-from verl.workers.reward import CustomRewardManager
+from ..single_controller.ray import RayWorkerGroup
+from ..utils.tokenizer import get_processor, get_tokenizer
+from ..workers.fsdp_workers import FSDPWorker
+from ..workers.reward import CustomRewardManager
+from .config import PPOConfig
+from .ray_trainer import RayPPOTrainer, ResourcePoolManager, Role
+
+
+# please make sure main_task is not scheduled on head
+@ray.remote(num_cpus=1)
+class Runner:
+    """A runner for RL training."""
+
+    def run(self, config: PPOConfig):
+        # print config
+        config.deep_post_init()
+        print(json.dumps(config.to_dict(), indent=2))
+
+        # instantiate tokenizer
+        tokenizer = get_tokenizer(
+            config.worker.actor.model.model_path,
+            trust_remote_code=config.worker.actor.model.trust_remote_code,
+            use_fast=True,
+        )
+        processor = get_processor(
+            config.worker.actor.model.model_path,
+            trust_remote_code=config.worker.actor.model.trust_remote_code,
+            use_fast=True,
+        )
+
+        # define worker classes
+        ray_worker_group_cls = RayWorkerGroup
+        role_worker_mapping = {
+            Role.ActorRollout: ray.remote(FSDPWorker),
+            Role.Critic: ray.remote(FSDPWorker),
+            Role.RefPolicy: ray.remote(FSDPWorker),
+        }
+        global_pool_id = "global_pool"
+        resource_pool_spec = {
+            global_pool_id: [config.trainer.n_gpus_per_node] * config.trainer.nnodes,
+        }
+        mapping = {
+            Role.ActorRollout: global_pool_id,
+            Role.Critic: global_pool_id,
+            Role.RefPolicy: global_pool_id,
+        }
+        resource_pool_manager = ResourcePoolManager(resource_pool_spec=resource_pool_spec, mapping=mapping)
+
+        reward_fn = CustomRewardManager(tokenizer=tokenizer, compute_score=config.worker.reward.compute_score)
+        val_reward_fn = CustomRewardManager(tokenizer=tokenizer, compute_score=config.worker.reward.compute_score)
+
+        trainer = RayPPOTrainer(
+            config=config,
+            tokenizer=tokenizer,
+            processor=processor,
+            role_worker_mapping=role_worker_mapping,
+            resource_pool_manager=resource_pool_manager,
+            ray_worker_group_cls=ray_worker_group_cls,
+            reward_fn=reward_fn,
+            val_reward_fn=val_reward_fn,
+        )
+        trainer.init_workers()
+        trainer.fit()
 
 
 def main():
     cli_args = OmegaConf.from_cli()
-    file_config = OmegaConf.load(cli_args.config)
-    del cli_args.config
-
     default_config = OmegaConf.structured(PPOConfig())
-    ppo_config = OmegaConf.merge(default_config, file_config, cli_args)
+
+    if hasattr(cli_args, "config"):
+        config_path = cli_args.pop("config", None)
+        file_config = OmegaConf.load(config_path)
+        default_config = OmegaConf.merge(default_config, file_config)
+
+    ppo_config = OmegaConf.merge(default_config, cli_args)
     ppo_config = OmegaConf.to_object(ppo_config)
 
+    # this is for local ray cluster
     if not ray.is_initialized():
-        # this is for local ray cluster
-        ray.init(runtime_env={"env_vars": {"TOKENIZERS_PARALLELISM": "true", "NCCL_DEBUG": "WARN"}})
-
-    ray.get(main_task.remote(ppo_config))
-
-
-@ray.remote(num_cpus=1)  # please make sure main_task is not scheduled on head
-def main_task(config: PPOConfig):
-    config.deep_post_init()
-    print(json.dumps(config.to_dict(), indent=2))
-    # instantiate tokenizer
-    tokenizer = get_tokenizer(config.worker.actor.model.model_path)
-    processor = get_processor(config.worker.actor.model.model_path, use_fast=True)
-
-    # define worker classes
-    ray_worker_group_cls = RayWorkerGroup
-    role_worker_mapping = {
-        Role.ActorRollout: ray.remote(FSDPWorker),
-        Role.Critic: ray.remote(FSDPWorker),
-        Role.RefPolicy: ray.remote(FSDPWorker),
-    }
-
-    global_pool_id = "global_pool"
-    resource_pool_spec = {
-        global_pool_id: [config.trainer.n_gpus_per_node] * config.trainer.nnodes,
-    }
-    mapping = {
-        Role.ActorRollout: global_pool_id,
-        Role.Critic: global_pool_id,
-        Role.RefPolicy: global_pool_id,
-    }
-
-    reward_fn = CustomRewardManager(
-        tokenizer=tokenizer, num_examine=1, compute_score=config.worker.reward.compute_score
-    )
-    val_reward_fn = CustomRewardManager(
-        tokenizer=tokenizer, num_examine=1, compute_score=config.worker.reward.compute_score
-    )
-
-    resource_pool_manager = ResourcePoolManager(resource_pool_spec=resource_pool_spec, mapping=mapping)
-
-    trainer = RayPPOTrainer(
-        config=config,
-        tokenizer=tokenizer,
-        processor=processor,
-        role_worker_mapping=role_worker_mapping,
-        resource_pool_manager=resource_pool_manager,
-        ray_worker_group_cls=ray_worker_group_cls,
-        reward_fn=reward_fn,
-        val_reward_fn=val_reward_fn,
-    )
-    trainer.init_workers()
-    trainer.fit()
+        # for rocm
+        if torch.version.hip is not None:
+            ray.init(num_gpus=torch.cuda.device_count(),
+                     ignore_reinit_error=True,
+                     runtime_env={"env_vars": {"TOKENIZERS_PARALLELISM": "true", "NCCL_DEBUG": "WARN"}})
+        else:
+            ray.init(runtime_env={"env_vars": {"TOKENIZERS_PARALLELISM": "true", "NCCL_DEBUG": "WARN"}})
+
+    runner = Runner.remote()
+    ray.get(runner.run.remote(ppo_config))
 
 
 if __name__ == "__main__":

verl/trainer/metrics.py

+# 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.
+
+from typing import Any, Dict, List
+
+import numpy as np
+import torch
+
+from ..protocol import DataProto
+
+
+def reduce_metrics(metrics: Dict[str, List[Any]]) -> Dict[str, Any]:
+    return {key: np.mean(value) for key, value in metrics.items()}
+
+
+def compute_data_metrics(batch: DataProto, use_critic: bool = False) -> Dict[str, Any]:
+    sequence_score = batch.batch["token_level_scores"].sum(-1)
+    sequence_reward = batch.batch["token_level_rewards"].sum(-1)
+
+    advantages = batch.batch["advantages"]
+    returns = batch.batch["returns"]
+
+    max_response_length = batch.batch["responses"].size(-1)
+
+    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()
+
+    valid_adv = torch.masked_select(advantages, response_mask)
+    valid_returns = torch.masked_select(returns, response_mask)
+
+    if use_critic:
+        values = batch.batch["values"]
+        valid_values = torch.masked_select(values, response_mask)
+        return_diff_var = torch.var(valid_returns - valid_values)
+        return_var = torch.var(valid_returns)
+
+    metrics = {
+        # score
+        "critic/score/mean": torch.mean(sequence_score).detach().item(),
+        "critic/score/max": torch.max(sequence_score).detach().item(),
+        "critic/score/min": torch.min(sequence_score).detach().item(),
+        # reward
+        "critic/rewards/mean": torch.mean(sequence_reward).detach().item(),
+        "critic/rewards/max": torch.max(sequence_reward).detach().item(),
+        "critic/rewards/min": torch.min(sequence_reward).detach().item(),
+        # adv
+        "critic/advantages/mean": torch.mean(valid_adv).detach().item(),
+        "critic/advantages/max": torch.max(valid_adv).detach().item(),
+        "critic/advantages/min": torch.min(valid_adv).detach().item(),
+        # returns
+        "critic/returns/mean": torch.mean(valid_returns).detach().item(),
+        "critic/returns/max": torch.max(valid_returns).detach().item(),
+        "critic/returns/min": torch.min(valid_returns).detach().item(),
+        **(
+            {
+                # values
+                "critic/values/mean": torch.mean(valid_values).detach().item(),
+                "critic/values/max": torch.max(valid_values).detach().item(),
+                "critic/values/min": torch.min(valid_values).detach().item(),
+                # vf explained var
+                "critic/vf_explained_var": (1.0 - return_diff_var / (return_var + 1e-5)).detach().item(),
+            }
+            if use_critic
+            else {}
+        ),
+        # 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(),
+        # 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(),
+    }
+    return metrics
+
+
+def compute_timing_metrics(batch: DataProto, timing_raw: Dict[str, float]) -> Dict[str, Any]:
+    num_response_tokens = torch.sum(batch.batch["response_mask"]).item()
+    num_overall_tokens = sum(batch.meta_info["global_token_num"])
+    num_tokens_of_section = {
+        **dict.fromkeys(["gen", "reward"], num_response_tokens),
+        **dict.fromkeys(["ref", "old", "values", "adv", "update_critic", "update_actor"], num_overall_tokens),
+    }
+    return {
+        **{f"timing_s/{name}": value for name, value in timing_raw.items()},
+        **{
+            f"timing_per_token_ms/{name}": timing_raw[name] * 1000 / num_tokens_of_section[name]
+            for name in set(num_tokens_of_section.keys()) & set(timing_raw.keys())
+        },
+    }
+
+
+def compute_throughout_metrics(batch: DataProto, timing_raw: Dict[str, float], n_gpus: int) -> Dict[str, Any]:
+    total_num_tokens = sum(batch.meta_info["global_token_num"])
+    time = timing_raw["step"]
+    return {
+        "perf/total_num_tokens": total_num_tokens,
+        "perf/time_per_step": time,
+        "perf/throughput": total_num_tokens / (time * n_gpus),
+    }

verl/trainer/ray_trainer.py

@@ -18,54 +18,71 @@ This trainer supports model-agonistic model initialization with huggingface
 
 import os
 import uuid
+from collections import defaultdict
 from contextlib import contextmanager
 from copy import deepcopy
 from dataclasses import dataclass, field
-from enum import Enum
-from pprint import pprint
-from typing import Any, Dict, Optional, Type
+from enum import Enum, IntEnum, auto
+from typing import Any, Callable, Dict, List, Optional, Type
 
 import numpy as np
+import ray
 import torch
 from codetiming import Timer
-from torch.utils.data import DataLoader, RandomSampler, SequentialSampler
+from ray.experimental.tqdm_ray import tqdm
+from torch.utils.data import RandomSampler, SequentialSampler
+from torchdata.stateful_dataloader import StatefulDataLoader
 from transformers import PreTrainedTokenizer, ProcessorMixin
 
-from verl import DataProto
-from verl.protocol import pad_dataproto_to_divisor, unpad_dataproto
-from verl.single_controller.base import Worker
-from verl.single_controller.ray import RayClassWithInitArgs, RayResourcePool, RayWorkerGroup
-from verl.single_controller.ray.base import create_colocated_worker_cls
-from verl.trainer import core_algos
-from verl.trainer.config import PPOConfig
-from verl.utils.rl_dataset import RLHFDataset, collate_fn
-from verl.utils.torch_functional import masked_mean
-from verl.utils.tracking import Tracking
-from verl.workers.fsdp_workers import FSDPWorker
+from ..protocol import DataProto, pad_dataproto_to_divisor, unpad_dataproto
+from ..single_controller.base import Worker
+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.dataset import RLHFDataset, collate_fn
+from ..utils.logger import Tracker
+from ..utils.py_functional import convert_dict_to_str
+from ..utils.seqlen_balancing import get_seqlen_balanced_partitions, log_seqlen_unbalance
+from ..workers.fsdp_workers import FSDPWorker
+from . import core_algos
+from .config import PPOConfig
+from .metrics import compute_data_metrics, compute_throughout_metrics, compute_timing_metrics, reduce_metrics
 
 
 WorkerType = Type[Worker]
 
 
-class Role(Enum):
+class Role(IntEnum):
     """
     To create more roles dynamically, you can subclass Role and add new members
     """
 
-    Actor = 0
-    Rollout = 1
-    ActorRollout = 2
-    Critic = 3
-    RefPolicy = 4
-    RewardModel = 5
-    ActorRolloutRef = 6
+    Actor = auto()
+    Rollout = auto()
+    ActorRollout = auto()
+    Critic = auto()
+    RefPolicy = auto()
+    RewardModel = auto()
+    ActorRolloutRef = auto()
+
+
+class AdvantageEstimator(str, Enum):
+    """
+    Using an enumeration class to avoid spelling errors in adv_estimator
+    """
+
+    GAE = "gae"
+    GRPO = "grpo"
+    REINFORCE_PLUS_PLUS = "reinforce_plus_plus"
+    REMAX = "remax"
+    RLOO = "rloo"
 
 
 @dataclass
 class ResourcePoolManager:
     """
     Define a resource pool specification. Resource pool will be initialized first.
-    Mapping
     """
 
     resource_pool_spec: dict[str, list[int]]
@@ -82,23 +99,41 @@ class ResourcePoolManager:
             )
             self.resource_pool_dict[resource_pool_name] = resource_pool
 
+        self._check_resource_available()
+
     def get_resource_pool(self, role: Role) -> RayResourcePool:
-        """Get the resource pool of the worker_cls"""
+        """Get the resource pool of the worker."""
         return self.resource_pool_dict[self.mapping[role]]
 
+    def get_n_gpus(self) -> int:
+        """Get the number of gpus in this cluster."""
+        return sum([n_gpus for process_on_nodes in self.resource_pool_spec.values() for n_gpus in process_on_nodes])
 
-def apply_kl_penalty(data: DataProto, kl_ctrl: core_algos.AdaptiveKLController, kl_penalty="kl"):
-    responses = data.batch["responses"]
-    response_length = responses.size(1)
+    def _check_resource_available(self):
+        """Check if the resource pool can be satisfied in this ray cluster."""
+        node_available_resources = ray.state.available_resources_per_node()
+        node_available_gpus = {node: node_info.get("GPU", 0) for node, node_info in node_available_resources.items()}
+
+        # check total required gpus can be satisfied
+        total_available_gpus = sum(node_available_gpus.values())
+        total_required_gpus = sum(
+            [n_gpus for process_on_nodes in self.resource_pool_spec.values() for n_gpus in process_on_nodes]
+        )
+        if total_available_gpus < total_required_gpus:
+            raise ValueError(
+                f"Total available GPUs {total_available_gpus} is less than total desired GPUs {total_required_gpus}."
+            )
+
+
+def apply_kl_penalty(data: DataProto, kl_ctrl: core_algos.KLController, kl_penalty="kl"):
     token_level_scores = data.batch["token_level_scores"]
     batch_size = data.batch.batch_size[0]
-    attention_mask = data.batch["attention_mask"]
-    response_mask = attention_mask[:, -response_length:]
+    response_mask = data.batch["response_mask"]
 
     # compute kl between ref_policy and current policy
-    if "ref_log_prob" in data.batch.keys():
+    if "ref_log_probs" in data.batch.keys():
         kld = core_algos.kl_penalty(
-            data.batch["old_log_probs"], data.batch["ref_log_prob"], kl_penalty=kl_penalty
+            data.batch["old_log_probs"], data.batch["ref_log_probs"], kl_penalty=kl_penalty
         )  # (batch_size, response_length)
         kld = kld * response_mask
         beta = kl_ctrl.value
@@ -108,191 +143,49 @@ def apply_kl_penalty(data: DataProto, kl_ctrl: core_algos.AdaptiveKLController,
 
     token_level_rewards = token_level_scores - beta * kld
 
-    current_kl = masked_mean(kld, mask=response_mask, axis=-1)  # average over sequence
+    current_kl = VF.masked_mean(kld, mask=response_mask, dim=-1)  # average over sequence
     current_kl = torch.mean(current_kl, dim=0).item()
 
-    # according to https://github.com/huggingface/trl/blob/951ca1841f29114b969b57b26c7d3e80a39f75a0/trl/trainer/ppo_trainer.py#L837
+    # According to https://github.com/huggingface/trl/blob/v0.11.0/trl/trainer/ppo_trainer.py#L880
     kl_ctrl.update(current_kl=current_kl, n_steps=batch_size)
     data.batch["token_level_rewards"] = token_level_rewards
 
-    metrics = {"critic/kl": current_kl, "critic/kl_coeff": beta}
-
+    metrics = {"critic/kl": current_kl, "critic/kl_coef": beta}
     return data, metrics
 
 
-def compute_advantage(data: DataProto, adv_estimator, gamma=1.0, lam=1.0, num_repeat=1):
-    # prepare response group
-    # TODO: add other ways to estimate advantages
-    if adv_estimator == "gae":
+def compute_advantage(data: DataProto, adv_estimator: AdvantageEstimator, gamma: float = 1.0, lam: float = 1.0):
+    token_level_rewards = data.batch["token_level_rewards"]
+    response_mask = data.batch["response_mask"]
+    index = data.non_tensor_batch["uid"]
+    if adv_estimator == AdvantageEstimator.GAE:
         values = data.batch["values"]
-        responses = data.batch["responses"]
-        response_length = responses.size(-1)
-        attention_mask = data.batch["attention_mask"]
-        response_mask = attention_mask[:, -response_length:]
-        token_level_rewards = data.batch["token_level_rewards"]
         advantages, returns = core_algos.compute_gae_advantage_return(
             token_level_rewards=token_level_rewards, values=values, eos_mask=response_mask, gamma=gamma, lam=lam
         )
-        data.batch["advantages"] = advantages
-        data.batch["returns"] = returns
-    elif adv_estimator == "grpo":
-        token_level_rewards = data.batch["token_level_rewards"]
-        index = data.non_tensor_batch["uid"]
-        responses = data.batch["responses"]
-        response_length = responses.size(-1)
-        attention_mask = data.batch["attention_mask"]
-        response_mask = attention_mask[:, -response_length:]
+    elif adv_estimator == AdvantageEstimator.GRPO:
         advantages, returns = core_algos.compute_grpo_outcome_advantage(
             token_level_rewards=token_level_rewards, eos_mask=response_mask, index=index
         )
-        data.batch["advantages"] = advantages
-        data.batch["returns"] = returns
-    elif adv_estimator == "reinforce_plus_plus":
-        token_level_rewards = data.batch["token_level_rewards"]
-        responses = data.batch["responses"]
-        response_length = responses.size(-1)
-        attention_mask = data.batch["attention_mask"]
-        response_mask = attention_mask[:, -response_length:]
+    elif adv_estimator == AdvantageEstimator.REINFORCE_PLUS_PLUS:
         advantages, returns = core_algos.compute_reinforce_plus_plus_outcome_advantage(
             token_level_rewards=token_level_rewards, eos_mask=response_mask, gamma=gamma
         )
-        data.batch["advantages"] = advantages
-        data.batch["returns"] = returns
-    elif adv_estimator == "remax":
-        token_level_rewards = data.batch["token_level_rewards"]
-        index = data.non_tensor_batch["uid"]
-        responses = data.batch["responses"]
-        response_length = responses.size(-1)
-        attention_mask = data.batch["attention_mask"]
-        response_mask = attention_mask[:, -response_length:]
-
+    elif adv_estimator == AdvantageEstimator.REMAX:
         reward_baselines = data.batch["reward_baselines"]
-
         advantages, returns = core_algos.compute_remax_outcome_advantage(
             token_level_rewards=token_level_rewards, reward_baselines=reward_baselines, eos_mask=response_mask
         )
-
-        data.batch["advantages"] = advantages
-        data.batch["returns"] = returns
+    elif adv_estimator == AdvantageEstimator.RLOO:
+        advantages, returns = core_algos.compute_rloo_outcome_advantage(
+            token_level_rewards=token_level_rewards, eos_mask=response_mask, index=index
+        )
     else:
         raise NotImplementedError
-    return data
 
-
-def reduce_metrics(metrics: Dict[str, Any]):
-    for key, val in metrics.items():
-        metrics[key] = np.mean(val)
-
-    return metrics
-
-
-def _compute_response_info(batch: DataProto):
-    response_length = batch.batch["responses"].shape[-1]
-
-    prompt_mask = batch.batch["attention_mask"][:, :-response_length]
-    response_mask = batch.batch["attention_mask"][:, -response_length:]
-
-    prompt_length = prompt_mask.sum(-1).float()
-    response_length = response_mask.sum(-1).float()  # (batch_size,)
-
-    return dict(
-        response_mask=response_mask,
-        prompt_length=prompt_length,
-        response_length=response_length,
-    )
-
-
-def compute_data_metrics(batch: DataProto, use_critic: bool = True):
-    # TODO: add response length
-    sequence_score = batch.batch["token_level_scores"].sum(-1)
-    sequence_reward = batch.batch["token_level_rewards"].sum(-1)
-
-    advantages = batch.batch["advantages"]
-    returns = batch.batch["returns"]
-
-    max_response_length = batch.batch["responses"].shape[-1]
-
-    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)
-
-    response_info = _compute_response_info(batch)
-    prompt_length = response_info["prompt_length"]
-    response_length = response_info["response_length"]
-
-    valid_adv = torch.masked_select(advantages, response_mask)
-    valid_returns = torch.masked_select(returns, response_mask)
-
-    if use_critic:
-        values = batch.batch["values"]
-        valid_values = torch.masked_select(values, response_mask)
-        return_diff_var = torch.var(valid_returns - valid_values)
-        return_var = torch.var(valid_returns)
-
-    metrics = {
-        # score
-        "critic/score/mean": torch.mean(sequence_score).detach().item(),
-        "critic/score/max": torch.max(sequence_score).detach().item(),
-        "critic/score/min": torch.min(sequence_score).detach().item(),
-        # reward
-        "critic/rewards/mean": torch.mean(sequence_reward).detach().item(),
-        "critic/rewards/max": torch.max(sequence_reward).detach().item(),
-        "critic/rewards/min": torch.min(sequence_reward).detach().item(),
-        # adv
-        "critic/advantages/mean": torch.mean(valid_adv).detach().item(),
-        "critic/advantages/max": torch.max(valid_adv).detach().item(),
-        "critic/advantages/min": torch.min(valid_adv).detach().item(),
-        # returns
-        "critic/returns/mean": torch.mean(valid_returns).detach().item(),
-        "critic/returns/max": torch.max(valid_returns).detach().item(),
-        "critic/returns/min": torch.min(valid_returns).detach().item(),
-        **(
-            {
-                # values
-                "critic/values/mean": torch.mean(valid_values).detach().item(),
-                "critic/values/max": torch.max(valid_values).detach().item(),
-                "critic/values/min": torch.min(valid_values).detach().item(),
-                # vf explained var
-                "critic/vf_explained_var": (1.0 - return_diff_var / (return_var + 1e-5)).detach().item(),
-            }
-            if use_critic
-            else {}
-        ),
-        # 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(),
-        # 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(),
-    }
-    return metrics
-
-
-def compute_timing_metrics(batch, timing_raw):
-    response_info = _compute_response_info(batch)
-    num_prompt_tokens = torch.sum(response_info["prompt_length"]).item()
-    num_response_tokens = torch.sum(response_info["response_length"]).item()
-    num_overall_tokens = num_prompt_tokens + num_response_tokens
-
-    num_tokens_of_section = {
-        "gen": num_response_tokens,
-        **{name: num_overall_tokens for name in ["ref", "values", "adv", "update_critic", "update_actor"]},
-    }
-
-    return {
-        **{f"timing_s/{name}": value for name, value in timing_raw.items()},
-        **{
-            f"timing_per_token_ms/{name}": timing_raw[name] * 1000 / num_tokens_of_section[name]
-            for name in set(num_tokens_of_section.keys()) & set(timing_raw.keys())
-        },
-    }
+    data.batch["advantages"] = advantages
+    data.batch["returns"] = returns
+    return data
 
 
 @contextmanager
@@ -308,8 +201,6 @@ class RayPPOTrainer:
     Note that this trainer runs on the driver process on a single CPU/GPU node.
     """
 
-    # TODO: support each role have individual ray_worker_group_cls,
-    # i.e., support different backend of different role
     def __init__(
         self,
         config: PPOConfig,
@@ -318,8 +209,8 @@ class RayPPOTrainer:
         role_worker_mapping: dict[Role, WorkerType],
         resource_pool_manager: ResourcePoolManager,
         ray_worker_group_cls: RayWorkerGroup = RayWorkerGroup,
-        reward_fn=None,
-        val_reward_fn=None,
+        reward_fn: Callable = None,
+        val_reward_fn: Callable = None,
     ):
         self.tokenizer = tokenizer
         self.processor = processor
@@ -328,42 +219,51 @@ class RayPPOTrainer:
         self.val_reward_fn = val_reward_fn
 
         self.hybrid_engine = config.worker.hybrid_engine
-        assert self.hybrid_engine, "Currently, only support hybrid engine"
-
         if self.hybrid_engine:
-            assert Role.ActorRollout in role_worker_mapping, f"{role_worker_mapping.keys()}"
+            assert Role.ActorRollout in role_worker_mapping, (
+                f"ActorRollout should be included in {role_worker_mapping.keys()}."
+            )
+        else:
+            raise NotImplementedError
 
         self.role_worker_mapping = role_worker_mapping
         self.resource_pool_manager = resource_pool_manager
-        self.use_reference_policy = Role.RefPolicy in role_worker_mapping
         self.use_reward_model = Role.RewardModel in role_worker_mapping
         self.ray_worker_group_cls = ray_worker_group_cls
 
         # define KL control
-        if self.use_reference_policy:
+        if Role.RefPolicy in role_worker_mapping and not config.algorithm.disable_kl:
+            self.use_reference_policy = True
             self.kl_ctrl = core_algos.get_kl_controller(config.algorithm)
         else:
-            self.kl_ctrl = core_algos.FixedKLController(kl_coef=0.0)
+            self.use_reference_policy = False
+            self.kl_ctrl = core_algos.FixedKLController(init_kl_coef=0.0)
+            print("KL is disabled, no KL metrics will be logged. Please set `kl_coef=0` to log KL metrics.")
 
-        if self.config.algorithm.adv_estimator == "gae":
+        if config.algorithm.adv_estimator == AdvantageEstimator.GAE:
             self.use_critic = True
-        elif self.config.algorithm.adv_estimator == "grpo":
-            self.use_critic = False
-        elif self.config.algorithm.adv_estimator == "reinforce_plus_plus":
-            self.use_critic = False
-        elif self.config.algorithm.adv_estimator == "remax":
-            self.use_critic = False
         else:
-            raise NotImplementedError
+            self.use_critic = False
+
+        if config.algorithm.adv_estimator not in list(AdvantageEstimator):
+            raise NotImplementedError(f"Unknown advantage estimator: {config.algorithm.adv_estimator}.")
+
+        if config.data.rollout_batch_size % config.worker.actor.global_batch_size != 0:
+            raise ValueError("Rollout batch size must be divisible by global batch size.")
+
+        if self.use_critic and config.data.rollout_batch_size % config.worker.critic.global_batch_size != 0:
+            raise ValueError("Rollout batch size must be divisible by global batch size.")
 
         self._create_dataloader()
 
-    def _create_dataloader(self):
+    def _create_dataloader(self) -> None:
         self.train_dataset = RLHFDataset(
             data_path=self.config.data.train_files,
             tokenizer=self.tokenizer,
             processor=self.processor,
             prompt_key=self.config.data.prompt_key,
+            answer_key=self.config.data.answer_key,
+            image_key=self.config.data.image_key,
             max_prompt_length=self.config.data.max_prompt_length,
             truncation="right",
             system_prompt=self.config.data.system_prompt,
@@ -378,13 +278,14 @@ class RayPPOTrainer:
         else:
             sampler = SequentialSampler(data_source=self.train_dataset)
 
-        self.train_dataloader = DataLoader(
+        self.train_dataloader = StatefulDataLoader(
             dataset=self.train_dataset,
             batch_size=self.config.data.rollout_batch_size,
+            sampler=sampler,
             num_workers=8,
-            drop_last=True,
             collate_fn=collate_fn,
-            sampler=sampler,
+            pin_memory=False,
+            drop_last=True,
         )
 
         self.val_dataset = RLHFDataset(
@@ -392,24 +293,28 @@ class RayPPOTrainer:
             tokenizer=self.tokenizer,
             processor=self.processor,
             prompt_key=self.config.data.prompt_key,
+            answer_key=self.config.data.answer_key,
+            image_key=self.config.data.image_key,
             max_prompt_length=self.config.data.max_prompt_length,
             truncation="right",
             system_prompt=self.config.data.system_prompt,
             min_pixels=self.config.data.min_pixels,
             max_pixels=self.config.data.max_pixels,
         )
-        self.val_dataloader = DataLoader(
+        self.val_dataloader = StatefulDataLoader(
             dataset=self.val_dataset,
-            batch_size=len(self.val_dataset),
-            num_workers=8,
+            batch_size=len(self.val_dataset)
+            if self.config.data.val_batch_size == -1
+            else self.config.data.val_batch_size,
             shuffle=False,
-            drop_last=False,
+            num_workers=8,
             collate_fn=collate_fn,
+            pin_memory=False,
+            drop_last=False,
         )
 
         assert len(self.train_dataloader) >= 1
         assert len(self.val_dataloader) >= 1
-
         print(f"Size of train dataloader: {len(self.train_dataloader)}")
         print(f"Size of val dataloader: {len(self.val_dataloader)}")
 
@@ -423,20 +328,11 @@ class RayPPOTrainer:
         self.config.worker.critic.optim.training_steps = training_steps
         print(f"Total training steps: {self.training_steps}")
 
-    def _maybe_log_val_generations_to_wandb(self, inputs, outputs, scores):
-        """Log a table of validation samples to wandb"""
-
-        generations_to_log = self.config.trainer.val_generations_to_log_to_wandb
-
-        if generations_to_log == 0:
+    def _maybe_log_val_generations(self, inputs: List[str], outputs: List[str], scores: List[float]) -> None:
+        """Log a table of validation samples"""
+        if self.config.trainer.val_generations_to_log <= 0:
             return
 
-        if generations_to_log > 0 and "wandb" not in self.config.trainer.logger:
-            print("WARNING: `val_generations_to_log_to_wandb` is set, but no wandb logger is found.")
-            return
-
-        import wandb
-
         # Create tuples of (input, output, score) and sort by input text
         samples = list(zip(inputs, outputs, scores))
         samples.sort(key=lambda x: x[0])  # Sort by input text
@@ -445,43 +341,14 @@ class RayPPOTrainer:
         rng = np.random.RandomState(42)
         rng.shuffle(samples)
 
-        # Take first N samples after shuffling
-        samples = samples[:generations_to_log]
+        samples = samples[: self.config.trainer.val_generations_to_log]
+        self.logger.log_generation(samples, self.global_step)
 
-        # Create column names for all samples
-        columns = ["step"] + sum(
-            [[f"input_{i + 1}", f"output_{i + 1}", f"score_{i + 1}"] for i in range(len(samples))], []
-        )
-
-        if not hasattr(self, "validation_table"):
-            # Initialize the table on first call
-            self.validation_table = wandb.Table(columns=columns)
-
-        # Create a new table with same columns and existing data
-        # Workaround for https://github.com/wandb/wandb/issues/2981#issuecomment-1997445737
-        new_table = wandb.Table(columns=columns, data=self.validation_table.data)
-
-        # Add new row with all data
-        row_data = []
-        row_data.append(self.global_steps)
-        for sample in samples:
-            row_data.extend(sample)
-
-        new_table.add_data(*row_data)
-
-        # Update reference and log
-        wandb.log({"val/generations": new_table}, step=self.global_steps)
-        self.validation_table = new_table
-
-    def _validate(self):
+    def _validate(self) -> Dict[str, Any]:
         reward_tensor_lst = []
-        data_source_lst = []
-
         # Lists to collect samples for the table
-        sample_inputs = []
-        sample_outputs = []
-        sample_scores = []
-
+        sample_inputs, sample_outputs, sample_scores = [], [], []
+        reward_metrics_lst = defaultdict(list)
         for test_data in self.val_dataloader:
             test_batch = DataProto.from_single_dict(test_data)
             # Store original inputs
@@ -489,10 +356,10 @@ class RayPPOTrainer:
             input_texts = [self.tokenizer.decode(ids, skip_special_tokens=True) for ids in input_ids]
             sample_inputs.extend(input_texts)
 
-            if "pixel_values" in test_batch.non_tensor_batch.keys():
+            if "multi_modal_inputs" in test_batch.non_tensor_batch.keys():
                 test_gen_batch = test_batch.pop(
                     batch_keys=["input_ids", "attention_mask", "position_ids"],
-                    non_tensor_batch_keys=["pixel_values", "image_grid_thw", "raw_prompt_ids", "images"],
+                    non_tensor_batch_keys=["raw_prompt_ids", "multi_modal_data", "multi_modal_inputs"],
                 )
             else:
                 test_gen_batch = test_batch.pop(
@@ -500,15 +367,10 @@ class RayPPOTrainer:
                     non_tensor_batch_keys=["raw_prompt_ids"],
                 )
 
-            test_gen_batch.meta_info = {"do_sample": False}
-
-            # pad to be divisible by dp_size
-            test_gen_batch_padded, pad_size = pad_dataproto_to_divisor(
-                test_gen_batch, self.actor_rollout_wg.world_size
-            )
-            test_output_gen_batch_padded = self.actor_rollout_wg.generate_sequences(test_gen_batch_padded)
-            # unpad
-            test_output_gen_batch = unpad_dataproto(test_output_gen_batch_padded, pad_size=pad_size)
+            test_gen_batch.meta_info = self.config.worker.rollout.val_override_config
+            test_gen_batch, pad_size = pad_dataproto_to_divisor(test_gen_batch, self.actor_rollout_wg.world_size)
+            test_output_gen_batch = self.actor_rollout_wg.generate_sequences(test_gen_batch)
+            test_output_gen_batch = unpad_dataproto(test_output_gen_batch, pad_size=pad_size)
             print("validation generation end")
 
             # Store generated outputs
@@ -519,40 +381,24 @@ class RayPPOTrainer:
             test_batch = test_batch.union(test_output_gen_batch)
 
             # evaluate using reward_function
-            reward_tensor = self.val_reward_fn(test_batch)
+            reward_tensor, reward_metrics = self.val_reward_fn(test_batch)
 
             # Store scores
             scores = reward_tensor.sum(-1).cpu().tolist()
             sample_scores.extend(scores)
 
             reward_tensor_lst.append(reward_tensor)
-            data_source_lst.append(
-                test_batch.non_tensor_batch.get("data_source", ["unknown"] * reward_tensor.shape[0])
-            )
+            for key, value in reward_metrics.items():
+                reward_metrics_lst[key].extend(value)
 
-        self._maybe_log_val_generations_to_wandb(inputs=sample_inputs, outputs=sample_outputs, scores=sample_scores)
+        self._maybe_log_val_generations(inputs=sample_inputs, outputs=sample_outputs, scores=sample_scores)
+        reward_score = torch.cat(reward_tensor_lst, dim=0).sum(-1).mean().item()
+        val_reward_metrics = {f"val/{key}_reward": value for key, value in reduce_metrics(reward_metrics_lst).items()}
+        return {"val/reward_score": reward_score, **val_reward_metrics}
 
-        reward_tensor = torch.cat(reward_tensor_lst, dim=0).sum(-1).cpu()  # (batch_size,)
-        data_sources = np.concatenate(data_source_lst, axis=0)
-
-        # evaluate test_score based on data source
-        data_source_reward = {}
-        for i in range(reward_tensor.shape[0]):
-            data_source = data_sources[i]
-            if data_source not in data_source_reward:
-                data_source_reward[data_source] = []
-            data_source_reward[data_source].append(reward_tensor[i].item())
-
-        metric_dict = {}
-        for data_source, rewards in data_source_reward.items():
-            metric_dict[f"val/test_score/{data_source}"] = np.mean(rewards)
-
-        return metric_dict
-
-    def init_workers(self):
+    def init_workers(self) -> None:
         """Init resource pool and worker group"""
         self.resource_pool_manager.create_resource_pool()
-
         self.resource_pool_to_cls = {pool: {} for pool in self.resource_pool_manager.resource_pool_dict.values()}
 
         # create actor and rollout
@@ -594,7 +440,7 @@ class RayPPOTrainer:
         # NOTE: if you want to use a different resource pool for each role, which can support different parallel size,
         # you should not use `create_colocated_worker_cls`. Instead, directly pass different resource pool to different worker groups.
         # See https://github.com/volcengine/verl/blob/master/examples/ray/tutorial.ipynb for more information.
-        all_wg = {}
+        all_wg: Dict[str, FSDPWorker] = {}
         self.wg_dicts = []
         for resource_pool, class_dict in self.resource_pool_to_cls.items():
             worker_dict_cls = create_colocated_worker_cls(class_dict=class_dict)
@@ -605,62 +451,80 @@ class RayPPOTrainer:
             self.wg_dicts.append(wg_dict)
 
         if self.use_critic:
-            self.critic_wg: FSDPWorker = all_wg["critic"]
+            self.critic_wg = all_wg["critic"]
             self.critic_wg.init_model()
 
         if self.use_reference_policy:
-            self.ref_policy_wg: FSDPWorker = all_wg["ref"]
+            self.ref_policy_wg = all_wg["ref"]
             self.ref_policy_wg.init_model()
 
         if self.use_reward_model:
-            self.rm_wg: FSDPWorker = all_wg["rm"]
+            self.rm_wg = all_wg["rm"]
             self.rm_wg.init_model()
 
         # we should create rollout at the end so that vllm can have a better estimation of kv cache memory
-        self.actor_rollout_wg: FSDPWorker = all_wg["actor_rollout"]
+        self.actor_rollout_wg = all_wg["actor_rollout"]
         self.actor_rollout_wg.init_model()
 
-    def _save_checkpoint(self):
-        # path: {save_checkpoint_path}/global_step_{global_steps}/actor
-        local_global_step_folder = os.path.join(
-            self.config.trainer.save_checkpoint_path, f"global_step_{self.global_steps}"
-        )
-        actor_local_path = os.path.join(local_global_step_folder, "actor")
-
-        self.actor_rollout_wg.save_checkpoint(
-            actor_local_path,
-            self.global_steps,
-            remove_previous_ckpt=self.config.trainer.remove_previous_ckpt,
+    def _save_checkpoint(self) -> None:
+        # path: {save_checkpoint_path}/global_step_{global_step}/{actor,critic}
+        remove_obsolete_ckpt(
+            self.config.trainer.save_checkpoint_path, self.global_step, self.config.trainer.save_limit
         )
+        folder_path = os.path.join(self.config.trainer.save_checkpoint_path, f"global_step_{self.global_step}")
+        actor_path = os.path.join(folder_path, "actor")
+        self.actor_rollout_wg.save_checkpoint(actor_path)
 
         if self.use_critic:
-            critic_local_path = os.path.join(local_global_step_folder, "critic")
-            self.critic_wg.save_checkpoint(
-                critic_local_path,
-                self.global_steps,
-                remove_previous_ckpt=self.config.trainer.remove_previous_ckpt,
-            )
+            critic_path = os.path.join(folder_path, "critic")
+            self.critic_wg.save_checkpoint(critic_path)
 
-        local_latest_checkpointed_iteration = os.path.join(
-            self.config.trainer.save_checkpoint_path, "latest_checkpointed_iteration.txt"
-        )
-        with open(local_latest_checkpointed_iteration, "w") as f:
-            f.write(str(self.global_steps))
+        dataloader_path = os.path.join(folder_path, "dataloader.pt")
+        dataloader_state_dict = self.train_dataloader.state_dict()
+        torch.save(dataloader_state_dict, dataloader_path)
 
-    def _load_checkpoint(self):
+        last_global_step_path = os.path.join(self.config.trainer.save_checkpoint_path, CHECKPOINT_TRACKER)
+        with open(last_global_step_path, "w") as f:
+            f.write(str(self.global_step))
+
+    def _load_checkpoint(self) -> None:
         if self.config.trainer.load_checkpoint_path is None:
             return
 
-        print(f"Load from checkpoint: {self.config.trainer.load_checkpoint_path}")
+        if "global_step_" not in self.config.trainer.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")
-        critic_path = os.path.join(self.config.trainer.load_checkpoint_path, "critic")
-        self.actor_rollout_wg.load_checkpoint(
-            actor_path, del_local_after_load=self.config.trainer.del_local_ckpt_after_load
-        )
+        self.actor_rollout_wg.load_checkpoint(actor_path)
         if self.use_critic:
-            self.critic_wg.load_checkpoint(
-                critic_path, del_local_after_load=self.config.trainer.del_local_ckpt_after_load
-            )
+            critic_path = os.path.join(self.config.trainer.load_checkpoint_path, "critic")
+            self.critic_wg.load_checkpoint(critic_path)
+
+        dataloader_path = os.path.join(self.config.trainer.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)
+        else:
+            print(f"No dataloader state found at {dataloader_path}, will start from scratch.")
+
+    def _balance_batch(self, batch: DataProto, metrics: Dict[str, Any], logging_prefix: str = "global_seqlen") -> None:
+        """Reorder the data on single controller such that each dp rank gets similar total tokens"""
+        attention_mask = batch.batch["attention_mask"]
+        batch_size = attention_mask.shape[0]
+        global_seqlen_lst = batch.batch["attention_mask"].view(batch_size, -1).sum(-1).tolist()  # (train_batch_size,)
+        world_size = self.actor_rollout_wg.world_size
+        global_partition_lst = get_seqlen_balanced_partitions(
+            global_seqlen_lst, k_partitions=world_size, equal_size=True
+        )
+        # reorder based on index. The data will be automatically equally partitioned by dispatch function
+        global_idx = torch.tensor([j for partition in global_partition_lst for j in partition])
+        batch.reorder(global_idx)
+        global_balance_stats = log_seqlen_unbalance(
+            seqlen_list=global_seqlen_lst, partitions=global_partition_lst, prefix=logging_prefix
+        )
+        metrics.update(global_balance_stats)
 
     def fit(self):
         """
@@ -668,13 +532,9 @@ class RayPPOTrainer:
         The driver process only need to call the compute functions of the worker group through RPC to construct the PPO dataflow.
         The light-weight advantage computation is done on the driver process.
         """
-        logger = Tracking(
-            project_name=self.config.trainer.project_name,
-            experiment_name=self.config.trainer.experiment_name,
-            default_backend=self.config.trainer.logger,
-            config=self.config.to_dict(),
-        )
-        self.global_steps = 0
+        self.logger = Tracker(loggers=self.config.trainer.logger, config=self.config.to_dict())
+        self.global_step = 0
+        val_metrics: Optional[Dict[str, Any]] = None
 
         # load checkpoint before doing anything
         self._load_checkpoint()
@@ -683,27 +543,24 @@ class RayPPOTrainer:
         # currently, we only support validation using the reward_function.
         if self.val_reward_fn is not None and self.config.trainer.val_before_train:
             val_metrics = self._validate()
-            pprint(f"Initial validation metrics: {val_metrics}")
-            logger.log(data=val_metrics, step=self.global_steps)
+            self.logger.log(data=val_metrics, step=self.global_step)
             if self.config.trainer.val_only:
                 return
 
-        for _ in range(self.config.trainer.total_episodes):
-            for batch_dict in self.train_dataloader:
-                self.global_steps += 1
-                if self.global_steps >= self.training_steps:
+        for _ in tqdm(range(self.config.trainer.total_episodes), desc="Episode", position=0):
+            for batch_dict in tqdm(self.train_dataloader, desc="Running step", position=1):
+                self.global_step += 1
+                if self.global_step > self.training_steps:
                     break
 
-                metrics = {}
-                timing_raw = {}
-
+                metrics, timing_raw = {}, {}
                 batch: DataProto = DataProto.from_single_dict(batch_dict)
 
                 # pop those keys for generation
-                if "pixel_values" in batch.non_tensor_batch.keys():
+                if "multi_modal_inputs" in batch.non_tensor_batch.keys():
                     gen_batch = batch.pop(
                         batch_keys=["input_ids", "attention_mask", "position_ids"],
-                        non_tensor_batch_keys=["pixel_values", "image_grid_thw", "raw_prompt_ids", "images"],
+                        non_tensor_batch_keys=["raw_prompt_ids", "multi_modal_data", "multi_modal_inputs"],
                     )
                 else:
                     gen_batch = batch.pop(
@@ -719,17 +576,15 @@ class RayPPOTrainer:
                     if self.config.algorithm.adv_estimator == "remax":
                         with _timer("gen_max", timing_raw):
                             gen_baseline_batch = deepcopy(gen_batch)
-                            gen_baseline_batch.meta_info["do_sample"] = False
+                            gen_baseline_batch.meta_info["temperature"] = 0.0
                             gen_baseline_output = self.actor_rollout_wg.generate_sequences(gen_baseline_batch)
 
                             batch = batch.union(gen_baseline_output)
-                            reward_baseline_tensor = self.reward_fn(batch)
+                            reward_baseline_tensor, _ = self.reward_fn(batch)
                             reward_baseline_tensor = reward_baseline_tensor.sum(dim=-1)
 
                             batch.pop(batch_keys=list(gen_baseline_output.batch.keys()))
-
                             batch.batch["reward_baselines"] = reward_baseline_tensor
-
                             del gen_baseline_batch, gen_baseline_output
 
                     batch.non_tensor_batch["uid"] = np.array(
@@ -739,24 +594,37 @@ class RayPPOTrainer:
                     batch = batch.repeat(repeat_times=self.config.worker.rollout.n, interleave=True)
                     batch = batch.union(gen_batch_output)
 
+                    # compute reward
+                    with _timer("reward", timing_raw):
+                        if self.use_reward_model:
+                            raise NotImplementedError("Reward model is not supported yet.")
+
+                        # we combine with rule-based rm
+                        reward_tensor, reward_metrics = self.reward_fn(batch)
+                        batch.batch["token_level_scores"] = reward_tensor
+                        reward_metrics = {
+                            f"reward/{key}": value for key, value in reduce_metrics(reward_metrics).items()
+                        }
+                        metrics.update(reward_metrics)
+
                     # balance the number of valid tokens on each dp rank.
                     # Note that this breaks the order of data inside the batch.
                     # Please take care when you implement group based adv computation such as GRPO and rloo
-                    # self._balance_batch(batch, metrics=metrics) # TODO: re-enable balance batch
+                    self._balance_batch(batch, metrics=metrics)
 
                     # compute global_valid tokens
                     batch.meta_info["global_token_num"] = torch.sum(batch.batch["attention_mask"], dim=-1).tolist()
 
                     # recompute old_log_probs
-                    with _timer("old_log_prob", timing_raw):
-                        old_log_prob = self.actor_rollout_wg.compute_log_prob(batch)
-                        batch = batch.union(old_log_prob)
+                    with _timer("old", timing_raw):
+                        old_log_probs = self.actor_rollout_wg.compute_log_probs(batch)
+                        batch = batch.union(old_log_probs)
 
+                    # compute ref_log_probs
                     if self.use_reference_policy:
-                        # compute reference log_prob
                         with _timer("ref", timing_raw):
-                            ref_log_prob = self.ref_policy_wg.compute_ref_log_prob(batch)
-                            batch = batch.union(ref_log_prob)
+                            ref_log_probs = self.ref_policy_wg.compute_ref_log_probs(batch)
+                            batch = batch.union(ref_log_probs)
 
                     # compute values
                     if self.use_critic:
@@ -765,18 +633,8 @@ class RayPPOTrainer:
                             batch = batch.union(values)
 
                     with _timer("adv", timing_raw):
-                        # compute scores. Support both model and function-based.
-                        # We first compute the scores using reward model. Then, we call reward_fn to combine
-                        # the results from reward model and rule-based results.
-                        if self.use_reward_model:
-                            raise NotImplementedError
-
-                        # we combine with rule-based rm
-                        reward_tensor = self.reward_fn(batch)
-                        batch.batch["token_level_scores"] = reward_tensor
-
-                        # compute rewards. apply_kl_penalty if available
-                        if not self.config.worker.actor.use_kl_loss:  # not grpo
+                        # apply kl penalty if available
+                        if not self.config.algorithm.use_kl_loss and self.use_reference_policy:  # apply kl penalty to reward
                             batch, kl_metrics = apply_kl_penalty(
                                 batch, kl_ctrl=self.kl_ctrl, kl_penalty=self.config.algorithm.kl_penalty
                             )
@@ -790,7 +648,6 @@ class RayPPOTrainer:
                             adv_estimator=self.config.algorithm.adv_estimator,
                             gamma=self.config.algorithm.gamma,
                             lam=self.config.algorithm.lam,
-                            num_repeat=self.config.worker.rollout.n,
                         )
 
                     # update critic
@@ -798,43 +655,51 @@ class RayPPOTrainer:
                         with _timer("update_critic", timing_raw):
                             critic_output = self.critic_wg.update_critic(batch)
 
-                        critic_output_metrics = reduce_metrics(critic_output.meta_info["metrics"])
-                        metrics.update(critic_output_metrics)
+                        critic_metrics = reduce_metrics(critic_output.non_tensor_batch)
+                        metrics.update(critic_metrics)
 
-                    # implement critic warmup
-                    if self.config.trainer.critic_warmup <= self.global_steps:
-                        # update actor
+                    # update actor
+                    if self.config.trainer.critic_warmup <= self.global_step:
                         with _timer("update_actor", timing_raw):
                             actor_output = self.actor_rollout_wg.update_actor(batch)
 
-                        actor_output_metrics = reduce_metrics(actor_output.meta_info["metrics"])
-                        metrics.update(actor_output_metrics)
+                        actor_metrics = reduce_metrics(actor_output.non_tensor_batch)
+                        metrics.update(actor_metrics)
 
                     # validate
                     if (
                         self.val_reward_fn is not None
-                        and self.config.trainer.test_freq > 0
-                        and self.global_steps % self.config.trainer.test_freq == 0
+                        and self.config.trainer.val_freq > 0
+                        and self.global_step % self.config.trainer.val_freq == 0
                     ):
-                        with _timer("testing", timing_raw):
-                            val_metrics: dict = self._validate()
+                        with _timer("validation", timing_raw):
+                            val_metrics = self._validate()
+
                         metrics.update(val_metrics)
 
-                    if self.config.trainer.save_freq > 0 and self.global_steps % self.config.trainer.save_freq == 0:
+                    if self.config.trainer.save_freq > 0 and self.global_step % self.config.trainer.save_freq == 0:
                         with _timer("save_checkpoint", timing_raw):
                             self._save_checkpoint()
 
                 # collect metrics
+                n_gpus = self.resource_pool_manager.get_n_gpus()
                 metrics.update(compute_data_metrics(batch=batch, use_critic=self.use_critic))
                 metrics.update(compute_timing_metrics(batch=batch, timing_raw=timing_raw))
+                metrics.update(compute_throughout_metrics(batch=batch, timing_raw=timing_raw, n_gpus=n_gpus))
 
-                # TODO: make a canonical logger that supports various backend
-                logger.log(data=metrics, step=self.global_steps)
+                self.logger.log(data=metrics, step=self.global_step)
 
         # perform validation after training
         if self.val_reward_fn is not None:
-            val_metrics = self._validate()
-            pprint(f"Final validation metrics: {val_metrics}")
-            logger.log(data=val_metrics, step=self.global_steps)
-
-        self._save_checkpoint()
+            if (
+                val_metrics is None
+                or self.config.trainer.val_freq <= 0
+                or self.global_step % self.config.trainer.val_freq != 0
+            ):
+                val_metrics = self._validate()
+                self.logger.log(data=val_metrics, step=self.global_step)
+
+            print(f"Final validation metrics: {convert_dict_to_str(val_metrics)}")
+
+        if self.config.trainer.save_freq <= 0 or self.global_step % self.config.trainer.save_freq != 0:
+            self._save_checkpoint()

verl/utils/__init__.py

@@ -11,8 +11,3 @@
 # 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.
-
-from .tokenizer import get_processor, get_tokenizer
-
-
-__all__ = ["get_processor", "get_tokenizer"]

verl/utils/checkpoint/__init__.py

@@ -11,3 +11,8 @@
 # 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.
+
+from .checkpoint_manager import CHECKPOINT_TRACKER, remove_obsolete_ckpt
+
+
+__all__ = ["CHECKPOINT_TRACKER", "remove_obsolete_ckpt"]

verl/utils/checkpoint/checkpoint_manager.py

@@ -11,20 +11,27 @@
 # 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 os
 import random
+import re
 import shutil
 import tempfile
+from abc import ABC, abstractmethod
+from typing import Any, Dict, Optional, Union
 
 import numpy as np
 import torch
-import torch.distributed
+import torch.distributed as dist
 from filelock import FileLock
 from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
 from transformers import PreTrainedTokenizer, ProcessorMixin
 
 
-class BaseCheckpointManager:
+CHECKPOINT_TRACKER = "latest_global_step.txt"
+
+
+class BaseCheckpointManager(ABC):
     """
     A checkpoint manager that saves and loads
     - model
@@ -44,42 +51,27 @@ class BaseCheckpointManager:
         model: FSDP,
         optimizer: torch.optim.Optimizer,
         lr_scheduler: torch.optim.lr_scheduler.LRScheduler,
-        tokenizer: PreTrainedTokenizer,
-        processor: ProcessorMixin
+        processing_class: Union[PreTrainedTokenizer, ProcessorMixin],
     ):
-        self.previous_global_step = None
-        self.previous_save_local_path = None
-
         self.model = model
         self.optimizer = optimizer
         self.lr_scheduler = lr_scheduler
-        self.tokenizer = tokenizer
-        self.processor = processor
+        self.processing_class = processing_class
 
         assert isinstance(self.model, FSDP)
-        self.rank = torch.distributed.get_rank()
-        self.world_size = torch.distributed.get_world_size()
+        self.rank = dist.get_rank()
+        self.world_size = dist.get_world_size()
 
+    @abstractmethod
     def load_checkpoint(self, *args, **kwargs):
         raise NotImplementedError
 
+    @abstractmethod
     def save_checkpoint(self, *args, **kwargs):
         raise NotImplementedError
 
-    def remove_previous_save_local_path(self):
-        if not self.previous_save_local_path:
-            return
-
-        abs_path = os.path.abspath(self.previous_save_local_path)
-        print(f"Checkpoint manager remove previous save local path: {abs_path}")
-        if not os.path.exists(abs_path):
-            return
-
-        # remove previous local_path
-        shutil.rmtree(abs_path, ignore_errors=True)
-
     @staticmethod
-    def local_mkdir(path):
+    def local_mkdir(path: str) -> str:
         if not os.path.isabs(path):
             working_dir = os.getcwd()
             path = os.path.join(working_dir, path)
@@ -89,18 +81,16 @@ class BaseCheckpointManager:
         lock_path = os.path.join(tempfile.gettempdir(), lock_filename)
 
         try:
-            with FileLock(lock_path, timeout=60):  # Add timeout
-                # make a new dir
+            with FileLock(lock_path, timeout=60):
                 os.makedirs(path, exist_ok=True)
         except Exception as e:
             print(f"Warning: Failed to acquire lock for {path}: {e}")
-            # Even if the lock is not acquired, try to create the directory
-            os.makedirs(path, exist_ok=True)
+            os.makedirs(path, exist_ok=True)  # even if the lock is not acquired, try to create the directory
 
         return path
 
     @staticmethod
-    def get_rng_state():
+    def get_rng_state() -> Dict[str, Any]:
         rng_state = {
             "cpu": torch.get_rng_state(),
             "cuda": torch.cuda.get_rng_state(),
@@ -110,14 +100,14 @@ class BaseCheckpointManager:
         return rng_state
 
     @staticmethod
-    def load_rng_state(rng_state):
+    def load_rng_state(rng_state: Dict[str, Any]):
         torch.set_rng_state(rng_state["cpu"])
         torch.cuda.set_rng_state(rng_state["cuda"])
         np.random.set_state(rng_state["numpy"])
         random.setstate(rng_state["random"])
 
 
-def find_latest_ckpt_path(path, directory_format="global_step_{}"):
+def find_latest_ckpt_path(path: Optional[str] = None, directory_format: str = "global_step_{}") -> Optional[str]:
     if path is None:
         return None
 
@@ -128,6 +118,7 @@ def find_latest_ckpt_path(path, directory_format="global_step_{}"):
 
     with open(tracker_file, "rb") as f:
         iteration = int(f.read().decode())
+
     ckpt_path = os.path.join(path, directory_format.format(iteration))
     if not os.path.exists(ckpt_path):
         print("Checkpoint does not exist: %s", ckpt_path)
@@ -137,8 +128,33 @@ def find_latest_ckpt_path(path, directory_format="global_step_{}"):
     return ckpt_path
 
 
-def get_checkpoint_tracker_filename(root_path: str):
+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, "latest_checkpointed_iteration.txt")
+    return os.path.join(root_path, CHECKPOINT_TRACKER)
+
+
+def remove_obsolete_ckpt(path: str, global_step: int, save_limit: int = -1, directory_format: str = "global_step_{}"):
+    """
+    Remove the obsolete checkpoints that exceed the save_limit.
+    """
+    if save_limit <= 0:
+        return
+
+    if not os.path.exists(path):
+        return
+
+    pattern = re.escape(directory_format).replace(r"\{\}", r"(\d+)")
+    ckpt_folders = []
+    for folder in os.listdir(path):
+        if match := re.match(pattern, folder):
+            step = int(match.group(1))
+            if step < global_step:
+                ckpt_folders.append((step, folder))
+
+    ckpt_folders.sort(reverse=True)
+    for _, folder in ckpt_folders[save_limit - 1 :]:
+        folder_path = os.path.join(path, folder)
+        shutil.rmtree(folder_path, ignore_errors=True)
+        print(f"Removed obsolete checkpoint: {folder_path}")