Seg fault

Author: joecummings

apps/grpo/main.py

@@ -14,9 +14,11 @@
 
 import torch
 import torch.nn.functional as F
+import torchstore as ts
 from datasets import load_dataset
 from forge.actors.policy import Policy
 from forge.actors.replay_buffer import ReplayBuffer
+from forge.actors.trainer import _qwen3_hf_to_vllm
 from forge.cli.config import parse
 from forge.controller.actor import ForgeActor
 from forge.controller.service import ServiceConfig, shutdown_service, spawn_service
@@ -26,8 +28,7 @@
 from omegaconf import DictConfig
 from src.forge.data.utils import exclude_service
 from torch import nn
-from torchstore import MultiProcessStore
-from torchstore._state_dict_utils import DELIM, push_state_dict
+from torchstore.state_dict_utils import DELIM, put_state_dict
 from transformers import AutoModelForCausalLM
 from vllm.transformers_utils.tokenizer import get_tokenizer
 
@@ -144,12 +145,11 @@ class Trainer(ForgeActor):
     learning_rate: float = 1e-5
     beta: float = 0.1
     device: torch.device | None = None
-    store: MultiProcessStore | None = None
     state_dict_key: str = "model_state_dict"
     dp_rank: int = 0  # TODO: support data parallelism, hard code it for now
 
     @endpoint
-    def setup(self):
+    async def setup(self):
         if self.device is None:
             self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
@@ -167,45 +167,9 @@ def setup(self):
 
         self.loss = SimpleGRPOLoss(self.beta)
 
-        self.logger.info(f"Trainer model initialized on {self.device}")
+        self.store = await ts.initialize()
 
-    def _qwen3_hf_to_vllm(self, saved_sd):
-        """Convert transformers state dict to vLLM format."""
-        load_sd = {}
-        num_layers = 28  # For Qwen3-1.7B
-
-        # Copy over directly mapped keys
-        for k in saved_sd:
-            if any(
-                x in k
-                for x in [
-                    "down_proj",
-                    "input_layernorm",
-                    "post_attention_layernorm",
-                    "o_proj",
-                    "norm.weight",
-                    "embed_tokens.weight",
-                    "lm_head.weight",
-                ]
-            ):
-                load_sd[k] = saved_sd[k]
-
-        # Fuse QKV and gate_up_proj
-        for i in range(num_layers):
-            prefix = f"model.layers.{i}."
-
-            # QKV fusion
-            q = saved_sd[prefix + "self_attn.q_proj.weight"]
-            k = saved_sd[prefix + "self_attn.k_proj.weight"]
-            v = saved_sd[prefix + "self_attn.v_proj.weight"]
-            load_sd[prefix + "self_attn.qkv_proj.weight"] = torch.cat([q, k, v], dim=0)
-
-            # MLP gate_up_proj fusion
-            gate = saved_sd[prefix + "mlp.gate_proj.weight"]
-            up = saved_sd[prefix + "mlp.up_proj.weight"]
-            load_sd[prefix + "mlp.gate_up_proj.weight"] = torch.cat([gate, up], dim=0)
-
-        return load_sd
+        self.logger.info(f"Trainer model initialized on {self.device}")
 
     @endpoint
     async def train_step(self, batch: list[list[Episode]]):
@@ -238,16 +202,16 @@ async def train_step(self, batch: list[list[Episode]]):
         loss.backward()
         self.optimizer.step()
 
-        return loss.detach()
+        return loss.item()
 
     @endpoint
     async def push_weights(self, version: int):
         """Update policy model weights with trainer's current weights."""
         start_time = time.time()
-        assert self.store is not None, "Store must be provided to save weights"
+        assert self.store is not None, "Store must be initialized to save weights"
         key = f"{self.state_dict_key}{DELIM}{version}"  # Use version as unique id
