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

swift/rlhf_trainers/utils.py

@@ -20,6 +20,7 @@
 from PIL import Image
 from pydantic import BaseModel, field_validator
 from torch import nn
+from torch.nn.utils.rnn import pad_sequence
 from torch.utils.data import DataLoader, RandomSampler
 from transformers.utils import is_torch_npu_available
 from types import MethodType
@@ -1873,45 +1874,44 @@ def pad_logps_back_to_batch(logps_rmpad: Optional[torch.Tensor],
         # Compute actual sequence lengths
         seq_lengths = cu_seqlens[1:] - cu_seqlens[:-1]
 
-    # Compute cumulative sequence lengths
-    cu_seqlens = torch.cumsum(torch.cat([torch.tensor([0], device=device), seq_lengths]), dim=0)
     max_seq_len = logits_to_keep  # All sequences will be padded to this length
-
-    # Initialize output tensors with padding value
-    logps_padded = torch.full((batch_size, max_seq_len), pad_value, dtype=dtype, device=device)
-    valid_mask = torch.zeros(batch_size, max_seq_len, dtype=torch.float32, device=device)
-
-    # Unflatten: assign each sequence's logps to the corresponding row
-    # Use LEFT PADDING (right-align the data) to match the standard padding convention
     logps_flat = logps_rmpad.squeeze(0)  # [total_nnz]
 
-    for i in range(batch_size):
-        start_idx = cu_seqlens[i].item()
-        end_idx = cu_seqlens[i + 1].item()
-        seq_len = int(seq_lengths[i].item())
-
-        actual_end_idx = min(end_idx, len(logps_flat))
-        actual_len = actual_end_idx - start_idx
-
-        if actual_len <= 0:
-            continue
-
-        # Left padding: place data at the RIGHT side of the row
-        # pad_len is the number of padding tokens at the beginning
-        pad_len = max_seq_len - seq_len
-
-        if actual_len < seq_len:
-            # Input data is shorter than expected seq_len
-            # This happens when logps_flat doesn't have enough data
-            # Place actual data at the rightmost positions
-            data_pad_len = max_seq_len - actual_len
-            logps_padded[i, data_pad_len:] = logps_flat[start_idx:actual_end_idx]
-            valid_mask[i, data_pad_len:] = 1.0
-        else:
-            # Normal case: seq_len tokens of data
-            logps_padded[i, pad_len:] = logps_flat[start_idx:end_idx]
+    if batch_size <= 2:
+        cu_seqlens = torch.cat((seq_lengths.new_zeros(1), seq_lengths.cumsum(0)))
+        logps_padded = torch.full((batch_size, max_seq_len), pad_value, dtype=dtype, device=device)
+        valid_mask = torch.zeros(batch_size, max_seq_len, dtype=torch.float32, device=device)
+        for i in range(batch_size):
+            start_idx = cu_seqlens[i].item()
+            end_idx = cu_seqlens[i + 1].item()
+            seq_len = int(seq_lengths[i].item())
+            actual_end_idx = min(end_idx, len(logps_flat))
+            actual_len = actual_end_idx - start_idx
+            if actual_len <= 0:
+                continue
+            pad_len = max_seq_len - actual_len if actual_len < seq_len else max_seq_len - seq_len
+            logps_padded[i, pad_len:] = logps_flat[start_idx:actual_end_idx]
             valid_mask[i, pad_len:] = 1.0
-
+        return logps_padded, valid_mask
+
+    lengths = seq_lengths.detach().tolist()
+    actual_lengths = []
+    remaining = logps_flat.numel()
+    for seq_len in lengths:
+        actual_lengths.append(min(max(remaining, 0), seq_len))
+        remaining -= seq_len
+
+    logps_flat = logps_flat.to(dtype=dtype)
+    sequences = torch.split(logps_flat[:sum(actual_lengths)], actual_lengths)
+    # Reverse before and after right-padding to support left-padding on older PyTorch versions.
+    logps_padded = pad_sequence([sequence.flip(0) for sequence in sequences], batch_first=True,
+                                padding_value=pad_value).flip(1)
+    if logps_padded.shape[1] < max_seq_len:
+        logps_padded = F.pad(logps_padded, (max_seq_len - logps_padded.shape[1], 0), value=pad_value)
+
+    actual_lengths = torch.tensor(actual_lengths, dtype=torch.long, device=device)
+    positions = torch.arange(max_seq_len, device=device)
+    valid_mask = (positions.unsqueeze(0) >= (max_seq_len - actual_lengths).unsqueeze(1)).to(torch.float32)
     return logps_padded, valid_mask
 
 

tests/test_align/test_rlhf_loss.py

@@ -9,6 +9,7 @@
 from swift.rlhf_trainers.dpo_trainer import DPOTrainer
 from swift.rlhf_trainers.kto_trainer import KTOTrainer
 from swift.rlhf_trainers.rlhf_mixin import RLHFTrainerMixin
+from swift.rlhf_trainers.utils import pad_logps_back_to_batch
 from swift.trainers.mixin import SwiftMixin
 from swift.utils import get_packed_seq_params
 
@@ -72,6 +73,27 @@ def _reference_segment_sum(values, lengths):
     return torch.stack(outputs)
 
 
+def _reference_pad_logps(logps_rmpad, seq_lengths, logits_to_keep, dtype=None, pad_value=-1e10):
+    if dtype is None:
+        dtype = logps_rmpad.dtype
+    lengths = seq_lengths.cpu().tolist()
+    device = logps_rmpad.device
+    output = torch.full((len(lengths), logits_to_keep), pad_value, dtype=dtype, device=device)
+    valid_mask = torch.zeros_like(output, dtype=torch.float32)
+    flat = logps_rmpad.flatten().to(dtype)
+    offset = 0
+    for i, seq_len in enumerate(lengths):
+        actual_len = min(max(flat.numel() - offset, 0), seq_len)
+        if actual_len <= 0:
+            offset += seq_len
+            continue
+        pad_len = logits_to_keep - (actual_len if actual_len < seq_len else seq_len)
+        output[i, pad_len:] = flat[offset:offset + actual_len]
+        valid_mask[i, pad_len:] = 1.0
+        offset += seq_len
+    return output, valid_mask
+
+
 def _reference_dpo_sum(values, lengths, num_examples, ld_alpha=None, is_ref_model=False):
     lengths_list = lengths.cpu().tolist()
     public_lengths = [min(lengths_list[i], lengths_list[i + num_examples]) for i in range(num_examples)]
@@ -226,6 +248,56 @@ def test_get_cu_seqlens(self):
                 actual = trainer.get_cu_seqlens(position_ids, 11)
                 torch.testing.assert_close(actual, expected, rtol=0, atol=0)
 
+    def test_pad_logps_back_to_batch(self):
+        cases = [
+            ('normal', [4, 3, 5, 2], 14, 16),
+            ('empty', [3, 0, 2, 5], 10, 8),
+            ('all_empty', [0, 0, 0], 0, 4),
+            ('truncated', [3, 5, 2], 5, 6),
+            ('extra', [3, 2, 1], 10, 5),
+            ('small_fast_path', [1, 2], 3, 4),
+        ]
+        for device in _test_devices():
+            for dtype in (torch.float32, torch.bfloat16):
+                for name, lengths_list, source_tokens, logits_to_keep in cases:
+                    with self.subTest(device=device, dtype=dtype, case=name):
+                        lengths = torch.tensor(lengths_list, dtype=torch.int32, device=device)
+                        logps = torch.arange(source_tokens, dtype=dtype, device=device).reshape(1, -1)
+                        expected = _reference_pad_logps(logps, lengths, logits_to_keep)
+                        actual = pad_logps_back_to_batch(
+                            logps, batch_size=len(lengths_list), seq_lengths=lengths, logits_to_keep=logits_to_keep)
+                        torch.testing.assert_close(actual[0].cpu(), expected[0].cpu())
+                        torch.testing.assert_close(actual[1].cpu(), expected[1].cpu())
+                        self.assertEqual(actual[0].dtype, dtype)
+                        self.assertEqual(actual[0].device, logps.device)
+
+                lengths = torch.tensor([4, 3, 5, 2], dtype=torch.int32, device=device)
+                logps = torch.arange(14, dtype=torch.bfloat16, device=device).reshape(1, -1)
+                expected = _reference_pad_logps(logps, lengths, 16, dtype=torch.float32)
+                actual = pad_logps_back_to_batch(
+                    logps, batch_size=4, seq_lengths=lengths, logits_to_keep=16, dtype=torch.float32)
+                torch.testing.assert_close(actual[0].cpu(), expected[0].cpu())
+                torch.testing.assert_close(actual[1].cpu(), expected[1].cpu())
+
+        position_ids = torch.cat([torch.arange(length) for length in [4, 3, 5, 2]]).unsqueeze(0)
+        logps = torch.arange(14, dtype=torch.float32).reshape(1, -1)
+        lengths = torch.tensor([4, 3, 5, 2], dtype=torch.int32)
+        expected = _reference_pad_logps(logps, lengths, 16)
+        actual = pad_logps_back_to_batch(logps, batch_size=4, position_ids=position_ids, logits_to_keep=16)
+        torch.testing.assert_close(actual[0], expected[0])
+        torch.testing.assert_close(actual[1], expected[1])
+
+        lengths = torch.tensor([4, 3, 5, 2], dtype=torch.int32)
+        logps = torch.randn(1, 14, requires_grad=True)
+        expected_input = logps.detach().clone().requires_grad_(True)
+        actual = pad_logps_back_to_batch(logps, batch_size=4, seq_lengths=lengths, logits_to_keep=8, pad_value=0.0)
+        expected = _reference_pad_logps(expected_input, lengths, 8, pad_value=0.0)
+        (actual[0].square().sum() + actual[1].sum()).backward()
+        (expected[0].square().sum() + expected[1].sum()).backward()
+        torch.testing.assert_close(actual[0], expected[0])
+        torch.testing.assert_close(actual[1], expected[1])
+        torch.testing.assert_close(logps.grad, expected_input.grad, rtol=0, atol=0)
+
     def test_packed_sequence_sum_forward_and_backward(self):
         lengths = [0, 3, 2, 0, 4]
         total_tokens = sum(lengths)