perf(rlhf): chunk selective_log_softmax instead of looping per token

[rrrxxx0510]

PR type

• Bug Fix
• New Feature
• Document Updates
• More Models or Datasets Support

PR information

Fixes #10357

swift imports selective_log_softmax from trl.trainer.utils. TRL's version loops over the leading dimension of logits "to reduce peak mem consumption". That is fine when the leading dimension is a batch dimension, but under padding_free/packing swift flattens the logits to [total_nnz, vocab_size] (grpo_trainer.py:1591), so the loop degenerates to one iteration per token — thousands of small log_softmax + gather launches per call instead of a few dozen.

This adds a chunked implementation in swift/rlhf_trainers/utils.py: the flattened rows are processed chunk_size at a time. Peak memory is still bounded by the chunk size, which was the point of the original loop, but the work gets batched again. The three call sites in grpo_trainer.py (sequence-parallel / padding_free / default) and the two in rlhf_mixin.py (DPO/CPO/KTO/ORPO) switch to it.

Correctness, checked against the naive gather(logits.log_softmax(-1), ...):

• bf16/fp16 are bit-for-bit identical (torch.equal) across 2-D and 3-D shapes, chunk_size 1/3/1024/10^6, top-K indices, temperature, row_mask, and gradients. bf16 is what training actually uses.
• fp32/fp64 keep TRL's x - logsumexp(x) form, so results can differ in the last bit; asserted within rtol=atol=1e-6.

chunk_size=1024 is close to optimal (2.6 ms vs 2.4 ms at 16384 on the 8192x152064 case), so the default is left alone.

Experiment results

Measured on an H200 (torch 2.9.1+cu130, bf16), output asserted bit-identical in every case:

rows	vocab	TRL loop	chunked	speedup
1024	32768	22.9 ms	0.1 ms	328x
4096	32768	89.7 ms	0.3 ms	345x
8192	32768	129.6 ms	0.5 ms	238x
2048	152064	64.0 ms	0.7 ms	95x
8192	152064	255.9 ms	2.6 ms	97x

The win is almost entirely launch overhead, which is why it is far larger on GPU than the ~5x measured on CPU.

Tests: tests/general/test_selective_log_softmax.py, 7 cases, no GPU or model needed.

Checklist

• I have linked the related issue above, if there is one.
• I have reviewed and understand every changed line, and I can explain how the change works and why it is correct.
• I have run the relevant tests or reproductions locally, and included the commands and results above (if applicable).
• I am responsible for the content of this PR, regardless of how it was produced.

[rrrxxx0510]

Hi — this is my first PR from a fork, so the workflows are parked waiting for a maintainer to approve them before CI runs. It replaces TRL's per-row selective_log_softmax with a chunked version (bf16 output is bit-for-bit identical; ~5x faster in my CPU measurements). Could you approve the workflow runs when you get a chance?

[rrrxxx0510]

GPU numbers are in — measured on an H200 (torch 2.9.1+cu130, bf16), output asserted bit-identical in every case:

rows	vocab	TRL loop	chunked	speedup
1024	32768	22.9 ms	0.1 ms	328x
4096	32768	89.7 ms	0.3 ms	345x
8192	32768	129.6 ms	0.5 ms	238x
2048	152064	64.0 ms	0.7 ms	95x
8192	152064	255.9 ms	2.6 ms	97x

rows=8192 is roughly 8 sequences x 1024 completion tokens, and this runs 2-3 times per training step (policy / old / ref). The win is almost entirely launch overhead, which is why it is much larger on GPU than the ~5x I measured on CPU.

I also swept chunk_size at rows=8192 / vocab=152064: 256 -> 2.9 ms, 1024 (the default) -> 2.6 ms, 16384 -> 2.4 ms. Left the default alone since it is already in the flat region.

Everything else in the description is unchanged: bf16/fp16 stay bit-for-bit identical, fp32/fp64 keep TRL's x - logsumexp(x) form.

[rrrxxx0510]

Hi @hjh0119 @tastelikefeet - sorry to ping directly. This one has been parked with every workflow run at action_required since yesterday, so CI never runs.

It replaces TRL's per-row selective_log_softmax with a chunked version, because under padding_free swift flattens logits to [total_nnz, vocab_size] and that loop degenerates to one iteration per token. bf16 output is bit-for-bit identical; measured 95-345x faster on an H200 (numbers in the comment above).

Could one of you approve the workflow runs when you get a chance? Happy to rebase or change anything once I can see the CI result.