Add --rollout-sample-filter-path (#961)

Author: zhuzilin

slime/rollout/sglang_rollout.py

@@ -410,6 +410,10 @@ async def generate_rollout_async(
 
     # reset the global state to prevent effects on the next rollout or eval.
     state.reset()
+    if args.rollout_sample_filter_path is not None:
+        filter_func = load_function(args.rollout_sample_filter_path)
+        filter_func(args, data)
+
     return RolloutFnTrainOutput(samples=data, metrics=metric_gatherer.collect()), aborted_samples
 
 

slime/utils/arguments.py

@@ -1094,6 +1094,18 @@ def add_rollout_buffer_arguments(parser):
                 default=128,
                 help="Multiplier for data padding size in data processing.",
             )
+            parser.add_argument(
+                "--rollout-sample-filter-path",
+                type=str,
+                default=None,
+                help=(
+                    "Path to the rollout sample filter function. "
+                    "This function determines whether a sample will participate in loss calculation. "
+                    "The function should take args and samples (list[Sample]) as input, and return None. "
+                    "Please directly modify the remove_sample attribute of Sample. "
+                    "Note: This attribute does not determine whether the sample participates in advantage normalization."
+                ),
+            )
             return parser
 
         def add_custom_megatron_plugins_arguments(parser):