actual run and add on-policy distillation example in GSM8K

Author: garyzhang99

docs/sphinx_doc/assets/opd_acc.png

@@ -0,0 +1,36 @@
+# Example: On-Policy Distillation on GSM8K dataset
+
+This example demonstrates On-Policy Distillation (OPD) algorithm training on the GSM8K dataset.
+
+On-Policy Distillation is a knowledge distillation method, where in this example:
+1. **Student model** (`Qwen/Qwen2.5-1.5B-Instruct`) generates trajectories with logprobs
+2. **Teacher model** (`Qwen/Qwen2.5-Math-7B-Instruct`) computes logprobs on the same trajectories
+3. The advantage is computed as: `advantages = kl_coef * (teacher_logprobs - student_logprobs)`
+4. The student model is trained to minimize this KL divergence, effectively learning from the teacher
+
+## Key Configuration
+
+- **Algorithm**: `on_policy_distill`
+- **Workflow**: `on_policy_distill_workflow`
+- **Student Model**: `Qwen/Qwen2.5-1.5B-Instruct`
+- **Teacher Model**: `Qwen/Qwen2.5-Math-7B-Instruct` (configured as auxiliary model)
+
+## Running the Example
+
+Download the model checkpoint and modify your config file, then run:
+```bash
+trinity run examples/opd_gsm8k/opd_gsm8k.yaml
+```
+
+Then you are all set! It should be pretty simple😄, and the training should converge very quick(Much quicker then RL).
+
+
+
+![](../../docs/sphinx_doc/assets/opd_acc.png)
+![](../../docs/sphinx_doc/assets/opd_kl.png)
+
+
+## References
+
+- https://arxiv.org/pdf/2306.13649
+- https://thinkingmachines.ai/blog/on-policy-distillation/

docs/sphinx_doc/assets/opd_kl.png

@@ -0,0 +1,74 @@
+project: "Trinity-RFT-gsm8k-opd"
+name: "qwen2.5-1.5B-distill-from-math-7B-lr1e-5"
+checkpoint_root_dir: ${oc.env:TRINITY_CHECKPOINT_ROOT_DIR,./checkpoints}
+algorithm:
+  algorithm_type: on_policy_distill
+  repeat_times: 8
+  optimizer:
+    lr: 1e-5
+  advantage_fn_args:
+    kl_coef: 1.0
+model:
+  # Student model
+  model_path: ${oc.env:TRINITY_MODEL_PATH,Qwen/Qwen2.5-1.5B-Instruct}
+  max_response_tokens: 1024
+  max_model_len: 2048
+cluster:
+  node_num: 1
+  gpu_per_node: 8
+buffer:
+  total_epochs: 1
+  batch_size: 96
+  explorer_input:
+    taskset:
+      name: gsm8k
+      storage_type: file
+      path: ${oc.env:TRINITY_TASKSET_PATH,openai/gsm8k}
+      subset_name: main
+      split: train
+      format:
+        prompt_key: 'question'
+        response_key: 'answer'
+      rollout_args:
+        temperature: 1.0
+    # Use on_policy_distill_math_workflow for Qwen2.5-Math style format with accuracy reward
+    default_workflow_type: 'on_policy_distill_math_workflow'
+  trainer_input:
+    experience_buffer:
+      name: gsm8k_opd_buffer
+      storage_type: queue
+explorer:
+  eval_interval: 50
+  runner_per_model: 8
+  rollout_model:
+    # Student model for rollout
+    engine_num: 4
+    tensor_parallel_size: 1
+    enable_prefix_caching: false
+    enforce_eager: true
+    dtype: bfloat16
+    seed: 42
+  auxiliary_models:
+    # Teacher model for distillation
+    - model_path: ${oc.env:TRINITY_MODEL_PATH,Qwen/Qwen2.5-Math-7B-Instruct}
+      engine_num: 1
+      tensor_parallel_size: 2
+      enable_prefix_caching: false
+      enforce_eager: true
+      dtype: bfloat16
+      seed: 42
+      max_model_len: 4096
+      max_prompt_tokens: 2048
+      max_response_tokens: 1024
+synchronizer:
+  sync_method: 'nccl'
+  sync_interval: 1
+  sync_timeout: 1200
+trainer:
+  save_interval: 100
+  grad_clip: 1.0
+  use_dynamic_bsz: true
+  max_token_len_per_gpu: 16384
+  ulysses_sequence_parallel_size: 1
+monitor:
+  monitor_type: wandb

