Skip to content

Latest commit

 

History

5 Commits

Folders and files

NameName
Last commit message
Last commit date
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 

Repository files navigation

LUFFY-POOL

Dual-Pool Hybrid Reinforcement Learning for LLM Mathematical Reasoning

LUFFY-POOL extends the LUFFY framework with a dual-pool training architecture that simultaneously leverages on-policy and off-policy data for more efficient reinforcement learning of large language models. The system is designed for training mathematical reasoning capabilities using GRPO (Group Relative Policy Optimization) with importance-weighted off-policy corrections.

Key Features

  • Dual-Pool Training: Combines on-policy rollouts (fresh samples from the current policy) with off-policy data (pre-collected trajectories) in a unified training loop, using importance sampling to correct distribution mismatch.
  • GRPO with Off-Policy Extension: Extends Group Relative Policy Optimization to support mixed-source advantage estimation, enabling more sample-efficient training.
  • Distributed Training Pipeline: Built on VERL (Volcano Engine RL) with FSDP + Ray + vLLM, supporting 3D parallelism (tensor/pipeline/data) for large-scale training.
  • Multi-Strategy Reward Functions: Supports rule-based symbolic verification, math-verify library grading, and LLM-as-Reward-Model (ORM) for nuanced evaluation of reasoning chains.

Architecture

┌─────────────────────────────────────────────────┐
│                  Mix Trainer                     │
│                                                  │
│  ┌──────────────┐       ┌──────────────────┐    │
│  │  On-Policy   │       │   Off-Policy     │    │
│  │    Pool       │       │     Pool          │    │
│  │  (vLLM       │       │  (Pre-collected   │    │
│  │   Rollout)   │       │   Trajectories)   │    │
│  └──────┬───────┘       └────────┬──────────┘    │
│         │                        │               │
│         └────────┬───────────────┘               │
│                  ▼                                │
│  ┌──────────────────────────────┐               │
│  │   GRPO Advantage Estimation  │               │
│  │  (importance-weighted mix)   │               │
│  └──────────────┬───────────────┘               │
│                  ▼                                │
│  ┌──────────────────────────────┐               │
│  │   FSDP Actor Update          │               │
│  │  (gradient checkpointing)    │               │
│  └──────────────────────────────┘               │
└─────────────────────────────────────────────────┘

Project Structure

LUFFY-POOL/
├── luffy/
│   ├── verl/                       # VERL framework (Volcano Engine RL)
│   │   └── verl/mix_src/           # Core dual-pool implementation
│   │       ├── main_mix_ppo.py     # Training entry point
│   │       ├── mix_trainer.py      # Dual-pool trainer with on/off-policy learning
│   │       ├── mix_core_alg.py     # GRPO advantage + on/off-policy loss functions
│   │       ├── mix_actor.py        # Actor model training logic
│   │       ├── mix_vllm_rollout.py # vLLM-based generation and rollout
│   │       ├── mix_fsdp_worker.py  # FSDP distributed training worker
│   │       └── rl_dataset_with_target.py  # Dataset with target action sampling
│   ├── deepscaler/                 # Reward module
│   │   ├── rewards/
│   │   │   ├── math_reward.py      # Rule-based + LLM-as-RM reward functions
│   │   │   └── math_utils/         # Symbolic/numeric answer grading
│   │   └── system_prompts.py       # ORM prompt templates
│   ├── requirements.txt
│   └── setup.py
├── data/                           # Data preparation scripts & datasets
│   ├── prepare_sft.py              # SFT data formatting
│   ├── prepare_train.py            # RL training data
│   └── prepare_train_sft_rl.py     # Mixed SFT+RL data
├── eval_scripts/                   # Evaluation tools
│   ├── generate_vllm.py            # vLLM batch generation
│   ├── collect_results.py          # Result aggregation
│   └── oat_math_grader.py          # Math grading utility
├── exp_scripts/                    # Experiment configs (Hydra YAML)
└── figures/                        # Visualizations

Technical Details

Dual-Pool Training Loop

  1. On-Policy Rollout: The current policy generates multiple candidate solutions (n=8) per math problem using vLLM, forming the on-policy pool.
  2. Off-Policy Sampling: Pre-collected trajectories (from stronger models or previous checkpoints) are sampled with importance weights to form the off-policy pool.
  3. Mixed Advantage Estimation: GRPO computes group-relative advantages across both pools, with importance sampling ratios correcting for the off-policy distribution shift.
  4. Policy Update: The actor is updated via clipped PPO-style objectives on the combined loss, with KL penalty against a reference policy.

Reward System

Reward Type Description
Rule-based verification Symbolic and numeric answer matching with tolerance
math-verify Library-based formal verification of mathematical expressions
LLM-as-RM (ORM) GPT/Gemini evaluates reasoning chain quality beyond final answer correctness
Format reward Checks for proper <think>...</think> chain-of-thought formatting

Training Configuration

  • Base Model: Qwen2.5-Math-7B (16k context, think variant)
  • Rollout: vLLM with sampling, 8 candidates per prompt, up to 8192 response tokens
  • Optimization: FSDP with gradient checkpointing, AdamW optimizer
  • Batch Size: 128 (train), 512 (validation)
  • Epochs: 30
  • Orchestration: Ray for multi-node coordination

Quick Start

# 1. Install dependencies
cd luffy && pip install -e . && pip install -r requirements.txt

# 2. Prepare training data
python data/prepare_train.py

# 3. Launch training (see exp_scripts/ for full configs)
python luffy/verl/verl/mix_src/main_mix_ppo.py \
  --config-path ../exp_scripts \
  --config-name config.yaml

# 4. Evaluate
python eval_scripts/generate_vllm.py \
  --model_path <checkpoint_path> \
  --data_path data/valid.parquet

Dependencies

  • Python 3.9+
  • PyTorch 2.4.0
  • Transformers 4.46.3
  • vLLM 0.6.3
  • Ray 2.12.0
  • DeepSpeed 0.15.0
  • Flash Attention 2.7.3

Acknowledgements

Built upon the LUFFY framework and VERL (Volcano Engine Reinforcement Learning).

About

带有双池子的混合策略强化学习,基于LUFFY开发

Resources

Stars

10 stars

Watchers

0 watching

Forks

Releases

Packages

Contributors

Languages