feat: add experimental native RL stack and arithmetic validation benchmark

examples/09_rl_training_methods.py

@@ -17,8 +17,9 @@
     ORPOTrainer, ORPOConfig,
     GRPOTrainer, GRPOConfig,
     # Utilities
-    prepare_preference_dataset,
+    prepare_rl_dataset,
     create_reward_function,
+    resume_from_checkpoint,
 )
 
 
@@ -60,6 +61,8 @@ def demo_dpo_training():
         },
     ]
 
+    prepared_dataset = prepare_rl_dataset(preference_data, mode="preference", tokenizer=tokenizer)
+
     # Configure DPO
     config = DPOConfig(
         beta=0.1,  # KL penalty coefficient
@@ -72,7 +75,7 @@ def demo_dpo_training():
     # Create trainer
     trainer = DPOTrainer(
         model=model,
-        train_dataset=preference_data,
+        train_dataset=prepared_dataset,
         tokenizer=tokenizer,
         args=config,
     )
@@ -122,12 +125,19 @@ def demo_grpo_training():
         },
     ]
 
-    # Create a math reward function
-    math_reward = create_reward_function("math")
+    prepared_dataset = prepare_rl_dataset(reasoning_data, mode="prompt", tokenizer=tokenizer)
+
+    # Compose rewards through the public RL API surface.
+    math_reward = create_reward_function(
+        rewards=[
+            {"name": "math", "source": "math", "weight": 1.0},
+            {"name": "length", "source": "length", "weight": 0.1},
+        ]
+    )
 
     # Configure GRPO
     config = GRPOConfig(
-        loss_type="grpo",  # grpo, dr_grpo, dapo, bnpo
+        loss_type="grpo",  # Phase 1 accepts grpo/dr_grpo/dapo/bnpo via one shared objective
         beta=0.04,
         num_generations=4,  # Multiple generations per prompt
         temperature=0.7,
@@ -139,7 +149,7 @@ def demo_grpo_training():
     # Create trainer with custom reward function
     trainer = GRPOTrainer(
         model=model,
-        train_dataset=reasoning_data,
+        train_dataset=prepared_dataset,
         tokenizer=tokenizer,
         reward_fn=math_reward,  # Custom reward!
         args=config,
@@ -153,6 +163,7 @@ def demo_grpo_training():
 
     # Would train with: trainer.train()
     print("\nTo train: trainer.train()")
+    print(f"To inspect a saved checkpoint first: {resume_from_checkpoint.__name__}('./grpo_output')")
 
 
 def demo_orpo_training():
@@ -221,14 +232,14 @@ def show_available_trainers():
         print(f"| {name} | {method} | {use_case} |")
 
     print("\n" + "=" * 70)
-    print("GRPO Loss Types (for reasoning models)")
+    print("GRPO Loss Types (accepted in Phase 1)")
     print("=" * 70)
 
     grpo_types = [
-        ("grpo", "Standard GRPO", "Default for reasoning"),
-        ("dr_grpo", "Dr. GRPO", "Distilled version"),
-        ("dapo", "DAPO", "Data-efficient variant"),
-        ("bnpo", "BNPO", "Batch-normalized variant"),
+        ("grpo", "Standard GRPO", "Primary Phase 1 name"),
+        ("dr_grpo", "Dr. GRPO", "Accepted alias; shared Phase 1 objective"),
+        ("dapo", "DAPO", "Accepted alias; shared Phase 1 objective"),
+        ("bnpo", "BNPO", "Accepted alias; shared Phase 1 objective"),
     ]
 
     print("\n| Loss Type | Name | Description |")

examples/10_qwen3_arithmetic_grpo_validation.py

@@ -0,0 +1,10 @@
+from pathlib import Path
+import sys
+
+sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
+
+from mlx_tune.arithmetic_grpo_validation import main
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

mlx_tune/__init__.py

@@ -15,7 +15,15 @@
 
 __version__ = "0.4.0"  # Renamed to mlx-tune (formerly unsloth-mlx)
 
-from mlx_tune.model import FastLanguageModel
+from mlx_tune.model import (
+    FastLanguageModel,
+    ReferencePolicy,
+    RLModelRoles,
+    RewardModel,
+    ValueModel,
+    build_value_model,
+    create_rl_model_roles,
+)
 from mlx_tune.trainer import (
     prepare_dataset,
     format_chat_template,
@@ -25,25 +33,44 @@
     get_training_config,
 )
 from mlx_tune.sft_trainer import SFTTrainer, SFTConfig, TrainingArguments
+from mlx_tune.rl_api import (
+    RLCheckpointBundle,
+    PreparedRLDataset,
+    prepare_rl_dataset,
+    build_reference_policy,
+    build_reward_model,
+    create_reward_function,
+    resume_from_checkpoint,
+)
 
 # RL Trainers
 from mlx_tune.rl_trainers import (
+    RewardTrainer,
+    RewardConfig,
     DPOTrainer,
     DPOConfig,
     ORPOTrainer,
     ORPOConfig,
     GRPOTrainer,
     GRPOConfig,
+    PPOTrainer,
+    PPOConfig,
+    OnlineDPOTrainer,
+    OnlineDPOConfig,
+    KTOConfig,
+    SimPOConfig,
     KTOTrainer,
     SimPOTrainer,
+    prepare_reward_dataset,
     prepare_preference_dataset,
-    create_reward_function,
+    score_reward_model,
 )
 
 # Loss functions for custom training
 from mlx_tune.losses import (
     compute_log_probs,
     compute_log_probs_with_lengths,
+    compute_completion_log_probs,
     dpo_loss,
     orpo_loss,
     kto_loss,
@@ -52,6 +79,16 @@
     grpo_loss,
     grpo_batch_loss,
     compute_reference_logprobs,
+    pairwise_reward_loss,
+    reward_model_pairwise_loss,
+    reward_model_regression_loss,
+    value_regression_loss,
+    value_model_regression_loss,
+    scalar_loss_metrics,
+    pairwise_ranking_accuracy,
+    precompute_preference_reference_logprobs,
+    precompute_kto_reference_logprobs,
+    ppo_sequence_loss,
 )
 
 # Vision Language Models
@@ -92,10 +129,24 @@
     HFDatasetConfig,
     load_dataset_with_config,
 )
+from mlx_tune.trl_compat import PatchFastRL
 
 __all__ = [
     # Core
     "FastLanguageModel",
+    "ReferencePolicy",
+    "RLModelRoles",
+    "RewardModel",
+    "ValueModel",
+    "build_reference_policy",
+    "build_reward_model",
+    "build_value_model",
+    "create_rl_model_roles",
+    "PreparedRLDataset",
+    "RLCheckpointBundle",
+    "prepare_rl_dataset",
+    "resume_from_checkpoint",
+    "PatchFastRL",
     "__version__",
     # SFT Training
     "SFTTrainer",
@@ -108,6 +159,14 @@
     "ORPOConfig",
     "GRPOTrainer",
     "GRPOConfig",
+    "RewardTrainer",
+    "RewardConfig",
+    "PPOTrainer",
+    "PPOConfig",
+    "OnlineDPOTrainer",
+    "OnlineDPOConfig",
+    "KTOConfig",
+    "SimPOConfig",
     "KTOTrainer",
     "SimPOTrainer",
     # Vision Models
@@ -116,6 +175,7 @@
     # Loss Functions
     "compute_log_probs",
     "compute_log_probs_with_lengths",
+    "compute_completion_log_probs",
     "dpo_loss",
     "orpo_loss",
     "kto_loss",
@@ -124,15 +184,27 @@
     "grpo_loss",
     "grpo_batch_loss",
     "compute_reference_logprobs",
+    "pairwise_reward_loss",
+    "reward_model_pairwise_loss",
+    "reward_model_regression_loss",
+    "value_regression_loss",
+    "value_model_regression_loss",
+    "scalar_loss_metrics",
+    "pairwise_ranking_accuracy",
+    "precompute_preference_reference_logprobs",
+    "precompute_kto_reference_logprobs",
+    "ppo_sequence_loss",
     # Utilities
     "prepare_dataset",
+    "prepare_reward_dataset",
     "prepare_preference_dataset",
     "format_chat_template",
     "create_training_data",
     "save_model_hf_format",
     "export_to_gguf",
     "get_training_config",
     "create_reward_function",
+    "score_reward_model",
     "load_vlm_dataset",
     # Chat Templates and Dataset Formatting
     "detect_dataset_format",