Merge branch 'main' into dev/fix_microbatch_loss_scale

Author: yanxi-chen

.github/workflows/unittest.yaml

@@ -97,6 +97,13 @@ jobs:
           fi
         fi
 
+    - name: Clean checkpoint dir
+      working-directory: trinity-${{ github.run_id }}/.github/workflows/docker
+      if: always()
+      run: |
+        docker compose exec trinity-node-1 rm -rf /mnt/checkpoints/*
+      continue-on-error: true
+
     - name: Upload test results
       if: env.tests_run == 'true' || failure()
       uses: actions/upload-artifact@v4

benchmark/config/countdown-template.yaml

@@ -35,19 +35,18 @@ buffer:
       rollout_args:
         temperature: 1.0
         logprobs: 0
+      default_workflow_type: math_workflow
+      default_reward_fn_type: countdown_reward
     eval_tasksets: []
-    default_workflow_type: math_workflow
-    default_reward_fn_type: countdown_reward
-    system_prompt: null
-    reply_prefix: null
   trainer_input:
     experience_buffer:
       name: experience_buffer
       storage_type: queue
-      use_priority_queue: true
-      replay_buffer_kwargs:
+      replay_buffer:
+        enable: true
         priority_fn: linear_decay
-        decay: 0.1
+        priority_fn_args:
+          decay: 0.1
 explorer:
   runner_per_model: 8
   max_timeout: 900

benchmark/config/gsm8k-template.yaml

@@ -40,19 +40,18 @@ buffer:
       rollout_args:
         temperature: 1.0
         logprobs: 0
+      default_workflow_type: math_workflow
+      default_reward_fn_type: math_reward
     eval_tasksets: []
-    default_workflow_type: math_workflow
-    default_reward_fn_type: math_reward
-    system_prompt: null
-    reply_prefix: null
   trainer_input:
     experience_buffer:
       name: experience_buffer
       storage_type: queue
-      use_priority_queue: true
-      replay_buffer_kwargs:
+      replay_buffer:
+        enable: true
         priority_fn: linear_decay
-        decay: 0.1
+        priority_fn_args:
+          decay: 0.1
 explorer:
   runner_per_model: 8
   max_timeout: 900
@@ -79,7 +78,7 @@ trainer:
   enable_preview: true
   grad_clip: 1.0
   use_dynamic_bsz: true
-  ppo_max_token_len_per_gpu: 10240
+  max_token_len_per_gpu: 10240
   ulysses_sequence_parallel_size: 1
 monitor:
   monitor_type: wandb

docs/sphinx_doc/source/index.rst

@@ -21,6 +21,7 @@ Welcome to Trinity-RFT's documentation!
    tutorial/develop_algorithm.md
    tutorial/example_mix_algo.md
    tutorial/develop_operator.md
+   tutorial/develop_selector.md
    tutorial/trinity_configs.md
    tutorial/synchronizer.md
 

docs/sphinx_doc/source/tutorial/develop_operator.md

@@ -1,13 +1,14 @@
+(Operators)=
 ## Operator Development Guide
 
 ### Step 0: Basic Concepts of Operator Module
 
 In Trinity-RFT, the operator module is responsible for processing experience data in the buffer module. It supports existing data processing capabilities from [Data-Juicer](https://github.com/modelscope/data-juicer) naturally, and allows developers to implement their own operators as well.
 By customizing operators, developers can implement various data processing functionalities, such as data augmentation, filtering, and transformation. You can even implement advantages/returns calculation as operators, as shown in {ref}`Algorithms <Algorithms>` section.
 
-- **DataJuicerOperator** ({class}`trinity.data.operators.DataJuicerOperator`): The operator that wraps the data processing operators from Data-Juicer. It provides a simple interface for developers to list the Data-Juicer operators they want to use. The full list of Data-Juicer operators can be found [here](https://modelscope.github.io/data-juicer/en/main/docs/Operators.html).
-- **ExperienceOperator** ({class}`trinity.data.operators.ExperienceOperator`): The base class for all operators used in experience data processing. It defines the interface and common functionalities that all operators should have. Each operator processes a batch of experience data and returns the processed data with metrics for logging.
-- **ExperiencePipeline** ({class}`trinity.data.pipelines.ExperiencePipeline`): The experience data processing pipeline that manages a sequence of operators. It takes raw experiences from the `Explorer`, passes them through each operator in the pipeline, and writes the final processed experiences into the input buffer of the `Trainer`.
+- **DataJuicerOperator** ({class}`trinity.buffer.operators.DataJuicerOperator`): The operator that wraps the data processing operators from Data-Juicer. It provides a simple interface for developers to list the Data-Juicer operators they want to use. The full list of Data-Juicer operators can be found [here](https://modelscope.github.io/data-juicer/en/main/docs/Operators.html).
+- **ExperienceOperator** ({class}`trinity.buffer.operators.ExperienceOperator`): The base class for all operators used in experience data processing. It defines the interface and common functionalities that all operators should have. Each operator processes a batch of experience data and returns the processed data with metrics for logging.
+- **ExperiencePipeline** ({class}`trinity.buffer.pipelines.ExperiencePipeline`): The experience data processing pipeline that manages a sequence of operators. It takes raw experiences from the `Explorer`, passes them through each operator in the pipeline, and writes the final processed experiences into the input buffer of the `Trainer`.
 
 ```{note}
 Except for `ExperiencePipeline`, Trinity-RFT also provides `TaskPipeline` for task data processing.
@@ -55,7 +56,7 @@ class RewardFilter(ExperienceOperator):
         return filtered_exps, metrics
 ```
 
-After implementation, you need to register this module through {class}`trinity.data.operators.EXPERIENCE_OPERATORS`. Once registered, the module can be configured in the configuration file using the registered name.
+After implementation, you need to register this module through {class}`trinity.buffer.operators.EXPERIENCE_OPERATORS`. Once registered, the module can be configured in the configuration file using the registered name.
 
 ### Step 2: Use Your Operator
 

docs/sphinx_doc/source/tutorial/develop_overview.md

@@ -11,6 +11,7 @@ The table below lists the main functions of each extension interface, its target
 | `Workflow`          | Agent Application Developers | Enhance agent's ability to complete tasks in a specified environment | [🔗](./develop_workflow.md) |
 | `Algorithm`         | RL Algorithm Researchers | Design new RL algorithms                 | [🔗](./develop_algorithm.md) |
 | `Operator`          | Data Engineers    | Design new data cleaning and augmentation strategies | [🔗](./develop_operator.md) |
+| `Selector`          | Data Engineers    | Design new task selection strategies           | [🔗](./develop_selector.md) |
 
 ```{tip}
 Trinity-RFT provides a modular development approach, allowing you to flexibly add custom modules without modifying the framework code.

docs/sphinx_doc/source/tutorial/develop_selector.md

@@ -0,0 +1,263 @@
+# 🧪 Experimental: Task Selection & Scheduling System
+
+```{note}
+This module is currently in **experimental status**. Interfaces may change in future versions.
+This document describes the functionality and intended usage of the system.
+```
+
+
+
+## Overview
+
+This system enables **intelligent, adaptive task sampling** from multiple datasets (called *tasksets*) during exploration. It consists of two core components:
+
+1. **`Selector`** – Controls how individual samples are selected *within* each taskset.
+2. **`TasksetScheduler`** – Manages *which* tasksets contribute to each batch and coordinates their sampling.
+
+Together, they support advanced training strategies such as:
+- Curriculum learning (easy → hard)
+- Multi-task interleaving or mixing
+- Difficulty-aware sampling
+- Adaptive data selection based on model performance
+
+These capabilities allow you to train models more efficiently by focusing on informative or challenging examples.
+
+
+
+## Module 1: Selector – Customizable Data Selection
+
+A `Selector` determines **which tasks (samples) to select** from its associated dataset (`Taskset`). Beyond basic strategies like sequential or random access, it supports **adaptive algorithms** that adjust sampling based on feedback—such as sample difficulty, model confidence, or reward signals.
+
+### Built-in Selectors
+
+| Selector Type | Description |
+|---------------|-------------|
+| `sequential` | Returns samples in fixed order (0, 1, ..., N). |
+| `shuffle` | Shuffles the dataset once per epoch; then iterates sequentially. |
+| `random` | Randomly samples without replacement within each batch. Independent across batches. |
+| `offline_easy2hard` | Sorts samples by pre-defined features (e.g., loss, length), serving easier ones first, progressing to harder ones. |
+| `difficulty_based` *(custom example)* | Dynamically selects samples near a target difficulty level using probabilistic modeling. |
+
+You can also **implement your own custom selector** to enable adaptive or curriculum-based learning.
+
+
+
+### ✅ Step 1: Implement a Custom Selector
+
+To create a new selector, inherit from `BaseSelector` and implement the following methods:
+
+#### Required Methods
+
+| Method | Purpose |
+|-------|--------|
+| `get_indices(batch_size: int, return_extra_info=False) -> List[int]` | Return a list of sample indices to read next. |
+| `update(indices: List[int], values: List[float])` | Update internal state using feedback (e.g., rewards, losses). |
+| `state_dict() -> Dict` | Serialize current state for checkpointing. |
+| `load_state_dict(state_dict: Dict)` | Restore state from a saved dictionary. |
+
+#### Example: `DifficultyBasedSelector`
+
+This selector focuses on samples whose predicted performance is closest to a target (e.g., 90% success rate), effectively choosing "just right" difficulty tasks.
+
+```python
+@SELECTORS.register_module("difficulty_based")
+class DifficultyBasedSelector(BaseSelector):
+    def __init__(self, data_source, config: TaskSelectorConfig) -> None:
+        super().__init__(data_source, config)
+        self.logger = get_logger("difficulty_based_selector")
+
+        # Build difficulty estimator using two input features (e.g., correctness, uncertainty)
+        self.diff_estimator = self.build_diff_estimator(
+            data_source.dataset, config.feature_keys, config.kwargs
+        )
+        self.current_index = 0
+        self.seed = config.seed
+
+        # Configuration parameters
+        self.do_sample = config.kwargs.get("do_sample", False)
+        self.target_reward = config.kwargs.get("target_reward", 1.0)
+        self.tau = config.kwargs.get("tau", 1.0)
+
+    # ... detailed implementation
+
+    def get_indices(self, batch_size, return_extra_info=False):
+        # Compute scores based on proximity to target reward
+        sampling_scores = self.get_scores()
+        sampling_scores = torch.from_numpy(sampling_scores)
+
+        if self.tau == 0:
+            # Greedy: take top-k highest scoring samples
+            selected_indices = torch.topk(sampling_scores, batch_size).indices
+        else:
+            # Stochastic: sample via softmax with temperature scaling
+            sampling_logits = sampling_scores / self.tau
+            sampling_logits -= sampling_logits.max()  # Stability
+            sampling_probabilities = torch.softmax(sampling_logits, dim=0)
+            rng = torch.Generator().manual_seed(self.seed + self.current_index)
+            selected_indices = torch.multinomial(
+                sampling_probabilities,
+                batch_size,
+                replacement=False,
+                generator=rng,
+            )
+
+        self.current_index += batch_size
+
+        if return_extra_info:
+            # Optional debugging info
+            extra_info = {
+                "indices": selected_indices.tolist(),
+                "scores": sampling_scores[selected_indices].tolist(),
+                # ... other metadata
+            }
+            return selected_indices, extra_info
+        else:
+            return selected_indices
+
+    def update(self, indices: List[int], values: List[float]) -> None:
+        # Update difficulty model with observed rewards
+        self.diff_estimator.update(indices, values)
+
+    def state_dict(self) -> Dict:
+        return {"current_index": self.current_index}
+
+    def load_state_dict(self, state_dict: Dict) -> None:
+        self.current_index = state_dict.get("current_index", 0)
+```
+
+> 🔁 After defining your class, use `@SELECTORS.register_module("your_name")` so it can be referenced by name in configs.
+
+
+
+### ✅ Step 2: Implement a Feedback Operator
+
+For adaptive selectors like `DifficultyBasedSelector`, you need to provide runtime feedback (e.g., task rewards). This is done via an **Experience Operator** that processes rollouts and computes metrics.
+
+> 📚 See the {ref}`Operator Development Guide<Operators>` for more on building custom experience processors.
+
+The operator must output a metric under the key `trinity.common.constants.SELECTOR_METRIC`, structured as:
+
+```python
+{
+    SELECTOR_METRIC: {
+        0: {  # taskset_id
+            "indices": [10, 25, 43],
+            "values": [0.8, 0.6, 0.9]  # e.g., average reward
+        },
+        1: { ... }
+    }
+}
+```
+
+#### Example: Pass Rate Calculator
+
+```python
+@EXPERIENCE_OPERATORS.register_module("pass_rate_calculator")
+class PassRateCalculator(ExperienceOperator):
+    def __init__(self, **kwargs):
+        pass
+
+    def process(self, exps: List[Experience]) -> Tuple[List[Experience], Dict]:
+        raw_metric = defaultdict(lambda: defaultdict(list))
+
+        for exp in exps:
+            task_index = exp.info["task_index"]
+            assert "taskset_id" in task_index and "index" in task_index
+            raw_metric[task_index["taskset_id"]][task_index["index"]].append(exp.reward)
+
+        metric = {}
+        for taskset_id, task_metrics in raw_metric.items():
+            indices = []
+            reward_means = []
+            for idx, rewards in task_metrics.items():
+                indices.append(idx)
+                reward_means.append(float(np.mean(rewards)))
+            metric[taskset_id] = {
+                "indices": indices,
+                "values": reward_means,
+            }
+
+        return exps, {SELECTOR_METRIC: metric}
+```
+
+This operator calculates the average reward per task and passes it back to the corresponding selector for updating difficulty estimates.
+
+
+
+### ✅ Step 3: Update Configuration
+
+After implementing your selector and operator, register them in the config file.
+
+#### Add the Operator to the Pipeline
+
+```yaml
+data_processor:
+  experience_pipeline:
+    operators:
+      - name: pass_rate_calculator  # Must match @register_module name
+```
+
+#### Configure the Taskset with Your Selector
+
+```yaml
+buffer:
+  explorer_input:
+    tasksets:
+      - name: my_taskset
+        storage_type: file
+        path: ./path/to/tasks
+        task_selector:
+          selector_type: difficulty_based   # Matches @register_module name
+          feature_keys: ["correct", "uncertainty"]
+          kwargs:
+            m: 16
+            lamb: 0.2
+            rho: 0.2
+            target_reward: 0.9
+            tau: 0.5
+            do_sample: true
+```
+
+> 💡 You can define multiple tasksets, each with its own selector type and configuration.
+
+
+
+## Module 2: TasksetScheduler – Multi-Taskset Orchestration
+
+The `TasksetScheduler` manages **how different tasksets are interleaved or mixed** during training.
+
+### Key Features
+
+- Supports **multiple tasksets** simultaneously.
+- Balances sampling proportionally to dataset sizes.
+- **Shuffles taskset access order** at the start of each epoch.
+- Enables **curriculum-style** or **interleaved multi-task training**.
+- Fully **checkpointable**: resumes exactly where it left off.
+- Integrates with any registered `Selector`.
+
+### How It Works
+
+At each training step:
+1. Determines which tasksets should contribute to the current batch.
+2. Queries each taskset’s selector to get specific sample indices.
+3. Reads the actual data asynchronously.
+4. Tags each task with `"taskset_id"` for downstream routing or analysis.
+
+Epochs are defined based on total data volume and batch size:
+```python
+steps_per_epoch = total_samples // batch_size
+```
+
+At the beginning of each epoch, the scheduler reshuffles the sequence of taskset accesses to introduce variability.
+
+
+
+## Summary
+
+With these components, you can:
+- Use simple strategies like random or sequential sampling.
+- Design **adaptive curricula** using custom selectors.
+- Combine multiple datasets intelligently.
+- Optimize training efficiency by focusing on high-value samples.
+
+By combining smart `Selectors` with the flexible `TasksetScheduler`, you gain fine-grained control over what your model sees—and when.