perf(rlhf): vectorize padding-free log-probability restore

[Ruihan11]

PR type

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

PR information

Summary

Optimize padding-free RLHF log-probability restoration for larger batches by replacing per-sequence NPU synchronization and slice assignments with a batched padding path.

Motivation

pad_logps_back_to_batch is shared by the GRPO, DPO, KTO, HF, Megatron, and Ray RLHF paths. The previous implementation performed a Python loop over every batch element. For each sequence it:

• read sequence boundaries from the device with .item();
• computed a device-to-host synchronization;
• performed individual output and mask slice assignments.

This overhead grows with the number of packed sequences and is visible in padding-free RLHF batches containing many short completions.

Changes

• Keep the original loop-based path for batch_size <= 2.
  • This avoids the fixed pad_sequence/temporary construction overhead for tiny batches.
  • The threshold is based on the Ascend 910B3 measurements below; it is not a hardware latency assertion in CI.
• For larger batches:
  • derive actual available lengths on the host with one tolist() transfer instead of one .item() transfer per boundary;
  • split the flat log-probability tensor into sequences;
  • reverse sequences before and after pad_sequence to preserve left-padding semantics on older PyTorch versions;
  • use F.pad to reach logits_to_keep when necessary;
  • construct valid_mask with vectorized device-side indexing.
• Preserve the existing function signature and caller behavior, including:
  • seq_lengths and position_ids fallback paths;
  • dtype and pad_value handling;
  • left-padded output layout;
  • empty, truncated, and extra-input behavior;
  • autograd propagation.
• Add focused correctness tests in tests/test_align/test_rlhf_loss.py.
• Do not change RLHF trainer APIs, callers, CUDA code, third-party dependencies, or default CI performance behavior.

Correctness validation

The focused test compares the optimized implementation with an independent loop reference and covers:

• normal variable-length packed sequences;
• empty and all-empty sequences;
• truncated flat input;
• extra flat input;
• batch-size 1/2 fast path;
• seq_lengths and position_ids input paths;
• FP32 and BF16;
• explicit output dtype conversion;
• output shape, dtype, device, values, and validity mask;
• backward gradients.

Results:

python -m pytest -q tests/test_align/test_rlhf_loss.py
10 passed, 32 subtests passed

A separate Ascend NPU correctness harness also compared the real production caller against the original implementation for FP32/BF16 normal, empty, and truncated inputs. Output tensors, masks, loss, gradients, and finite-state checks all passed.

Experiment results

Environment

Hardware: Ascend 910B3
Visible devices: 1 NPU
Python: 3.10.12
PyTorch: 2.10.0+cpu
torch_npu: 2.10.0
CANN: 8.5.0
Warmup: 20 iterations
Measurement: 100 iterations
Timing: torch.npu.synchronize() before and after every measured iteration

Baseline and optimized measurements used the same inputs, sequence lengths, dtype, and benchmark settings.

Helper microbenchmark

The helper benchmark covered FP32/BF16, uniform and skewed sequence distributions, 4096 and 16384 total tokens, and representative batch sizes up to 256. The following table shows synchronized p50/p95 results for the 4096-token uniform distribution:

FP32

Batch size	Baseline p50	Optimized p50	p50 speedup	Baseline p95	Optimized p95	p95 speedup
1	0.491 ms	0.482 ms	1.02x	0.503 ms	0.514 ms	0.98x
2	0.730 ms	0.674 ms	1.08x	0.744 ms	0.689 ms	1.08x
4	1.212 ms	0.579 ms	2.09x	1.236 ms	0.592 ms	2.09x
8	2.163 ms	0.723 ms	2.99x	2.189 ms	0.740 ms	2.96x
16	4.069 ms	1.036 ms	3.93x	4.115 ms	1.062 ms	3.87x
32	7.888 ms	1.658 ms	4.76x	7.998 ms	1.692 ms	4.73x
64	15.268 ms	2.888 ms	5.29x	15.470 ms	2.945 ms	5.25x
128	29.856 ms	5.460 ms	5.47x	30.136 ms	5.625 ms	5.36x
256	59.722 ms	10.633 ms	5.62x	60.974 ms	10.835 ms	5.63x

