Repository navigation
fix(dschat): avoid NaN from masked PPO log-probability padding - #1022
Open
ot4562-glitch wants to merge 1 commit into
Open
ot4562-glitch wants to merge 1 commit into
ot4562-glitch wants to merge 1 commit into
Conversation
Signed-off-by: ot4562-glitch <266149313+ot4562-glitch@users.noreply.github.com>
ot4562-glitch
force-pushed
the
fix/dschat-ppo-mask-before-exp
branch
from
October 26, 2026 07:02
eb72230 to
5fd77ac
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
DeepSpeed-Chat's PPO actor loss computes log_ratio = (logprobs - old_logprobs) * mask. Padded positions may carry -inf sentinel values or nonfinite logprobs. When (-inf - -inf) becomes NaN, multiplying by zero does not remove it. The subsequent exp and PPO loss then become NaN despite the valid response tokens being finite.
Select zero with torch.where(mask.bool(), ..., 0) before exponentiation, preserving the valid-token policy ratio and clipping objective.
Regression
New CPU tests invoke the real DeepSpeedPPOTrainer.actor_loss_fn directly using an uninitialized trainer stub with cliprange=0.2. They exercise -inf/-inf and NaN padded logprobs, assert finite loss, zero excluded-token gradient, and unchanged valid-token gradients. A normal case is included.
The source-equivalent PyTorch 2.10 CPU screen reproduced NaN before and finite values after the change. The project's full DeepSpeed environment, GPU training, and CI suite were not run here. This PR does not claim to handle an all-empty action mask.
QuotientGuard
This issue was triaged during a 10-repository numerical investigation using QuotientGuard (https://github.com/ot4562-glitch/QuotientGuard), including explicit synthetic coefficient controls. The controls were source-equivalent, not DeepSpeed-Chat runtime capture, and QuotientGuard did not automatically certify a bug (reconstruction_validation=false). The independent PyTorch regression is the primary evidence; no new dependency is added.