[WIP]Support ray

examples/sampler/ray/sample.sh

@@ -0,0 +1 @@
+swift sample --config sampling.yaml

examples/sampler/ray/sampling.yaml

@@ -0,0 +1,34 @@
+ray_exp_name: sampling
+
+use_ray: true
+
+model: Qwen/Qwen2.5-VL-3B-Instruct 
+dataset: tastelikefeet/competition_math#16
+num_return_sequences: 2
+max_length: 2048
+system: "You are a math model, you should **think step by step** carefully, and always consider the basic math principles to avoid making calculating mistakes. Give the final answer wrapped with \\boxed{{}}"
+load_args: false
+sampler_engine: vllm
+max_new_tokens: 768
+orm_model: math
+prm_model: Qwen/Qwen2.5-Math-PRM-7B
+override_exist_file: true
+num_sampling_per_gpu_batch_size: 4
+top_p: 1.0
+temperature: 1.0
+prm_threshold: 0.8
+output_file: sampling.jsonl
+
+device_groups:
+  nproc_per_node: 4
+  sample_group:
+    device: GPU
+    ranks: list(range(0, 2))
+    workers:
+      - sampler
+  rm_group:
+    device: GPU
+    ranks: list(range(2, 4))
+    workers:
+      - prm
+      - orm

requirements/framework.txt

@@ -27,6 +27,7 @@ requests
 rouge
 safetensors
 scipy
+omegaconf
 sentencepiece
 simplejson>=3.3.0
 sortedcontainers>=1.5.9

swift/cli/main.py

@@ -3,7 +3,8 @@
 import os
 import subprocess
 import sys
-from typing import Dict, List, Optional
+import json
+from typing import Dict, List, Optional, Any
 
 from swift.utils import get_logger
 
@@ -45,6 +46,44 @@ def get_torchrun_args() -> Optional[List[str]]:
     return torchrun_args
 
 
+def prepare_config_args(argv):
+    for i in range(0, len(argv[1:]), 2):
+        arg_name = argv[i]
+        arg_value = argv[i + 1]
+        if arg_name == '--config':
+            from omegaconf import OmegaConf, DictConfig
+            from swift.ray import RayHelper
+            config = OmegaConf.load(arg_value)
+
+            def parse_dict_config(cfg: DictConfig) -> Dict[str, Any]:
+                result = {}
+                def _traverse(config: Any, parent_key: str = ""):
+                    if isinstance(config, DictConfig):
+                        for key, value in config.items():
+                            if key == 'device_groups':
+                                result[key] = json.dumps(OmegaConf.to_container(value))
+                            else:
+                                current_path = f"{parent_key}.{key}" if parent_key else key
+                                _traverse(value, current_path)
+                    else:
+                        last_key = parent_key.split('.')[-1] if parent_key else ""
+                        result[last_key] = config
+
+                _traverse(cfg)
+                return result
+
+            cfg = parse_dict_config(config)
+            for key, value in cfg.items():
+                argv.append(f'--{key}')
+                if not isinstance(value, str):
+                    value = str(value)
+                argv.append(value)
+
+            argv.pop(i)
+            argv.pop(i)
+            break
+
+
 def _compat_web_ui(argv):
     # [compat]
     method_name = argv[0]
@@ -56,11 +95,14 @@ def _compat_web_ui(argv):
 def cli_main(route_mapping: Optional[Dict[str, str]] = None) -> None:
     route_mapping = route_mapping or ROUTE_MAPPING
     argv = sys.argv[1:]
+    if 'local-rank' in argv[0]:
+        argv = argv[1:]
     _compat_web_ui(argv)
     method_name = argv[0].replace('_', '-')
     argv = argv[1:]
     file_path = importlib.util.find_spec(route_mapping[method_name]).origin
     torchrun_args = get_torchrun_args()
