[skyrl-train] Implement loss reduction via advantage normalization and fix token_mean reduction strategy

[justinvyu]

This is a breaking change for the default token_mean loss behavior, as well as observed policy_loss metrics! See the "Differences in reported loss metric magnitudes" section below.

Summary

• Change reduce_loss() to always returns a simple masked sum ((loss * mask).sum()). To achieve different reduction strategies, we pre-scale the advantages before they enter the loss function, which also aligns with how Tinker's API handles it.
  • Scales the loss by the DP size before calling backward() to counteract the default data parallel mean gradient all-reduce across workers to do a sum instead.
• Fixes the token_mean loss reduction method to take a mean across all tokens in the minibatch rather than averaging across microbatches. Allows running with the old loss reduction with the token_mean_legacy config.

Loss reduction strategies

• Option 1: token_mean

  • Average loss per token across the entire mini-batch.
  • This is the fixed version where the denominator is the total token count across the full mini-batch, so the gradient is independent of how the minibbatch is split into micro-batches.

• Option 1b: token_mean_legacy

  • Compute token-mean loss within each micro-batch, then average across micro-batches.
  • This reproduces the token_mean behavior before this PR.
  • The problem: if micro-batches have different token counts, the effective weighting differs from a true global token mean. This is also less usable since changing micro batch size affects the loss and the training dynamics.
  • Kept as a fallback in case of performance regressions — we should remove this down the line.

• Option 2: sequence_mean

  • Compute per-token loss within each sequence, average across sequences.
  • This is unchanged and is just implemented via advantage normalization instead.

• Option 3: seq_mean_token_sum_norm

  • Dr. GRPO style — normalize by a fixed constant to avoid any length-dependent weighting.
  • This is unchanged and is just implemented via advantage normalization instead.

Mean all-reduce -> sum all-reduce

We need the loss to be summed across microbatches and data parallel workers:

• DDP/FSDP defaults to a mean all-reduce for gradients across workers. This PR counteracts this by multiplying by the DP world size in order to keep the loss sum across data parallel groups.
• Megatron also does a similar mean reduction across microbatches and workers, so we counteract this by multiplying by num microbatches and DP size to achieve the sum.

Difference in reported loss metric magnitudes

You will observe that the loss metric reported has a different magnitude compared to your older runs. This is beacuse the old token_mean implementation was somewhere between a true token mean and a sequence mean due to per-micro-batch normalization (ex: micro_batch_size=1 was equivalent to sequence mean).

The new token_mean is a proper minibatch token mean, while sequence_mean weights every sequence equally regardless of length. When comparing the loss produced by different reduction methods computed on the same advantages, from a real run:

  token_mean (new):  0.322   — every token weighted equally across the mini-batch
  token_mean (old):  0.065   — mean of per-micro-batch token means, where micro_batch_size=4
  sequence_mean:     0.00098 — every sequence weighted equally

The old token_mean gave each micro-batch equal weight rather than each token, so its scale depended on how advantages were distributed across micro-batches. The new implementation is invariant to micro-batch size.

Note that token_mean_legacy reports the old metrics still, and the sequence_mean and seq_mean_token_sum_norm modes also match exactly. See this comment for more details.

Tinker compatibility

Here was the first attempt at fixing the loss reduction across microbatches: #909

This method was to track total tokens and then do one big normalization at the optim_step in order to get an average per-token loss. But, we decided to align with Tinker's way of just summing up the loss at the end, and pushing any loss normalization to the user's advantage calculation.

The benefit is that users have full control of customizing their loss reduction strategy, rather than having it happen in our opaque forward_backward, optim_step implementation which would require some configuration argument that diverges from tinker's API. For example, we would need to add a config somewhere to determine how to average/sum the loss:

client.forward_backward(...)
client.optim_step(..., loss_reduction="token_mean")  # no longer tinker compatible

The current PR aligns with Tinker semantics:

Notice that for all objectives we sum the token-level losses over the sequence length unlike some other loss implementations. If you would like to explore different aggregation schemes, you can include that in the advantage tensor computation.

Example for loss_reduction="token_mean":

