Merge branch 'main' of github.com:modelscope/Trinity-RFT into dev/config_manager

Author: chenyushuo

examples/grpo_gsm8k/gsm8k.yaml

@@ -1,8 +1,9 @@
 data:
   # basic info
-  dataset_path: '/PATH/TO/DATASET/'
+  dataset_path: 'openai/gsm8k'
+  subset_name: "main"
   train_split: 'train'
-  eval_split: ''
+  eval_split: 'test'
   format_config:
     prompt_key: 'question'
     response_key: 'answer'
@@ -24,7 +25,7 @@ model:
   model_path: '/PATH/TO/MODEL/'
   max_prompt_tokens: 256
   max_response_tokens: 1024
-  checkpoint_path: '/PATH/TO/CHECKPOINT/'
+  checkpoint_path: ""
 cluster:
   node_num: 1
   gpu_per_node: 8
@@ -34,7 +35,8 @@ buffer:
   train_dataset:
     name: gsm8k_buffer
     storage_type: queue
-    path: 'sqlite:////gsm8k.db'
+    algorithm_type: ppo
+    path: 'sqlite:///gsm8k.db'
   # sft_warmup_dataset: # Uncomment these to enable sft warmup
   #   name: warmup_data
   #   storage_type: file

trinity/cli/launcher.py

@@ -129,7 +129,8 @@ def run(config_path: str):
         data_config.dj_config_path or data_config.dj_process_desc
     ):
         activate_data_module(data_config.data_workflow_url, config_path)
-    ray.init()
+    if not ray.is_initialized():
+        ray.init()
     if config.mode == "explore":
         explore(config)
     elif config.mode == "train":

trinity/common/config.py

@@ -45,6 +45,7 @@ class DataConfig:
 
     dataset_path: str = ""
     train_split: str = "train"
+    subset_name: Optional[str] = None
     eval_split: Optional[str] = None  # TODO: check data format
     format_config: FormatConfig = field(default_factory=FormatConfig)
 

trinity/common/models/vllm_worker.py

@@ -5,7 +5,7 @@
 import torch.distributed
 from vllm.worker.worker import Worker
 
-from trinity.utils.distributed import init_process_group
+from trinity.utils.distributed import init_process_group, is_ipv6_address
 from trinity.utils.log import get_logger
 
 logger = get_logger(__name__)
@@ -43,9 +43,15 @@ def init_process_group(
             )
             self._weight_update_rank = torch.distributed.get_rank() + rank_offset
 
+        if is_ipv6_address(master_address):
+            # using tcp://ipv6:port will lead to ValueError
+            init_method = f"tcp://[{master_address}]:{master_port}"
+        else:
+            init_method = f"tcp://{master_address}:{master_port}"
+
         self._model_update_group = init_process_group(
             backend=backend,
-            init_method=f"tcp://{master_address}:{master_port}",
+            init_method=init_method,
             world_size=world_size,
             rank=self._weight_update_rank,
             group_name=group_name,

trinity/common/task.py

@@ -101,6 +101,15 @@ def task_generator(
         yield task
 
 
+def load_hf_dataset(config: DataConfig, split: str):
+    """Load a Hugging Face dataset with optional configuration name."""
+    if config.subset_name is not None:
+        hf_dataset = load_dataset(config.dataset_path, config.subset_name, split=split)
+    else:
+        hf_dataset = load_dataset(config.dataset_path, split=split)
+    return hf_dataset
+
+
 @dataclass
 class TaskSet:
     """A TaskSet class that defines a set of tasks and their associated reward functions."""
@@ -125,7 +134,8 @@ def load(
         # disable datasets caching to avoid reuse old-version dataset
         datasets.disable_caching()
         if task_type == TaskType.EVAL:
-            dataset = load_dataset(config.dataset_path)[config.eval_split]
+            assert config.eval_split is not None, "eval_split must be provided for eval taskset."
+            dataset = load_hf_dataset(config, config.eval_split)
         else:  # default
             if task_type != TaskType.EVAL and config.db_url != "":
                 logger.info(f"Loading dataset from database with url: {config.db_url}")
@@ -134,7 +144,7 @@ def load(
                 dataset = Dataset.from_sql(RftDatasetModel.__tablename__, f"{db_type}:///{db_name}")
             elif config.dataset_path != "":
                 logger.info(f"Loading dataset from local file with path: {config.dataset_path}.")
-                dataset = load_dataset(config.dataset_path)[config.train_split]
+                dataset = load_hf_dataset(config, config.train_split)
             else:
                 raise ValueError("No dataset path or db url provided.")
         datasets.enable_caching()

trinity/trainer/verl/fsdp_workers.py

@@ -51,7 +51,7 @@
 from verl.workers.sharding_manager.fsdp_ulysses import FSDPUlyssesShardingManager
 
 from trinity.common.constants import AlgorithmType
-from trinity.utils.distributed import init_process_group
+from trinity.utils.distributed import init_process_group, is_ipv6_address
 
 logger = logging.getLogger(__file__)
 logger.setLevel(os.getenv("VERL_PPO_LOGGING_LEVEL", "WARN"))
@@ -592,9 +592,15 @@ def setup_weight_sync_group(self):
                 setup_ref = explorer.setup_weight_sync_group.remote(
                     master_address, master_port, self.state_dict_meta
                 )
+                if is_ipv6_address(master_address):
+                    # using tcp://ipv6:port will lead to ValueError
+                    init_method = f"tcp://[{master_address}]:{master_port}"
+                else:
+                    init_method = f"tcp://{master_address}:{master_port}"
+
                 self._model_update_group = init_process_group(
                     backend=backend,
-                    init_method=f"tcp://{master_address}:{master_port}",
+                    init_method=init_method,
                     world_size=world_size,
                     rank=0,
                     group_name=group_name,

trinity/utils/distributed.py

@@ -1,5 +1,6 @@
 # -*- coding: utf-8 -*-
 """For distributed training with multiple process groups."""
+import ipaddress
 from datetime import timedelta
 from typing import Any, Optional, Union
 
@@ -15,6 +16,14 @@
 )
 
 
+def is_ipv6_address(ip_str: str) -> bool:
+    try:
+        ip = ipaddress.ip_address(ip_str)
+        return isinstance(ip, ipaddress.IPv6Address)
+    except ValueError:
+        return False
+
+
 def init_process_group(
     backend: Union[str, Backend] = None,
     init_method: Optional[str] = None,