Merge remote-tracking branch 'origin/main' into ref-actor

Author: Jack-Khuu

apps/grpo/main.py

@@ -5,6 +5,7 @@
 # LICENSE file in the root directory of this source tree.
 
 import asyncio
+import logging
 import time
 from dataclasses import dataclass
 from typing import Callable
@@ -15,12 +16,15 @@
 from forge.actors.reference_actor import compute_sequence_logprobs, RefModel
 from forge.actors.replay_buffer import ReplayBuffer
 from forge.controller.actor import ForgeActor
-from forge.controller.service import ServiceConfig, spawn_service
+from forge.controller.service import ServiceConfig, shutdown_service, spawn_service
 from forge.data.rewards import MathReward, ThinkingReward
 from forge.util.metric_logging import get_metric_logger
 from monarch.actor import endpoint
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
+logger = logging.getLogger(__name__)
+logger.setLevel(logging.DEBUG)
+
 
 @dataclass
 class Group:
@@ -242,18 +246,18 @@ async def __call__(self, groups: list[Group]) -> list[float]:
 class DatasetActor(ForgeActor):
     """Actor wrapper for HuggingFace dataset to provide async interface."""
 
-    def __init__(self, *args, **kwargs):
+    def __init__(
+        self, path: str, config_name: str, split: str, streaming: bool, **kwargs
+    ):
         super().__init__()
-        self._setup_dataset(*args, **kwargs)
 
-    def _setup_dataset(self, *args, **kwargs):
         def gsm8k_to_messages(sample):
             question = sample["question"]
             full_answer: str = sample["answer"]
             answer = full_answer.split("#### ")[1]
             return {"question": question, "answer": answer}
 
-        ds = load_dataset(*args, **kwargs)
+        ds = load_dataset(path, config_name, split=split, streaming=streaming)
         ds = ds.map(gsm8k_to_messages)
         ds = ds.shuffle()
         self._iterator = iter(ds)
@@ -279,66 +283,69 @@ async def main():
     )
 
     # ---- Setup services ---- #
