refactor(sunjx): refactor dataset and reward module

[Jiaxuan-Sun]

1. Dataset Module Refactoring (lightrft/datasets/)

Modified:

• __init__.py: Refactored imports with unified interfaces and improved optional dependency handling

Added:

• config.py: DatasetConfig class

  • Unified configuration for train/eval/pretrain datasets
  • Auto-normalization of data_path and data_probs (supports string/list)
  • Factory methods: for_train(), for_eval(), for_pretrain()
  • Parameter validation

• loader.py: DatasetLoader class

  • Unified loading interface for train/eval/pretrain datasets
  • Automatic handling of blending_datasets parameters
  • Support for PromptDatasetVL and SFTDatasetVL
  • Consistent logging

2. Reward Module (lightrft/reward/)

Added:

• __init__.py: Module entry point with unified exports

• base.py: BaseReward abstract base class

  • Unified compute() method signature
  • Consistent return format: (rewards, metrics)

• rule.py: RuleReward class

  • Rule-based reward implementation
  • Format checking (e.g., <think> tags, \boxed{} notation)
  • Accuracy verification using mathruler grader
  • Registry pattern for custom rule types
  • Built-in rules: default, geo3k_*, gsm8k_*

• model.py: Reward model implementations

  • SingleRewardModel: Single reward model wrapper with auto load/offload
  • MultiRewardModel: Multiple reward model ensemble with recipe-based aggregation
  • Supports standard PyTorch models and custom engines (e.g., SGLang)

• manager.py: RewardManager class

  • Unified manager for all reward types
  • Auto-selection of reward implementation (rule/single/multi)
  • from_config() factory method

[puyuan1996]

Please remove the Author line.

[puyuan1996]

We need to update the examples in gsm8k and geo3k (examples/gsm8k_geo3k/train_colocate.py) to use the refactored dataset and reward APIs introduced here.

[puyuan1996]

LightRFT