Arbor GRPO Sync Update

docs/docs/tutorials/rl_multihop/index.ipynb

@@ -227,7 +227,6 @@
    "outputs": [],
    "source": [
     "from dspy.teleprompt.grpo import GRPO\n",
-    "from dspy.clients.utils_finetune import MultiGPUConfig\n",
     "\n",
     "program = ResearchHop(num_docs=4, num_hops=2)\n",
     "program.set_lm(local_lm)\n",
@@ -261,7 +260,6 @@
     "    num_steps_for_val=10,\n",
     "    train_kwargs=train_kwargs,\n",
     "    report_train_scores=False,\n",
-    "    gpu_config=MultiGPUConfig(num_inference_gpus=1, num_training_gpus=1),\n",
     ")\n",
     "\n",
     "optimized_program = compiler.compile(\n",

docs/docs/tutorials/rl_papillon/index.ipynb

@@ -256,7 +256,6 @@
    "outputs": [],
    "source": [
     "from dspy.teleprompt.grpo import GRPO\n",
-    "from dspy.clients.utils_finetune import MultiGPUConfig\n",
     "\n",
     "papillon = PAPILLON(untrusted_model=openai_lm)\n",
     "papillon.set_lm(local_lm)\n",
@@ -291,7 +290,6 @@
     "    num_steps_for_val=10,\n",
     "    train_kwargs=train_kwargs,\n",
     "    report_train_scores=False,\n",
-    "    gpu_config=MultiGPUConfig(num_inference_gpus=2, num_training_gpus=2),\n",
     ")\n",
     "\n",
     "optimized_papillon = compiler.compile(\n",

dspy/clients/lm.py

@@ -13,7 +13,7 @@
 from dspy.clients.cache import request_cache
 from dspy.clients.openai import OpenAIProvider
 from dspy.clients.provider import Provider, ReinforceJob, TrainingJob
-from dspy.clients.utils_finetune import MultiGPUConfig, TrainDataFormat
+from dspy.clients.utils_finetune import TrainDataFormat
 from dspy.dsp.utils.settings import settings
 from dspy.utils.callback import BaseCallback
 
@@ -229,15 +229,15 @@ def thread_function_wrapper():
         return job
 
     def reinforce(
-        self, train_kwargs, gpu_config: MultiGPUConfig = MultiGPUConfig(num_inference_gpus=1, num_training_gpus=1)
+        self, train_kwargs
     ) -> ReinforceJob:
         # TODO(GRPO Team): Should we return an initialized job here?
         from dspy import settings as settings
 
         err = f"Provider {self.provider} does not implement the reinforcement learning interface."
         assert self.provider.reinforceable, err
 
-        job = self.provider.ReinforceJob(lm=self, train_kwargs=train_kwargs, gpu_config=gpu_config)
+        job = self.provider.ReinforceJob(lm=self, train_kwargs=train_kwargs)
         job.initialize()
         return job
 

dspy/clients/lm_local_arbor.py

@@ -1,23 +1,20 @@
+import json
 import time
 from datetime import datetime
-from typing import TYPE_CHECKING, Any, TypedDict
+from typing import TYPE_CHECKING, Any
 from urllib.parse import urljoin
 
 import openai
 import requests
 
 import dspy
 from dspy.clients.provider import Provider, ReinforceJob, TrainingJob
-from dspy.clients.utils_finetune import GRPOGroup, MultiGPUConfig, TrainDataFormat, TrainingStatus, save_data
+from dspy.clients.utils_finetune import GRPOGroup, GRPOStatus, TrainDataFormat, TrainingStatus, save_data
 
 if TYPE_CHECKING:
     from dspy.clients.lm import LM
 
 
-class GRPOTrainKwargs(TypedDict):
-    num_generations: int
-
-
 class ArborTrainingJob(TrainingJob):
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
@@ -48,7 +45,7 @@ def status(self) -> TrainingStatus:
 
 class ArborReinforceJob(ReinforceJob):
     DEFAULT_TRAIN_KWARGS = {  # noqa: RUF012
-        "temperature": 0.9,
+        "temperature": 1.0,
         "beta": 0.04,
         "num_iterations": 1,
         "per_device_train_batch_size": 8,
@@ -57,6 +54,7 @@ class ArborReinforceJob(ReinforceJob):
         # This is false by default in TRL, but I think it makes sense to be true for us
         "gradient_checkpointing": True,
         "lr_scheduler_type": "constant_with_warmup",
+        "warmup_steps": 10,
         "max_prompt_length": None,
         "max_completion_length": None,
         "gradient_checkpointing_kwargs": None,
@@ -65,13 +63,13 @@ class ArborReinforceJob(ReinforceJob):
         "max_grad_norm": 1.0,
         "report_to": "none",
         "log_completions": True,
-        "logging_steps": 100,
+        "logging_steps": 10,
         # By default, none is the model's max context length
         "max_context_length": None,
         "lora": False,
     }
 
-    def __init__(self, lm: "LM", train_kwargs: GRPOTrainKwargs, gpu_config: MultiGPUConfig = MultiGPUConfig(num_inference_gpus=1, num_training_gpus=1)):
+    def __init__(self, lm: "LM", train_kwargs: dict[str, Any]):
         # The teleprompter must ensure that this is set
         if "num_generations" not in train_kwargs:
             raise ValueError("num_generations must be set in the training kwargs")
@@ -81,7 +79,6 @@ def __init__(self, lm: "LM", train_kwargs: GRPOTrainKwargs, gpu_config: MultiGPU
         self.provider_job_id = None
         self.checkpoints = {}
         self.last_checkpoint = None
-        self.gpu_config = gpu_config
 
     def initialize(self):
         # TODO(GRPO Team): Set provider job ID
@@ -100,6 +97,7 @@ def initialize(self):
             "gradient_checkpointing", self.DEFAULT_TRAIN_KWARGS["gradient_checkpointing"]
         )
         lr_scheduler_type = self.train_kwargs.get("lr_scheduler_type", self.DEFAULT_TRAIN_KWARGS["lr_scheduler_type"])
+        warmup_steps = self.train_kwargs.get("warmup_steps", self.DEFAULT_TRAIN_KWARGS["warmup_steps"])
         max_prompt_length = self.train_kwargs.get("max_prompt_length", self.DEFAULT_TRAIN_KWARGS["max_prompt_length"])
         max_completion_length = self.train_kwargs.get(
             "max_completion_length", self.DEFAULT_TRAIN_KWARGS["max_completion_length"]
@@ -116,43 +114,54 @@ def initialize(self):
         max_context_length = self.train_kwargs.get(
             "max_context_length", self.DEFAULT_TRAIN_KWARGS["max_context_length"]
         )
-        lora = self.train_kwargs.get("lora", self.DEFAULT_TRAIN_KWARGS["lora"])
+        max_steps = self.train_kwargs.get("max_steps",500)
+        # lora = self.train_kwargs.get("lora", self.DEFAULT_TRAIN_KWARGS["lora"])
         api_base = self.lm.kwargs["api_base"]
 
         finetune_model = ArborProvider._remove_provider_prefix(self.lm.model)
-        # Only multi-GPU is supported for now
-        gpu_config_type = "multi"
         data = {
             "model": finetune_model,
-            "num_generations": num_generations,
-            "temperature": temperature,
-            "beta": beta,
-            "num_iterations": num_iterations,
-            "per_device_train_batch_size": per_device_train_batch_size,
-            "learning_rate": learning_rate,
-            "gradient_accumulation_steps": gradient_accumulation_steps,
-            "gradient_checkpointing": gradient_checkpointing,
-            "lr_scheduler_type": lr_scheduler_type,
-            "max_prompt_length": max_prompt_length,
-            "max_completion_length": max_completion_length,
-            "bf16": bf16,
-            "scale_rewards": scale_rewards,
-            "gradient_checkpointing_kwargs": gradient_checkpointing_kwargs,
-            "max_grad_norm": max_grad_norm,
-            "report_to": report_to,
-            "log_completions": log_completions,
-            "logging_steps": logging_steps,
-            "max_context_length": max_context_length,
-            "lora": lora,
+            "trainer_config": {
+                "num_generations": num_generations,
+                "temperature": temperature,
+                "beta": beta,
+                "per_device_train_batch_size": per_device_train_batch_size,
+                "learning_rate": learning_rate,
+                "gradient_accumulation_steps": gradient_accumulation_steps,
+                "gradient_checkpointing": gradient_checkpointing,
+                "lr_scheduler_type": lr_scheduler_type,
+                "warmup_steps": warmup_steps,
+                "max_prompt_length": max_prompt_length,
+                "max_completion_length": max_completion_length,
+                "bf16": bf16,
+                "scale_rewards": scale_rewards,
+                "gradient_checkpointing_kwargs": gradient_checkpointing_kwargs,
+                "max_grad_norm": max_grad_norm,
+                "report_to": report_to,
+                "log_completions": log_completions,
+                "logging_steps": logging_steps,
+                 # "max_context_length": max_context_length,
+                # "max_seq_len": max_context_length,
+                "max_steps": max_steps,
+                # "lora": lora,
+            },
+            "inference_config": {
+                "model": finetune_model,
+                "max_context_length": max_context_length,
+            },
             "gpu_config": {
-                "type": gpu_config_type,
-                gpu_config_type: self.gpu_config,
+                "type": "multi",
+                "multi": {
+                    "num_inference_gpus": 1,
+                    "num_training_gpus": 1,
+                },
             },
         }
-        url = urljoin(api_base, "fine_tuning/grpo/initialize")
+        url = f"{api_base}fine_tuning/grpo/initialize"
         headers = {"Content-Type": "application/json"}
         response = requests.post(url=url, headers=headers, json=data)
-        assert response.status_code == 200, f"Failed to initialize GRPO: {response}"
+        print(json.dumps(response.json(), indent=2))
+        response.raise_for_status()
         response = response.json()
         self.lm.model = ArborProvider._add_provider_prefix(response["current_model"])
         self.provider_job_id = response.get("job_id")
@@ -165,8 +174,8 @@ def _run_grpo_step_one_group(
         # api_key = self.lm.kwargs["api_key"]
 
         finetune_model = ArborProvider._remove_provider_prefix(self.lm.model)
-        data = {"job_id": self.provider_job_id, "model": finetune_model, "batch": train_group}
-        url = urljoin(api_base, f"fine_tuning/grpo/{self.provider_job_id}/step")
+        data = {"job_id": self.provider_job_id, "model": finetune_model, "batch": train_group["group"], "batch_id": train_group["batch_id"]}
+        url = urljoin(api_base, "fine_tuning/grpo/step")
         headers = {"Content-Type": "application/json"}
         response = requests.post(url, headers=headers, json=data)
         assert response.status_code == 200, f"Failed to run a GRPO step: {response.text}"
@@ -190,6 +199,15 @@ def step(self, train_data: list[GRPOGroup], train_data_format: TrainDataFormat |
         for group in train_data:
             self._run_grpo_step_one_group(group, train_data_format)
 
+    def get_status(self) -> GRPOStatus:
+        api_base = self.lm.kwargs["api_base"]
+        url = f"{api_base}fine_tuning/grpo/status"
+        headers = {"Content-Type": "application/json"}
+        body = {"job_id": self.provider_job_id}
+        response = requests.post(url, headers=headers, json=body)
+        assert response.status_code == 200, f"Failed to get GRPO status: {response.text}"
+        return GRPOStatus(**response.json())
+
     def save_checkpoint(self, checkpoint_name: str, score: float | None = None):
         api_base = self.lm.kwargs["api_base"]
         url = urljoin(api_base, f"fine_tuning/grpo/{self.provider_job_id}/checkpoint")
@@ -222,15 +240,14 @@ def terminate(self):
         self.lm.model = ArborProvider._add_provider_prefix(current_model)
 
     def cancel(self):
-        if self.provider_job_id:
-            api_base = self.lm.kwargs["api_base"]
-            url = urljoin(api_base, f"fine_tuning/grpo/{self.provider_job_id}/cancel")
-            headers = {"Content-Type": "application/json"}
-            response = requests.post(url, headers=headers)
-            if response.status_code == 200:
-                self.provider_job_id = None
-            else:
-                raise Exception(f"Failed to cancel GRPO job: {response.text}")
+        if ArborProvider.does_job_exist(self.provider_job_id):
+            status = self.status()
+            if ArborProvider.is_terminal_training_status(status):
+                err_msg = "Jobs that are complete cannot be canceled."
+                err_msg += f" Job with ID {self.provider_job_id} is done."
+                raise Exception(err_msg)
+            openai.fine_tuning.jobs.cancel(self.provider_job_id)
+            self.provider_job_id = None
 
     def status(self) -> TrainingStatus:
         status = ArborProvider.get_training_status(self.provider_job_id)

dspy/clients/provider.py

@@ -3,7 +3,7 @@
 from threading import Thread
 from typing import TYPE_CHECKING, Any
 
-from dspy.clients.utils_finetune import MultiGPUConfig, TrainDataFormat
+from dspy.clients.utils_finetune import TrainDataFormat
 
 if TYPE_CHECKING:
     from dspy.clients.lm import LM
@@ -36,13 +36,11 @@ def status(self):
 
 
 class ReinforceJob:
-    def __init__(self, lm: "LM", train_kwargs: dict[str, Any] | None = None, gpu_config: MultiGPUConfig = MultiGPUConfig(num_inference_gpus=1, num_training_gpus=1)):
+    def __init__(self, lm: "LM", train_kwargs: dict[str, Any] | None = None):
         self.lm = lm
         self.train_kwargs = train_kwargs or {}
-        self.gpu_config = gpu_config
         self.checkpoints = {}
         self.last_checkpoint = None
-        self.gpu_config = gpu_config
 
 
     @abstractmethod

dspy/clients/utils_finetune.py

@@ -40,14 +40,17 @@ class GRPOChatData(TypedDict):
     reward: float
 
 
-GRPOGroup = list[GRPOChatData]
-
-
-class MultiGPUConfig(TypedDict):
-    # Number of GPUs to use for inference
-    num_inference_gpus: int
-    # Number of GPUs to use for training
-    num_training_gpus: int
+class GRPOGroup(TypedDict):
+    batch_id: int | None
+    group: list[GRPOChatData]
+
+class GRPOStatus(TypedDict):
+    job_id: str
+    status: str | None = None
+    current_model: str
+    checkpoints: dict[str, str]
+    last_checkpoint: str | None = None
+    pending_batch_ids: list[int] = []
 
 
 def infer_data_format(adapter: Adapter) -> str: