github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-06-21 05:04:47 +00:00
Code Issues Packages Projects Releases Wiki Activity

update loss

Browse Source
This commit is contained in:
Tong Li 2025-03-06 10:53:03 +08:00
parent b96d69055e
commit 678f5a9eca
1 changed files with 5 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
applications/ColossalChat/coati/distributed/loss.py
View File

@ -10,16 +10,18 @@ class PolicyLoss(nn.Module):
Policy Loss for PPO Policy Loss for PPO
""" """
def __init__(self, clip_eps: float = 0.2, skip_threshold: float = 20.0) -> None: def __init__(self, clip_eps: float = 0.2, skip_threshold: float = 20.0, beta: float = 0.01) -> None:
super().__init__() super().__init__()
self.clip_eps = clip_eps self.clip_eps = clip_eps
self.skip_threshold = skip_threshold self.skip_threshold = skip_threshold
self.beta = beta
def forward( def forward(
self, self,
log_probs: torch.Tensor, log_probs: torch.Tensor,
old_log_probs: torch.Tensor, old_log_probs: torch.Tensor,
advantages: torch.Tensor, advantages: torch.Tensor,
per_token_kl: torch.Tensor,
action_mask: Optional[torch.Tensor] = None, action_mask: Optional[torch.Tensor] = None,
) -> torch.Tensor: ) -> torch.Tensor:
skip = False skip = False
@ -35,7 +37,8 @@ class PolicyLoss(nn.Module):
ratio = ratio_.clamp(0.0, 10.0) ratio = ratio_.clamp(0.0, 10.0)
surr1 = ratio * advantages surr1 = ratio * advantages
surr2 = ratio.clamp(1 - self.clip_eps, 1 + self.clip_eps) * advantages surr2 = ratio.clamp(1 - self.clip_eps, 1 + self.clip_eps) * advantages
loss = -torch.min(surr1, surr2) loss = -torch.min(surr1, surr2) + self.beta * per_token_kl
if action_mask is not None: if action_mask is not None:
loss = masked_mean(loss, action_mask) loss = masked_mean(loss, action_mask)
else: else: