From d6b21f2afaa1c280837ee4b115dece3859b5283d Mon Sep 17 00:00:00 2001
From: RerankerGuo <121015044+RerankerGuo@users.noreply.github.com>
Date: Fri, 7 Aug 2026 13:07:15 +0800
Subject: [PATCH] fix(sampling): preserve remainder samples across shards

Distribute dataset remainders across data_range shards so parallel sampling covers every input exactly once, including datasets smaller than the shard count.

Test: .venv/bin/python tests/run.py --test_dir tests/sample --pattern test_sampling.py
---
 swift/pipelines/sampling/sampling.py |  6 ++--
 tests/sample/test_sampling.py        | 47 ++++++++++++++++++++++++++++
 2 files changed, 51 insertions(+), 2 deletions(-)
 create mode 100644 tests/sample/test_sampling.py

diff --git a/swift/pipelines/sampling/sampling.py b/swift/pipelines/sampling/sampling.py
index 6ec443acd0..d3fbfee667 100644
--- a/swift/pipelines/sampling/sampling.py
+++ b/swift/pipelines/sampling/sampling.py
@@ -43,8 +43,10 @@ def _get_dataset(self):
             args.dataset, split_dataset_ratio=0., shuffle=args.dataset_shuffle, **dataset_kwargs)
         logger.info(f'Sampling_dataset: {sampling_dataset}')
         dataset_len = len(sampling_dataset)
-        piece_len = dataset_len // self.total_piece
-        sampling_dataset = sampling_dataset.select(range(piece_len * self.cur_piece, piece_len * (self.cur_piece + 1)))
+        piece_len, remainder = divmod(dataset_len, self.total_piece)
+        start = piece_len * self.cur_piece + min(self.cur_piece, remainder)
+        end = start + piece_len + int(self.cur_piece < remainder)
+        sampling_dataset = sampling_dataset.select(range(start, end))
         return sampling_dataset
 
     def run(self):
diff --git a/tests/sample/test_sampling.py b/tests/sample/test_sampling.py
new file mode 100644
index 0000000000..3cebce25df
--- /dev/null
+++ b/tests/sample/test_sampling.py
@@ -0,0 +1,47 @@
+import unittest
+from types import SimpleNamespace
+from unittest.mock import patch
+
+from swift.pipelines.sampling.sampling import SwiftSampling
+
+
+class FakeDataset:
+
+    def __init__(self, values):
+        self.values = values
+
+    def __len__(self):
+        return len(self.values)
+
+    def select(self, indices):
+        return FakeDataset([self.values[i] for i in indices])
+
+
+class TestSampling(unittest.TestCase):
+
+    @patch('swift.pipelines.sampling.sampling.load_dataset')
+    def test_data_range_partitions_cover_dataset(self, mock_load_dataset):
+        args = SimpleNamespace(
+            dataset=['test-dataset'],
+            dataset_shuffle=False,
+            get_dataset_kwargs=lambda: {},
+        )
+
+        for dataset_size, expected_sizes in [(10, [4, 3, 3]), (2, [1, 1, 0])]:
+            with self.subTest(dataset_size=dataset_size):
+                dataset = FakeDataset(list(range(dataset_size)))
+                mock_load_dataset.return_value = dataset, None
+                shards = []
+                for shard_index in range(3):
+                    sampling = SwiftSampling.__new__(SwiftSampling)
+                    sampling.args = args
+                    sampling.cur_piece = shard_index
+                    sampling.total_piece = 3
+                    shards.append(sampling._get_dataset().values)
+
+                self.assertEqual([len(shard) for shard in shards], expected_sizes)
+                self.assertEqual([item for shard in shards for item in shard], list(range(dataset_size)))
+
+
+if __name__ == '__main__':
+    unittest.main()