[bugfix] fix infer pt dp

swift/llm/infer/infer.py

@@ -247,7 +247,10 @@ def infer_dataset(self) -> List[Dict[str, Any]]:
             while idx < len(val_dataset):
                 shard_size = min(args.write_batch_size, len(val_dataset) - idx)
                 shard_dataset = val_dataset.select(range(idx, idx + shard_size))
-                result_list += self._batch_infer(shard_dataset, request_config)
+                result = self._batch_infer(shard_dataset, request_config)
+                if self.jsonl_writer:
+                    self.jsonl_writer.append(result, gather_obj=True)
+                result_list += result
                 idx += shard_size
                 prog_bar.update(shard_size)
             prog_bar.close()
@@ -267,6 +270,11 @@ def _batch_infer(self, val_dataset, request_config):
             data_parallel_size = args.global_world_size // args.vllm_tensor_parallel_size
         else:
             rank, data_parallel_size = args.rank, args.global_world_size
+        # The dataset is insufficient for DP partitioning
+        if len(val_dataset) < data_parallel_size:
+            if rank >= len(val_dataset):
+                return []
+            data_parallel_size = len(val_dataset)
         if rank >= 0 and data_parallel_size > 1:
             val_dataset = val_dataset.shard(data_parallel_size, rank, contiguous=True)
         val_dataset = list(val_dataset)
@@ -279,14 +287,13 @@ def _batch_infer(self, val_dataset, request_config):
             labels_list.append(labels)
 
         resp_list = self.infer(val_dataset, request_config, template=self.template, use_tqdm=True, **self.infer_kwargs)
-        if not (args.infer_backend == 'vllm' and rank >= 0 and args.rank % args.vllm_tensor_parallel_size != 0):
+        if not (args.infer_backend == 'vllm' and rank >= 0
+                and args.rank % args.vllm_tensor_parallel_size != 0):  # DP & TP
             for data, resp, labels in zip(val_dataset, resp_list, labels_list):
                 response = resp.choices[0].message.content
                 data['messages'].append({'role': 'assistant', 'content': response})
                 data = {'response': response, 'labels': labels, 'logprobs': resp.choices[0].logprobs, **data}
                 result_list.append(data)
-        if self.jsonl_writer:
-            self.jsonl_writer.append(result_list, gather_obj=True)
         return result_list