-    policy = await spawn_service(
-        ServiceConfig(procs_per_replica=1, with_gpus=True, num_replicas=1),
-        Policy,
-        PolicyConfig(
-            num_workers=1,
-            worker_params=WorkerConfig(model=model),
-            sampling_params=SamplingOverrides(num_samples=group_size, max_tokens=16),
+    (
+        dataloader,
+        policy,
+        trainer,
+        replay_buffer,
+        compute_advantages,
+        ref_model,
+        reward_actor,
+    ) = await asyncio.gather(
+        spawn_service(
+            ServiceConfig(procs_per_replica=1, num_replicas=1),
+            DatasetActor,
+            path="openai/gsm8k",
+            config_name="main",
+            split="train",
+            streaming=True,
+        ),
+        spawn_service(
+            ServiceConfig(procs_per_replica=1, with_gpus=True, num_replicas=1),
+            Policy,
+            config=PolicyConfig(
+                worker_params=WorkerConfig(model=model),
+                sampling_params=SamplingOverrides(
+                    num_samples=group_size, max_tokens=16
+                ),
+            ),
+        ),
+        spawn_service(
+            ServiceConfig(procs_per_replica=1, with_gpus=True, num_replicas=1),
+            Trainer,
+            learning_rate=1e-5,
+            beta=0.1,
+            model_name=model,
+        ),
+        spawn_service(
+            ServiceConfig(procs_per_replica=1, num_replicas=1),
+            ReplayBuffer,
+            batch_size=4,
+            max_policy_age=1,
+        ),
+        spawn_service(
+            ServiceConfig(procs_per_replica=1, num_replicas=1),
+            ComputeAdvantages,
+            gamma=0.99,
+            lambda_=0.95,
+        ),
+        spawn_service(
+            ServiceConfig(procs_per_replica=1, num_replicas=1, with_gpus=True),
+            RefModel,
+            model_name=model,
+        ),
+        spawn_service(
+            ServiceConfig(procs_per_replica=1, num_replicas=1),
+            RewardActor,
+            reward_functions=[MathReward(), ThinkingReward()],
         ),
-    )
-
-    trainer = await spawn_service(
-        ServiceConfig(procs_per_replica=1, with_gpus=True, num_replicas=1),
-        Trainer,
-        learning_rate=1e-5,
-        beta=0.1,
-        model_name=model,
-    )
-
-    replay_buffer = await spawn_service(
-        ServiceConfig(procs_per_replica=1, num_replicas=1),
-        ReplayBuffer,
-        batch_size=4,
-        max_policy_age=1,
-    )
-
-    dataloader = await spawn_service(
-        ServiceConfig(procs_per_replica=1, num_replicas=1),
-        DatasetActor,
-        "openai/gsm8k",
-        "main",
-        split="train",
-        streaming=True,
-    )
-
-    compute_advantages = await spawn_service(
-        ServiceConfig(procs_per_replica=1, num_replicas=1),
-        ComputeAdvantages,
-        gamma=0.99,
-        lambda_=0.95,
-    )
-
-    ref_model = await spawn_service(
-        ServiceConfig(procs_per_replica=1, num_replicas=1, with_gpus=True),
-        RefModel,
-        model_name=model,
-    )
-
-    reward_actor = await spawn_service(
-        ServiceConfig(procs_per_replica=1, num_replicas=1),
-        RewardActor,
-        reward_functions=[MathReward(), ThinkingReward()],
     )
 
     print("All services initialized successfully!")
 
     # ---- Core RL loops ---- #
     async def continuous_rollouts():
         rollout_count = 0
-        # TODO: Move this into setup
-        asyncio.create_task(policy.run_processing.call())
         while True:
             sample = await dataloader.__next__.choose()
             if sample is None:
@@ -409,6 +416,17 @@ async def continuous_training():
         print("Training interrupted by user")
         rollout_task.cancel()
         training_task.cancel()
+    finally:
+        print("Shutting down...")
+        await asyncio.gather(
+            shutdown_service(policy),
+            shutdown_service(trainer),
+            shutdown_service(replay_buffer),
+            shutdown_service(dataloader),
+            shutdown_service(compute_advantages),
+            shutdown_service(ref_model),
+            shutdown_service(reward_actor),
+        )
 
 
 if __name__ == "__main__":

apps/vllm/main.py

@@ -16,7 +16,7 @@
 from typing import List
 
 from forge.actors.policy import Policy, PolicyConfig, SamplingOverrides, WorkerConfig
-from forge.controller.service import ServiceConfig, spawn_service
+from forge.controller.service import ServiceConfig, shutdown_service, spawn_service
 from vllm.outputs import CompletionOutput
 
 
@@ -58,9 +58,11 @@ def parse_args() -> Namespace:
 
 
 def get_configs(args: Namespace) -> (PolicyConfig, ServiceConfig):
+
+    worker_size = 2
     worker_params = WorkerConfig(
         model=args.model,
-        tensor_parallel_size=2,
+        tensor_parallel_size=worker_size,
         pipeline_parallel_size=1,
         enforce_eager=True,
         vllm_args=None,
@@ -72,35 +74,34 @@ def get_configs(args: Namespace) -> (PolicyConfig, ServiceConfig):
     )
 
     policy_config = PolicyConfig(
-        num_workers=2, worker_params=worker_params, sampling_params=sampling_params
+        worker_params=worker_params, sampling_params=sampling_params
+    )
+    service_config = ServiceConfig(
+        procs_per_replica=worker_size, num_replicas=1, with_gpus=True
     )
-    service_config = ServiceConfig(procs_per_replica=1, num_replicas=1)
 
     return policy_config, service_config
 
 
 async def run_vllm(service_config: ServiceConfig, config: PolicyConfig, prompt: str):
     print("Spawning service...")
     policy = await spawn_service(service_config, Policy, config=config)
-    session_id = await policy.start_session()
 
-    print("Starting background processing...")
-    processing_task = asyncio.create_task(policy.run_processing.call())
+    async with policy.session():
+        print("Requesting generation...")
+        responses: List[CompletionOutput] = await policy.generate.choose(prompt=prompt)
 
-    print("Requesting generation...")
-    responses: List[CompletionOutput] = await policy.generate.choose(prompt=prompt)
+        print("\nGeneration Results:")
+        print("=" * 80)
+        for batch, response in enumerate(responses):
+            print(f"Sample {batch + 1}:")
+            print(f"User: {prompt}")
+            print(f"Assistant: {response.text}")
+            print("-" * 80)
 
-    print("\nGeneration Results:")
-    print("=" * 80)
-    for batch, response in enumerate(responses):
-        print(f"Sample {batch + 1}:")
-        print(f"User: {prompt}")
-        print(f"Assistant: {response.text}")
-        print("-" * 80)
+        print("\nShutting down...")
 
-    print("\nShutting down...")
-    await policy.shutdown.call()
-    await policy.terminate_session(session_id)
+    await shutdown_service(policy)
 
 
 if __name__ == "__main__":