BF16

Batch size	Baseline p50	Optimized p50	p50 speedup	Baseline p95	Optimized p95	p95 speedup
1	0.489 ms	0.456 ms	1.07x	0.502 ms	0.470 ms	1.07x
2	0.730 ms	0.695 ms	1.05x	0.744 ms	0.705 ms	1.06x
4	1.197 ms	0.699 ms	1.71x	1.219 ms	0.714 ms	1.71x
8	2.152 ms	0.938 ms	2.29x	2.171 ms	0.960 ms	2.26x
16	4.038 ms	1.369 ms	2.95x	4.107 ms	1.390 ms	2.96x
32	7.775 ms	2.270 ms	3.43x	8.104 ms	2.302 ms	3.52x
64	15.524 ms	4.038 ms	3.84x	17.076 ms	4.081 ms	4.18x
128	30.155 ms	7.707 ms	3.91x	30.719 ms	7.791 ms	3.94x
256	60.589 ms	14.782 ms	4.10x	61.687 ms	14.942 ms	4.13x

Skewed distributions and 16384-token inputs showed the same trend. For example, at 4096 tokens and batch size 256:

FP32 skewed: 62.562 -> 10.884 ms, 5.75x p50; 5.85x p95
BF16 skewed: 60.083 -> 16.657 ms, 3.61x p50; 3.65x p95

Production-caller synthetic NPU E2E

The benchmark invokes the actual production method:

GRPOTrainer._unpad_logps_and_entropies

It covers both the logps-only path (compute_entropy=False) and the logps-plus-entropy path (compute_entropy=True). Results below are FP32, 4096 total tokens, synchronized p50 latency:

Logps only

Batch size	Baseline	Optimized	Speedup
1	0.510 ms	0.459 ms	1.11x
2	0.780 ms	0.717 ms	1.09x
4	1.304 ms	0.620 ms	2.10x
8	2.358 ms	0.819 ms	2.88x
16	4.412 ms	1.212 ms	3.64x
32	8.539 ms	1.967 ms	4.34x
64	16.865 ms	3.506 ms	4.81x
128	32.166 ms	6.536 ms	4.92x

Logps plus entropy

Batch size	Baseline	Optimized	Speedup
1	0.972 ms	0.888 ms	1.09x
2	1.502 ms	1.407 ms	1.07x
4	2.532 ms	1.268 ms	2.00x
8	4.483 ms	1.683 ms	2.66x
16	8.424 ms	2.501 ms	3.37x
32	20.045 ms	4.081 ms	4.91x
64	37.299 ms	7.318 ms	5.10x
128	73.109 ms	13.842 ms	5.28x

Peak allocation increases slightly on the optimized path because it uses temporary batched padding buffers. For the 4096-token uniform FP32 helper case at batch 256, peak allocation was approximately 54 KiB baseline versus 163 KiB optimized. Skewed/longer inputs reached approximately 0.1--1.7 MiB additional peak allocation depending on the sequence distribution, with no persistent growth observed. Reserved memory remained allocator-dependent and was not used as a correctness criterion.

Validation commands

python -m pytest -q tests/test_align/test_rlhf_loss.py
python -m compileall -q swift
flake8
isort --check-only .
yapf --diff --recursive swift tests
git diff --check

All commands passed.

Limitations and scope

• This change targets the shared padding-free RLHF log-probability restoration helper; no public API or trainer caller changes are required.
• CUDA behavior is unchanged.
• The small-batch loop path is intentionally retained because the vectorized path does not provide a reliable benefit for batch sizes 1--2.
• Peak temporary allocation is slightly higher on the optimized path; the measured increase is small relative to model memory and does not grow across repeated iterations.