Policy uses the generic completion data model (#207)

* using the completions data model

* working changes

* few mode fixes

* lints

* review comments

---------

Co-authored-by: rithesh <[email protected]>

Author: Ritesh1905

apps/grpo/main.py

@@ -305,10 +305,8 @@ async def continuous_rollouts():
                 device="cuda",
             )
             # Populate episode info and calculate rewards
-            for i, (episode, response) in enumerate(
-                zip(group.episodes, responses.outputs)
-            ):
-                episode.request_tokens = responses.prompt_token_ids
+            for i, (episode, response) in enumerate(zip(group.episodes, responses)):
+                episode.request_tokens = response.prompt_ids
                 episode.response_tokens = response.token_ids
                 episode.response = response.text
                 input_ids[i, :max_req_tokens] = episode.request_tensor

apps/toy_rl/sumdigits.py

@@ -358,7 +358,7 @@ class RewardActor(ForgeActor):
 
     @endpoint
     async def evaluate_response(self, prompt: str, response: str, target: str) -> float:
-        reward = 1.0 if response.strip() == "10" else 0.0
+        reward = 1.0 if response.strip() == target else 0.0
         return reward
 
 
@@ -396,18 +396,7 @@ def generate_sample(self, step: int) -> dict[str, str]:
 
     def generate_one(self, step: int) -> str:
         """Generate number based on training step for curriculum learning."""
-        if step < 200:
-            # Early training: 2-digit numbers (10-99)
-            min_val, max_val = 10, 99
-        elif step < 1000:
-            # Later training: 1-4 digit numbers (0-1000)
-            min_val, max_val = 0, 1000
-        elif step < 3000:
-            # Later training: 1-6 digit numbers (0-100000)
-            min_val, max_val = 0, 100000
-        else:
-            # Later training: 1-8 digit numbers (0-10000000)
-            min_val, max_val = 0, 10000000
+        min_val, max_val = 10, 100
 
         number = random.randint(min_val, max_val)
         return str(number)
@@ -497,19 +486,11 @@ async def continuous_rollouts():
             )
 
             # TODO: Parallelize the following calculation
-            for episode, response in zip(group.episodes, responses.outputs):
-                episode.request_tokens = responses.prompt_token_ids
+            for episode, response in zip(group.episodes, responses):
+                episode.request_tokens = response.prompt_ids
                 episode.response_tokens = response.token_ids
                 episode.response = response.text
-                episode.response_logprobs = torch.tensor(
-                    [
-                        top_k_dict[token].logprob
-                        for token, top_k_dict in zip(
-                            response.token_ids,
-                            response.logprobs,
-                        )
-                    ]
-                )
+                episode.response_logprobs = response.logprobs
                 episode.ref_logprobs = await ref_model.forward.choose(episode)
                 episode.reward = await reward_actor.evaluate_response.choose(
                     prompt=prompt, response=response.text, target=target

apps/vllm/main.py

@@ -17,8 +17,8 @@
 from forge.cli.config import parse
 from forge.controller.provisioner import shutdown
 
+from forge.data_models.completion import Completion
 from omegaconf import DictConfig
-from vllm.outputs import RequestOutput
 
 os.environ["HYPERACTOR_MESSAGE_DELIVERY_TIMEOUT_SECS"] = "600"
 os.environ["HYPERACTOR_CODE_MAX_FRAME_LENGTH"] = "1073741824"
@@ -36,11 +36,13 @@ async def run(cfg: DictConfig):
     try:
         async with policy.session():
             print("Requesting generation...")
-            response_output: RequestOutput = await policy.generate.choose(prompt=prompt)
+            response_output: list[Completion] = await policy.generate.choose(
+                prompt=prompt
+            )
 
             print("\nGeneration Results:")
             print("=" * 80)
-            for batch, response in enumerate(response_output.outputs):
+            for batch, response in enumerate(response_output):
                 print(f"Sample {batch + 1}:")
                 print(f"User: {prompt}")
                 print(f"Assistant: {response.text}")

src/forge/actors/policy.py

@@ -25,7 +25,7 @@
 from vllm.entrypoints.utils import _validate_truncation_size
 from vllm.executor.multiproc_worker_utils import set_multiprocessing_worker_envs
 from vllm.lora.request import LoRARequest
-from vllm.outputs import RequestOutput
+from vllm.outputs import CompletionOutput, RequestOutput
 from vllm.sampling_params import GuidedDecodingParams, RequestOutputKind, SamplingParams
 from vllm.transformers_utils.tokenizer_group import init_tokenizer_from_configs
 from vllm.usage.usage_lib import UsageContext
@@ -44,6 +44,9 @@
 from forge.controller import ForgeActor, get_proc_mesh, stop_proc_mesh
 
 from forge.data.sharding import VLLMSharding
+from forge.data_models.completion import Completion
+from forge.data_models.prompt import to_prompt
+
 from forge.interfaces import Policy as PolicyInterface
 from forge.types import ProcessConfig
 
@@ -258,7 +261,7 @@ def start_processing(self):
             self._run_task = asyncio.create_task(self.run())
 
     @endpoint
-    async def generate(self, prompt: str, priority: int = 0) -> RequestOutput:
+    async def generate(self, prompt: str, priority: int = 0) -> list[Completion]:
         """Generate a response for the given prompt
 
         Args:
@@ -362,8 +365,9 @@ async def run(self):
 
             for request_output in processed_outputs.request_outputs:
                 if request_output.finished:
+                    completions = self._to_completions(request_output)
                     _, fut = self.requests.pop(request_output.request_id)
-                    fut.set_result(request_output)
+                    fut.set_result(completions)
 
     @endpoint
     async def update_weights(self, policy_version: int):
@@ -396,6 +400,42 @@ async def get_version(self) -> int:
     async def stop(self):
         self.running = False
 
+    def _to_completions(self, request_output: RequestOutput) -> list[Completion]:
+        """Convert a RequestOutput to a list of Completion objects."""
+        completions = []
+        original_prompt = request_output.prompt
+        prompt_token_ids = request_output.prompt_token_ids
+        for output in request_output.outputs:
+            completions.append(
+                Completion(
+                    # TODO: the to_prompt encoding will be different from the original.
+                    # This is okay for now, since I don't see any direct usage of prompt using completion object.
+                    prompt=to_prompt(original_prompt),
+                    stop_reason=output.finish_reason,
+                    text=output.text,
+                    prompt_ids=torch.tensor(prompt_token_ids),
+                    token_ids=torch.tensor(output.token_ids),
+                    logprobs=self._extract_logprobs(output),
+                )
+            )
+
+        return completions
+
+    def _extract_logprobs(self, one_sample: CompletionOutput) -> torch.Tensor | None:
+        """
+        Extract log probabilities from a sample, if available.
+        """
+        if one_sample.logprobs is not None:
+            return torch.tensor(
+                [
+                    top_k_dict[token].logprob
+                    for token, top_k_dict in zip(
+                        one_sample.token_ids, one_sample.logprobs
+                    )
+                ]
+            )
+        return None
+
 
 @dataclass
 class PolicyWorker(ForgeActor):

src/forge/data_models/completion.py

@@ -8,6 +8,7 @@
 from typing import Optional
 
 import torch
+
 from forge.data_models.prompt import Prompt
 
 
@@ -28,4 +29,7 @@ class Completion:
     token_ids: torch.Tensor
 
     # the log probabilities of the target tokens
-    log_probs: Optional[torch.Tensor] = None
+    logprobs: Optional[torch.Tensor] = None
+
+    # the reason for stopping the generation
+    stop_reason: str | None = None