Feat: multi-threads data fetching for sft data

slime/ray/placement_group.py

@@ -1,5 +1,6 @@
 import logging
 import socket
+import os
 
 import ray
 from ray.util.placement_group import placement_group
@@ -10,6 +11,7 @@
 
 logger = logging.getLogger(__name__)
 
+ROLLOUT_NUM_CPUS = os.environ.get("ROLLOUT_NUM_CPUS", 8)
 
 @ray.remote(num_gpus=1)
 class InfoActor:
@@ -168,7 +170,7 @@ def create_training_models(args, pgs, rollout_manager):
 
 def create_rollout_manager(args, pg):
     rollout_manager = RolloutManager.options(
-        num_cpus=1,
+        num_cpus=ROLLOUT_NUM_CPUS,
         num_gpus=0,
     ).remote(args, pg)
 

slime/rollout/sft_rollout.py

@@ -1,4 +1,7 @@
 import logging
+import concurrent.futures
+import time
+import os
 
 from slime.utils.mask_utils import MultiTurnLossMaskGenerator
 from slime.utils.processing_utils import load_processor, load_tokenizer
@@ -7,11 +10,25 @@
 
 logger = logging.getLogger(__name__)
 
-
+ROLLOUT_NUM_CPUS = os.environ.get("ROLLOUT_NUM_CPUS", 8)
 TOKENIZER = None
 PROCESSOR = None
 MASK_GENERATOR = None
 SAMPLE_PRINTED = False
+_EXECUTOR = None
+
+
+def _init_worker(hf_checkpoint, loss_mask_type):
+    global TOKENIZER, PROCESSOR, MASK_GENERATOR
+    TOKENIZER = load_tokenizer(hf_checkpoint, trust_remote_code=True)
+    PROCESSOR = load_processor(hf_checkpoint, trust_remote_code=True)
+    MASK_GENERATOR = MultiTurnLossMaskGenerator(TOKENIZER, tokenizer_type=loss_mask_type)
+
+
+def _process_single_sample(messages, tools):
+    token_ids, loss_mask = MASK_GENERATOR.get_loss_mask(messages, tools=tools)
+    response_length = MASK_GENERATOR.get_response_lengths([loss_mask])[0]
+    return token_ids, loss_mask, response_length
 
 
 def generate_rollout(args, rollout_id, data_buffer, evaluation=False):
@@ -29,26 +46,28 @@ def generate_rollout(args, rollout_id, data_buffer, evaluation=False):
     assert not evaluation
     assert args.rollout_global_dataset
 
-    global TOKENIZER, PROCESSOR, MASK_GENERATOR, SAMPLE_PRINTED
-    if TOKENIZER is None:
-        TOKENIZER = load_tokenizer(args.hf_checkpoint, trust_remote_code=True)
-
-    if PROCESSOR is None:
-        PROCESSOR = load_processor(args.hf_checkpoint, trust_remote_code=True)
-
-    if MASK_GENERATOR is None:
-        MASK_GENERATOR = MultiTurnLossMaskGenerator(TOKENIZER, tokenizer_type=args.loss_mask_type)
-
+    global TOKENIZER, PROCESSOR, MASK_GENERATOR, SAMPLE_PRINTED, _EXECUTOR
+
     samples = data_buffer.get_samples(args.rollout_batch_size)
+
+
+    if _EXECUTOR is None:
+        _EXECUTOR = concurrent.futures.ProcessPoolExecutor(
+            max_workers=ROLLOUT_NUM_CPUS,
+            initializer=_init_worker,
+            initargs=(args.hf_checkpoint, args.loss_mask_type)
+        )
 
-    for i, sample in enumerate(samples):
-        (sample,) = sample
-        messages = sample.prompt
-        tools = sample.metadata.get("tools", None)
+    tasks = []
+    for sample in samples:
+        (s,) = sample
+        tasks.append((s.prompt, s.metadata.get("tools", None)))
 
-        token_ids, loss_mask = MASK_GENERATOR.get_loss_mask(messages, tools=tools)
+    results = list(_EXECUTOR.map(lambda x: _process_single_sample(*x), tasks))
 
-        response_length = MASK_GENERATOR.get_response_lengths([loss_mask])[0]
+    for i, (sample_wrapper, result) in enumerate(zip(samples, results)):
+        (sample,) = sample_wrapper
+        token_ids, loss_mask, response_length = result
 
         sample.tokens = token_ids
         sample.response_length = response_length

slime/utils/arguments.py

@@ -1224,7 +1224,7 @@ def add_rollout_buffer_arguments(parser):
                 "--loss-mask-type",
                 type=str,
                 default="qwen",
-                choices=["qwen", "qwen3", "distill_qwen"],
+                choices=["qwen", "qwen3"],
                 help="Loss mask type",
             )
             parser.add_argument(

slime/utils/mask_utils.py

@@ -108,34 +108,11 @@ def gen_multi_turn_loss_mask_qwen3(
 
         return all_token_ids, all_loss_masks
 
-    def gen_multi_turn_loss_mask_distill_qwen(
-        self, messages: list[dict], tools: list[dict] = None
-    ) -> tuple[list[int], list[int]]:
-        prompt = self.tokenizer.apply_chat_template(
-            messages[:1], tokenize=False, add_generation_prompt=True, tools=tools
-        )
-        response = messages[-1]["content"]
-        prompt_tokens = self.tokenizer(prompt, add_special_tokens=False)["input_ids"]
-        response_tokens = self.tokenizer(response, add_special_tokens=False)["input_ids"]
-
-        response_length = len(response_tokens)
-        token_ids = prompt_tokens + response_tokens
-        loss_mask = [0] * len(prompt_tokens) + [1] * response_length
-
-        if messages[-1].get("step_loss_mask", 1) != 1:
-            loss_mask = [0] * len(token_ids)
-        return token_ids, loss_mask
-
     def get_loss_mask(self, messages: list[dict], tools: list[dict] = None) -> tuple[list[int], list[int]]:
         if self.tokenizer_type == "qwen":
-            if "<｜Assistant｜>" in self.tokenizer.get_added_vocab():
-                return self.gen_multi_turn_loss_mask_distill_qwen(messages, tools)
-
             return self.gen_multi_turn_loss_mask_qwen(messages, tools)
         elif self.tokenizer_type == "qwen3":
             return self.gen_multi_turn_loss_mask_qwen3(messages, tools)
-        elif self.tokenizer_type == "distill_qwen":
-            return self.gen_multi_turn_loss_mask_distill_qwen(messages, tools)
         else:
             raise ValueError(f"Unsupported tokenizer type: {self.tokenizer_type}")