resolving merge conflicts

Author: pradeepfn

apps/grpo/main.py

@@ -4,6 +4,8 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
+# Usage: python -m apps.grpo.main --config apps/grpo/qwen3_1_7b.yaml
+
 import asyncio
 import logging
 import uuid
@@ -13,13 +15,16 @@
 import torch
 import torch.nn.functional as F
 from datasets import load_dataset
-from forge.actors.policy import Policy, PolicyConfig, SamplingOverrides, WorkerConfig
+from forge.actors.policy import Policy
 from forge.actors.replay_buffer import ReplayBuffer
+from forge.cli.config import parse
 from forge.controller.actor import ForgeActor
 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 omegaconf import DictConfig
+from src.forge.data.utils import exclude_service
 from torch import nn
 from transformers import AutoModelForCausalLM
 from vllm.transformers_utils.tokenizer import get_tokenizer
@@ -286,11 +291,11 @@ async def forward(self, episode: Episode) -> torch.Tensor:
 class DatasetActor(ForgeActor):
     """Actor wrapper for HuggingFace dataset to provide async interface."""
 
-    path: str
-    revision: str
-    data_split: str
-    streaming: bool
-    model: str
+    path: str = "openai/gsm8k"
+    revision: str = "main"
+    data_split: str = "train"
+    streaming: bool = True
+    model: str = "Qwen/Qwen3-1.7B-Base"
 
     @endpoint
     def setup(self):
@@ -326,12 +331,13 @@ async def pad_token(self):
         return self.tokenizer.pad_token_id
 
 
-async def main():
+async def main(cfg: DictConfig):
     """Main GRPO training loop with rollout and training processes."""
-    group_size = 4
-    model = "Qwen/Qwen3-1.7B-Base"
-    max_req_tokens = 512
-    max_res_tokens = 128
+    # Get parameters from config with fallbacks
+    group_size = cfg.group_size
+    model = cfg.model
+    max_req_tokens = cfg.max_req_tokens
+    max_res_tokens = cfg.max_res_tokens
 
     # ---- Setup WandB Logger ---- #
     logger = get_metric_logger(
@@ -351,47 +357,37 @@ async def main():
         reward_actor,
     ) = await asyncio.gather(
         spawn_service(
-            ServiceConfig(procs_per_replica=1, num_replicas=1),
+            ServiceConfig(**cfg.dataset.service),
             DatasetActor,
-            path="openai/gsm8k",
-            revision="main",
-            data_split="train",
-            streaming=True,
-            model=model,
+            **exclude_service(cfg.dataset),
         ),
         spawn_service(
-            ServiceConfig(procs_per_replica=1, with_gpus=True, num_replicas=1),
+            ServiceConfig(**cfg.policy.service),
             Policy,
-            config=PolicyConfig(
-                worker_params=WorkerConfig(model=model),
-                sampling_params=SamplingOverrides(
-                    n=group_size, max_tokens=max_res_tokens
-                ),
-            ),
+            **exclude_service(cfg.policy),
         ),
         spawn_service(
-            ServiceConfig(procs_per_replica=1, with_gpus=True, num_replicas=1),
+            ServiceConfig(**cfg.trainer.service),
             Trainer,
-            learning_rate=1e-5,
             model_name=model,
+            **exclude_service(cfg.trainer),
         ),
         spawn_service(
-            ServiceConfig(procs_per_replica=1, num_replicas=1),
+            ServiceConfig(**cfg.replay_buffer.service),
             ReplayBuffer,
-            batch_size=4,
-            max_policy_age=1,
+            **exclude_service(cfg.replay_buffer),
         ),
         spawn_service(
-            ServiceConfig(procs_per_replica=1, num_replicas=1),
+            ServiceConfig(**cfg.compute_advantages.service),
             ComputeAdvantages,
         ),
         spawn_service(
-            ServiceConfig(procs_per_replica=1, num_replicas=1, with_gpus=True),
+            ServiceConfig(**cfg.ref_model.service),
             RefModel,
             model_name=model,
         ),
         spawn_service(
-            ServiceConfig(procs_per_replica=1, num_replicas=1),
+            ServiceConfig(**cfg.reward_actor.service),
             RewardActor,
             reward_functions=[MathReward(), ThinkingReward()],
         ),
@@ -485,5 +481,10 @@ async def continuous_training():
         )
 
 
+@parse
+def recipe_main(cfg: DictConfig) -> None:
+    asyncio.run(main(cfg))
+
+
 if __name__ == "__main__":
-    asyncio.run(main())
+    recipe_main()

apps/grpo/qwen3_1_7b.yaml

@@ -0,0 +1,74 @@
+# GRPO Training Configuration
+
+# Global configuration
+group_size: 4
+batch_size: 4
+max_req_tokens: 512
+max_res_tokens: 128
+model: "Qwen/Qwen3-1.7B-Base"
+
+# Dataset configuration
+dataset:
+  path: "openai/gsm8k"
+  revision: "main"
+  data_split: "train"
+  streaming: true
+  service:
+    procs_per_replica: 1
+    num_replicas: 1
+    with_gpus: false
+
+# Policy configuration
+policy:
+  engine_config:
+    model: ${model}
+    tensor_parallel_size: 1
+    pipeline_parallel_size: 1
+    enforce_eager: true
+  sampling_config:
+    n: 4
+    max_tokens: 128
+    temperature: 1.0
+    top_p: 1.0
+  service:
+    procs_per_replica: 1
+    num_replicas: 1
+    with_gpus: true
+
+# Trainer configuration
+trainer:
+  learning_rate: 1e-5
+  service:
+    procs_per_replica: 1
+    num_replicas: 1
+    with_gpus: true
+
+# Replay buffer configuration
+replay_buffer:
+  batch_size: ${batch_size}
+  max_policy_age: 0
+  service:
+    procs_per_replica: 1
+    num_replicas: 1
+    with_gpus: false
+
+# Compute advantages configuration
+compute_advantages:
+  service:
+    procs_per_replica: 1
+    num_replicas: 1
+    with_gpus: false
+
+# Reference model configuration
+ref_model:
+  service:
+    procs_per_replica: 1
+    num_replicas: 1
+    with_gpus: true
+
+# Reward actor configuration
+reward_actor:
+  service:
+    procs_per_replica: 1
+    num_replicas: 1
+    with_gpus: false

