Add YAML-based configuration support for vLLM main

apps/grpo/main.py

@@ -13,7 +13,7 @@
 import torch
 import torch.nn.functional as F
 from datasets import load_dataset
-from forge.actors.policy import EngineConfig, Policy, SamplingOverrides
+from forge.actors.policy import EngineArgOverrides, Policy, SamplingOverrides
 from forge.actors.replay_buffer import ReplayBuffer
 from forge.controller.actor import ForgeActor
 from forge.controller.service import ServiceConfig, shutdown_service, spawn_service
@@ -362,8 +362,10 @@ async def main():
         spawn_service(
             ServiceConfig(procs_per_replica=1, with_gpus=True, num_replicas=1),
             Policy,
-            engine_params=EngineConfig(model=model),
-            sampling_overrides=SamplingOverrides(n=group_size, max_tokens=16),
+            engine_params=EngineArgOverrides(model=model),
+            sampling_overrides=SamplingOverrides(
+                n=group_size, max_tokens=max_res_tokens
+            ),
         ),
         spawn_service(
             ServiceConfig(procs_per_replica=1, with_gpus=True, num_replicas=1),

apps/vllm/main.py

@@ -22,7 +22,9 @@
 
 async def run(cfg: DictConfig):
 
-    if "prompt" in cfg and cfg["prompt"] is not None:
+    if (prompt := cfg.get("prompt")) is None:
+        gd = cfg.policy.get("sampling_overrides", {}).get("guided_decoding", False)
+        prompt = "What is 3+5?" if gd else "Tell me a joke"
         prompt = cfg["prompt"]
     else:
         gd = cfg.policy.get("sampling_overrides", {}).get("guided_decoding", False)

src/forge/actors/policy.py

@@ -82,9 +82,9 @@ def from_dict(cls, d: Mapping):
 
 
 @dataclass
-class EngineConfig(EngineArgs):
+class EngineArgOverrides(EngineArgs):
     """
-    EngineConfig extends EngineArgs with worker-specific fields.
+    EngineArgOverrides extends EngineArgs with worker-specific fields.
     Overlapping keys in input dict will override EngineArgs defaults.
     """
 
@@ -103,8 +103,12 @@ def from_dict(cls, d: Mapping):
 
 @dataclass
 class Policy(PolicyInterface):
-    engine_params: EngineConfig | Mapping = field(default_factory=EngineConfig)
-    sampling_overrides: SamplingOverrides = field(default_factory=SamplingOverrides)
+    engine_params: EngineArgOverrides | Mapping = field(
+        default_factory=EngineArgOverrides
+    )
+    sampling_overrides: SamplingOverrides | Mapping = field(
+        default_factory=SamplingOverrides
+    )
     available_devices: str | None = None
     # Gets set up by setup
     sampling_params: SamplingParams | None = None
@@ -119,7 +123,7 @@ def __post_init__(self):
         self._worker_procs: ProcMesh | None = None
         self.weights_version: int = 0
         if isinstance(self.engine_params, Mapping):
-            self.engine_params = EngineConfig.from_dict(self.engine_params)
+            self.engine_params = EngineArgOverrides.from_dict(self.engine_params)
         if isinstance(self.sampling_overrides, Mapping):
             self.sampling_overrides = SamplingOverrides.from_dict(
                 self.sampling_overrides
@@ -130,7 +134,7 @@ async def launch(  # pyright: ignore[reportIncompatibleMethodOverride]
         cls: type["Policy"],
         *,
         process_config: ProcessConfig,
-        engine_params: EngineConfig | Mapping = EngineConfig(),
+        engine_params: EngineArgOverrides | Mapping = EngineArgOverrides(),
         sampling_overrides: SamplingOverrides | Mapping = SamplingOverrides(),
         available_devices: str | None = None,
         store: MultiProcessStore | None = None,
@@ -148,7 +152,7 @@ async def launch(  # pyright: ignore[reportIncompatibleMethodOverride]
         policy_proc = await get_proc_mesh(process_config=policy_proc_config)
 
         if isinstance(engine_params, Mapping):
-            engine_params = EngineConfig.from_dict(engine_params)
+            engine_params = EngineArgOverrides.from_dict(engine_params)
 
         if isinstance(engine_params, Mapping):
             sampling_overrides = SamplingOverrides(**sampling_overrides)
@@ -368,7 +372,7 @@ async def stop(self):
 
 @dataclass
 class PolicyWorker(ForgeActor):
-    vllm_args: EngineConfig | dict = EngineConfig()
+    vllm_args: EngineArgOverrides | Mapping = EngineArgOverrides()
     state_dict_key: str = "model_state_dict"
 
     def __post_init__(self):
@@ -384,12 +388,9 @@ def __post_init__(self):
         - all LLM generate methods, verify against LLM inputs
         - all executor methods verify no changes
         """
-        if isinstance(self.vllm_args, dict):
-            self.vllm_args = EngineConfig.from_dict(self.vllm_args)
-        elif not isinstance(self.vllm_args, EngineConfig):
-            raise TypeError(
-                f"vllm_args must be a EngineConfig or dict, got {type(self.vllm_args)}"
-            )
+        if isinstance(self.vllm_args, Mapping):
+            self.vllm_args = EngineArgOverrides.from_dict(self.vllm_args)
+
         # Original method returns False when not run in the main thread
         self.vllm_args._is_v1_supported_oracle = lambda *_: True
         # Build Config

tests/integration_tests/test_policy_update.py

@@ -11,7 +11,7 @@
 
 import torch
 
-from forge.actors.policy import EngineConfig, Policy, SamplingOverrides
+from forge.actors.policy import EngineArgOverrides, Policy, SamplingOverrides
 from forge.controller.service import ServiceConfig, spawn_service
 from forge.data.sharding import VLLMSharding
 from torchstore import MultiProcessStore
@@ -168,7 +168,7 @@ def validate_loaded_tensors_equals_original(
 
 def get_configs(worker_size: int, model_name: str) -> Tuple[Dict, ServiceConfig]:
 
-    engine_params = EngineConfig(
+    engine_params = EngineArgOverrides(
         model=model_name,
         tensor_parallel_size=worker_size,
         pipeline_parallel_size=1,

tests/unit_tests/test_policy_config.py

@@ -9,7 +9,7 @@
 
 import yaml
 
-from forge.actors.policy import EngineConfig, Policy, SamplingOverrides
+from forge.actors.policy import EngineArgOverrides, Policy, SamplingOverrides
 
 
 class TestPolicyConfig(unittest.TestCase):
@@ -20,7 +20,7 @@ def test_policy_default_initialization(self):
         policy = Policy()
 
         # Default factories
-        self.assertIsInstance(policy.engine_params, EngineConfig)
+        self.assertIsInstance(policy.engine_params, EngineArgOverrides)
         self.assertIsInstance(policy.sampling_overrides, SamplingOverrides)
         self.assertIsNone(policy.available_devices)
 
@@ -62,7 +62,7 @@ def test_policy_with_dict_configs(self):
             available_devices="test-gpu-device-abcd",
         )
 
-        self.assertIsInstance(policy.engine_params, EngineConfig)
+        self.assertIsInstance(policy.engine_params, EngineArgOverrides)
         self.assertIsInstance(policy.sampling_overrides, SamplingOverrides)
 
         # Test basic fields
@@ -124,15 +124,15 @@ def test_policy_yaml_config_loading(self):
 
             self.assertEqual(policy.available_devices, "yaml-test-device-xyz")
 
-    def test_engineconfig_ignores_invalid_keys(self):
-        """EngineConfig.from_dict ignores unexpected keys."""
+    def test_engineargoverrides_ignores_invalid_keys(self):
+        """EngineArgOverrides.from_dict ignores unexpected keys."""
         engine_params = {
             "model": "custom-model",
             "tensor_parallel_size": 2,
             "invalid_key_123": "should be ignored",
         }
 
-        config = EngineConfig.from_dict(engine_params)
+        config = EngineArgOverrides.from_dict(engine_params)
 
         self.assertEqual(config.model, "custom-model")
         self.assertEqual(config.tensor_parallel_size, 2)