Add files using upload-large-folder tool

64b9802 verified 17 days ago

8.21 kB

# 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.
# 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 numpy as np
from collections import defaultdict
from types import SimpleNamespace

def masked_mean(values, mask, axis=None):
    """Compute mean of tensor with a masked values."""
    return (values * mask).sum(axis=axis) / (mask.sum(axis=axis) + 1e-8)

verl_F=SimpleNamespace(masked_mean=masked_mean)

def compute_grpo_outcome_advantage(token_level_rewards: torch.Tensor,
                                   response_mask: torch.Tensor,
                                   index: np.ndarray,
                                   epsilon: float = 1e-6,
                                   uniform_std=False):
    """
    Compute advantage for GRPO, operating only on Outcome reward 
    (with only one scalar reward for each response).
    Args:
        token_level_rewards: `(torch.Tensor)`
            shape: (bs, response_length)
        response_mask: `(torch.Tensor)`
            shape: (bs, response_length)
    Returns:
        advantages: `(torch.Tensor)`
            shape: (bs, response_length)
        Returns: `(torch.Tensor)`
            shape: (bs, response_length)
    """
    scores = token_level_rewards.sum(dim=-1)

    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}")
        if uniform_std:
            std = torch.tensor(list(id2std.values())).mean()
            for i in range(bsz):
                scores[i] = (scores[i] - id2mean[index[i]]) / (std + epsilon)
        else:
            for i in range(bsz):
                scores[i] = (scores[i] - id2mean[index[i]]) / (id2std[index[i]] + epsilon)
        scores = scores.unsqueeze(-1) * response_mask

    return scores, scores

def agg_loss(loss_mat: torch.Tensor, loss_mask: torch.Tensor, loss_agg_mode: str):
    """
    Aggregate the loss matrix into a scalar.
    Args:
        loss_mat: `(torch.Tensor)`
            shape: (bs, response_length)
        loss_mask: `(torch.Tensor)`
            shape: (bs, response_length)
        loss_agg_mode: (str) choices: "token-mean" / "seq-mean-token-sum" / "seq-mean-token-mean"
            "token-mean" is the default behavior
    Returns:
        loss: `a scalar torch.Tensor`
            aggregated loss
    """
    if loss_agg_mode == "token-mean":
        loss = verl_F.masked_mean(loss_mat, loss_mask)
    elif loss_agg_mode == "seq-mean-token-sum":
        seq_losses = torch.sum(loss_mat * loss_mask, dim=-1)  # token-sum
        loss = torch.mean(seq_losses)  # seq-mean
    elif loss_agg_mode == "seq-mean-token-mean":
        seq_losses = torch.sum(loss_mat * loss_mask, dim=-1) / torch.sum(loss_mask, dim=-1)  # token-mean
        loss = torch.mean(seq_losses)  # seq-mean
    else:
        raise ValueError(f"Invalid loss_agg_mode: {loss_agg_mode}")

    return loss

def compute_policy_loss(old_log_prob,
                        log_prob,
                        advantages,
                        response_mask,
                        cliprange=None,
                        cliprange_low=None,
                        cliprange_high=None,
                        clip_ratio_c=3.0,
                        loss_agg_mode="token-mean"):
    """Adapted from https://github.com/huggingface/trl/blob/main/trl/trainer/ppo_trainer.py#L1122
    Args:
        old_log_prob: `(torch.Tensor)`
            shape: (bs, response_length)
        log_prob: `(torch.Tensor)`
            shape: (bs, response_length)
        advantages: `(torch.Tensor)`
            shape: (bs, response_length)
        response_mask: `(torch.Tensor)`
            shape: (bs, response_length)
        cliprange: (float)
            The clip range used in PPO. See https://arxiv.org/abs/1707.06347
        cliprange_low: (float)
            The lower clip range used in PPO.
        cliprange_high: (float)
            The higher clip range used in PPO.
        clip_ratio_c: (float) default: 3.0
            The lower bound of the ratio for dual-clip PPO, See https://arxiv.org/pdf/1912.09729
        loss_agg_mode: (str) choices: "token-mean" / "seq-mean-token-sum" / "seq-mean-token-mean"
            "token-mean" is the default behavior        
    Returns:
        pg_loss: `a scalar torch.Tensor`
            policy gradient loss computed via PPO
        pg_clipfrac: (float)
            the fraction of policy gradient loss being clipped
        ppo_kl: (float)
            the estimated KL divergence between the latest updating policy and the old sampling policy
        pg_clipfrac_lower: (float)
            the fraction of policy gradient loss being clipped when the advantage is negative
    """
    assert clip_ratio_c > 1.0, f"The lower bound of the clip_ratio_c for dual-clip PPO should be greater than 1.0, but get the value: {clip_ratio_c}."

    negative_approx_kl = log_prob - old_log_prob
    ratio = torch.exp(negative_approx_kl)
    ppo_kl = verl_F.masked_mean(-negative_approx_kl, response_mask)

    pg_losses1 = -advantages * ratio
    if cliprange_low is None:
        cliprange_low = cliprange
    if cliprange_high is None:
        cliprange_high = cliprange
    pg_losses2 = -advantages * torch.clamp(ratio, 1 - cliprange_low,
                                           1 + cliprange_high)  # - clip(ratio, 1-cliprange, 1+cliprange) * A
    clip_pg_losses1 = torch.maximum(pg_losses1,
                                    pg_losses2)  # max(-ratio * A, -clip(ratio, 1-cliprange, 1+cliprange) * A)
    pg_clipfrac = verl_F.masked_mean(torch.gt(pg_losses2, pg_losses1).float(), response_mask)

    pg_losses3 = -advantages * clip_ratio_c
    clip_pg_losses2 = torch.min(pg_losses3, clip_pg_losses1)
    pg_clipfrac_lower = verl_F.masked_mean(
        torch.gt(clip_pg_losses2, pg_losses3) * (advantages < 0).float(), response_mask)

    pg_losses = torch.where(advantages < 0, clip_pg_losses2, clip_pg_losses1)
    pg_loss = agg_loss(loss_mat=pg_losses, loss_mask=response_mask, loss_agg_mode=loss_agg_mode)

    return pg_loss, pg_clipfrac, ppo_kl, pg_clipfrac_lower