• Move the 1/num_minibatch_tokens normalization into the advantage: loss = sum( -advantage_i * ratio_i for i in range(num_minibatch_tokens) ) / num_minibatch_tokens
• -> sum( -(advantage_i / num_minibatch_tokens) * ratio_i for i in range(num_minibatch_tokens) )

Learning curve comparisons before/after the PR

FSDP (wandb)

Megatron (wandb)

1.7B:

30B lora:

master baseline:

token_mean_legacy + fixed token_mean:

---

[devin-ai-integration]

✅ Devin Review: No Issues Found

Devin Review analyzed this PR and found no bugs or issues to report.

[erictang000]

this looks almost good to merge, super clean thanks for adding the token_mean_legacy path

just want to check my understanding + 1 question about the metrics code that I think I probably wrote on the old PR...

[justinvyu]

Metrics fix 1: remove dp_size multiplier in reported metrics, since there's no average that we need to correct for, since reduce_microbatch_metrics and all_reduce_metrics both do sums for *_loss metrics.

[justinvyu]

Metrics fix 2: Take an average across minibatches instead of still summing. This is because the loss reduction normalization happens at the minibatch level. Across different minibatches we should just average, otherwise we'll increase the reported loss scale by ~num_minibatches

[justinvyu]

Q: should this affect the SFT case? SFT doesn't look at the normalized advantages either, similar to the critic loss case.

Before the PR, the SFT case does a sum across the negative log likelihoods within a microbatch, but still averaged over microbatches and dp workers.
Now, we are summing negative log likelihood across the entire minibatch. What's the desired behavior here?

[justinvyu]

Reverting to the old behavior for now and we can tackle it in a follow-up. SFT loss reduction is broken already due to taking a sum within the microbatch but then a mean across microbatches/workers. The current behavior does not align with this comment: https://github.com/justinvyu/SkyRL/blob/c5feb83b38f4635c7fc705c2bb192a7d6ad16947/skyrl/backends/skyrl_train/utils/ppo_utils.py#L917

[justinvyu]

To sanity check the difference in loss metric magnitudes, I dumped the raw advantages on the first step and manually calculated the loss with the different reduction methods on the same dumped data.

Using dumped advantage tensors from a real GRPO run to compare old vs. new:

With micro_batch_size=4 (128 micro-batches), old and new differ by ~5x — matching what's observed in the actual run:

Average: old=0.065  new=0.322  ratio=4.92

With micro_batch_size=1, the old token_mean reduces to sequence_mean (each sequence weighted equally). The old values match sequence_mean exactly:

token_mean old:  Average=-0.024
sequence_mean:   Average=-0.024  ratio=1.0000

With micro_batch_size=512 (1 micro-batch = full mini-batch), old and new converge:

Average: old=0.322  new=0.322  ratio=1.0000

The new token_mean value (0.322) is the same regardless of micro_batch_size — which is the correct behavior. The old value varied between -0.024 (at micro_batch_size=1, i.e. sequence_mean) and 0.322 (at micro_batch_size=512) depending on how micro-batches were formed.

token_mean_legacy reproduces the old behavior. Runs using token_mean won't be directly comparable to before, but the difference is analogous to comparing token_mean vs. sequence_mean — a different weighting, not a bug.

[devin-ai-integration]

Devin Review found 1 new potential issue.

View 14 additional findings in Devin Review.

[devin-ai-integration]

🟡 microbatch_weight uses configured size instead of actual micro-batch size, incorrect for uneven splits

The microbatch_weight is always computed as micro_batch_size / len(data) using the configured micro_batch_size, rather than the actual size of the current micro-batch. When BatchIterator is used with drop_last=False and the data doesn't divide evenly, the last micro-batch will be smaller but receive the same weight as full-sized ones. This over-weights the KL/entropy loss terms (RL path, line 890) and the entire loss (SFT path, line 818) for that final micro-batch. In practice, validate_batch_sizes currently ensures even divisibility, so this is unlikely to trigger, but the code is still incorrect if validation constraints change.

---

Was this helpful? React with 👍 or 👎 to provide feedback.