-        new_sd = self._qwen3_hf_to_vllm(self.model.state_dict())
-        await push_state_dict(self.store, new_sd, key)
+        new_sd = _qwen3_hf_to_vllm(self.model.state_dict(), num_layers=28)
+        await put_state_dict(self.store, new_sd, key)
         end_time = time.time()
         self.logger.debug(
             f"Pushed weights to {key} in {end_time - start_time:.2f} seconds"
@@ -322,11 +286,11 @@ class DatasetActor(ForgeActor):
     revision: str = "main"
     data_split: str = "train"
     streaming: bool = True
-    tokenizer: str = "Qwen/Qwen3-1.7B"
+    model: str = "Qwen/Qwen3-1.7B"
 
     @endpoint
     def setup(self):
-        self._tokenizer = get_tokenizer(self.tokenizer)
+        self._tokenizer = get_tokenizer(self.model)
 
         def gsm8k_transform(sample):
             system_prompt = """
@@ -380,7 +344,6 @@ async def main(cfg: DictConfig):
     )
 
     # ---- Setup services ---- #
-    store = await MultiProcessStore.create_store()
     (
         dataloader,
         policy,
@@ -399,13 +362,11 @@ async def main(cfg: DictConfig):
             ServiceConfig(**cfg.policy.service),
             Policy,
             **exclude_service(cfg.policy),
-            store=store,
         ),
         spawn_service(
             ServiceConfig(**cfg.trainer.service),
             Trainer,
             **exclude_service(cfg.trainer),
-            store=store,
         ),
         spawn_service(
             ServiceConfig(**cfg.replay_buffer.service),

apps/grpo/qwen3_1_7b.yaml

@@ -13,7 +13,7 @@ dataset:
   revision: "main"
   data_split: "train"
   streaming: true
-  tokenizer: ${model}
+  model: ${model}
   service:
     procs_per_replica: 1
     num_replicas: 1

src/forge/actors/policy.py

@@ -13,9 +13,9 @@
 from dataclasses import asdict, dataclass, field, fields
 
 import torch
+import torchstore as ts
 from monarch.actor import current_rank, endpoint, ProcMesh
-from torchstore import MultiProcessStore
-from torchstore._state_dict_utils import DELIM
+from torchstore.state_dict_utils import DELIM
 
 from vllm.engine.arg_utils import EngineArgs
 from vllm.entrypoints.utils import _validate_truncation_size
@@ -107,14 +107,13 @@ class Policy(PolicyInterface):
     lora_request: LoRARequest | None = None
     tokenization_kwargs: dict = field(default_factory=dict)
     policy_worker: "PolicyWorker" = None
-    store: MultiProcessStore | None = None
 
     def __post_init__(self):
         self._run_task: asyncio.Task | None = None
         self._policy_proc: ProcMesh | None = None
         self._worker_procs: ProcMesh | None = None
         self.weights_version: int = 0
-        self.running: bool = False
+        self.running = False
         if isinstance(self.engine_config, Mapping):
             self.engine_config = EngineConfig.from_dict(self.engine_config)
         if isinstance(self.sampling_config, Mapping):
@@ -128,7 +127,6 @@ async def launch(  # pyright: ignore[reportIncompatibleMethodOverride]
         engine_config: EngineConfig | Mapping = EngineConfig(),
         sampling_config: SamplingConfig | Mapping = SamplingConfig(),
         available_devices: str | None = None,
-        store: MultiProcessStore | None = None,
         **kwargs,
     ) -> "Policy":
         # Note - get_proc_mesh will set MASTER_ADDR, MASTER_PORT and CUDA_VISIBLE_DEVICES
@@ -161,7 +159,6 @@ async def launch(  # pyright: ignore[reportIncompatibleMethodOverride]
             sampling_config=sampling_config,
             available_devices=available_devices,
             policy_worker=workers,
-            store=store,
         )
         policy._policy_proc = policy_proc
         policy._worker_procs = worker_procs
@@ -189,7 +186,7 @@ async def shutdown(  # pyright: ignore[reportIncompatibleMethodOverride]
     async def setup(self):
         # Set up policy_worker
         assert self.policy_worker is not None, "Policy worker should not be None"
-        await self.policy_worker.setup.call(store=self.store)
+        await self.policy_worker.setup.call()
 
         self.request_id = 0
         self.requests: dict[str, tuple[None | ParentRequest, asyncio.Future]] = {}
@@ -343,9 +340,8 @@ async def run(self):
 
             for request_output in processed_outputs.request_outputs:
                 if request_output.finished:
-                    if request_output.request_id in self.requests:
-                        _, fut = self.requests.pop(request_output.request_id)
-                        fut.set_result(request_output)
+                    _, fut = self.requests.pop(request_output.request_id)
+                    fut.set_result(request_output)
 
     @endpoint
     async def update_weights(self):
@@ -403,8 +399,8 @@ def __post_init__(self):
         self.vllm_args = self.vllm_args.create_engine_config(UsageContext.LLM_CLASS)
 
     @endpoint
-    async def setup(self, store: MultiProcessStore = None):
-        self.torchstore = store
+    async def setup(self):
+        self.store = await ts.initialize()
         # TODO: remove ["gpus"] when monarch implements a flat rank
         self.rank = current_rank()["gpus"]
         self.worker = self.setup_worker()
@@ -428,7 +424,7 @@ async def _load_tensor_parallel_state_dict(
 
             # Load the full tensor from torchstore
             # TODO: only get the part of the tensor that is needed
-            stored_tensor = await self.torchstore.get(
+            stored_tensor = await self.store.get(
                 f"{self.state_dict_key}{DELIM}{version}{DELIM}{param_name}"
             )
             sharding.load_from_source_to_target(
@@ -440,7 +436,7 @@ async def _load_tensor_parallel_state_dict(
     @endpoint
     async def update(self, version: int):
         """Update model weights by reading state dict from torchstore"""
-        if self.torchstore is None:
+        if self.store is None:
             raise Exception("No torchstore configured, skipping model update")
         key = f"{self.state_dict_key}{DELIM}{version}"
         model = self.worker.model_runner.model

src/forge/actors/trainer.py

@@ -268,3 +268,49 @@ def push_weights(self) -> None:
     async def cleanup(self) -> None:
         if self.engine.checkpointer:
             self.engine.checkpointer.close()
+
+
+def _qwen3_hf_to_vllm(
+    sd: dict[str, torch.Tensor], num_layers: int
+) -> dict[str, torch.Tensor]:
+    """Convert transformers state dict to vLLM format. Specifically, this fuses
+    QKV projection and MLP gate_up_proj layers.
+
+    Args:
+        sd (dict): State dict from HF model.
+        num_layers (int): Number of layers in the model.
+
+    Returns:
+        dict: State dict in vLLM format.
+    """
+    load_sd = {}
+
+    # Copy over directly mapped keys
+    for k in sd:
+        if any(
+            x in k
+            for x in [
+                "down_proj",
+                "input_layernorm",
+                "post_attention_layernorm",
+                "o_proj",
+                "norm.weight",
+                "embed_tokens.weight",
+                "lm_head.weight",
+            ]
+        ):
+            load_sd[k] = sd[k]
+
+    for i in range(num_layers):
+        prefix = f"model.layers.{i}."
+        # QKV fusion
+        q = sd[prefix + "self_attn.q_proj.weight"]
+        k = sd[prefix + "self_attn.k_proj.weight"]
+        v = sd[prefix + "self_attn.v_proj.weight"]
+        load_sd[prefix + "self_attn.qkv_proj.weight"] = torch.cat([q, k, v], dim=0)
+        # MLP gate_up_proj fusion
+        gate = sd[prefix + "mlp.gate_proj.weight"]
+        up = sd[prefix + "mlp.up_proj.weight"]
+        load_sd[prefix + "mlp.gate_up_proj.weight"] = torch.cat([gate, up], dim=0)
+
+    return load_sd