Skip to content

GitLab

Commit ff7fb65e authored by chenych's avatar chenych
Browse files

Update

parent c132cbcb
Hide whitespace changes
Inline Side-by-side
Showing with 253 additions and 1030 deletions
scripts/model_merger.py
View file @ ff7fb65e
...@@ -66,7 +66,7 @@ if __name__ == "__main__": ...@@ -66,7 +66,7 @@ if __name__ == "__main__":
print(f"Got device mesh {mesh}, mesh_dim_names {mesh_dim_names}") print(f"Got device mesh {mesh}, mesh_dim_names {mesh_dim_names}")
assert mesh_dim_names in (("fsdp",),), f"Unsupported mesh_dim_names {mesh_dim_names}" assert mesh_dim_names in (("fsdp",), ("ddp", "fsdp")), f"Unsupported mesh_dim_names {mesh_dim_names}"
if "tp" in mesh_dim_names: if "tp" in mesh_dim_names:
# fsdp * tp # fsdp * tp
...@@ -106,9 +106,10 @@ if __name__ == "__main__": ...@@ -106,9 +106,10 @@ if __name__ == "__main__":
if isinstance(tensor, DTensor): if isinstance(tensor, DTensor):
state_dict[key].append(tensor._local_tensor.bfloat16()) state_dict[key].append(tensor._local_tensor.bfloat16())
placements = tuple(tensor.placements) placements = tuple(tensor.placements)
# replicated placement at dp dimension can be discarded # replicated placement at ddp dimension can be discarded
if mesh_dim_names[0] == "dp": if mesh_dim_names[0] == "ddp":
placements = placements[1:] placements = placements[1:]
if key not in param_placements: if key not in param_placements:
param_placements[key] = placements param_placements[key] = placements
else: else:
......
verl/models/registry.py deleted 100644 → 0
View file @ c132cbcb
# 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.
# Supported models using HF Rmpad
# TODO(sgm): HF may supported more than listed here, we should add more after testing
from transformers import GemmaConfig, LlamaConfig, MistralConfig, Qwen2Config
_REOVEPAD_MODELS = {"llama": LlamaConfig, "mistral": MistralConfig, "gemma": GemmaConfig, "qwen2": Qwen2Config}
def check_model_support_rmpad(model_type: str):
assert isinstance(model_type, str)
if model_type not in _REOVEPAD_MODELS.keys():
raise ValueError(
f"Model architecture {model_type} is not supported for now. "
f"RMPad supported architectures: {_REOVEPAD_MODELS.keys()}."
f"Please set `use_remove_padding=False` in the model config."
)
verl/models/transformers/llama.py deleted 100644 → 0
View file @ c132cbcb
# Copyright 2024 Bytedance Ltd. and/or its affiliates
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import sys
from typing import Optional, Tuple
import torch
if sys.version_info >= (3, 11):
pass
else:
pass
from transformers.cache_utils import Cache
from transformers.modeling_flash_attention_utils import _flash_attention_forward
from transformers.models.llama.modeling_llama import apply_rotary_pos_emb
from transformers.utils import logging
from verl.utils.ulysses import (
gather_heads_scatter_seq,
gather_seq_scatter_heads,
get_ulysses_sequence_parallel_world_size,
)
logger = logging.get_logger(__name__)
def llama_flash_attn_forward(
self,
hidden_states: torch.Tensor,
attention_mask: Optional[torch.LongTensor] = None,
position_ids: Optional[torch.LongTensor] = None,
past_key_value: Optional[Cache] = None,
output_attentions: bool = False,
use_cache: bool = False,
cache_position: Optional[torch.LongTensor] = None,
position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, # will become mandatory in v4.46
**kwargs,
) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
"""
adapt from transformers 4.47.1
"""
output_attentions = False
bsz, q_len, _ = hidden_states.size()
query_states = self.q_proj(hidden_states)
key_states = self.k_proj(hidden_states)
value_states = self.v_proj(hidden_states)
# Flash attention requires the input to have the shape
# batch_size x seq_length x head_dim x hidden_dim
# therefore we just need to keep the original shape
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)
# trade off: repeat first and then all to all
# key_states = repeat_kv(key_states, self.num_key_value_groups)
# value_states = repeat_kv(value_states, self.num_key_value_groups)
########## AlltoAll for Ulysses ##########
ulysses_sp_size = get_ulysses_sequence_parallel_world_size()
if ulysses_sp_size > 1:
# (bsz, n_head, seq_len/n, head_dim) -> (bsz, n_head/n, seq_len, head_dim)
query_states = gather_seq_scatter_heads(query_states, seq_dim=2, head_dim=1)
key_states = gather_seq_scatter_heads(key_states, seq_dim=2, head_dim=1)
value_states = gather_seq_scatter_heads(value_states, seq_dim=2, head_dim=1)
full_q_len = query_states.size(2) # full seq length
if position_embeddings is None:
logger.warning_once(
"The attention layers in this model are transitioning from computing the RoPE embeddings internally "
"through `position_ids` (2D tensor with the indexes of the tokens), to using externally computed "
"`position_embeddings` (Tuple of tensors, containing cos and sin). In v4.46 `position_ids` will be "
"removed and `position_embeddings` will be mandatory."
)
cos, sin = self.rotary_emb(value_states, position_ids)
else:
cos, sin = position_embeddings
query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
if past_key_value is not None:
# sin and cos are specific to RoPE models; cache_position needed for the static cache
cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position}
key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)
# TODO: These transpose are quite inefficient but Flash Attention requires the layout [batch_size, sequence_length, num_heads, head_dim]. We would need to refactor the KV cache
# to be able to avoid many of these transpose/reshape/view.
query_states = query_states.transpose(1, 2)
key_states = key_states.transpose(1, 2)
value_states = value_states.transpose(1, 2)
dropout_rate = self.attention_dropout if self.training else 0.0
# In PEFT, usually we cast the layer norms in float32 for training stability reasons
# therefore the input hidden states gets silently casted in float32. Hence, we need
# cast them back in the correct dtype just to be sure everything works as expected.
# This might slowdown training & inference so it is recommended to not cast the LayerNorms
# in fp32. (LlamaRMSNorm handles it correctly)
input_dtype = query_states.dtype
if input_dtype == torch.float32:
if torch.is_autocast_enabled():
target_dtype = torch.get_autocast_gpu_dtype()
# Handle the case where the model is quantized
elif hasattr(self.config, "_pre_quantization_dtype"):
target_dtype = self.config._pre_quantization_dtype
else:
target_dtype = self.q_proj.weight.dtype
logger.warning_once(
f"The input hidden states seems to be silently casted in float32, this might be related to"
f" the fact you have upcasted embedding or layer norm layers in float32. We will cast back the input in"
f" {target_dtype}."
)
query_states = query_states.to(target_dtype)
key_states = key_states.to(target_dtype)
value_states = value_states.to(target_dtype)
attn_output = _flash_attention_forward(
query_states,
key_states,
value_states,
attention_mask,
full_q_len,
position_ids=position_ids,
dropout=dropout_rate,
sliding_window=getattr(self, "sliding_window", None),
use_top_left_mask=self._flash_attn_uses_top_left_mask,
is_causal=self.is_causal,
**kwargs,
)
attn_output = attn_output.reshape(bsz, full_q_len, -1, self.head_dim).contiguous()
########## AlltoAll for Ulysses ##########
if ulysses_sp_size > 1:
attn_output = gather_heads_scatter_seq(attn_output, seq_dim=1, head_dim=2)
attn_output = attn_output.reshape(bsz, q_len, -1).contiguous()
attn_output = self.o_proj(attn_output)
if not output_attentions:
attn_weights = None
return attn_output, attn_weights, past_key_value
verl/models/transformers/monkey_patch.py deleted 100644 → 0
View file @ c132cbcb
# 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.
"""
Apply monkey-patch function to models
"""
#### Open Source Models
#### transformers version < 4.48
import importlib.metadata
from functools import lru_cache
from packaging import version
from transformers import PretrainedConfig
def apply_monkey_patch_to_llama():
from transformers.models.llama.modeling_llama import LlamaFlashAttention2
from verl.models.transformers.llama import llama_flash_attn_forward
LlamaFlashAttention2.forward = llama_flash_attn_forward
def apply_monkey_patch_to_qwen2():
from transformers.models.qwen2.modeling_qwen2 import Qwen2FlashAttention2
from verl.models.transformers.qwen2 import qwen2_flash_attn_forward
Qwen2FlashAttention2.forward = qwen2_flash_attn_forward
_PATCH_NAME_TO_FUNC = {
"llama": apply_monkey_patch_to_llama,
"qwen2": apply_monkey_patch_to_qwen2,
}
def apply_monkey_patch(config: PretrainedConfig, verbose=True):
if not is_transformers_version_in_range("4.45.0", "4.47.1"):
raise AssertionError(
"The installed `transformers` version doesn't support ulysses patch. "
"Please install a version between 4.45.0 and 4.47.1 to use this ulysses feature."
)
success_apply_monkey_patch = False
if config.model_type in _PATCH_NAME_TO_FUNC:
_PATCH_NAME_TO_FUNC[config.model_type]()
success_apply_monkey_patch = True
if success_apply_monkey_patch and verbose:
print(f"Applying monkey patch to model {config.model_type}")
elif not success_apply_monkey_patch:
raise NotImplementedError(
f"Ulysses for model {config.model_type} is not implemented, \
please set `ulysses_sequence_parallel_size=1`"
)
return success_apply_monkey_patch
@lru_cache
def is_transformers_version_in_range(min_version: str, max_version: str) -> bool:
try:
# Get the installed version of the transformers library
transformers_version = importlib.metadata.version("transformers")
except importlib.metadata.PackageNotFoundError:
raise ModuleNotFoundError("The `transformers` package is not installed.")
# Check if the version is within the specified range
return version.parse(min_version) <= version.parse(transformers_version) <= version.parse(max_version)
verl/models/transformers/qwen2.py deleted 100644 → 0
View file @ c132cbcb
# 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 Optional, Tuple
import torch
from transformers.cache_utils import Cache
from transformers.modeling_flash_attention_utils import _flash_attention_forward
from transformers.models.llama.modeling_llama import apply_rotary_pos_emb, repeat_kv
from transformers.utils import logging
from verl.utils.ulysses import (
gather_heads_scatter_seq,
gather_seq_scatter_heads,
get_ulysses_sequence_parallel_world_size,
)
logger = logging.get_logger(__name__)
def qwen2_flash_attn_forward(
self,
hidden_states: torch.Tensor,
attention_mask: Optional[torch.Tensor] = None,
position_ids: Optional[torch.LongTensor] = None,
past_key_value: Optional[Cache] = None,
output_attentions: bool = False,
use_cache: bool = False,
cache_position: Optional[torch.LongTensor] = None,
position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, # will become mandatory in v4.46
):
bsz, q_len, _ = hidden_states.size()
query_states = self.q_proj(hidden_states)
key_states = self.k_proj(hidden_states)
value_states = self.v_proj(hidden_states)
query_states = query_states.view(bsz, q_len, -1, self.head_dim).transpose(1, 2)
key_states = key_states.view(bsz, q_len, -1, self.head_dim).transpose(1, 2)
value_states = value_states.view(bsz, q_len, -1, self.head_dim).transpose(1, 2)
########## AlltoAll for Ulysses ##########
ulysses_sp_size = get_ulysses_sequence_parallel_world_size()
if ulysses_sp_size > 1:
# (bsz, n_head, seq_len/n, head_dim) -> (bsz, n_head/n, seq_len, head_dim)
query_states = gather_seq_scatter_heads(query_states, seq_dim=2, head_dim=1)
key_states = gather_seq_scatter_heads(key_states, seq_dim=2, head_dim=1)
value_states = gather_seq_scatter_heads(value_states, seq_dim=2, head_dim=1)
full_q_len = query_states.size(2) # full seq length
if position_embeddings is None:
logger.warning_once(
"The attention layers in this model are transitioning from computing the RoPE embeddings internally "
"through `position_ids` (2D tensor with the indexes of the tokens), to using externally computed "
"`position_embeddings` (Tuple of tensors, containing cos and sin). In v4.46 `position_ids` will be "
"removed and `position_embeddings` will be mandatory."
)
cos, sin = self.rotary_emb(value_states, position_ids)
else:
cos, sin = position_embeddings
query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
if past_key_value is not None:
cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position} # Specific to RoPE models
key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)
# repeat k/v heads if n_kv_heads < n_heads
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
# In PEFT, usually we cast the layer norms in float32 for training stability reasons
# therefore the input hidden states gets silently casted in float32. Hence, we need
# cast them back in float16 just to be sure everything works as expected.
input_dtype = query_states.dtype
if input_dtype == torch.float32:
if torch.is_autocast_enabled():
target_dtype = torch.get_autocast_gpu_dtype()
# Handle the case where the model is quantized
elif hasattr(self.config, "_pre_quantization_dtype"):
target_dtype = self.config._pre_quantization_dtype
else:
target_dtype = self.q_proj.weight.dtype
logger.warning_once(
f"The input hidden states seems to be silently casted in float32, this might be related to"
f" the fact you have upcasted embedding or layer norm layers in float32. We will cast back the input in"
f" {target_dtype}."
)
query_states = query_states.to(target_dtype)
key_states = key_states.to(target_dtype)
value_states = value_states.to(target_dtype)
# Reashape to the expected shape for Flash Attention
query_states = query_states.transpose(1, 2)
key_states = key_states.transpose(1, 2)
value_states = value_states.transpose(1, 2)
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
else:
sliding_window = None
attn_output = _flash_attention_forward(
query_states,
key_states,
value_states,
attention_mask,
full_q_len,
position_ids=position_ids,
dropout=dropout_rate,
sliding_window=sliding_window,
is_causal=self.is_causal,
use_top_left_mask=self._flash_attn_uses_top_left_mask,
)
# use full_q_len to reshape
attn_output = attn_output.reshape(bsz, full_q_len, -1, self.head_dim).contiguous()
########## AlltoAll for Ulysses ##########
if ulysses_sp_size > 1:
attn_output = gather_heads_scatter_seq(attn_output, seq_dim=1, head_dim=2)
attn_output = attn_output.reshape(bsz, q_len, -1).contiguous()
attn_output = self.o_proj(attn_output)
if not output_attentions:
attn_weights = None
return attn_output, attn_weights, past_key_value
verl/models/transformers/qwen2_5_vl.py deleted 100644 → 0
View file @ c132cbcb
from typing import Optional
import torch
from transformers.models.qwen2_5_vl.processing_qwen2_5_vl import Qwen2_5_VLProcessor
def get_rope_index(
processor: Qwen2_5_VLProcessor,
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
verl/protocol.py
View file @ ff7fb65e
...@@ -122,7 +122,7 @@ def batch_collate(features: List[Dict[str, Any]]) -> Dict[str, List[Any]]: ...@@ -122,7 +122,7 @@ def batch_collate(features: List[Dict[str, Any]]) -> Dict[str, List[Any]]:
return batch_features return batch_features
def fold_batch_dim(data: "DataProto", new_batch_size): def fold_batch_dim(data: "DataProto", new_batch_size: int):
""" """
Fold a batch dim from [bsz, xxx] into [new_bsz, bsz // new_bsz, xxx] Fold a batch dim from [bsz, xxx] into [new_bsz, bsz // new_bsz, xxx]
""" """
...@@ -158,8 +158,8 @@ def collate_fn(data_items: list["DataProtoItem"]): ...@@ -158,8 +158,8 @@ def collate_fn(data_items: list["DataProtoItem"]):
@dataclass @dataclass
class DataProtoItem: class DataProtoItem:
batch: Optional[TensorDict] = None batch: Optional[TensorDict] = None
non_tensor_batch: Dict = field(default_factory=dict) non_tensor_batch: Dict[str, NDArray] = field(default_factory=dict)
meta_info: Dict = field(default_factory=dict) meta_info: Dict[str, Any] = field(default_factory=dict)
@dataclass @dataclass
...@@ -172,13 +172,13 @@ class DataProto: ...@@ -172,13 +172,13 @@ class DataProto:
""" """
batch: Optional[TensorDict] = None batch: Optional[TensorDict] = None
non_tensor_batch: Dict[str, Any] = field(default_factory=dict) non_tensor_batch: Dict[str, NDArray] = field(default_factory=dict)
meta_info: Dict[str, Any] = field(default_factory=dict) meta_info: Dict[str, Any] = field(default_factory=dict)
def __post_init__(self): def __post_init__(self):
self.check_consistency() # perform necessary checking self.check_consistency() # perform necessary checking
def __len__(self): def __len__(self) -> int:
if self.batch is not None: if self.batch is not None:
return self.batch.batch_size[0] return self.batch.batch_size[0]
elif self.non_tensor_batch is not None and len(self.non_tensor_batch) > 0: elif self.non_tensor_batch is not None and len(self.non_tensor_batch) > 0:
...@@ -187,43 +187,41 @@ class DataProto: ...@@ -187,43 +187,41 @@ class DataProto:
else: else:
return 0 return 0
def __getitem__(self, item): def __getitem__(self, item: Union[int, slice]) -> Union["DataProto", "DataProtoItem"]:
tensor_data = self.batch[item] tensor_data = self.batch[item]
non_tensor_data = {key: val[item] for key, val in self.non_tensor_batch.items()} non_tensor_data = {key: val[item] for key, val in self.non_tensor_batch.items()}
return_type = DataProto if isinstance(item, slice) else DataProtoItem 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) return return_type(batch=tensor_data, non_tensor_batch=non_tensor_data, meta_info=self.meta_info)
def __getstate__(self): def __getstate__(self) -> Tuple[bytes, Dict[str, NDArray], Dict[str, Any]]:
buffer = io.BytesIO() buffer = io.BytesIO()
if self.batch is not None: if self.batch is not None:
self.batch = self.batch.contiguous() self.batch: TensorDict = self.batch.contiguous()
self.batch = self.batch.consolidate() self.batch: TensorDict = self.batch.consolidate()
torch.save(self.batch, buffer) torch.save(self.batch, buffer)
buffer_bytes = buffer.getvalue() buffer_bytes = buffer.getvalue()
return buffer_bytes, self.non_tensor_batch, self.meta_info return buffer_bytes, self.non_tensor_batch, self.meta_info
def __setstate__(self, data): def __setstate__(self, data: Tuple[bytes, Dict[str, NDArray], Dict[str, Any]]) -> None:
batch_deserialized_bytes, non_tensor_batch, meta_info = data batch_deserialized_bytes, non_tensor_batch, meta_info = data
batch_deserialized = io.BytesIO(initial_bytes=batch_deserialized_bytes) batch_deserialized = io.BytesIO(batch_deserialized_bytes)
batch = torch.load( batch = torch.load(batch_deserialized, weights_only=False, map_location="cpu")
batch_deserialized, weights_only=False, map_location="cpu" if not torch.cuda.is_available() else None
)
self.batch = batch self.batch = batch
self.non_tensor_batch = non_tensor_batch self.non_tensor_batch = non_tensor_batch
self.meta_info = meta_info self.meta_info = meta_info
def save_to_disk(self, filepath): def save_to_disk(self, filepath: str) -> None:
with open(filepath, "wb") as f: with open(filepath, "wb") as f:
pickle.dump(self, f) pickle.dump(self, f)
@staticmethod @staticmethod
def load_from_disk(filepath) -> "DataProto": def load_from_disk(filepath: str) -> "DataProto":
with open(filepath, "rb") as f: with open(filepath, "rb") as f:
data = pickle.load(f) data = pickle.load(f)
return data return data
def print_size(self, prefix=""): def print_size(self, prefix: str = "") -> None:
size_of_tensordict = 0 size_of_tensordict = 0
for tensor in self.batch.values(): for tensor in self.batch.values():
if isinstance(tensor, torch.Tensor): if isinstance(tensor, torch.Tensor):
...@@ -255,7 +253,11 @@ class DataProto: ...@@ -255,7 +253,11 @@ class DataProto:
assert len(val) == batch_size, f"key {key} length {len(val)} is not equal to batch size {batch_size}." assert len(val) == batch_size, f"key {key} length {len(val)} is not equal to batch size {batch_size}."
@classmethod @classmethod
def from_single_dict(cls, data: Dict[str, Union[torch.Tensor, NDArray]], meta_info=None): def from_single_dict(
cls,
data: Dict[str, Union[torch.Tensor, NDArray]],
meta_info: Optional[Dict[str, Any]] = None,
) -> "DataProto":
tensors = {} tensors = {}
non_tensors = {} non_tensors = {}
for key, value in data.items(): for key, value in data.items():
...@@ -269,7 +271,13 @@ class DataProto: ...@@ -269,7 +271,13 @@ class DataProto:
return DataProto.from_dict(tensors=tensors, non_tensors=non_tensors, meta_info=meta_info) return DataProto.from_dict(tensors=tensors, non_tensors=non_tensors, meta_info=meta_info)
@classmethod @classmethod
def from_dict(cls, tensors: Dict[str, torch.Tensor], non_tensors=None, meta_info=None, num_batch_dims=1): def from_dict(
cls,
tensors: Dict[str, torch.Tensor],
non_tensors: Dict[str, NDArray] = None,
meta_info: Optional[Dict[str, Any]] = None,
num_batch_dims: int = 1,
) -> "DataProto":
"""Create a DataProto from a dict of tensors. This assumes that """Create a DataProto from a dict of tensors. This assumes that
1. All the tensor in tensors have the same dim0 1. All the tensor in tensors have the same dim0
2. Only dim0 is the batch dim 2. Only dim0 is the batch dim
...@@ -293,13 +301,14 @@ class DataProto: ...@@ -293,13 +301,14 @@ class DataProto:
else: else:
current_batch = tensor.shape[:num_batch_dims] current_batch = tensor.shape[:num_batch_dims]
assert batch_size == current_batch, ( assert batch_size == current_batch, (
f"Not all the tensor in tensors have the same batch size with batch_dims={num_batch_dims}. Got {pivot_key} has {batch_size}, {key} has {current_batch}" f"Not all the tensor in tensors have the same batch size with batch_dims={num_batch_dims}. "
f"Got {pivot_key} has {batch_size}, {key} has {current_batch}"
) )
tensor_dict = TensorDict(source=tensors, batch_size=batch_size) tensor_dict = TensorDict(source=tensors, batch_size=batch_size)
return cls(batch=tensor_dict, non_tensor_batch=non_tensors, meta_info=meta_info) return cls(batch=tensor_dict, non_tensor_batch=non_tensors, meta_info=meta_info)
def to(self, device) -> "DataProto": def to(self, device: torch.device) -> "DataProto":
"""move the batch to device """move the batch to device
Args: Args:
...@@ -314,7 +323,13 @@ class DataProto: ...@@ -314,7 +323,13 @@ class DataProto:
return self return self
def select(self, batch_keys=None, non_tensor_batch_keys=None, meta_info_keys=None, deepcopy=False) -> "DataProto": def select(
self,
batch_keys: Optional[List[str]] = None,
non_tensor_batch_keys: Optional[List[str]] = None,
meta_info_keys: Optional[List[str]] = None,
deepcopy: bool = False,
) -> "DataProto":
"""Select a subset of the DataProto via batch_keys and meta_info_keys """Select a subset of the DataProto via batch_keys and meta_info_keys
Args: Args:
...@@ -332,7 +347,7 @@ class DataProto: ...@@ -332,7 +347,7 @@ class DataProto:
sub_batch = self.batch sub_batch = self.batch
if non_tensor_batch_keys is not None: if non_tensor_batch_keys is not None:
non_tensor_batch = {key: val for key, val in self.non_tensor_batch.items() if key in non_tensor_batch_keys} non_tensor_batch = {k: v for k, v in self.non_tensor_batch.items() if k in non_tensor_batch_keys}
else: else:
non_tensor_batch = self.non_tensor_batch non_tensor_batch = self.non_tensor_batch
...@@ -340,7 +355,7 @@ class DataProto: ...@@ -340,7 +355,7 @@ class DataProto:
non_tensor_batch = copy.deepcopy(non_tensor_batch) non_tensor_batch = copy.deepcopy(non_tensor_batch)
if meta_info_keys is not None: if meta_info_keys is not None:
sub_meta_info = {key: val for key, val in self.meta_info.items() if key in meta_info_keys} sub_meta_info = {k: v for k, v in self.meta_info.items() if k in meta_info_keys}
else: else:
sub_meta_info = self.meta_info sub_meta_info = self.meta_info
...@@ -349,7 +364,12 @@ class DataProto: ...@@ -349,7 +364,12 @@ class DataProto:
return DataProto(batch=sub_batch, non_tensor_batch=non_tensor_batch, meta_info=sub_meta_info) return DataProto(batch=sub_batch, non_tensor_batch=non_tensor_batch, meta_info=sub_meta_info)
def pop(self, batch_keys=None, non_tensor_batch_keys=None, meta_info_keys=None) -> "DataProto": def pop(
self,
batch_keys: Optional[List[str]] = None,
non_tensor_batch_keys: Optional[List[str]] = None,
meta_info_keys: Optional[List[str]] = None,
) -> "DataProto":
"""Pop a subset of the DataProto via `batch_keys` and `meta_info_keys` """Pop a subset of the DataProto via `batch_keys` and `meta_info_keys`
Args: Args:
...@@ -377,7 +397,9 @@ class DataProto: ...@@ -377,7 +397,9 @@ class DataProto:
return DataProto.from_dict(tensors=tensors, non_tensors=non_tensors, meta_info=meta_info) return DataProto.from_dict(tensors=tensors, non_tensors=non_tensors, meta_info=meta_info)
def rename(self, old_keys=None, new_keys=None) -> "DataProto": def rename(
self, old_keys: Optional[Union[str, List[str]]] = None, new_keys: Optional[Union[str, List[str]]] = None
) -> "DataProto":
""" """
Note that this function only rename the key in the batch Note that this function only rename the key in the batch
""" """
...@@ -422,7 +444,9 @@ class DataProto: ...@@ -422,7 +444,9 @@ class DataProto:
self.meta_info = union_two_dict(self.meta_info, other.meta_info) self.meta_info = union_two_dict(self.meta_info, other.meta_info)
return self return self
def make_iterator(self, mini_batch_size, epochs, seed=None, dataloader_kwargs=None): def make_iterator(
self, mini_batch_size: int, epochs: int, seed: int = None, dataloader_kwargs: Dict[str, Any] = None
):
"""Make an iterator from the DataProto. This is built upon that TensorDict can be used as a normal Pytorch """Make an iterator from the DataProto. This is built upon that TensorDict can be used as a normal Pytorch
dataset. See https://pytorch.org/tensordict/tutorials/data_fashion for more details. dataset. See https://pytorch.org/tensordict/tutorials/data_fashion for more details.
...@@ -509,9 +533,7 @@ class DataProto: ...@@ -509,9 +533,7 @@ class DataProto:
Returns: Returns:
DataProto: concatenated DataProto DataProto: concatenated DataProto
""" """
batch_lst = [] batch_lst = [batch.batch for batch in data]
for batch in data:
batch_lst.append(batch.batch)
if batch_lst[0] is not None: if batch_lst[0] is not None:
new_batch = torch.cat(batch_lst, dim=0) new_batch = torch.cat(batch_lst, dim=0)
else: else:
...@@ -523,7 +545,7 @@ class DataProto: ...@@ -523,7 +545,7 @@ class DataProto:
return DataProto(batch=new_batch, non_tensor_batch=non_tensor_batch, meta_info=data[0].meta_info) return DataProto(batch=new_batch, non_tensor_batch=non_tensor_batch, meta_info=data[0].meta_info)
def reorder(self, indices): def reorder(self, indices: torch.Tensor) -> None:
""" """
Note that this operation is in-place Note that this operation is in-place
""" """
...@@ -531,7 +553,7 @@ class DataProto: ...@@ -531,7 +553,7 @@ class DataProto:
self.batch = self.batch[indices] self.batch = self.batch[indices]
self.non_tensor_batch = {key: val[indices_np] for key, val in self.non_tensor_batch.items()} self.non_tensor_batch = {key: val[indices_np] for key, val in self.non_tensor_batch.items()}
def repeat(self, repeat_times=2, interleave=True): def repeat(self, repeat_times: int = 2, interleave: bool = True) -> "DataProto":
""" """
Repeat the batch data a specified number of times. Repeat the batch data a specified number of times.
......
verl/trainer/config.py
View file @ ff7fb65e
...@@ -42,7 +42,7 @@ class DataConfig: ...@@ -42,7 +42,7 @@ class DataConfig:
max_response_length: int = 512 max_response_length: int = 512
rollout_batch_size: int = 512 rollout_batch_size: int = 512
val_batch_size: int = -1 val_batch_size: int = -1
system_prompt: Optional[str] = None format_prompt: Optional[str] = None
shuffle: bool = True shuffle: bool = True
seed: int = 1 seed: int = 1
max_pixels: int = 4194304 max_pixels: int = 4194304
......
verl/trainer/core_algos.py
View file @ ff7fb65e
...@@ -34,15 +34,22 @@ if TYPE_CHECKING: ...@@ -34,15 +34,22 @@ if TYPE_CHECKING:
class KLController(ABC): class KLController(ABC):
kl_coef: float
"""KL coefficient."""
@abstractmethod @abstractmethod
def update(self, current_kl: float, n_steps: int) -> None: ... def update(self, current_kl: float, n_steps: int) -> None:
"""Update kl_coef according to current KL."""
...
class AdaptiveKLController(KLController): class AdaptiveKLController(KLController):
"""Adaptive KL controller described in: https://arxiv.org/pdf/1909.08593.pdf""" """Adaptive KL controller described in: https://arxiv.org/pdf/1909.08593.pdf
Copied from https://github.com/huggingface/trl/blob/v0.11.0/trl/trainer/utils.py#L54"""
def __init__(self, init_kl_coef: float, target_kl: float, horizon: float): def __init__(self, init_kl_coef: float, target_kl: float, horizon: float):
self.value = init_kl_coef self.kl_coef = init_kl_coef
self.target = target_kl self.target = target_kl
self.horizon = horizon self.horizon = horizon
...@@ -50,14 +57,16 @@ class AdaptiveKLController(KLController): ...@@ -50,14 +57,16 @@ class AdaptiveKLController(KLController):
target = self.target target = self.target
proportional_error = np.clip(current_kl / target - 1, -0.2, 0.2) proportional_error = np.clip(current_kl / target - 1, -0.2, 0.2)
mult = 1 + proportional_error * n_steps / self.horizon mult = 1 + proportional_error * n_steps / self.horizon
self.value *= mult self.kl_coef *= mult
class FixedKLController(KLController): class FixedKLController(KLController):
"""Fixed KL controller.""" """Fixed KL controller.
Copeid from https://github.com/huggingface/trl/blob/v0.11.0/trl/trainer/utils.py#L72"""
def __init__(self, init_kl_coef: float): def __init__(self, init_kl_coef: float):
self.value = init_kl_coef self.kl_coef = init_kl_coef
def update(self, current_kl: float, n_steps: int) -> None: def update(self, current_kl: float, n_steps: int) -> None:
pass pass
...@@ -84,7 +93,7 @@ def get_kl_controller(algorithm_config: "AlgorithmConfig") -> KLController: ...@@ -84,7 +93,7 @@ def get_kl_controller(algorithm_config: "AlgorithmConfig") -> KLController:
def compute_gae_advantage_return( def compute_gae_advantage_return(
token_level_rewards: torch.Tensor, token_level_rewards: torch.Tensor,
values: torch.Tensor, values: torch.Tensor,
eos_mask: torch.Tensor, response_mask: torch.Tensor,
gamma: torch.Tensor, gamma: torch.Tensor,
lam: torch.Tensor, lam: torch.Tensor,
) -> Tuple[torch.Tensor, torch.Tensor]: ) -> Tuple[torch.Tensor, torch.Tensor]:
...@@ -95,8 +104,8 @@ def compute_gae_advantage_return( ...@@ -95,8 +104,8 @@ def compute_gae_advantage_return(
shape: (bs, response_length) shape: (bs, response_length)
values: `(torch.Tensor)` values: `(torch.Tensor)`
shape: (bs, response_length) shape: (bs, response_length)
eos_mask: `(torch.Tensor)` response_mask: `(torch.Tensor)`
shape: (bs, response_length). [EOS] mask. The token after [EOS] have mask zero. shape: (bs, response_length). The token after eos tokens have mask zero.
gamma: `(float)` gamma: `(float)`
discounted factor used in RL discounted factor used in RL
lam: `(float)` lam: `(float)`
...@@ -105,7 +114,7 @@ def compute_gae_advantage_return( ...@@ -105,7 +114,7 @@ def compute_gae_advantage_return(
Returns: Returns:
advantages: `(torch.Tensor)` advantages: `(torch.Tensor)`
shape: (bs, response_length) shape: (bs, response_length)
Returns: `(torch.Tensor)` returns: `(torch.Tensor)`
shape: (bs, response_length) shape: (bs, response_length)
""" """
...@@ -119,32 +128,33 @@ def compute_gae_advantage_return( ...@@ -119,32 +128,33 @@ def compute_gae_advantage_return(
advantages_reversed.append(lastgaelam) advantages_reversed.append(lastgaelam)
advantages = torch.stack(advantages_reversed[::-1], dim=1) advantages = torch.stack(advantages_reversed[::-1], dim=1)
returns = (advantages + values) * eos_mask returns = advantages + values
advantages = VF.masked_whiten(advantages, eos_mask) * eos_mask advantages = VF.masked_whiten(advantages, response_mask)
return advantages, returns return advantages, returns
# NOTE(sgm): this implementation only consider outcome supervision, where the reward is a scalar. # NOTE(sgm): this implementation only consider outcome supervision, where the reward is a scalar.
@torch.no_grad() @torch.no_grad()
def compute_grpo_outcome_advantage( def compute_grpo_outcome_advantage(
token_level_rewards: torch.Tensor, eos_mask: torch.Tensor, index: torch.Tensor, eps: float = 1e-6 token_level_rewards: torch.Tensor, response_mask: torch.Tensor, index: torch.Tensor, eps: float = 1e-6
) -> Tuple[torch.Tensor, torch.Tensor]: ) -> Tuple[torch.Tensor, torch.Tensor]:
""" """
Compute advantage for GRPO, operating only on Outcome reward Compute advantage for GRPO, operating only on Outcome reward
(with only one scalar reward for each response). (with only one scalar reward for each response).
Args: Args:
token_level_rewards: `(torch.Tensor)` token_level_rewards: `(torch.Tensor)`
shape: (bs, response_length) shape: (bs, response_length)
eos_mask: `(torch.Tensor)` response_mask: `(torch.Tensor)`
shape: (bs, response_length) shape: (bs, response_length)
Returns: Returns:
advantages: `(torch.Tensor)` advantages: `(torch.Tensor)`
shape: (bs, response_length) shape: (bs, response_length)
Returns: `(torch.Tensor)` returns: `(torch.Tensor)`
shape: (bs, response_length) shape: (bs, response_length)
""" """
response_length = token_level_rewards.shape[-1]
scores = token_level_rewards.sum(dim=-1) scores = token_level_rewards.sum(dim=-1)
id2score = defaultdict(list) id2score = defaultdict(list)
id2mean, id2std = {}, {} id2mean, id2std = {}, {}
...@@ -154,86 +164,78 @@ def compute_grpo_outcome_advantage( ...@@ -154,86 +164,78 @@ def compute_grpo_outcome_advantage(
id2score[index[i]].append(scores[i]) id2score[index[i]].append(scores[i])
for idx in id2score: for idx in id2score:
if len(id2score[idx]) == 1: assert len(id2score[idx]) > 1, "GRPO needs rollout.n > 1."
id2mean[idx] = torch.tensor(0.0) id2mean[idx] = torch.mean(torch.tensor(id2score[idx]))
id2std[idx] = torch.tensor(1.0) id2std[idx] = torch.std(torch.tensor(id2score[idx]))
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): for i in range(bsz):
scores[i] = (scores[i] - id2mean[index[i]]) / (id2std[index[i]] + eps) scores[i] = (scores[i] - id2mean[index[i]]) / (id2std[index[i]] + eps)
scores = scores.unsqueeze(-1).tile([1, response_length]) * eos_mask returns = scores.unsqueeze(-1) * response_mask
return scores, scores return returns, returns
@torch.no_grad() @torch.no_grad()
def compute_rloo_outcome_advantage( def compute_rloo_outcome_advantage(
token_level_rewards: torch.Tensor, eos_mask: torch.Tensor, index: torch.Tensor token_level_rewards: torch.Tensor, response_mask: torch.Tensor, index: torch.Tensor
) -> Tuple[torch.Tensor, torch.Tensor]: ) -> Tuple[torch.Tensor, torch.Tensor]:
""" """
Compute advantage for RLOO based on https://arxiv.org/abs/2402.14740 Compute advantage for RLOO based on https://arxiv.org/abs/2402.14740
Args: Args:
token_level_rewards: `(torch.Tensor)` token_level_rewards: `(torch.Tensor)`
shape: (bs, response_length) shape: (bs, response_length)
eos_mask: `(torch.Tensor)` response_mask: `(torch.Tensor)`
shape: (bs, response_length) shape: (bs, response_length)
Returns: Returns:
advantages: `(torch.Tensor)` advantages: `(torch.Tensor)`
shape: (bs, response_length) shape: (bs, response_length)
Returns: `(torch.Tensor)` returns: `(torch.Tensor)`
shape: (bs, response_length) shape: (bs, response_length)
""" """
response_length = token_level_rewards.shape[-1]
scores = token_level_rewards.sum(dim=-1) scores = token_level_rewards.sum(dim=-1)
id2score = defaultdict(list) id2score = defaultdict(list)
id2mean = {} id2sum = {}
bsz = scores.shape[0] bsz = scores.shape[0]
for i in range(bsz): for i in range(bsz):
id2score[index[i]].append(scores[i]) id2score[index[i]].append(scores[i])
for idx in id2score: for idx in id2score:
if len(id2score[idx]) == 1: id2sum[idx] = torch.sum(torch.tensor(id2score[idx]))
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): for i in range(bsz):
response_num = len(id2score[index[i]]) sample_num = len(id2score[index[i]])
if response_num > 1: assert sample_num > 1, "RLOO needs rollout.n > 1."
scores[i] = scores[i] * response_num / (response_num - 1) - id2mean[index[i]] * response_num / ( baseline = (id2sum[index[i]] - scores[i]) / (sample_num - 1)
response_num - 1 scores[i] = scores[i] - baseline
)
scores = scores.unsqueeze(-1).tile([1, response_length]) * eos_mask returns = scores.unsqueeze(-1) * response_mask
return scores, scores return returns, returns
@torch.no_grad() @torch.no_grad()
def compute_reinforce_plus_plus_outcome_advantage( def compute_reinforce_plus_plus_outcome_advantage(
token_level_rewards: torch.Tensor, eos_mask: torch.Tensor, gamma: torch.Tensor token_level_rewards: torch.Tensor, response_mask: torch.Tensor, gamma: torch.Tensor
) -> Tuple[torch.Tensor, torch.Tensor]: ) -> Tuple[torch.Tensor, torch.Tensor]:
""" """
Compute advantage for REINFORCE++. Compute advantage for REINFORCE++.
This implementation is based on the paper: https://arxiv.org/abs/2501.03262 This implementation is based on the paper: https://arxiv.org/abs/2501.03262
Args: Args:
token_level_rewards: `(torch.Tensor)` token_level_rewards: `(torch.Tensor)`
shape: (bs, response_length) shape: (bs, response_length)
eos_mask: `(torch.Tensor)` response_mask: `(torch.Tensor)`
shape: (bs, response_length) shape: (bs, response_length)
Returns: Returns:
advantages: `(torch.Tensor)` advantages: `(torch.Tensor)`
shape: (bs, response_length) shape: (bs, response_length)
Returns: `(torch.Tensor)` returns: `(torch.Tensor)`
shape: (bs, response_length) shape: (bs, response_length)
""" """
returns = torch.zeros_like(token_level_rewards) returns = torch.zeros_like(token_level_rewards)
running_return = 0 running_return = 0
...@@ -241,17 +243,15 @@ def compute_reinforce_plus_plus_outcome_advantage( ...@@ -241,17 +243,15 @@ def compute_reinforce_plus_plus_outcome_advantage(
running_return = token_level_rewards[:, t] + gamma * running_return running_return = token_level_rewards[:, t] + gamma * running_return
returns[:, t] = running_return returns[:, t] = running_return
# Reset after EOS # Reset after EOS
running_return = running_return * eos_mask[:, t] running_return = running_return * response_mask[:, t]
advantages = VF.masked_whiten(returns, eos_mask) advantages = VF.masked_whiten(returns, response_mask)
advantages *= eos_mask
returns *= eos_mask
return advantages, returns return advantages, returns
@torch.no_grad() @torch.no_grad()
def compute_remax_outcome_advantage( def compute_remax_outcome_advantage(
token_level_rewards: torch.Tensor, reward_baselines: torch.Tensor, eos_mask: torch.Tensor token_level_rewards: torch.Tensor, reward_baselines: torch.Tensor, response_mask: torch.Tensor
) -> Tuple[torch.Tensor, torch.Tensor]: ) -> Tuple[torch.Tensor, torch.Tensor]:
""" """
Compute advantage for ReMax, operating only on Outcome reward Compute advantage for ReMax, operating only on Outcome reward
...@@ -263,20 +263,19 @@ def compute_remax_outcome_advantage( ...@@ -263,20 +263,19 @@ def compute_remax_outcome_advantage(
shape: (bs, response_length) shape: (bs, response_length)
reward_baselines: `(torch.Tensor)` reward_baselines: `(torch.Tensor)`
shape: (bs,) shape: (bs,)
eos_mask: `(torch.Tensor)` response_mask: `(torch.Tensor)`
shape: (bs, response_length) shape: (bs, response_length)
Returns: Returns:
advantages: `(torch.Tensor)` advantages: `(torch.Tensor)`
shape: (bs, response_length) shape: (bs, response_length)
Returns: `(torch.Tensor)` returns: `(torch.Tensor)`
shape: (bs, response_length) shape: (bs, response_length)
""" """
response_length = token_level_rewards.shape[-1] scores = token_level_rewards.sum(dim=-1) - reward_baselines
# scores = token_level_rewards.sum(dim=-1) returns = scores.unsqueeze(-1) * response_mask
returns = (token_level_rewards * eos_mask).flip(dims=[-1]).cumsum(dim=-1).flip(dims=[-1]) * eos_mask return returns, returns
advantages = returns - reward_baselines.unsqueeze(-1).tile([1, response_length]) * eos_mask
return advantages, returns
def compute_rewards( def compute_rewards(
...@@ -293,9 +292,11 @@ def compute_policy_loss( ...@@ -293,9 +292,11 @@ def compute_policy_loss(
old_log_probs: torch.Tensor, old_log_probs: torch.Tensor,
log_probs: torch.Tensor, log_probs: torch.Tensor,
advantages: torch.Tensor, advantages: torch.Tensor,
eos_mask: torch.Tensor, response_mask: torch.Tensor,
cliprange: float, clip_ratio_low: float,
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: clip_ratio_high: float,
clip_ratio_dual: float,
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
"""Compute the policy loss. """Compute the policy loss.
Adapted from https://github.com/huggingface/trl/blob/v0.15.0/trl/trainer/ppo_trainer.py#L568 Adapted from https://github.com/huggingface/trl/blob/v0.15.0/trl/trainer/ppo_trainer.py#L568
...@@ -307,42 +308,61 @@ def compute_policy_loss( ...@@ -307,42 +308,61 @@ def compute_policy_loss(
shape: (bs, response_length) shape: (bs, response_length)
advantages: `(torch.Tensor)` advantages: `(torch.Tensor)`
shape: (bs, response_length) shape: (bs, response_length)
eos_mask: `(torch.Tensor)` response_mask: `(torch.Tensor)`
shape: (bs, response_length) shape: (bs, response_length)
cliprange: (float) clip_ratio_low: (float)
The clip range used in PPO. See https://arxiv.org/abs/1707.06347 The lower clip range used in PPO. See https://arxiv.org/abs/1707.06347
clip_ratio_high: (float)
The higher clip range used in DAPO. See https://arxiv.org/pdf/2503.14476
clip_ratio_dual: (float)
The dual clip range used in Dual-clip PPO. See https://arxiv.org/pdf/1912.09729
Returns: Returns:
pg_loss: `a scalar torch.Tensor` pg_loss: `a scalar torch.Tensor`
policy gradient loss computed via PPO policy gradient loss computed via PPO
pg_clipfrac: (float) pg_clipfrac_higher: (float)
a float number indicating the fraction of policy gradient loss being clipped a float number indicating the fraction of policy gradient loss being clipped to a higher value
pg_clipfrac_lower: (float)
a float number indicating the fraction of policy gradient loss being clipped to a lower value
ppo_kl: (float)
a float number indicating the mean KL divergence between the old policy and the new policy
""" """
negative_approx_kl = log_probs - old_log_probs negative_approx_kl = log_probs - old_log_probs
# clamp the ratio before exp to avoid nan # clamp the ratio before exp to avoid nan
# see: https://github.com/pytorch/pytorch/issues/10729 # see: https://github.com/pytorch/pytorch/issues/10729
ratio = torch.exp(negative_approx_kl) ratio = torch.exp(negative_approx_kl)
clipped_ratio = torch.exp(torch.clamp(negative_approx_kl, np.log(1.0 - cliprange), np.log(1.0 + cliprange))) clipped_ratio = torch.exp(
ppo_kl = VF.masked_mean(-negative_approx_kl, eos_mask) torch.clamp(negative_approx_kl, np.log(1.0 - clip_ratio_low), np.log(1.0 + clip_ratio_high))
)
pg_losses = -advantages * ratio pg_loss = -advantages * ratio
pg_losses2 = -advantages * clipped_ratio pg_loss2 = -advantages * clipped_ratio
pg_loss3 = -advantages * clip_ratio_dual
pg_loss = VF.masked_mean(torch.max(pg_losses, pg_losses2), eos_mask) clipped_pg_loss_higher = torch.max(pg_loss, pg_loss2) # clip if pg_loss < pg_loss2
pg_clipfrac = VF.masked_mean(torch.gt(pg_losses2, pg_losses).float(), eos_mask) pg_clipfrac_higher = (pg_loss < pg_loss2).float()
return pg_loss, pg_clipfrac, ppo_kl clipped_pg_loss_lower = torch.min(clipped_pg_loss_higher, pg_loss3) # clip if pg_loss > pg_loss3 and adv < 0
final_pg_loss = torch.where(advantages < 0, clipped_pg_loss_lower, clipped_pg_loss_higher)
pg_clipfrac_lower = (clipped_pg_loss_higher > pg_loss3).float() * (advantages < 0).float()
final_pg_loss = VF.masked_mean(final_pg_loss, response_mask)
pg_clipfrac_higher = VF.masked_mean(pg_clipfrac_higher, response_mask)
pg_clipfrac_lower = VF.masked_mean(pg_clipfrac_lower, response_mask)
ppo_kl = VF.masked_mean(-negative_approx_kl, response_mask)
return final_pg_loss, pg_clipfrac_higher, pg_clipfrac_lower, ppo_kl
def compute_value_loss( def compute_value_loss(
vpreds: torch.Tensor, vpreds: torch.Tensor,
returns: torch.Tensor, returns: torch.Tensor,
values: torch.Tensor, values: torch.Tensor,
eos_mask: torch.Tensor, action_mask: torch.Tensor,
cliprange_value: float, cliprange_value: float,
) -> Tuple[torch.Tensor, float]: ) -> Tuple[torch.Tensor, float]:
"""Compute the value loss. """Compute the value loss.
Copied from https://github.com/huggingface/trl/blob/v0.15.0/trl/trainer/ppo_trainer.py#L556 Adapted from https://github.com/huggingface/trl/blob/v0.15.0/trl/trainer/ppo_trainer.py#L556
Args: Args:
vpreds (`torch.FloatTensor`): vpreds (`torch.FloatTensor`):
...@@ -351,7 +371,7 @@ def compute_value_loss( ...@@ -351,7 +371,7 @@ def compute_value_loss(
Ground truth returns, shape (`batch_size`, `response_length`) Ground truth returns, shape (`batch_size`, `response_length`)
values (`torch.FloatTensor`): values (`torch.FloatTensor`):
Old values of value head, shape (`batch_size`, `response_length`) Old values of value head, shape (`batch_size`, `response_length`)
eos_mask: `(torch.Tensor)` action_mask: `(torch.Tensor)`
shape: (bs, response_length) shape: (bs, response_length)
cliprange_value: (float) cliprange_value: (float)
The clip range for value net used in PPO. See https://arxiv.org/abs/1707.06347 The clip range for value net used in PPO. See https://arxiv.org/abs/1707.06347
...@@ -361,25 +381,29 @@ def compute_value_loss( ...@@ -361,25 +381,29 @@ def compute_value_loss(
value function loss value function loss
vf_clipfrac: a float vf_clipfrac: a float
The ratio of vf being clipped The ratio of vf being clipped
""" """
vpredclipped = torch.clamp(vpreds, values - cliprange_value, values + cliprange_value) vpredclipped = torch.clamp(vpreds, values - cliprange_value, values + cliprange_value)
vf_losses1 = torch.square(vpreds - returns) vf_loss1 = torch.square(vpreds - returns)
vf_losses2 = torch.square(vpredclipped - returns) vf_loss2 = torch.square(vpredclipped - returns)
vf_loss = 0.5 * VF.masked_mean(torch.max(vf_losses1, vf_losses2), eos_mask) vf_loss = 0.5 * VF.masked_mean(torch.max(vf_loss1, vf_loss2), action_mask) # clip if vf_loss1 < vf_loss2
vf_clipfrac = VF.masked_mean(torch.gt(vf_losses2, vf_losses1).float(), eos_mask) vf_clipfrac = VF.masked_mean((vf_loss1 < vf_loss2).float(), action_mask)
return vf_loss, vf_clipfrac return vf_loss, vf_clipfrac
def kl_penalty(log_probs: torch.FloatTensor, ref_log_probs: torch.FloatTensor, kl_penalty: str) -> torch.Tensor: def compute_kl(log_probs: torch.FloatTensor, ref_log_probs: torch.FloatTensor, kl_penalty: str) -> torch.Tensor:
"""Compute KL divergence given log_probs and ref_log_probs. """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
Adapted from https://github.com/huggingface/trl/blob/v0.11.0/trl/trainer/ppo_trainer.py#L1150
Args: Args:
log_probs: torch.Tensor log_probs: torch.Tensor
ref_log_probs: torch.Tensor ref_log_probs: torch.Tensor
kl_penalty: str
Returns: Returns:
kl_div: torch.Tensor kl_div: torch.Tensor
""" """
log_probs, ref_log_probs = log_probs.float(), ref_log_probs.float() log_probs, ref_log_probs = log_probs.float(), ref_log_probs.float()
if kl_penalty == "kl": if kl_penalty == "kl":
......
verl/trainer/main.py
View file @ ff7fb65e
...@@ -69,8 +69,8 @@ class Runner: ...@@ -69,8 +69,8 @@ class Runner:
} }
resource_pool_manager = ResourcePoolManager(resource_pool_spec=resource_pool_spec, mapping=mapping) resource_pool_manager = ResourcePoolManager(resource_pool_spec=resource_pool_spec, mapping=mapping)
reward_fn = CustomRewardManager(tokenizer=tokenizer, compute_score=config.worker.reward.compute_score) reward_fn = CustomRewardManager(tokenizer=tokenizer, config=config.worker.reward)
val_reward_fn = CustomRewardManager(tokenizer=tokenizer, compute_score=config.worker.reward.compute_score) val_reward_fn = CustomRewardManager(tokenizer=tokenizer, config=config.worker.reward)
trainer = RayPPOTrainer( trainer = RayPPOTrainer(
config=config, config=config,
...@@ -98,9 +98,8 @@ def main(): ...@@ -98,9 +98,8 @@ def main():
ppo_config = OmegaConf.merge(default_config, cli_args) ppo_config = OmegaConf.merge(default_config, cli_args)
ppo_config = OmegaConf.to_object(ppo_config) ppo_config = OmegaConf.to_object(ppo_config)
# this is for local ray cluster
if not ray.is_initialized(): if not ray.is_initialized():
# for rocm # this is for local ray cluster
if torch.version.hip is not None: if torch.version.hip is not None:
ray.init(num_gpus=torch.cuda.device_count(), ray.init(num_gpus=torch.cuda.device_count(),
ignore_reinit_error=True, ignore_reinit_error=True,
...@@ -108,6 +107,7 @@ def main(): ...@@ -108,6 +107,7 @@ def main():
else: else:
ray.init(runtime_env={"env_vars": {"TOKENIZERS_PARALLELISM": "true", "NCCL_DEBUG": "WARN"}}) ray.init(runtime_env={"env_vars": {"TOKENIZERS_PARALLELISM": "true", "NCCL_DEBUG": "WARN"}})
runner = Runner.remote() runner = Runner.remote()
ray.get(runner.run.remote(ppo_config)) ray.get(runner.run.remote(ppo_config))
......
verl/trainer/ray_trainer.py
View file @ ff7fb65e
...@@ -23,7 +23,7 @@ from contextlib import contextmanager ...@@ -23,7 +23,7 @@ from contextlib import contextmanager
from copy import deepcopy from copy import deepcopy
from dataclasses import dataclass, field from dataclasses import dataclass, field
from enum import Enum, IntEnum, auto from enum import Enum, IntEnum, auto
from typing import Any, Callable, Dict, List, Optional, Type from typing import Any, Callable, Dict, List, Optional, Tuple, Type
import numpy as np import numpy as np
import ray import ray
...@@ -50,9 +50,6 @@ from .config import PPOConfig ...@@ -50,9 +50,6 @@ from .config import PPOConfig
from .metrics import compute_data_metrics, compute_throughout_metrics, compute_timing_metrics, reduce_metrics from .metrics import compute_data_metrics, compute_throughout_metrics, compute_timing_metrics, reduce_metrics
WorkerType = Type[Worker]
class Role(IntEnum): class Role(IntEnum):
""" """
To create more roles dynamically, you can subclass Role and add new members To create more roles dynamically, you can subclass Role and add new members
...@@ -132,25 +129,19 @@ def apply_kl_penalty(data: DataProto, kl_ctrl: core_algos.KLController, kl_penal ...@@ -132,25 +129,19 @@ def apply_kl_penalty(data: DataProto, kl_ctrl: core_algos.KLController, kl_penal
# compute kl between ref_policy and current policy # compute kl between ref_policy and current policy
if "ref_log_probs" in data.batch.keys(): if "ref_log_probs" in data.batch.keys():
kld = core_algos.kl_penalty( kld = core_algos.compute_kl(data.batch["old_log_probs"], data.batch["ref_log_probs"], kl_penalty=kl_penalty)
data.batch["old_log_probs"], data.batch["ref_log_probs"], kl_penalty=kl_penalty kld = kld * response_mask # (batch_size, response_length)
) # (batch_size, response_length)
kld = kld * response_mask
beta = kl_ctrl.value
else: else:
beta = 0
kld = torch.zeros_like(response_mask, dtype=torch.float32) kld = torch.zeros_like(response_mask, dtype=torch.float32)
token_level_rewards = token_level_scores - beta * kld data.batch["token_level_rewards"] = token_level_scores - kl_ctrl.kl_coef * kld
current_kl = VF.masked_mean(kld, mask=response_mask, dim=-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() current_kl = torch.mean(current_kl, dim=0).item()
metrics = {"critic/kl": current_kl, "critic/kl_coef": kl_ctrl.kl_coef}
# According to https://github.com/huggingface/trl/blob/v0.11.0/trl/trainer/ppo_trainer.py#L880 # 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) 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_coef": beta}
return data, metrics return data, metrics
...@@ -161,25 +152,21 @@ def compute_advantage(data: DataProto, adv_estimator: AdvantageEstimator, gamma: ...@@ -161,25 +152,21 @@ def compute_advantage(data: DataProto, adv_estimator: AdvantageEstimator, gamma:
if adv_estimator == AdvantageEstimator.GAE: if adv_estimator == AdvantageEstimator.GAE:
values = data.batch["values"] values = data.batch["values"]
advantages, returns = core_algos.compute_gae_advantage_return( advantages, returns = core_algos.compute_gae_advantage_return(
token_level_rewards=token_level_rewards, values=values, eos_mask=response_mask, gamma=gamma, lam=lam token_level_rewards, values, response_mask, gamma, lam
) )
elif adv_estimator == AdvantageEstimator.GRPO: elif adv_estimator == AdvantageEstimator.GRPO:
advantages, returns = core_algos.compute_grpo_outcome_advantage( advantages, returns = core_algos.compute_grpo_outcome_advantage(token_level_rewards, response_mask, index)
token_level_rewards=token_level_rewards, eos_mask=response_mask, index=index
)
elif adv_estimator == AdvantageEstimator.REINFORCE_PLUS_PLUS: elif adv_estimator == AdvantageEstimator.REINFORCE_PLUS_PLUS:
advantages, returns = core_algos.compute_reinforce_plus_plus_outcome_advantage( advantages, returns = core_algos.compute_reinforce_plus_plus_outcome_advantage(
token_level_rewards=token_level_rewards, eos_mask=response_mask, gamma=gamma token_level_rewards, response_mask, gamma
) )
elif adv_estimator == AdvantageEstimator.REMAX: elif adv_estimator == AdvantageEstimator.REMAX:
reward_baselines = data.batch["reward_baselines"] reward_baselines = data.batch["reward_baselines"]
advantages, returns = core_algos.compute_remax_outcome_advantage( advantages, returns = core_algos.compute_remax_outcome_advantage(
token_level_rewards=token_level_rewards, reward_baselines=reward_baselines, eos_mask=response_mask token_level_rewards, reward_baselines, response_mask
) )
elif adv_estimator == AdvantageEstimator.RLOO: elif adv_estimator == AdvantageEstimator.RLOO:
advantages, returns = core_algos.compute_rloo_outcome_advantage( advantages, returns = core_algos.compute_rloo_outcome_advantage(token_level_rewards, response_mask, index)
token_level_rewards=token_level_rewards, eos_mask=response_mask, index=index
)
else: else:
raise NotImplementedError raise NotImplementedError
...@@ -206,11 +193,11 @@ class RayPPOTrainer: ...@@ -206,11 +193,11 @@ class RayPPOTrainer:
config: PPOConfig, config: PPOConfig,
tokenizer: PreTrainedTokenizer, tokenizer: PreTrainedTokenizer,
processor: Optional[ProcessorMixin], processor: Optional[ProcessorMixin],
role_worker_mapping: dict[Role, WorkerType], role_worker_mapping: dict[Role, Type[Worker]],
resource_pool_manager: ResourcePoolManager, resource_pool_manager: ResourcePoolManager,
ray_worker_group_cls: RayWorkerGroup = RayWorkerGroup, ray_worker_group_cls: Type[RayWorkerGroup] = RayWorkerGroup,
reward_fn: Callable = None, reward_fn: Optional[Callable[[DataProto], Tuple[torch.Tensor, Dict[str, List[float]]]]] = None,
val_reward_fn: Callable = None, val_reward_fn: Optional[Callable[[DataProto], Tuple[torch.Tensor, Dict[str, List[float]]]]] = None,
): ):
self.tokenizer = tokenizer self.tokenizer = tokenizer
self.processor = processor self.processor = processor
...@@ -249,10 +236,31 @@ class RayPPOTrainer: ...@@ -249,10 +236,31 @@ class RayPPOTrainer:
raise NotImplementedError(f"Unknown advantage estimator: {config.algorithm.adv_estimator}.") raise NotImplementedError(f"Unknown advantage estimator: {config.algorithm.adv_estimator}.")
if config.data.rollout_batch_size % config.worker.actor.global_batch_size != 0: 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.") raise ValueError("Rollout batch size must be divisible by actor global batch size.")
if (
config.data.rollout_batch_size * config.worker.rollout.n
) % config.worker.actor.micro_batch_size_per_device_for_experience != 0:
raise ValueError(
"Rollout batch size * rollout.n must be divisible by actor micro batch size for experience."
)
if self.use_critic and config.data.rollout_batch_size % config.worker.critic.global_batch_size != 0: if self.use_critic:
raise ValueError("Rollout batch size must be divisible by global batch size.") if config.data.rollout_batch_size % config.worker.critic.global_batch_size != 0:
raise ValueError("Rollout batch size must be divisible by critic global batch size.")
if (
config.data.rollout_batch_size * config.worker.rollout.n
) % config.worker.critic.micro_batch_size_per_device_for_experience != 0:
raise ValueError(
"Rollout batch size * rollout.n must be divisible by critic micro batch size for experience."
)
if (
config.algorithm.adv_estimator in (AdvantageEstimator.GRPO, AdvantageEstimator.RLOO)
and config.worker.rollout.n == 1
):
raise ValueError("GRPO and RLOO algorithm need `config.worker.rollout.n > 1`.")
self._create_dataloader() self._create_dataloader()
...@@ -266,7 +274,7 @@ class RayPPOTrainer: ...@@ -266,7 +274,7 @@ class RayPPOTrainer:
image_key=self.config.data.image_key, image_key=self.config.data.image_key,
max_prompt_length=self.config.data.max_prompt_length, max_prompt_length=self.config.data.max_prompt_length,
truncation="right", truncation="right",
system_prompt=self.config.data.system_prompt, format_prompt=self.config.data.format_prompt,
min_pixels=self.config.data.min_pixels, min_pixels=self.config.data.min_pixels,
max_pixels=self.config.data.max_pixels, max_pixels=self.config.data.max_pixels,
) )
...@@ -297,7 +305,7 @@ class RayPPOTrainer: ...@@ -297,7 +305,7 @@ class RayPPOTrainer:
image_key=self.config.data.image_key, image_key=self.config.data.image_key,
max_prompt_length=self.config.data.max_prompt_length, max_prompt_length=self.config.data.max_prompt_length,
truncation="right", truncation="right",
system_prompt=self.config.data.system_prompt, format_prompt=self.config.data.format_prompt,
min_pixels=self.config.data.min_pixels, min_pixels=self.config.data.min_pixels,
max_pixels=self.config.data.max_pixels, max_pixels=self.config.data.max_pixels,
) )
...@@ -328,13 +336,15 @@ class RayPPOTrainer: ...@@ -328,13 +336,15 @@ class RayPPOTrainer:
self.config.worker.critic.optim.training_steps = training_steps self.config.worker.critic.optim.training_steps = training_steps
print(f"Total training steps: {self.training_steps}") print(f"Total training steps: {self.training_steps}")
def _maybe_log_val_generations(self, inputs: List[str], outputs: List[str], scores: List[float]) -> None: def _maybe_log_val_generations(
self, inputs: List[str], outputs: List[str], labels: List[str], scores: List[float]
) -> None:
"""Log a table of validation samples""" """Log a table of validation samples"""
if self.config.trainer.val_generations_to_log <= 0: if self.config.trainer.val_generations_to_log <= 0:
return return
# Create tuples of (input, output, score) and sort by input text # Create tuples of (input, output, score) and sort by input text
samples = list(zip(inputs, outputs, scores)) samples = list(zip(inputs, outputs, labels, scores))
samples.sort(key=lambda x: x[0]) # Sort by input text samples.sort(key=lambda x: x[0]) # Sort by input text
# Use fixed random seed for deterministic shuffling # Use fixed random seed for deterministic shuffling
...@@ -347,10 +357,10 @@ class RayPPOTrainer: ...@@ -347,10 +357,10 @@ class RayPPOTrainer:
def _validate(self) -> Dict[str, Any]: def _validate(self) -> Dict[str, Any]:
reward_tensor_lst = [] reward_tensor_lst = []
# Lists to collect samples for the table # Lists to collect samples for the table
sample_inputs, sample_outputs, sample_scores = [], [], [] sample_inputs, sample_outputs, sample_labels, sample_scores = [], [], [], []
reward_metrics_lst = defaultdict(list) reward_metrics_lst = defaultdict(list)
for test_data in self.val_dataloader: for batch_dict in self.val_dataloader:
test_batch = DataProto.from_single_dict(test_data) test_batch = DataProto.from_single_dict(batch_dict)
# Store original inputs # Store original inputs
input_ids = test_batch.batch["input_ids"] input_ids = test_batch.batch["input_ids"]
input_texts = [self.tokenizer.decode(ids, skip_special_tokens=True) for ids in input_ids] input_texts = [self.tokenizer.decode(ids, skip_special_tokens=True) for ids in input_ids]
...@@ -377,7 +387,7 @@ class RayPPOTrainer: ...@@ -377,7 +387,7 @@ class RayPPOTrainer:
output_ids = test_output_gen_batch.batch["responses"] output_ids = test_output_gen_batch.batch["responses"]
output_texts = [self.tokenizer.decode(ids, skip_special_tokens=True) for ids in output_ids] output_texts = [self.tokenizer.decode(ids, skip_special_tokens=True) for ids in output_ids]
sample_outputs.extend(output_texts) sample_outputs.extend(output_texts)
sample_labels.extend(test_batch.non_tensor_batch["ground_truth"].tolist())
test_batch = test_batch.union(test_output_gen_batch) test_batch = test_batch.union(test_output_gen_batch)
# evaluate using reward_function # evaluate using reward_function
...@@ -391,7 +401,7 @@ class RayPPOTrainer: ...@@ -391,7 +401,7 @@ class RayPPOTrainer:
for key, value in reward_metrics.items(): for key, value in reward_metrics.items():
reward_metrics_lst[key].extend(value) reward_metrics_lst[key].extend(value)
self._maybe_log_val_generations(inputs=sample_inputs, outputs=sample_outputs, scores=sample_scores) self._maybe_log_val_generations(sample_inputs, sample_outputs, sample_labels, sample_scores)
reward_score = torch.cat(reward_tensor_lst, dim=0).sum(-1).mean().item() 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()} 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} return {"val/reward_score": reward_score, **val_reward_metrics}
...@@ -576,7 +586,8 @@ class RayPPOTrainer: ...@@ -576,7 +586,8 @@ class RayPPOTrainer:
if self.config.algorithm.adv_estimator == "remax": if self.config.algorithm.adv_estimator == "remax":
with _timer("gen_max", timing_raw): with _timer("gen_max", timing_raw):
gen_baseline_batch = deepcopy(gen_batch) gen_baseline_batch = deepcopy(gen_batch)
gen_baseline_batch.meta_info["temperature"] = 0.0 gen_baseline_batch.meta_info["temperature"] = 0
gen_baseline_batch.meta_info["n"] = 1
gen_baseline_output = self.actor_rollout_wg.generate_sequences(gen_baseline_batch) gen_baseline_output = self.actor_rollout_wg.generate_sequences(gen_baseline_batch)
batch = batch.union(gen_baseline_output) batch = batch.union(gen_baseline_output)
...@@ -634,7 +645,8 @@ class RayPPOTrainer: ...@@ -634,7 +645,8 @@ class RayPPOTrainer:
with _timer("adv", timing_raw): with _timer("adv", timing_raw):
# apply kl penalty if available # apply kl penalty if available
if not self.config.algorithm.use_kl_loss and self.use_reference_policy: # apply kl penalty to reward 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_metrics = apply_kl_penalty(
batch, kl_ctrl=self.kl_ctrl, kl_penalty=self.config.algorithm.kl_penalty batch, kl_ctrl=self.kl_ctrl, kl_penalty=self.config.algorithm.kl_penalty
) )
......
verl/utils/dataset.py
View file @ ff7fb65e
...@@ -90,7 +90,7 @@ class RLHFDataset(Dataset, ImageProcessMixin): ...@@ -90,7 +90,7 @@ class RLHFDataset(Dataset, ImageProcessMixin):
image_key: str = "images", image_key: str = "images",
max_prompt_length: int = 1024, max_prompt_length: int = 1024,
truncation: str = "error", truncation: str = "error",
system_prompt: str = None, format_prompt: str = None,
max_pixels: int = None, max_pixels: int = None,
min_pixels: int = None, min_pixels: int = None,
): ):
...@@ -101,7 +101,7 @@ class RLHFDataset(Dataset, ImageProcessMixin): ...@@ -101,7 +101,7 @@ class RLHFDataset(Dataset, ImageProcessMixin):
self.image_key = image_key self.image_key = image_key
self.max_prompt_length = max_prompt_length self.max_prompt_length = max_prompt_length
self.truncation = truncation self.truncation = truncation
self.system_prompt = system_prompt self.format_prompt = format_prompt
self.max_pixels = max_pixels self.max_pixels = max_pixels
self.min_pixels = min_pixels self.min_pixels = min_pixels
...@@ -123,8 +123,8 @@ class RLHFDataset(Dataset, ImageProcessMixin): ...@@ -123,8 +123,8 @@ class RLHFDataset(Dataset, ImageProcessMixin):
def __getitem__(self, index): def __getitem__(self, index):
row_dict: dict = self.dataset[index] row_dict: dict = self.dataset[index]
prompt_str: str = row_dict[self.prompt_key] prompt_str: str = row_dict[self.prompt_key]
if self.system_prompt: if self.format_prompt:
prompt_str = " ".join((self.system_prompt.strip(), prompt_str)) prompt_str = prompt_str + " " + self.format_prompt.strip()
if self.image_key in row_dict: if self.image_key in row_dict:
# https://huggingface.co/docs/transformers/en/tasks/image_text_to_text # https://huggingface.co/docs/transformers/en/tasks/image_text_to_text
......
verl/utils/logger/aggregate_logger.py deleted 100644 → 0
View file @ c132cbcb
# 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.
"""
A Ray logger will receive logging info from different processes.
"""
import numbers
from typing import Dict
def concat_dict_to_str(dict: Dict, step):
output = [f"step {step}:"]
for k, v in dict.items():
if isinstance(v, numbers.Number):
output.append(f"{k}:{v:.3f}")
output_str = " - ".join(output)
return output_str
class LocalLogger:
def __init__(self, remote_logger=None, enable_wandb=False, print_to_console=False):
self.print_to_console = print_to_console
if print_to_console:
print("Using LocalLogger is deprecated. The constructor API will change.")
def flush(self):
pass
def log(self, data, step):
if self.print_to_console:
print(concat_dict_to_str(data, step=step), flush=True)
verl/utils/logger/gen_logger.py
View file @ ff7fb65e
...@@ -31,22 +31,23 @@ if is_package_available("swanlab"): ...@@ -31,22 +31,23 @@ if is_package_available("swanlab"):
@dataclass @dataclass
class GenerationLogger(ABC): class GenerationLogger(ABC):
@abstractmethod @abstractmethod
def log(self, samples: List[Tuple[str, str, float]], step: int) -> None: ... def log(self, samples: List[Tuple[str, str, str, float]], step: int) -> None: ...
@dataclass @dataclass
class ConsoleGenerationLogger(GenerationLogger): class ConsoleGenerationLogger(GenerationLogger):
def log(self, samples: List[Tuple[str, str, float]], step: int) -> None: def log(self, samples: List[Tuple[str, str, str, float]], step: int) -> None:
for inp, out, score in samples: for inp, out, lab, score in samples:
print(f"[prompt] {inp}\n[output] {out}\n[score] {score}\n") print(f"[prompt] {inp}\n[output] {out}\n[ground_truth] {lab}\n[score] {score}\n")
@dataclass @dataclass
class WandbGenerationLogger(GenerationLogger): class WandbGenerationLogger(GenerationLogger):
def log(self, samples: List[Tuple[str, str, float]], step: int) -> None: def log(self, samples: List[Tuple[str, str, str, float]], step: int) -> None:
# Create column names for all samples # Create column names for all samples
columns = ["step"] + sum( columns = ["step"] + sum(
[[f"input_{i + 1}", f"output_{i + 1}", f"score_{i + 1}"] for i in range(len(samples))], [] [[f"input_{i + 1}", f"output_{i + 1}", f"label_{i + 1}", f"score_{i + 1}"] for i in range(len(samples))],
[],
) )
if not hasattr(self, "validation_table"): if not hasattr(self, "validation_table"):
...@@ -69,10 +70,12 @@ class WandbGenerationLogger(GenerationLogger): ...@@ -69,10 +70,12 @@ class WandbGenerationLogger(GenerationLogger):
@dataclass @dataclass
class SwanlabGenerationLogger(GenerationLogger): class SwanlabGenerationLogger(GenerationLogger):
def log(self, samples: List[Tuple[str, str, float]], step: int) -> None: def log(self, samples: List[Tuple[str, str, str, float]], step: int) -> None:
swanlab_text_list = [] swanlab_text_list = []
for i, sample in enumerate(samples): for i, sample in enumerate(samples):
row_text = f"input: {sample[0]}\n\n---\n\noutput: {sample[1]}\n\n---\n\nscore: {sample[2]}" row_text = "\n\n---\n\n".join(
(f"input: {sample[0]}", f"output: {sample[1]}", f"label: {sample[2]}", f"score: {sample[3]}")
)
swanlab_text_list.append(swanlab.Text(row_text, caption=f"sample {i + 1}")) swanlab_text_list.append(swanlab.Text(row_text, caption=f"sample {i + 1}"))
swanlab.log({"val/generations": swanlab_text_list}, step=step) swanlab.log({"val/generations": swanlab_text_list}, step=step)
...@@ -94,6 +97,6 @@ class AggregateGenerationsLogger: ...@@ -94,6 +97,6 @@ class AggregateGenerationsLogger:
if logger in GEN_LOGGERS: if logger in GEN_LOGGERS:
self.loggers.append(GEN_LOGGERS[logger]()) self.loggers.append(GEN_LOGGERS[logger]())
def log(self, samples: List[Tuple[str, str, float]], step: int) -> None: def log(self, samples: List[Tuple[str, str, str, float]], step: int) -> None:
for logger in self.loggers: for logger in self.loggers:
logger.log(samples, step) logger.log(samples, step)
verl/utils/logger/logger.py
View file @ ff7fb65e
...@@ -71,7 +71,7 @@ class TensorBoardLogger(Logger): ...@@ -71,7 +71,7 @@ class TensorBoardLogger(Logger):
os.makedirs(tensorboard_dir, exist_ok=True) os.makedirs(tensorboard_dir, exist_ok=True)
print(f"Saving tensorboard log to {tensorboard_dir}.") print(f"Saving tensorboard log to {tensorboard_dir}.")
self.writer = SummaryWriter(tensorboard_dir) self.writer = SummaryWriter(tensorboard_dir)
self.writer.add_hparams(flatten_dict(config)) self.writer.add_hparams(hparam_dict=flatten_dict(config), metric_dict={})
def log(self, data: Dict[str, Any], step: int) -> None: def log(self, data: Dict[str, Any], step: int) -> None:
for key, value in data.items(): for key, value in data.items():
...@@ -146,7 +146,7 @@ class Tracker: ...@@ -146,7 +146,7 @@ class Tracker:
for logger in self.loggers: for logger in self.loggers:
logger.log(data=data, step=step) logger.log(data=data, step=step)
def log_generation(self, samples: List[Tuple[str, str, float]], step: int) -> None: def log_generation(self, samples: List[Tuple[str, str, str, float]], step: int) -> None:
self.gen_logger.log(samples, step) self.gen_logger.log(samples, step)
def __del__(self): def __del__(self):
......
verl/utils/performance.py deleted 100644 → 0
View file @ c132cbcb
# Copyright 2024 Bytedance Ltd. and/or its affiliates
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import torch
import torch.distributed as dist
def log_gpu_memory_usage(head: str, rank: int = 0):
if (not dist.is_initialized()) or (rank is None) or (dist.get_rank() == rank):
memory_allocated = torch.cuda.memory_allocated() / 1024**3
memory_reserved = torch.cuda.memory_reserved() / 1024**3
print(f"{head}, memory allocated (GB): {memory_allocated}, memory reserved (GB): {memory_reserved}.")
verl/utils/reward_score/math.py
View file @ ff7fb65e
...@@ -30,6 +30,7 @@ def math_acc_reward(predict_str: str, ground_truth: str) -> float: ...@@ -30,6 +30,7 @@ def math_acc_reward(predict_str: str, ground_truth: str) -> float:
def math_compute_score(predict_str: str, ground_truth: str) -> Dict[str, float]: def math_compute_score(predict_str: str, ground_truth: str) -> Dict[str, float]:
predict_str = re.sub(r"\s*(<|>|/)\s*", r"\1", predict_str) # handle qwen2.5vl-32b format
format = math_format_reward(predict_str) format = math_format_reward(predict_str)
accuracy = math_acc_reward(predict_str, ground_truth) accuracy = math_acc_reward(predict_str, ground_truth)
return { return {
......
verl/utils/rl_dataset.py deleted 100644 → 0
View file @ c132cbcb
# Copyright 2024 Bytedance Ltd. and/or its affiliates
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import math
from collections import defaultdict
from typing import Any, Dict, List, Optional
import torch
from datasets import load_dataset
from PIL import Image
from PIL.Image import Image as ImageObject
from torch.utils.data import Dataset
from transformers import PreTrainedTokenizer, ProcessorMixin
import verl.utils.torch_functional as verl_F
from verl.models.transformers.qwen2_5_vl import get_rope_index
def collate_fn(features: List[Dict[str, Any]]) -> Dict[str, Any]:
tensors = defaultdict(list)
non_tensors = defaultdict(list)
for feature in features:
for key, value in feature.items():
if isinstance(value, torch.Tensor):
tensors[key].append(value)
else:
non_tensors[key].append(value)
for key, value in tensors.items():
if key not in ["pixel_values", "image_grid_thw"]:
tensors[key] = torch.stack(value, dim=0)
return {**tensors, **non_tensors}
def process_image(image: ImageObject, max_pixels: int, min_pixels: int) -> ImageObject:
if (image.width * image.height) > max_pixels:
resize_factor = math.sqrt(max_pixels / (image.width * image.height))
width, height = int(image.width * resize_factor), int(image.height * resize_factor)
image = image.resize((width, height), resample=Image.Resampling.NEAREST)
if (image.width * image.height) < min_pixels:
resize_factor = math.sqrt(min_pixels / (image.width * image.height))
width, height = int(image.width * resize_factor), int(image.height * resize_factor)
image = image.resize((width, height), resample=Image.Resampling.NEAREST)
if image.mode != "RGB":
image = image.convert("RGB")
return image
class RLHFDataset(Dataset):
"""
We assume the dataset contains a column that contains prompts and other information
"""
def __init__(
self,
data_path: str,
tokenizer: PreTrainedTokenizer,
processor: Optional[ProcessorMixin],
prompt_key="prompt",
max_prompt_length=1024,
truncation="error",
system_prompt=None,
max_pixels=None,
min_pixels=None,
):
self.tokenizer = tokenizer
self.processor = processor
self.prompt_key = prompt_key
self.max_prompt_length = max_prompt_length
self.truncation = truncation
self.system_prompt = system_prompt
self.max_pixels = max_pixels
self.min_pixels = min_pixels
if "@" in data_path:
data_path, data_split = data_path.split("@")
else:
data_split = "train"
self.dataset = load_dataset(data_path, split=data_split)
def __len__(self):
return len(self.dataset)
def __getitem__(self, index):
"""
Note that we also return the raw_input_ids so that it can be combined with other chat template
"""
row_dict = self.dataset[index]
messages = [
{"role": "system", "content": self.system_prompt},
{"role": "user", "content": row_dict[self.prompt_key]},
]
prompt = self.tokenizer.apply_chat_template(messages, add_generation_prompt=True, tokenize=False)
if "images" in row_dict: # expand image token
raw_prompt = prompt.replace("<image>", "<|vision_start|><|image_pad|><|vision_end|>")
row_dict["images"] = [
process_image(image, self.max_pixels, self.min_pixels) for image in row_dict["images"]
]
image_inputs = self.processor.image_processor(row_dict["images"], return_tensors="pt")
image_grid_thw = image_inputs["image_grid_thw"]
row_dict.update(image_inputs)
if image_grid_thw is not None:
merge_length = self.processor.image_processor.merge_size**2
index = 0
while "<image>" in prompt:
prompt = prompt.replace(
"<image>",
"<|vision_start|>"
+ "<|placeholder|>" * (image_grid_thw[index].prod() // merge_length)
+ "<|vision_end|>",
1,
)
index += 1
prompt = prompt.replace("<|placeholder|>", self.processor.image_token)
else:
raw_prompt = prompt
input_ids, attention_mask = verl_F.tokenize_and_postprocess_data(
prompt=prompt,
tokenizer=self.tokenizer,
max_length=self.max_prompt_length,
pad_token_id=self.tokenizer.pad_token_id,
left_pad=True,
truncation=self.truncation,
)
if "images" in row_dict:
position_ids = get_rope_index(
self.processor,
input_ids=input_ids,
image_grid_thw=image_grid_thw,
attention_mask=attention_mask,
) # (3, seq_len)
else:
position_ids = torch.clip(attention_mask.cumsum(dim=0) - 1, min=0, max=None) # (seqlen,)
row_dict["input_ids"] = input_ids
row_dict["attention_mask"] = attention_mask
row_dict["position_ids"] = position_ids
row_dict["raw_prompt_ids"] = self.tokenizer.encode(raw_prompt, add_special_tokens=False)
return row_dict
verl/utils/torch_functional.py
View file @ ff7fb65e
...@@ -94,27 +94,29 @@ def masked_whiten(values: torch.Tensor, mask: torch.Tensor, eps: float = 1e-8) - ...@@ -94,27 +94,29 @@ def masked_whiten(values: torch.Tensor, mask: torch.Tensor, eps: float = 1e-8) -
return (values - mean) * torch.rsqrt(var + eps) return (values - mean) * torch.rsqrt(var + eps)
def get_eos_mask(response_ids: torch.Tensor, eos_token_id: Union[int, List[int]] = 2, dtype: torch.dtype = torch.long): def get_response_mask(
response_ids: torch.Tensor, eos_token_id: Union[int, List[int]] = 2, dtype: torch.dtype = torch.long
):
"""Get the mask for the response ids, the mask will be 0 after the first eos token. """Get the mask for the response ids, the mask will be 0 after the first eos token.
eos_token_id can be int or list: 1 or [1, 2]. eos_token_id can be int or list: 1 or [1, 2].
``` ```
e.g. eos_token = 1 e.g. eos_token = 1
response_ids: [0, 0, 2, 4, 3, 5, 1, 0, 0] response_ids: [0, 0, 2, 4, 3, 5, 1, 0, 0]
eos_mask: [1, 1, 1, 1, 1, 1, 1, 0, 0] response_mask: [1, 1, 1, 1, 1, 1, 1, 0, 0]
``` ```
""" """
if isinstance(eos_token_id, int): if isinstance(eos_token_id, int):
eos_token_id = [eos_token_id] eos_token_id = [eos_token_id]
eos_mask = torch.zeros_like(response_ids, dtype=torch.bool) response_mask = torch.zeros_like(response_ids, dtype=torch.bool)
for token_id in eos_token_id: for token_id in eos_token_id:
eos_mask |= response_ids.eq(token_id) response_mask |= response_ids.eq(token_id)
eos_mask = eos_mask.long() response_mask = response_mask.long()
eos_mask = (torch.cumsum(eos_mask, dim=1) - eos_mask).bool() response_mask = (torch.cumsum(response_mask, dim=1) - response_mask).bool()
eos_mask = torch.logical_not(eos_mask).to(dtype) response_mask = torch.logical_not(response_mask).to(dtype)
return eos_mask return response_mask
def pad_2d_list_to_length( def pad_2d_list_to_length(
......
verl/utils/tracking.py deleted 100644 → 0
View file @ c132cbcb
# 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.
"""
A unified tracking interface that supports logging data to different backend
"""
import os
from typing import List, Union
from verl.utils.logger.aggregate_logger import LocalLogger
class Tracking:
supported_backend = ["wandb", "mlflow", "swanlab", "console"]
def __init__(self, project_name, experiment_name, default_backend: Union[str, List[str]] = "console", config=None):
if isinstance(default_backend, str):
default_backend = [default_backend]
for backend in default_backend:
assert backend in self.supported_backend, f"{backend} is not supported"
self.logger = {}
if "wandb" in default_backend:
import wandb # type: ignore
wandb.init(project=project_name, name=experiment_name, config=config)
self.logger["wandb"] = wandb
if "mlflow" in default_backend:
import mlflow # type: ignore
mlflow.start_run(run_name=experiment_name)
mlflow.log_params(config)
self.logger["mlflow"] = _MlflowLoggingAdapter()
if "swanlab" in default_backend:
import swanlab # type: ignore
SWANLAB_API_KEY = os.environ.get("SWANLAB_API_KEY", None)
SWANLAB_LOG_DIR = os.environ.get("SWANLAB_LOG_DIR", "swanlog")
SWANLAB_MODE = os.environ.get("SWANLAB_MODE", "cloud")
if SWANLAB_API_KEY:
swanlab.login(SWANLAB_API_KEY) # NOTE: previous login information will be overwritten
swanlab.init(
project=project_name,
experiment_name=experiment_name,
config=config,
logdir=SWANLAB_LOG_DIR,
mode=SWANLAB_MODE,
)
self.logger["swanlab"] = swanlab
if "console" in default_backend:
self.console_logger = LocalLogger(print_to_console=True)
self.logger["console"] = self.console_logger
def log(self, data, step, backend=None):
for default_backend, logger_instance in self.logger.items():
if backend is None or default_backend in backend:
logger_instance.log(data=data, step=step)
def __del__(self):
if "wandb" in self.logger:
self.logger["wandb"].finish(exit_code=0)
if "swanlab" in self.logger:
self.logger["swanlab"].finish()
class _MlflowLoggingAdapter:
def log(self, data, step):
import mlflow # type: ignore
mlflow.log_metrics(metrics=data, step=step)
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment