fix(dschat): avoid NaN from masked PPO log-probability padding

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,

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]]))