Merge branch 'main' into ungroup

Author: pbontrager

apps/grpo/main.py

@@ -6,6 +6,7 @@
 
 import asyncio
 import copy
+import logging
 import time
 import uuid
 from dataclasses import dataclass
@@ -15,15 +16,18 @@
 from datasets import load_dataset
 from forge.actors.policy import Policy, PolicyConfig, SamplingOverrides, WorkerConfig
 from forge.actors.replay_buffer import ReplayBuffer
-from forge.controller import ServiceConfig, spawn_service
 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 torch import nn
 from transformers import AutoModelForCausalLM
 from vllm.transformers_utils.tokenizer import get_tokenizer
 
+logger = logging.getLogger(__name__)
+logger.setLevel(logging.DEBUG)
+
 
 def compute_logprobs(
     logits: torch.Tensor, input_ids: torch.Tensor, temperature: float = 1.0
@@ -365,66 +369,60 @@ async def main():
     )
 
     # ---- Setup services ---- #
-    default_service_cfg = ServiceConfig(
-        procs_per_replica=1,
-        num_replicas=1,
-    )
-
-    policy = await spawn_service(
-        default_service_cfg,
-        Policy,
-        PolicyConfig(
-            num_workers=1,
-            worker_params=WorkerConfig(model=model),
-            sampling_params=SamplingOverrides(n=group_size, max_tokens=max_res_tokens),
-            available_devices="3",
+    (
+        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",
+            name="main",
+            data_split="train",
+            streaming=True,
+            model=model,
+        ),
+        spawn_service(
+            ServiceConfig(procs_per_replica=1, with_gpus=True, num_replicas=1),
+            Policy,
+            config=PolicyConfig(
+                worker_params=WorkerConfig(model=model),
+                sampling_params=SamplingOverrides(
+                    n=group_size, max_tokens=max_res_tokens
+                ),
+            ),
+        ),
+        spawn_service(
+            ServiceConfig(procs_per_replica=1, with_gpus=True, num_replicas=1),
+            Trainer,
+            learning_rate=1e-5,
+            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,
+        ),
+        spawn_service(
+            ServiceConfig(procs_per_replica=1, num_replicas=1, with_gpus=True),
+            RefModel,
+            model=titan_model,
+        ),
+        spawn_service(
+            ServiceConfig(procs_per_replica=1, num_replicas=1),
+            RewardActor,
+            reward_functions=[MathReward(), ThinkingReward()],
         ),
-    )
-
-    trainer = await spawn_service(
-        default_service_cfg,
-        Trainer,
-        learning_rate=1e-5,
-        beta=0.1,
-        model_name=model,
-        device=torch.device("cuda:1"),
-    )
-
-    replay_buffer = await spawn_service(
-        default_service_cfg,
-        ReplayBuffer,
-        batch_size=4,
-        max_policy_age=1,
-    )
-
-    dataloader = await spawn_service(
-        default_service_cfg,
-        DatasetActor,
-        "openai/gsm8k",
-        "main",
-        data_split="train",
-        streaming=True,
-        model=model,
-    )
-
-    compute_advantages = await spawn_service(
-        default_service_cfg,
-        ComputeAdvantages,
-        gamma=0.99,
-        lambda_=0.95,
-    )
-
-    ref_model = await spawn_service(
-        default_service_cfg,
-        RefModel,
-        model_name=model,
-        device=torch.device("cuda:2"),
-    )
-
-    reward_actor = await spawn_service(
-        default_service_cfg,
-        RewardActor,
-        reward_functions=[MathReward(), ThinkingReward()],
     )
 
     print("All services initialized successfully!")
@@ -433,8 +431,6 @@ async def main():
     async def continuous_rollouts():
         rollout_count = 0
         pad_id = dataloader.pad_token.choose()
-        # TODO: Move this into setup
-        asyncio.create_task(policy.run_processing.call())
         while True:
             sample = await dataloader.sample.choose()
             if sample is None:
@@ -501,6 +497,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/rl/llama3_8b.yaml

@@ -18,6 +18,7 @@ trainer:
   processes:
     scheduler: local # local | mast (not supported yet)
     num_hosts: 1
+    with_gpus: True
     num_procs: 4
 
   optimizer:
@@ -33,9 +34,11 @@ trainer:
     seq_len: 2048
     max_norm: 1.0
     steps: 5
-    compile: false
     dataset: "c4"
 
+  compile:
+    enable: false
+
   parallelism:
     data_parallel_replicate_degree: 1
     data_parallel_shard_degree: -1
@@ -65,6 +68,7 @@ replay_buffer:
   processes:
     scheduler: local # local | mast (not supported yet)
     num_hosts: 1
+    with_gpus: False
     num_procs: 1
 
 # policy:

apps/rl/main.py

@@ -20,7 +20,6 @@
 from forge.controller import spawn_actors
 from omegaconf import DictConfig
 
-
 logger = logging.getLogger(__name__)
 logger.setLevel(logging.INFO)
 
@@ -30,7 +29,7 @@ async def run(cfg: DictConfig):
         spawn_actors(
             name="trainer",
             actor_cls=RLTrainer,
-            cfg={"config": cfg.trainer},
+            cfg=cfg.trainer,
             processes=cfg.trainer.pop("processes"),
             set_address=True,
         ),

apps/sft_v2/llama3_8b.yaml

@@ -14,12 +14,13 @@ comm:
 model:
   name: llama3
   flavor: 8B
-  tokenizer_path: /tmp/Meta-Llama-3.1-8B-Instruct
+  tokenizer_path: /tmp/Llama-3.1-8B-Instruct
 
 processes:
   scheduler: local # local | mast (not supported yet)
   num_hosts: 1
   num_procs: 8
+  num_gpus: 8
 
 optimizer:
   name: AdamW

apps/sft_v2/main.py

@@ -6,7 +6,7 @@
 
 """To run:
 
-python -m apps.sft.main --config apps/sft/llama3_8b.yaml
+python -m apps.sft_v2.main --config apps/sft_v2/llama3_8b.yaml
 
 """
 

apps/vllm/main.py

@@ -5,7 +5,7 @@
 # LICENSE file in the root directory of this source tree.
 
 """To run:
-
+export HF_HUB_DISABLE_XET=1
 python -m apps.vllm.main --guided-decoding --num-samples 3
 
 """
@@ -16,8 +16,7 @@
 from typing import List
 
 from forge.actors.policy import Policy, PolicyConfig, SamplingOverrides, WorkerConfig
-from forge.controller.service import ServiceConfig
-from forge.controller.spawn import spawn_service
+from forge.controller.service import ServiceConfig, shutdown_service, spawn_service
 from vllm.outputs import CompletionOutput, RequestOutput
 
 
@@ -66,9 +65,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,
@@ -81,36 +82,35 @@ 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())
-
-    print("Requesting generation...")
-    request_output: RequestOutput = await policy.generate.choose(prompt=prompt)
-    responses: List[CompletionOutput] = request_output.outputs
-
-    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...")
-    await policy.shutdown.call()
-    await policy.terminate_session(session_id)
+
+    async with policy.session():
+        print("Requesting generation...")
+        request_output: RequestOutput = await policy.generate.choose(prompt=prompt)
+        responses: List[CompletionOutput] = request_output.outputs
+
+        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...")
+
+    await shutdown_service(policy)
 
 
 if __name__ == "__main__":

pyproject.toml

@@ -22,6 +22,7 @@ dependencies = [
     "tokenizers",
     # Miscellaneous
     "omegaconf",
+    "wandb",
 ]
 dynamic = ["version"]
 

src/forge/actors/__init__.py

@@ -4,7 +4,7 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
-__all__ = ["Policy", "PolicyRouter", "RLTrainer", "ReplayBuffer"]
+__all__ = ["Policy", "PolicyRouter", "RLTrainer", "ReplayBuffer", "TitanRefModel"]
 
 
 def __getattr__(name):
@@ -24,5 +24,9 @@ def __getattr__(name):
         from .replay_buffer import ReplayBuffer
 
         return ReplayBuffer
+    elif name == "TitanRefModel":
+        from .reference_actor import TitanRefModel
+
+        return TitanRefModel
     else:
         raise AttributeError(f"module {__name__} has no attribute {name}")