feat(loss): add --pg-loss-divisor: first-class constant-divisor pg_loss normalization (Dr.GRPO)

[EazyReal]

Problem

The default pg_loss reducer normalizes each sample by its own active-token count (a sum of per-sample token means). Dr.GRPO (arXiv:2503.20783, Eq. 2) identifies this data-dependent denominator as a source of length bias and replaces it with a constant divisor (e.g. the max context length); DeepSWE trains with the same rule.

slime had no way to express this directly. docs/en/get_started/customization.md §11 pointed at examples/DrGRPO/custom_reducer.py:get_pg_loss_reducer as the way to get it, but that file never existed (the doc link 404s). An earlier revision of this PR filled the gap with a ~265-line example reducer — and writing it exposed the real problem: the --custom-pg-loss-reducer-function-path factory is invoked with only four positional arguments (total_lengths, response_lengths, loss_masks, calculate_per_token_loss) — no args, no max_seq_lens. So a reducer implementing the one-line Dr.GRPO rule had to:

• reach into Megatron's global args for its divisor and qkv_format,
• re-implement the per-CP-rank loss-mask chunking that the default reducer (get_logits_and_tokens_offset_with_cp) already does, and
• still fail loud on qkv_format=bshd under context parallelism, because the call site does not forward max_seq_lens.

The example also had to own a core CLI argument (--drgrpo-divisor in slime/utils/arguments.py) — an example owning a global arg.

Before / After

Before (this PR's earlier shape — or any out-of-tree reducer today):

--custom-pg-loss-reducer-function-path examples.DrGRPO.custom_reducer.get_pg_loss_reducer \
--drgrpo-divisor 40960

~265 lines of example code reading Megatron globals, duplicating CP mask slicing, raising on bshd.

After:

--pg-loss-divisor 40960   # a constant, e.g. the max context length

pg_loss is aggregated as sum(token_loss * loss_mask) / 40960 instead of the default sum of per-sample active-token means. Concretely, for two responses with masked sums s1 (10 active tokens) and s2 (1000 active tokens): default gives s1/10 + s2/1000 (the short response dominates per token); constant mode gives (s1 + s2) / 40960 (every token weighs the same regardless of response length — Dr.GRPO's point).

Flag unset → byte-identical to today. Other metrics (pg_clipfrac, ppo_kl, entropy_loss, kl_loss) keep the default reducer, matching the scope of the custom hook.

Why first-class instead of an example

The core reducer site already has all the layout context the hook lacks, so the whole feature is ~15 lines there:

• cp_utils.get_sum_of_sample_mean grows a constant_divisor option that divides the masked token-sum closure it already builds — reusing the existing CP mask chunking, so it is correct for thd and bshd under context parallelism.
• No extra allreduce: the divisor is identical on every CP rank, so Megatron's gradient sum-allreduce already yields the full-batch value.
• No global-args reach-around: the divisor arrives through args at the selection site in loss.py (policy_loss_function), which also forwards qkv_format and max_seq_lens.
• --calculate-per-token-loss: the divisor is intentionally not applied — Megatron already divides by the all-reduced token count (same contract as the default reducer).
• Fail at startup, not mid-train: slime_validate_args rejects non-positive (or NaN) divisors.

The --custom-pg-loss-reducer-function-path hook is untouched for normalizations that need more than a constant. Docs §11 (en + zh) now shows the flag for the Dr.GRPO use case — fixing the dangling example reference that originally motivated this PR — and examples/DrGRPO/README.md is a thin pointer to the flag.

Tests

• tests/test_cp_utils.py: constant mode = masked token sum / constant (contrasted with the default per-sample means on the same inputs); divisor not applied under calculate_per_token_loss; summing per-CP-rank reducer outputs reproduces the cp=1 value for both thd and bshd.
• tests/test_megatron_argument_validation.py: startup validation rejects 0 / negative / NaN divisors, accepts positive ones.

Both files are already in the cpu-unittest matrix, so no CI registration changes. New tests were mutation-checked: ignoring the divisor, applying it in per-token mode, or removing the validation each makes the corresponding test fail.

🤖 Generated with Claude Code

[EazyReal]

Superseded by #2090, which adds --loss-aggregation {sample_mean, prompt_mean, token_mean, constant} — the Dr.GRPO constant divisor becomes the first-class --loss-aggregation constant --loss-aggregation-divisor L mode (same constant_divisor design as here), alongside the DAPO prompt-average, GRPO sample-average, and token-average. Closing in favor of the unified flag.