move chunk_list

Author: hjh0119

swift/pipelines/infer/rollout.py

@@ -40,8 +40,8 @@
                                          UpdateWeightsRequest)
 from swift.rlhf_trainers.utils import (FlattenedTensorBucket, FlattenedTensorMetadata, TensorLoRARequest,
                                        UpdateAdapterRequest, UpdateFlattenedAdapterRequest,
-                                       UpdateFlattenedParamsRequest, check_vllm_version_ge, patch_vllm_load_adapter,
-                                       patch_vllm_moe_model_weight_loader)
+                                       UpdateFlattenedParamsRequest, check_vllm_version_ge, chunk_list,
+                                       patch_vllm_load_adapter, patch_vllm_moe_model_weight_loader)
 from swift.rollout import RolloutScheduler, multi_turns
 from swift.utils import get_logger, get_seed, is_vllm_ascend_available
 from ..base import SwiftPipeline
@@ -57,8 +57,6 @@
     if is_vllm_ascend_available():
         from vllm_ascend.distributed.device_communicators.pyhccl import PyHcclCommunicator as PyNcclCommunicator  # noqa
 
-    from trl.scripts.vllm_serve import chunk_list
-
 except ImportError:
     pass
 """

swift/rlhf_trainers/utils.py

@@ -57,6 +57,26 @@ def embeddings(self):
             return self.lora_embeddings
 
 
+def chunk_list(lst: list, n: int) -> list[list]:
+    """
+    Split list `lst` into `n` evenly distributed sublists.
+
+    Example:
+    ```python
+    >>> chunk_list([1, 2, 3, 4, 5, 6], 2)
+    [[1, 2, 3], [4, 5, 6]]
+
+    >>> chunk_list([1, 2, 3, 4, 5, 6], 4)
+    [[1, 2], [3, 4], [5], [6]]
+
+    >>> chunk_list([1, 2, 3, 4, 5, 6], 8)
+    [[1], [2], [3], [4], [5], [6], [], []]
+    ```
+    """
+    k, r = divmod(len(lst), n)
+    return [lst[i * k + min(i, r):(i + 1) * k + min(i + 1, r)] for i in range(n)]
+
+
 def is_valid_ipv6_address(address: str) -> bool:
     """Check if the given address is a valid IPv6 address."""
     try: