Merge pull request #288 from thwu1/nightly

[feat] Tinker Workflow Trainer

Author: jeffreysijuntan

examples/math_tinker/train_math_tinker.py

@@ -16,23 +16,6 @@
 from rllm.trainer.tinker.tinker_agent_trainer import TinkerAgentTrainer
 
 
-class SimpleDataLoader:
-    """Simple reusable dataloader."""
-
-    def __init__(self, dataset, batch_size):
-        self.dataset = dataset
-        self.batch_size = batch_size
-
-    def __iter__(self):
-        for i in range(0, len(self.dataset), self.batch_size):
-            yield self.dataset[i : i + self.batch_size]
-
-
-def create_dataloader(dataset, batch_size):
-    """Create a simple reusable dataloader from dataset."""
-    return SimpleDataLoader(dataset, batch_size)
-
-
 @hydra.main(version_base=None, config_path="../../rllm/trainer/config", config_name="tinker_agent_trainer")
 def main(config: DictConfig):
     """
@@ -48,19 +31,15 @@ def main(config: DictConfig):
     if train_dataset is None or test_dataset is None:
         raise ValueError("Datasets not found! Please run prepare_tinker_math_dataset.py first:\n  python -m examples.math_tinker.prepare_tinker_math_dataset")
 
-    # Create dataloaders
-    train_dataloader = create_dataloader(train_dataset, config.data.train_batch_size)
-    test_dataloader = create_dataloader(test_dataset, config.data.val_batch_size)
-
     # Create trainer (uses separated components internally)
     trainer = TinkerAgentTrainer(
         config=config,
         agent_class=MathAgentWithFewshot,
         env_class=SingleTurnEnvironment,
         agent_args={"use_fewshot": True},
         env_args={"reward_fn": math_reward_fn},
-        train_dataloader=train_dataloader,
-        val_dataloader=test_dataloader,
+        train_dataset=train_dataset,
+        val_dataset=test_dataset,
     )
 
     # Train (all orchestration handled internally by TinkerAgentTrainer)

examples/solver_judge/train_solver_judge_flow.sh

@@ -10,7 +10,7 @@ python3 -m examples.solver_judge.train_solver_judge_flow \
     data.train_batch_size=64 \
     data.max_prompt_length=2048 \
     data.max_response_length=1024 \
-    actor_rollout_ref.model.path=Qwen/Qwen3-0.6B \
+    actor_rollout_ref.model.path=Qwen/Qwen3-4B-Instruct-2507 \
     actor_rollout_ref.actor.optim.lr=1e-6 \
     actor_rollout_ref.model.use_remove_padding=True \
     actor_rollout_ref.actor.loss_agg_mode=seq-mean-token-mean \
@@ -31,12 +31,13 @@ python3 -m examples.solver_judge.train_solver_judge_flow \
     actor_rollout_ref.rollout.name=vllm \
     actor_rollout_ref.rollout.mode="async" \
     actor_rollout_ref.rollout.enforce_eager=False \
-    actor_rollout_ref.rollout.temperature=0.6 \
+    actor_rollout_ref.rollout.temperature=1.0 \
+    actor_rollout_ref.rollout.top_p=1.0 \
     actor_rollout_ref.rollout.gpu_memory_utilization=0.8 \
     actor_rollout_ref.rollout.n=4 \
     actor_rollout_ref.rollout.val_kwargs.n=1 \
-    actor_rollout_ref.rollout.val_kwargs.temperature=0.6 \
-    actor_rollout_ref.rollout.val_kwargs.top_p=0.95 \
+    actor_rollout_ref.rollout.val_kwargs.temperature=1.0 \
+    actor_rollout_ref.rollout.val_kwargs.top_p=1.0 \
     actor_rollout_ref.ref.fsdp_config.param_offload=True \
     algorithm.adv_estimator=grpo \
     rllm.compact_filtering.enable=False \
@@ -59,6 +60,7 @@ python3 -m examples.solver_judge.train_solver_judge_flow \
     trainer.test_freq=10 \
     trainer.default_hdfs_dir=null \
     trainer.total_epochs=100 \
-    rllm.workflow.use_workflow=True
+    rllm.workflow.use_workflow=True \
+    +ray_init._temp_dir=/home/tianhao/tmp
 
 pkill -9 -f 'ray::WorkerDict' 

examples/solver_judge_tinker/run_solver_judge_flow_tinker.py

@@ -0,0 +1,149 @@
+import asyncio
+import json
+import os
+
+# Import countdown-specific modules
+import sys
+from copy import deepcopy
+
+import tinker
+from solver_judge_flow import SolverJudgeWorkflow
+from transformers import AutoTokenizer
+
+from rllm.data.dataset import DatasetRegistry
+from rllm.engine.agent_workflow_engine import AgentWorkflowEngine
+from rllm.engine.rollout.tinker_engine import TinkerEngine
+from rllm.rewards.countdown_reward import countdown_reward_fn
+
+sys.path.append(os.path.join(os.path.dirname(__file__), "..", "countdown"))
+
+
+def load_data(n=1):
+    """Load countdown data using the Dataset interface."""
+    dataset = DatasetRegistry.load_dataset("countdown", "test")
+    if dataset is None:
+        print("Dataset not found, preparing dataset...")
+        from prepare_countdown_data import prepare_countdown_data
+
+        _, dataset, _, _ = prepare_countdown_data()
+
+    data = []
+    for idx, example in enumerate(dataset):
+        processed = process_countdown_fn(example, idx)
+        for i in range(n):
+            data.append(deepcopy(processed))
+    return data
+
+
+def process_countdown_fn(example, idx):
+    """Process countdown example into the expected format."""
+    question = example["question"]
+    target = example["target"]
+    nums = example["nums"]
+
+    # Create ground truth in the format expected by countdown_reward_fn
+    ground_truth = {"target": target, "numbers": nums}
+
+    task = {"question": question, "ground_truth": ground_truth, "idx": idx, "data_source": "countdown", "target": target, "nums": nums}
+    return task
+
+
+def evaluate_results(results):
+    """Evaluate the results and compute pass@k metrics."""
+    from collections import defaultdict
+
+    # Create a map to store correct answers per problem
+    problem_correct_map = defaultdict(int)
+    problem_total_map = defaultdict(int)
+
+    # Count correct answers for each problem
+    for episode in results:
+        problem = episode.task["question"]
+
+        # Use the episode-level is_correct flag set by the workflow
+        is_correct = episode.is_correct
+
+        problem_correct_map[problem] += int(is_correct)
+        problem_total_map[problem] += 1
+
+    # Calculate pass@1 and pass@k
+    k = max(problem_total_map.values()) if problem_total_map else 1
+    total_problems = len(problem_correct_map)
+
+    if total_problems > 0:
+        pass_at_1 = sum(problem_correct_map.values()) / sum(problem_total_map.values())
+        pass_at_k = sum(1 for problem, correct in problem_correct_map.items() if correct > 0) / total_problems
+    else:
+        pass_at_1 = 0.0
+        pass_at_k = 0.0
+
+    print("Total unique problems:", total_problems)
+    print("Average Pass@1 Accuracy:", pass_at_1)
+    print(f"Average Pass@{k} Accuracy:", pass_at_k)
+
+
+if __name__ == "__main__":
+    import os
+
+    os.environ["TOKENIZERS_PARALLELISM"] = "true"
+
+    # Configuration
+    n_parallel_tasks = 4
+    n_solutions = 2  # Number of solutions to generate per problem
+
+    model_name = "Qwen/Qwen3-8B"
+    service_client = tinker.ServiceClient(base_url=None)
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    rollout_engine = TinkerEngine(
+        base_url=None,
+        model_name=model_name,
+        tokenizer=tokenizer,
+        service_client=service_client,
+        max_prompt_length=2048,
+        max_response_length=1024,
+        sampling_params={"temperature": 0.6, "top_p": 0.95},
+    )
+    training_client = service_client.create_lora_training_client(
+        base_model=model_name,
+        rank=4,
+    )
+    sampler_future = training_client.save_weights_for_sampler(name="000000")
+    sampler_result = sampler_future.result()
+    sampling_client = training_client.create_sampling_client(sampler_result.path)
+
+    rollout_engine.set_sampling_client(sampling_client)
+
+    engine = AgentWorkflowEngine(
+        workflow_cls=SolverJudgeWorkflow,
+        workflow_args={
+            "n_solutions": n_solutions,
+            "reward_function": countdown_reward_fn,
+        },
+        rollout_engine=rollout_engine,
+        config=None,
+        n_parallel_tasks=n_parallel_tasks,
+        retry_limit=1,
+    )
+
+    # Load countdown tasks
+    tasks = load_data(n=1)
+    print(f"Loaded {len(tasks)} countdown tasks")
+    tasks = tasks[:4]
+
+    results = asyncio.run(engine.execute_tasks(tasks))
+    import pdb
+
+    pdb.set_trace()
+
+    print(results[1])
+
+    # Evaluate results (rewards are already assigned in the workflow)
+    print("Evaluating results...")
+    evaluate_results(results)
+
+    # Save results
+    os.makedirs("logs", exist_ok=True)
+    with open("logs/solver_judge_countdown.json", "w") as f:
+        json.dump([episode.to_dict() for episode in results], f, indent=4)
+
+    print("\nResults saved to logs/solver_judge_countdown.json")

examples/solver_judge_tinker/train_solver_judge_flow_tinker.py

@@ -0,0 +1,28 @@
+import hydra
+
+from examples.solver_judge.solver_judge_flow import SolverJudgeWorkflow
+from rllm.data.dataset import DatasetRegistry
+from rllm.rewards.countdown_reward import countdown_reward_fn
+from rllm.trainer.tinker.tinker_workflow_trainer import TinkerWorkflowTrainer
+
+
+@hydra.main(config_path="pkg://rllm.trainer.config", config_name="tinker_workflow_trainer", version_base=None)
+def main(config):
+    train_dataset = DatasetRegistry.load_dataset("countdown", "train")
+    test_dataset = DatasetRegistry.load_dataset("countdown", "test")
+
+    trainer = TinkerWorkflowTrainer(
+        workflow_class=SolverJudgeWorkflow,
+        workflow_args={
+            "n_solutions": 2,
+            "reward_function": countdown_reward_fn,
+        },
+        config=config,
+        train_dataset=train_dataset,
+        val_dataset=test_dataset,
+    )
+    trainer.fit_agent()
+
+
+if __name__ == "__main__":
+    main()

examples/solver_judge_tinker/train_solver_judge_flow_tinker.sh

@@ -0,0 +1,25 @@
+set -x
+
+MODEL_PATH=Qwen/Qwen3-4B-Instruct-2507
+
+python3 -m examples.solver_judge_tinker.train_solver_judge_flow_tinker \
+    model.name=$MODEL_PATH \
+    model.lora_rank=32 \
+    training.group_size=4 \
+    training.learning_rate=4e-5 \
+    sampling.temperature=1.0 \
+    sampling.top_p=1.0 \
+    algorithm.adv_estimator=grpo \
+    algorithm.norm_adv_by_std_in_grpo=true \
+    data.max_prompt_length=2048 \
+    data.max_response_length=1024 \
+    data.train_batch_size=64 \
+    data.val_batch_size=512 \
+    trainer.total_epochs=100 \
+    trainer.logger=['wandb'] \
+    trainer.project_name='solver-judge-workflow' \
+    trainer.experiment_name='countdown-solver-judge-tinker-norm-by-std' \
+    trainer.val_before_train=False \
+    trainer.test_freq=10 \
+    trainer.save_freq=20 \
+    trainer.default_local_dir='/tmp/countdown-solver-judge-tinker-norm-by-std'

rllm/trainer/config/tinker_workflow_trainer.yaml

@@ -0,0 +1,66 @@
+# Tinker Backend Configuration for rLLM
+# This config is used when training agents with Tinker backend
+# Default settings match tinker_cookbook.recipes.math_rl for MATH dataset
+
+# Tinker-specific settings
+tinker_base_url: null  # Tinker service URL (null for local)
+  
+# Model Configuration
+model:
+  name: "Qwen/Qwen3-8B"  # Default model for MATH dataset
+  lora_rank: 32
+  train_unembed: true  # Train LoRA on output embedding layer (set to false for Fireworks compatibility)
+  train_attn: true     # Train LoRA on attention layers
+  train_mlp: true      # Train LoRA on MLP layers
+
+# Training Configuration
+training:
+  group_size: 16  # Number of rollouts per prompt (for GRPO)
+  learning_rate: 2e-5  # 2e-5 for MATH dataset
+  beta1: 0.9
+  beta2: 0.95
+  eps: 1e-8
+  max_length: 32768
+  num_minibatches: 1
+
+# Sampling Configuration
+sampling:
+  temperature: 0.6
+  top_p: 0.95
+
+# Algorithm Configuration (compatible with verl)
+algorithm:
+  adv_estimator: grpo  # REINFORCE, GRPO
+  gamma: 1.0
+  lam: 0.95
+  norm_adv_by_std_in_grpo: false  # math_rl doesn't normalize by std
+
+workflow:
+  n_parallel_tasks: 256
+  retry_limit: 3 
+
+# Data Configuration
+data:
+  train_files: null
+  val_files: null
+  max_prompt_length: 2048
+  max_response_length: 2048
+  train_batch_size: 64
+  val_batch_size: 32
+
+# Trainer Configuration
+trainer:
+  total_epochs: 10
+  logger: ['console']  # Options: 'console', 'wandb', 'tensorboard'
+  project_name: 'rllm-tinker'
+  experiment_name: 'default'
+  test_freq: 5
+  save_freq: 
+  reward_broadcast: 'step'
+  val_before_train: true
+  default_local_dir: '/tmp/rllm-tinker-checkpoints'
+
+# Hydra configuration
+hydra:
+  run:
+    dir: outputs/${now:%Y-%m-%d}/${now:%H-%M-%S}