examples/opd_gsm8k/README.md

@@ -80,7 +80,7 @@ class FormatConfig:
 class GenerationConfig:
     temperature: Optional[float] = None  # 1.0
     top_p: Optional[float] = None  # 1.0
-    top_k: Optional[int] = None  # -1
+    top_k: int = -1  # -1 means disabled
     logprobs: Optional[int] = None  # 0  # vLLM return `logprobs + 1` elements
     max_tokens: Optional[int] = None  # if None, use model.max_response_tokens
     # repeat each task for `n` times

examples/opd_gsm8k/opd_gsm8k.yaml

@@ -48,7 +48,9 @@
     MathTrainableRULERWorkflow,
 )
 from trinity.common.workflows.on_policy_distill_workflow import (
+    AsyncOnPolicyDistillMathWorkflow,
     AsyncOnPolicyDistillWorkflow,
+    OnPolicyDistillMathWorkflow,
     OnPolicyDistillWorkflow,
 )
 from trinity.common.workflows.rubric_judge_workflow import RubricJudgeWorkflow
@@ -103,4 +105,6 @@
     # On-policy distillation workflows
     "OnPolicyDistillWorkflow",
     "AsyncOnPolicyDistillWorkflow",
+    "OnPolicyDistillMathWorkflow",
+    "AsyncOnPolicyDistillMathWorkflow",
 ]

trinity/common/config.py

@@ -11,22 +11,28 @@
 5. Train with importance_sampling loss
 """
 
+from dataclasses import asdict
 from typing import List, Optional
 
 import openai
 
 from trinity.common.experience import Experience
 from trinity.common.models.model import ModelWrapper
-from trinity.common.workflows.workflow import WORKFLOWS, BaseSimpleWorkflow, Task
+from trinity.common.rewards.qwen25_eval import verify_math_answer
+from trinity.common.workflows.workflow import WORKFLOWS, Task, Workflow
 
 
 @WORKFLOWS.register_module("on_policy_distill_workflow")
-class OnPolicyDistillWorkflow(BaseSimpleWorkflow):
+class OnPolicyDistillWorkflow(Workflow):
     """On-policy distillation workflow.
 
     Computes and stores teacher_logprobs in experience.info.
     The advantage_fn in trainer will compute:
         advantages = teacher_logprobs - student_logprobs
