Make torchstore actually work!

Author: joecummings

apps/grpo/main.py

@@ -28,7 +28,7 @@
 from omegaconf import DictConfig
 from src.forge.data.utils import exclude_service
 from torch import nn
-from torchstore.state_dict_utils import DELIM, put_state_dict
+from torchstore.state_dict_utils import DELIM
 from transformers import AutoModelForCausalLM
 from vllm.transformers_utils.tokenizer import get_tokenizer
 
@@ -167,8 +167,6 @@ async def setup(self):
 
         self.loss = SimpleGRPOLoss(self.beta)
 
-        self.store = await ts.initialize()
-
         self.logger.info(f"Trainer model initialized on {self.device}")
 
     @endpoint
@@ -207,11 +205,10 @@ async def train_step(self, batch: list[list[Episode]]):
     @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 initialized to save weights"
         key = f"{self.state_dict_key}{DELIM}{version}"  # Use version as unique id
         new_sd = _qwen3_hf_to_vllm(self.model.state_dict(), num_layers=28)
-        await put_state_dict(self.store, new_sd, key)
+        start_time = time.time()
+        await ts.put_state_dict(new_sd, key)
         end_time = time.time()
         self.logger.debug(
             f"Pushed weights to {key} in {end_time - start_time:.2f} seconds"
@@ -344,6 +341,7 @@ async def main(cfg: DictConfig):
     )
 
     # ---- Setup services ---- #
+    await ts.initialize()
     (
         dataloader,
         policy,

src/forge/actors/policy.py

@@ -400,7 +400,6 @@ def __post_init__(self):
 
     @endpoint
     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()
@@ -424,7 +423,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.store.get(
+            stored_tensor = await ts.get(
                 f"{self.state_dict_key}{DELIM}{version}{DELIM}{param_name}"
             )
             sharding.load_from_source_to_target(
@@ -436,8 +435,6 @@ 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.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
         current_state_dict = model.state_dict()