def kl_penalty(logprob: torch.FloatTensor, ref_logprob: torch.FloatTensor, kl_penalty) -> torch.FloatTensor:
    """Compute KL divergence given logprob and ref_logprob.
    Copied from https://github.com/huggingface/trl/blob/main/trl/trainer/ppo_trainer.py#L1104
    Args:
        logprob:
        ref_logprob:
    Returns:
    """
    if kl_penalty == "kl":
        return logprob - ref_logprob

    if kl_penalty == "abs":
        return (logprob - ref_logprob).abs()

    if kl_penalty == "mse":
        return 0.5 * (logprob - ref_logprob).square()

    # J. Schulman. Approximating kl divergence, 2020.
    # # 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()
        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

    raise NotImplementedError

	# 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.
	# 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 numpy as np
	from collections import defaultdict
	from types import SimpleNamespace

	def masked_mean(values, mask, axis=None):
	"""Compute mean of tensor with a masked values."""
	return (values * mask).sum(axis=axis) / (mask.sum(axis=axis) + 1e-8)

	verl_F=SimpleNamespace(masked_mean=masked_mean)

	def compute_grpo_outcome_advantage(token_level_rewards: torch.Tensor,
	response_mask: torch.Tensor,
	index: np.ndarray,
	epsilon: float = 1e-6,
	uniform_std=False):
	"""
	Compute advantage for GRPO, operating only on Outcome reward
	(with only one scalar reward for each response).
	Args:
	token_level_rewards: `(torch.Tensor)`
	shape: (bs, response_length)
	response_mask: `(torch.Tensor)`
	shape: (bs, response_length)

	Returns:
	advantages: `(torch.Tensor)`
	shape: (bs, response_length)
	Returns: `(torch.Tensor)`
	shape: (bs, response_length)
	"""
	scores = token_level_rewards.sum(dim=-1)

	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}")
	if uniform_std:
	std = torch.tensor(list(id2std.values())).mean()
	for i in range(bsz):
	scores[i] = (scores[i] - id2mean[index[i]]) / (std + epsilon)
	else:
	for i in range(bsz):
	scores[i] = (scores[i] - id2mean[index[i]]) / (id2std[index[i]] + epsilon)
	scores = scores.unsqueeze(-1) * response_mask

	return scores, scores

	def agg_loss(loss_mat: torch.Tensor, loss_mask: torch.Tensor, loss_agg_mode: str):
	"""
	Aggregate the loss matrix into a scalar.
	Args:
	loss_mat: `(torch.Tensor)`
	shape: (bs, response_length)
	loss_mask: `(torch.Tensor)`
	shape: (bs, response_length)
	loss_agg_mode: (str) choices: "token-mean" / "seq-mean-token-sum" / "seq-mean-token-mean"
	"token-mean" is the default behavior
	Returns:
	loss: `a scalar torch.Tensor`
	aggregated loss
	"""
	if loss_agg_mode == "token-mean":
	loss = verl_F.masked_mean(loss_mat, loss_mask)
	elif loss_agg_mode == "seq-mean-token-sum":
	seq_losses = torch.sum(loss_mat * loss_mask, dim=-1) # token-sum
	loss = torch.mean(seq_losses) # seq-mean
	elif loss_agg_mode == "seq-mean-token-mean":
	seq_losses = torch.sum(loss_mat * loss_mask, dim=-1) / torch.sum(loss_mask, dim=-1) # token-mean
	loss = torch.mean(seq_losses) # seq-mean
	else:
	raise ValueError(f"Invalid loss_agg_mode: {loss_agg_mode}")

	return loss

	def compute_policy_loss(old_log_prob,
	log_prob,
	advantages,
	response_mask,
	cliprange=None,
	cliprange_low=None,
	cliprange_high=None,
	clip_ratio_c=3.0,
	loss_agg_mode="token-mean"):
	"""Adapted from https://github.com/huggingface/trl/blob/main/trl/trainer/ppo_trainer.py#L1122
	Args:
	old_log_prob: `(torch.Tensor)`
	shape: (bs, response_length)
	log_prob: `(torch.Tensor)`
	shape: (bs, response_length)
	advantages: `(torch.Tensor)`
	shape: (bs, response_length)
	response_mask: `(torch.Tensor)`
	shape: (bs, response_length)
	cliprange: (float)
	The clip range used in PPO. See https://arxiv.org/abs/1707.06347
	cliprange_low: (float)
	The lower clip range used in PPO.
	cliprange_high: (float)
	The higher clip range used in PPO.
	clip_ratio_c: (float) default: 3.0
	The lower bound of the ratio for dual-clip PPO, See https://arxiv.org/pdf/1912.09729
	loss_agg_mode: (str) choices: "token-mean" / "seq-mean-token-sum" / "seq-mean-token-mean"
	"token-mean" is the default behavior

	Returns:
	pg_loss: `a scalar torch.Tensor`
	policy gradient loss computed via PPO
	pg_clipfrac: (float)
	the fraction of policy gradient loss being clipped
	ppo_kl: (float)
	the estimated KL divergence between the latest updating policy and the old sampling policy
	pg_clipfrac_lower: (float)
	the fraction of policy gradient loss being clipped when the advantage is negative
	"""
	assert clip_ratio_c > 1.0, f"The lower bound of the clip_ratio_c for dual-clip PPO should be greater than 1.0, but get the value: {clip_ratio_c}."

	negative_approx_kl = log_prob - old_log_prob
	ratio = torch.exp(negative_approx_kl)
	ppo_kl = verl_F.masked_mean(-negative_approx_kl, response_mask)

	pg_losses1 = -advantages * ratio
	if cliprange_low is None:
	cliprange_low = cliprange
	if cliprange_high is None:
	cliprange_high = cliprange
	pg_losses2 = -advantages * torch.clamp(ratio, 1 - cliprange_low,
	1 + cliprange_high) # - clip(ratio, 1-cliprange, 1+cliprange) * A
	clip_pg_losses1 = torch.maximum(pg_losses1,
	pg_losses2) # max(-ratio * A, -clip(ratio, 1-cliprange, 1+cliprange) * A)
	pg_clipfrac = verl_F.masked_mean(torch.gt(pg_losses2, pg_losses1).float(), response_mask)

	pg_losses3 = -advantages * clip_ratio_c
	clip_pg_losses2 = torch.min(pg_losses3, clip_pg_losses1)
	pg_clipfrac_lower = verl_F.masked_mean(
	torch.gt(clip_pg_losses2, pg_losses3) * (advantages < 0).float(), response_mask)

	pg_losses = torch.where(advantages < 0, clip_pg_losses2, clip_pg_losses1)
	pg_loss = agg_loss(loss_mat=pg_losses, loss_mask=response_mask, loss_agg_mode=loss_agg_mode)

	return pg_loss, pg_clipfrac, ppo_kl, pg_clipfrac_lower

	def kl_penalty(logprob: torch.FloatTensor, ref_logprob: torch.FloatTensor, kl_penalty) -> torch.FloatTensor:
	"""Compute KL divergence given logprob and ref_logprob.
	Copied from https://github.com/huggingface/trl/blob/main/trl/trainer/ppo_trainer.py#L1104

	Args:
	logprob:
	ref_logprob:

	Returns:

	"""
	if kl_penalty == "kl":
	return logprob - ref_logprob

	if kl_penalty == "abs":
	return (logprob - ref_logprob).abs()

	if kl_penalty == "mse":
	return 0.5 * (logprob - ref_logprob).square()

	# J. Schulman. Approximating kl divergence, 2020.
	# # 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()
	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

	raise NotImplementedError