Debugging Cuda issue

Author: Jack-Khuu

apps/grpo/main.py

@@ -13,13 +13,18 @@
 import torch
 from datasets import load_dataset
 from forge.actors.policy import Policy, PolicyConfig, SamplingOverrides, WorkerConfig
-from forge.actors.reference_actor import compute_sequence_logprobs, RefModel
+from forge.actors.reference_actor import (
+    compute_sequence_logprobs,
+    HuggingFaceRefModel,
+    TitanRefModel,
+)
 from forge.actors.replay_buffer import ReplayBuffer
 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 torchtitan.config.job_config import Model as TitanJobModelConfig
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
 logger = logging.getLogger(__name__)
@@ -329,10 +334,16 @@ async def main():
             gamma=0.99,
             lambda_=0.95,
         ),
+        # spawn_service(
+        #     ServiceConfig(procs_per_replica=1, num_replicas=1, with_gpus=True),
+        #     RefModel,
+        #     model_name=model,
+        # ),
+        # GOAL: Swap this in and everything should just "work"
         spawn_service(
             ServiceConfig(procs_per_replica=1, num_replicas=1, with_gpus=True),
-            RefModel,
-            model_name=model,
+            TitanRefModel,
+            model=TitanJobModelConfig(name=model),
         ),
         spawn_service(
             ServiceConfig(procs_per_replica=1, num_replicas=1),

apps/grpo/test.py

@@ -0,0 +1,101 @@
+import asyncio
+
+from datasets import load_dataset
+
+from forge.actors.policy import Policy, PolicyConfig, SamplingOverrides, WorkerConfig
+from forge.actors.reference_actor import HuggingFaceRefModel, TitanRefModel
+
+from forge.controller.actor import ForgeActor
+from forge.controller.service import ServiceConfig, shutdown_service, spawn_service
+from monarch.actor import endpoint
+
+
+class DatasetActor(ForgeActor):
+    """Actor wrapper for HuggingFace dataset to provide async interface."""
+
+    def __init__(
+        self, path: str, config_name: str, split: str, streaming: bool, **kwargs
+    ):
+        super().__init__()
+
+        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(path, config_name, split=split, streaming=streaming)
+        ds = ds.map(gsm8k_to_messages)
+        ds = ds.shuffle()
+        self._iterator = iter(ds)
+
+    @endpoint
+    async def __next__(self) -> dict[str, str] | None:
+        return next(self._iterator)
+
+
+# Sandbox; will be removed
+async def main():
+    group_size = 1
+
+    # For Torchtitan
+    model = "Qwen/Qwen3-1.7B"
+
+    # Spawn Reference "Agents"
+    hf_model = await spawn_service(
+        ServiceConfig(procs_per_replica=1, num_replicas=1, with_gpus=True),
+        HuggingFaceRefModel,
+        model_name=model,
+    )
+    titan_model = await spawn_service(
+        ServiceConfig(procs_per_replica=1, num_replicas=1, with_gpus=True),
+        TitanRefModel,
+    )
+
+    # Spawn Policy for getting responses
+    policy = await 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),
+        ),
+    )
+
+    # Load Dataset
+    dataloader = await spawn_service(
+        ServiceConfig(procs_per_replica=1, num_replicas=1),
+        DatasetActor,
+        path="openai/gsm8k",
+        config_name="main",
+        split="train",
+        streaming=True,
+    )
+    sample = await dataloader.__next__.choose()
+    prompt, target = sample["question"], sample["answer"]
+    print("Sample: ", sample)
+
+    # Generate output from policy, then pass to reference agents
+    actions = await policy.generate.choose(prompt)
+    for action in actions:
+        print("Generated Action tok_ids: ", action.token_ids)
+
+        print("HuggingFace Results")
+        hf_logprobs: float = await hf_model.forward.choose(action.token_ids)
+        print("HF logprob: ", hf_logprobs)
+
+        print("Titan Results")
+        titan_logprobs: float = await titan_model.forward.choose(action.token_ids)
+        print("Titan logprob: ", titan_logprobs)
+        # TODO: Update forward to convert probs (logits) to logprobs
+
+    await asyncio.gather(
+        shutdown_service(policy),
+        shutdown_service(dataloader),
+        shutdown_service(hf_model),
+        shutdown_service(titan_model),
+    )
+
+
+if __name__ == "__main__":
+    asyncio.run(main())

apps/rl/main.py

@@ -26,7 +26,7 @@
 
 
 async def run(cfg: DictConfig):
-    trainer, buffer, reference = await asyncio.gather(
+    trainer, buffer = await asyncio.gather(
         spawn_actors(
             name="trainer",
             actor_cls=RLTrainer,
@@ -40,24 +40,18 @@ async def run(cfg: DictConfig):
             cfg=cfg.replay_buffer,
             processes=cfg.replay_buffer.pop("processes"),
         ),
-        spawn_actors(
-            name="reference_actor",
-            actor_cls=ReferenceActor,
-        ),
     )
     print("Actors spawned")
 
     # Initialize everything
     await asyncio.gather(
         buffer.setup.call(),
         trainer.setup.call(),
-        reference.setup.call(),
     )
     print("Setup done")
 
     print("shutting down...")
     await asyncio.gather(*[a.mesh.stop() for a in [trainer]])
-    await reference.cleanup.call()
 
 
 @parse

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}")

src/forge/actors/reference_actor.py

@@ -37,6 +37,13 @@
 
 @dataclass
 class ReferenceActor(ForgeActor):
+    """
+    Original idea (not updated); On second throught this might be overkill
+    if we can rely on the Service Replicas to handle the queue since there's no
+    real pre/post proc or host management (maybe later for DP?). For now just
+    directly spin up services of the reference models
+    """
+
     model: Model = field(default_factory=Model)
     # parallelism: Parallelism = field(default_factory=Parallelism)
     # comm: Comm = field(default_factory=Comm)