+    prepare_config_args(argv)
     python_cmd = sys.executable
     if torchrun_args is None or method_name not in {'pt', 'sft', 'rlhf', 'infer'}:
         args = [python_cmd, file_path, *argv]

swift/llm/argument/base_args/base_args.py

@@ -16,6 +16,7 @@
 from .model_args import ModelArguments
 from .quant_args import QuantizeArguments
 from .template_args import TemplateArguments
+from .ray_args import RayArguments
 
 logger = get_logger()
 
@@ -52,7 +53,7 @@ def __post_init__(self: 'BaseArguments'):
 
 @dataclass
 class BaseArguments(CompatArguments, GenerationArguments, QuantizeArguments, DataArguments, TemplateArguments,
-                    ModelArguments):
+                    ModelArguments, RayArguments):
     """
     BaseArguments class is a dataclass that inherits from multiple argument classes:
     GenerationArguments, QuantizeArguments, DataArguments, TemplateArguments, ModelArguments.
@@ -173,6 +174,7 @@ def __post_init__(self):
         QuantizeArguments.__post_init__(self)
         TemplateArguments.__post_init__(self)
         DataArguments.__post_init__(self)
+        RayArguments.__post_init__(self)
         if self.max_length is None and self.model_info is not None:
             self.max_length = self.model_info.max_model_len
         if self.packing and self.packing_length is None:

swift/llm/argument/base_args/ray_args.py

@@ -0,0 +1,17 @@
+import json
+from dataclasses import dataclass
+from typing import Optional
+
+
+@dataclass
+class RayArguments:
+
+    use_ray: bool = False
+
+    ray_exp_name: Optional[str] = None
+
+    device_groups: Optional[str] = None
+
+    def __post_init__(self):
+        if isinstance(self.device_groups, str):
+            self.device_groups = json.loads(self.device_groups)

swift/llm/base.py

@@ -8,6 +8,7 @@
 from swift.utils import get_logger, parse_args, seed_everything
 from .argument import BaseArguments
 from .utils import ProcessorMixin
+from swift.ray.base import RayHelper
 
 logger = get_logger()
 
@@ -18,6 +19,9 @@ class SwiftPipeline(ABC, ProcessorMixin):
     def __init__(self, args: Optional[Union[List[str], args_class]] = None):
         self.args = self._parse_args(args)
         args = self.args
+        if self.args.use_ray:
+            from swift.ray import RayHelper
+            RayHelper.initialize(self.args.device_groups)
         if hasattr(args, 'seed'):
             seed = args.seed + max(getattr(args, 'rank', -1), 0)
             seed_everything(seed)

swift/llm/sampling/base.py

@@ -2,6 +2,7 @@
 
 from swift.llm import SamplingArguments
 from swift.plugin import orms, prms
+from swift.ray.base import RayHelper
 from swift.utils import get_logger
 
 logger = get_logger()
@@ -17,13 +18,15 @@ def __init__(self, input_args: SamplingArguments):
         self.orm_model = None
         self._prepare_model_tokenizer()
         self._prepare_template()
-        self._prepare_rm()
+        self._prepare_prm()
+        self._prepare_orm()
 
     def _prepare_model_tokenizer(self):
         args = self.args
         _, self.processor = args.get_model_processor(load_model=False)
 
-    def _prepare_rm(self):
+    @RayHelper.function(group='prm')
+    def _prepare_prm(self):
         if self.args.prm_model is None:
             self.prm_model = None
             logger.warning('prm_model is None.')
@@ -33,6 +36,8 @@ def _prepare_rm(self):
             from swift.llm import PtEngine
             self.prm_model = PtEngine(self.args.prm_model, max_batch_size=64)
 
+    @RayHelper.function(group='orm')
+    def _prepare_orm(self):
         if self.args.orm_model is None:
             self.orm_model = None
             logger.warning('orm_model is None.')

swift/llm/sampling/vanilla_sampler.py

@@ -1,25 +1,29 @@
 # Copyright (c) Alibaba, Inc. and its affiliates.
+import json
 import os
 from copy import deepcopy
-from typing import Any, Dict, List
 
-import json
 import numpy as np
 
 from swift.llm import RequestConfig
 from swift.llm.sampling.base import Sampler
-from swift.llm.template.template_inputs import InferRequest
+from swift.ray.base import RayHelper
 from swift.utils import get_logger
 from .utils import get_messages_md5, get_reward
 
 logger = get_logger()
 
 
+@RayHelper.worker(group=['sampler', 'prm', 'orm'])
 class VanillaSampler(Sampler):
 
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
+        self.prepare_sampler()
+        self.caches = self.read_cache()
 
+    @RayHelper.function(group='sampler')
+    def prepare_sampler(self):
         if self.args.sampler_engine == 'pt':
             from swift.llm import PtEngine
             _Engine = PtEngine
@@ -38,8 +42,8 @@ def __init__(self, *args, **kwargs):
             self.infer_engine = _Engine(
                 self.args.model, model_type=self.args.model_type, template=self.template, **self.args.engine_kwargs)
             self.infer_engine.strict = False
-        self.caches = self.read_cache()
 
+    @RayHelper.function(group='sampler')
     def read_cache(self):
         cache_files = self.args.cache_files
         caches = {}
@@ -82,6 +86,7 @@ def check_row_valid(rows):
             assert not row.get('videos') or all([isinstance(video, str) and video for video in row['videos']])
             assert not row.get('audios') or all([isinstance(audio, str) and audio for audio in row['audios']])
 
+    @RayHelper.function(group='sampler', dispatch=lambda n, i, data: ([{'messages': data['messages'][i * len(data['messages']) // n : (i + 1) * len(data['messages']) // n]}], {}), collect='flatten')
     def generate(self, data):
         resp_all = []
         infer_requests = []
@@ -141,6 +146,20 @@ def generate(self, data):
             _cur += 1
         return resp_all
 
+    @RayHelper.function(group='orm', execute='first')
+    def get_orm_score(self, infer_requests, ground_truth):
+        return get_reward(
+            self.orm_model, infer_requests, ground_truths=[ground_truth] * len(infer_requests),
+            threshold=0.0)
+
+    @RayHelper.function(group='prm', execute='first')
+    def get_prm_score(self, infer_requests, ground_truth):
+        return get_reward(
+            self.prm_model,
+            infer_requests,
+            ground_truths=[ground_truth] * len(infer_requests),
+            threshold=self.args.prm_threshold)
+
     def do_sample(self, data):
         generated = []
         resp_all = self.generate(data)
@@ -160,18 +179,13 @@ def do_sample(self, data):
             _resps = deepcopy(resps)
             _resps['messages'][-1]['content'] = ground_truth
             infer_requests.append(_resps)
-            if self.orm_model is not None:
-                orm_score, _orm_mask = get_reward(
-                    self.orm_model, infer_requests, ground_truths=[ground_truth] * len(infer_requests), threshold=0.0)
+            if self.args.orm_model is not None:
+                orm_score, _orm_mask = self.get_orm_score(infer_requests, ground_truth)
             else:
                 orm_score = np.array([1.0] * len(infer_requests))
                 _orm_mask = np.array([True] * len(infer_requests))
-            if self.prm_model is not None:
-                prm_score, _prm_mask = get_reward(
-                    self.prm_model,
-                    infer_requests,
-                    ground_truths=[ground_truth] * len(infer_requests),
-                    threshold=self.args.prm_threshold)
+            if self.args.prm_model is not None:
+                prm_score, _prm_mask = self.get_prm_score(infer_requests, ground_truth)
             else:
                 prm_score = np.array([1.0] * len(infer_requests))
                 _prm_mask = np.array([True] * len(infer_requests))

swift/ray/__init__.py

@@ -0,0 +1 @@
+from .base import RayHelper