diff --git a/applications/DeepSpeed-Chat/dschat/rlhf/ppo_trainer.py b/applications/DeepSpeed-Chat/dschat/rlhf/ppo_trainer.py index 22cba6be0..64515dfaa 100644 --- a/applications/DeepSpeed-Chat/dschat/rlhf/ppo_trainer.py +++ b/applications/DeepSpeed-Chat/dschat/rlhf/ppo_trainer.py @@ -281,7 +281,8 @@ def get_overflow(self): def actor_loss_fn(self, logprobs, old_logprobs, advantages, mask): ## policy gradient loss - log_ratio = (logprobs - old_logprobs) * mask + # Mask by selection: (-inf - -inf) * 0 would remain NaN. + log_ratio = torch.where(mask.bool(), logprobs - old_logprobs, torch.zeros_like(logprobs)) ratio = torch.exp(log_ratio) pg_loss1 = -advantages * ratio pg_loss2 = -advantages * torch.clamp(ratio, 1.0 - self.cliprange, diff --git a/applications/DeepSpeed-Chat/tests/test_actor_loss_masking.py b/applications/DeepSpeed-Chat/tests/test_actor_loss_masking.py new file mode 100644 index 000000000..79c79efd8 --- /dev/null +++ b/applications/DeepSpeed-Chat/tests/test_actor_loss_masking.py @@ -0,0 +1,48 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +"""CPU regression checks for masked PPO policy log-ratios.""" + +import pytest +import torch + +from dschat.rlhf.ppo_trainer import DeepSpeedPPOTrainer + + +@pytest.mark.parametrize("padding_case", ["negative_infinity", "nan"]) +def test_actor_loss_ignores_nonfinite_masked_logprobs(padding_case): + trainer = DeepSpeedPPOTrainer.__new__(DeepSpeedPPOTrainer) + trainer.cliprange = 0.2 + + old = torch.tensor([[-1.0, -1.0, -1.0]]) + logprobs = old.clone() + if padding_case == "negative_infinity": + old[0, 0] = float("-inf") + logprobs[0, 0] = float("-inf") + else: + logprobs[0, 0] = float("nan") + logprobs.requires_grad_() + + mask = torch.tensor([[0.0, 1.0, 1.0]]) + loss = trainer.actor_loss_fn(logprobs, old, torch.ones_like(logprobs), mask) + loss.backward() + + torch.testing.assert_close(loss, torch.tensor(-1.0)) + torch.testing.assert_close(logprobs.grad, torch.tensor([[0.0, -0.5, -0.5]])) + assert torch.isfinite(logprobs.grad).all() + + +def test_actor_loss_keeps_normal_unmasked_gradient(): + trainer = DeepSpeedPPOTrainer.__new__(DeepSpeedPPOTrainer) + trainer.cliprange = 0.2 + old = torch.tensor([[-1.0, -1.0]]) + logprobs = old.clone().requires_grad_(True) + loss = trainer.actor_loss_fn( + logprobs, + old, + torch.ones_like(logprobs), + torch.ones_like(logprobs), + ) + loss.backward() + torch.testing.assert_close(loss, torch.tensor(-1.0)) + torch.testing.assert_close(logprobs.grad, torch.tensor([[-0.5, -0.5]]))