@@ -95,13 +102,18 @@ async def setup(self):
         # Spawn the RefModel
         self.ref_model = await spawn_service(
             default_service_cfg,
-            RefModel,
+            HuggingFaceRefModel,
             model_name=self.model.name,
             device=self.device,
         )
 
         # Kick off background processing
-        asyncio.create_task(self.run_processing.call())
+        self.start_processing()
+
+    def start_processing(self):
+        """Start the replica's processing loop if not already running."""
+        if self._run_task is None or self._run_task.done():
+            self._run_task = asyncio.create_task(self.run())
 
     @endpoint
     async def forward(self, token_ids: list[int]) -> torch.Tensor:
@@ -112,8 +124,7 @@ async def forward(self, token_ids: list[int]) -> torch.Tensor:
         self.queue.append((token_ids, fut))
         return await fut
 
-    @endpoint
-    async def run_processing(self):
+    async def run(self):
         """
         Simple loop to pass things along to the ref model
         """
@@ -127,11 +138,105 @@ async def run_processing(self):
             fut.set_result(model_output)
 
     @endpoint
-    async def cleanup(self) -> None:
+    async def stop(self) -> None:
         self.running = False
 
 
-class RefModel(ForgeActor):
+@dataclass
+class TitanRefModel(ForgeActor):
+    """
+    Represents a reference actor leveraging a torchtitan model for execution
+    """
+
+    # Refer to titan JobConfig for enabling more ForgeEngine configuration
+    model: Model = field(default_factory=Model)
+    parallelism: Parallelism = field(default_factory=Parallelism)
+
+    # Populated in setup (commented out for now for engine_config parsing)
+    # engine: ForgeEngine | None = None
+
+    def __post_init__(self):
+        """Initializes config types and env variables."""
+        # Instantiate dict fields
+        for f in fields(self):
+            attr = getattr(self, f.name)
+            if isinstance(attr, Mapping):
+                setattr(self, f.name, f.type(**attr))
+            elif not isinstance(attr, f.type):
+                raise TypeError(
+                    f"{f.name} should be a {f.type} type or a dict like object"
+                )
+
+        """
+        torchrun normally hands env variables, but we need to do it ourselves
+        in monarch for now.
+        """
+        self.rank = current_rank().rank
+        self.size = math.prod(current_size().values())
+
+        env = {
+            "RANK": str(self.rank),
+            "LOCAL_RANK": str(self.rank),
+            "LOCAL_WORLD_SIZE": str(self.size),
+            "GROUP_RANK": str(self.size),
+            "GROUP_WORLD_SIZE": str(self.size),
+            "ROLE_RANK": str(self.rank),
+            "ROLE_WORLD_SIZE": str(self.size),
+            "ROLE_NAME": "rank",
+            "WORLD_SIZE": str(self.size),
+            "PYTORCH_CUDA_ALLOC_CONF": "expandable_segments:True",
+        }
+        os.environ.update(env)
+
+    @endpoint
+    async def setup(self):
+        engine_config = {f.name: getattr(self, f.name) for f in fields(self)}
+        self.engine = ForgeEngine(ForgeJobConfig(**engine_config))
+
+    @endpoint
+    async def forward(self, token_ids: list[int]) -> torch.Tensor:
+        """
+        Given a return the log_probability of the token_ids
+        (Used as the reference_logprobs for KL Divergence)
+        """
+        model_parts = self.engine.model_parts
+        parallel_dims = self.engine.parallel_dims
+
+        # Use provided token_ids directly
+        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        input_ids = torch.tensor(token_ids, dtype=torch.long).unsqueeze(0).to(device)
+
+        optional_context_parallel_ctx = (
+            dist_utils.create_context_parallel_ctx(
+                cp_mesh=parallel_dims.world_mesh["cp"],
+                cp_buffers=[inputs, labels] + [m.freqs_cis for m in model_parts],
+                cp_seq_dims=[1, 1] + [0 for _ in model_parts],
+                cp_no_restore_buffers={inputs, labels},
+                cp_rotate_method=self.job_config.parallelism.context_parallel_rotate_method,
+            )
+            if parallel_dims.cp_enabled
+            else None
+        )
+
+        if parallel_dims.pp_enabled:
+            raise NotImplementedError("PP not implemented yet")
+        else:
+            # Non-PP forward / backward
+            with self.engine.train_context(optional_context_parallel_ctx):
+                assert len(model_parts) == 1
+                with self.engine.maybe_enable_amp:
+                    pred = model_parts[0](input_ids)
+
+        # TODO: Update compute_sequence_logprobs to convert probs (logits) to logprobs
+        return pred
+
+
+# Maintained to keep GRPO app prior to migration
+class HuggingFaceRefModel(ForgeActor):
+    """
+    Represents a reference actor leveraging HuggingFace for execution
+    """
+
     def __init__(self, model_name, device: torch.device | None = None):
         super().__init__()
         self.model_name = model_name

src/forge/actors/trainer.py

@@ -12,6 +12,8 @@
 from dataclasses import dataclass, field, fields
 
 import torch
+
+from forge.controller import ForgeActor
 from monarch.actor import current_rank, current_size, endpoint
 from torchtitan.config.job_config import (
     ActivationCheckpoint,
@@ -30,8 +32,6 @@
 from torchtitan.experiments.forge.engine import ForgeEngine
 from torchtitan.experiments.forge.job_config import ForgeJobConfig
 
-from forge.controller import ForgeActor
-
 logger = logging.getLogger(__name__)
 logger.setLevel(logging.INFO)