apps/vllm/llama3_8b.yaml

@@ -0,0 +1,19 @@
+policy:
+  engine_config:
+    model: "meta-llama/Llama-3.1-8B-Instruct"
+    tensor_parallel_size: 2
+    pipeline_parallel_size: 1
+    enforce_eager: true
+  sampling_config:
+    n: 2
+    guided_decoding: false
+    max_tokens: 512
+  available_devices: null
+  service:
+    procs_per_replica: 2
+    num_replicas: 1
+    with_gpus: true
+
+
+# Optional, otherwise argparse fallback kicks in
+prompt: "Tell me a joke"

apps/vllm/main.py

@@ -6,95 +6,34 @@
 
 """To run:
 export HF_HUB_DISABLE_XET=1
-python -m apps.vllm.main --guided-decoding --num-samples 3
-
+python -m apps.vllm.main --config apps/vllm/llama3_8b.yaml
 """
 
-import argparse
 import asyncio
-from argparse import Namespace
 
-from forge.actors.policy import Policy, PolicyConfig, SamplingOverrides, WorkerConfig
+from forge.actors.policy import Policy
+from forge.cli.config import parse
 from forge.controller.service import ServiceConfig, shutdown_service, spawn_service
-from vllm.outputs import RequestOutput
-from vllm.transformers_utils.tokenizer import get_tokenizer
-
-
-async def main():
-    """Main application for running vLLM policy inference."""
-    args = parse_args()
 
-    # Create configuration objects
-    policy_config, service_config = get_configs(args)
-
-    # Resolve the Prompts
-    if args.prompt is None:
-        prompt = "What is 3+5?" if args.guided_decoding else "Tell me a joke"
-    else:
-        prompt = args.prompt
-
-    # format prompt
-    tokenizer = get_tokenizer(policy_config.worker_params.model)
-    messages = [{"role": "user", "content": prompt}]
-    prompt = tokenizer.apply_chat_template(
-        messages, tokenize=False, add_generation_prompt=True
-    )
-
-    # Run the policy
-    await run_vllm(service_config, policy_config, prompt)
-
-
-def parse_args() -> Namespace:
-    parser = argparse.ArgumentParser(description="VLLM Policy Inference Application")
-    parser.add_argument(
-        "--model",
-        type=str,
-        default="Qwen/Qwen3-1.7B",  # "meta-llama/Llama-3.1-8B-Instruct",
-        help="Model to use",
-    )
-    parser.add_argument(
-        "--num-samples", type=int, default=2, help="Number of samples to generate"
-    )
-    parser.add_argument(
-        "--guided-decoding", action="store_true", help="Enable guided decoding"
-    )
-    parser.add_argument(
-        "--prompt", type=str, default=None, help="Custom prompt to use for generation"
-    )
-    return parser.parse_args()
+from omegaconf import DictConfig
+from src.forge.data.utils import exclude_service
+from vllm.outputs import RequestOutput
 
 
-def get_configs(args: Namespace) -> (PolicyConfig, ServiceConfig):
+async def run(cfg: DictConfig):
 
-    worker_size = 2
-    worker_params = WorkerConfig(
-        model=args.model,
-        tensor_parallel_size=worker_size,
-        pipeline_parallel_size=1,
-        enforce_eager=True,
-        vllm_args=None,
-    )
+    if (prompt := cfg.get("prompt")) is None:
+        gd = cfg.policy.get("sampling_config", {}).get("guided_decoding", False)
+        prompt = "What is 3+5?" if gd else "Tell me a joke"
 
-    sampling_params = SamplingOverrides(
-        n=args.num_samples,
-        guided_decoding=args.guided_decoding,
-        max_tokens=16,
-    )
+    print("Spawning service...")
 
-    policy_config = PolicyConfig(
-        worker_params=worker_params, sampling_params=sampling_params
-    )
-    service_config = ServiceConfig(
-        procs_per_replica=worker_size, num_replicas=1, with_gpus=True
+    policy = await spawn_service(
+        ServiceConfig(**cfg.policy.service),
+        Policy,
+        **exclude_service(cfg.policy),
     )
 
-    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)
-
     async with policy.session():
         print("Requesting generation...")
         response_output: RequestOutput = await policy.generate.choose(prompt=prompt)
@@ -112,5 +51,10 @@ async def run_vllm(service_config: ServiceConfig, config: PolicyConfig, prompt:
     await shutdown_service(policy)
 
 
+@parse
+def recipe_main(cfg: DictConfig) -> None:
+    asyncio.run(run(cfg))
+
+
 if __name__ == "__main__":
-    asyncio.run(main())
+    recipe_main()