+
+    Note: This workflow does NOT use reward_fn because:
+    - Advantage is computed from teacher-student logprobs difference
+    - No external reward signal is needed
     """
 
     is_async: bool = True
@@ -41,8 +47,13 @@ def __init__(
         auxiliary_models: Optional[List[openai.OpenAI]] = None,
         auxiliary_model_wrappers: Optional[List[ModelWrapper]] = None,
     ):
-        super().__init__(task=task, model=model, auxiliary_models=auxiliary_models)
-        self.auxiliary_model_wrappers = auxiliary_model_wrappers
+        super().__init__(
+            task=task,
+            model=model,
+            auxiliary_models=auxiliary_models,
+            auxiliary_model_wrappers=auxiliary_model_wrappers,
+        )
+        self.reset(task)
 
         assert (
             auxiliary_model_wrappers is not None and len(auxiliary_model_wrappers) >= 1
@@ -51,6 +62,49 @@ def __init__(
 
         self.temperature = task.workflow_args.get("temperature", 1.0)
 
+    def reset(self, task: Task):
+        """Reset the workflow with a new task.
+
+        Unlike BaseSimpleWorkflow, this does NOT require reward_fn.
+        """
+        self.task = task
+        self.format_args = task.format_args
+        self.system_prompt = task.format_args.system_prompt
+        self.reply_prefix = task.format_args.reply_prefix
+        self.raw_task = task.raw_task
+        self.task_desc = task.task_desc
+        self.truth = task.truth
+
+    def set_repeat_times(self, repeat_times, run_id_base):
+        self.repeat_times = repeat_times
+        self.task.rollout_args.n = repeat_times
+        self.run_id_base = run_id_base
+
+    @property
+    def rollout_args(self):
+        return asdict(self.task.rollout_args)
+
+    def format_messages(self):
+        """Format messages for the instruct model.
+
+        Default format: system_prompt (optional) + task_desc + reply_prefix (optional)
+        """
+        messages = []
+        if self.system_prompt:
+            messages.append({"role": "system", "content": self.system_prompt})
+        messages.append({"role": "user", "content": self.task_desc})
+        if self.reply_prefix:
+            messages.append({"role": "assistant", "content": self.reply_prefix})
+        return messages
+
+    def compute_reward(self, response: Experience) -> float:
+        """Compute reward for a response.
+
+        In base class, returns 0.0 as advantage is computed from teacher-student logprobs.
+        Subclasses can override this to compute actual rewards.
+        """
+        return 0.0
+
     async def run_async(self) -> List[Experience]:
         messages = self.format_messages()
 
@@ -79,13 +133,16 @@ async def run_async(self) -> List[Experience]:
             # Step 3: Store teacher_logprobs for advantage_fn
             response.teacher_logprobs = teacher_resp_logprobs
 
-            # Set a dummy reward (actual advantage computed by advantage_fn)
-            response.reward = 0.0
-            response.eid.run = i + self.run_id_base
-
-            # Metrics for monitoring
+            # Initialize metrics
             if response.metrics is None:
                 response.metrics = {}
+
+            # Compute reward (subclasses can override compute_reward)
+            response.reward = self.compute_reward(response)
+
+            response.eid.run = i + self.run_id_base
+
+            # KL divergence for monitoring
             kl = (student_resp_logprobs - teacher_resp_logprobs).sum().item()
             response.metrics["kl_divergence"] = kl
 
@@ -94,4 +151,53 @@ async def run_async(self) -> List[Experience]:
 
 @WORKFLOWS.register_module("async_on_policy_distill_workflow")
 class AsyncOnPolicyDistillWorkflow(OnPolicyDistillWorkflow):
+    """Alias for OnPolicyDistillWorkflow (already async)."""
+
+    pass
+
+
+@WORKFLOWS.register_module("on_policy_distill_math_workflow")
+class OnPolicyDistillMathWorkflow(OnPolicyDistillWorkflow):
+    """On-policy distillation workflow with Qwen2.5-Math style format.
+
+    This workflow:
+    - Uses Qwen2.5-Math style prompt format (same as math_eval_workflow)
+    - Computes accuracy using verify_math_answer as reward
+    - Suitable for math reasoning tasks like GSM8K, MATH, etc.
+    """
+
+    def format_messages(self):
+        """Format messages using Qwen2.5-Math style.
+
+        System prompt: "You are a helpful assistant."
+        User prompt: "{question}\nPlease reason step by step, and put your final answer within \\boxed{}."
+        """
+        system_prompt = "You are a helpful assistant."
+        user_prompt = f"{self.task_desc}\nPlease reason step by step, and put your final answer within \\boxed{{}}."
+        return [
+            {"role": "system", "content": system_prompt},
+            {"role": "user", "content": user_prompt},
+        ]
+
+    def compute_reward(self, response: Experience) -> float:
+        """Compute accuracy as reward using Qwen2.5-Math evaluation.
+
+        Returns 1.0 if answer is correct, 0.0 otherwise.
+        """
+        if response.response_text and self.truth:
+            accuracy, _ = verify_math_answer(
+                response_text=response.response_text, ground_truth=self.truth
+            )
+            # Store accuracy in metrics
+            if response.metrics is None:
+                response.metrics = {}
+            response.metrics["accuracy"] = accuracy
+            return float(accuracy)
+        return 0.0
+
+
+@WORKFLOWS.register_module("async_on_policy_distill_math_workflow")
+class AsyncOnPolicyDistillMathWorkflow(OnPolicyDistillMathWorkflow):
+    """Alias for OnPolicyDistillMathWorkflow (already async)."""
+
     pass