Weight loading working correctly with tp: use vllm builtin load_weights()

apps/grpo/main.py

@@ -19,6 +19,7 @@
 from forge.actors.policy import Policy
 from forge.actors.reference_model import ReferenceModel  # noqa: F401
 from forge.actors.replay_buffer import ReplayBuffer
+from forge.actors.torchstore_utils import get_param_key
 from forge.actors.trainer import _qwen3_hf_to_vllm
 from forge.cli.config import parse
 from forge.controller.actor import ForgeActor
@@ -185,14 +186,13 @@ 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."""
-        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)
-        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"
-        )
+        start_time = time.perf_counter()
+        hf_state_dict = self.model.state_dict()
+        for name, param in hf_state_dict.items():
+            key = get_param_key(version, name)
+            await ts.put(key, param)
+        end_time = time.perf_counter()
+        self.logger.debug(f"Pushed weights in {end_time - start_time:.2f} seconds")
 
 
 @dataclass
@@ -318,7 +318,7 @@ async def main(cfg: DictConfig):
     mlogger = get_metric_logger(
         "wandb",
         freq=1,
-        project="grpo-training",
+        project="yuxuanh-grpo-training-debug",
     )
 
     # ---- Setup services ---- #
@@ -397,20 +397,28 @@ async def continuous_rollouts():
 
     async def continuous_training():
         training_step = 0
-        policy_version = 0
         while True:
-            batch = await replay_buffer.sample.choose(
-                curr_policy_version=policy_version
-            )
+            batch = await replay_buffer.sample.choose(curr_policy_version=training_step)
             if batch is None:
                 await asyncio.sleep(0.1)
             else:
                 loss = await trainer.train_step.choose(batch)
                 training_step += 1
                 mlogger.log("loss/training_step", loss, training_step)
-                await trainer.push_weights.call(policy_version)
-                policy_version += 1
-                await policy.update_weights.call()
+                start_time = time.perf_counter()
+                await trainer.push_weights.call(training_step)
+                mlogger.log(
+                    "push_weights_time/training_step",
+                    time.perf_counter() - start_time,
+                    training_step,
+                )
+                start_time = time.perf_counter()
+                await policy.update_weights.call(training_step)
+                mlogger.log(
+                    "update_weights_time/training_step",
+                    time.perf_counter() - start_time,
+                    training_step,
+                )
 
     print("Starting GRPO training loops...")
     # TODO: Start multiple rollouts once all serivces support it

src/forge/actors/policy.py

@@ -17,7 +17,6 @@
 import torch
 import torchstore as ts
 from monarch.actor import current_rank, endpoint, ProcMesh
-from torchstore.state_dict_utils import DELIM
 from vllm.config import VllmConfig
 
 from vllm.engine.arg_utils import EngineArgs
@@ -40,11 +39,17 @@
 from vllm.v1.structured_output import StructuredOutputManager
 from vllm.worker.worker_base import WorkerWrapperBase
 
-from forge.controller import ForgeActor, get_proc_mesh, stop_proc_mesh
+from forge.actors.torchstore_utils import (
+    extract_param_name,
+    get_param_key,
+    get_param_prefix,
+)
 
+from forge.controller import ForgeActor, get_proc_mesh, stop_proc_mesh
 from forge.data.sharding import VLLMSharding
 from forge.interfaces import Policy as PolicyInterface
 from forge.types import ProcessConfig
+from forge.util.async_utils import make_sync_generator
 
 
 @dataclass
@@ -364,16 +369,16 @@ async def run(self):
                     fut.set_result(request_output)
 
     @endpoint
-    async def update_weights(self):
+    async def update_weights(self, policy_version: int):
         # TODO: If generating long sequences, this might be long and will block policy weight updates
         curr_requests = [fut for _, fut in self.requests.values()]
         if curr_requests:
             self.logger.debug(f"Waiting for {len(curr_requests)} pending requests")
             await asyncio.gather(*curr_requests)
 
         self.logger.debug(f"Starting weight update on {self.__class__.__name__}")
-        await self.policy_worker.update.call(version=self.weights_version)
-        self.weights_version += 1
+        await self.policy_worker.update.call(version=policy_version)
+        self.weights_version = policy_version
         self.logger.info(f"Weight update completed (now v{self.weights_version})")
 
     @endpoint
@@ -395,7 +400,6 @@ async def stop(self):
 @dataclass
 class PolicyWorker(ForgeActor):
     vllm_config: VllmConfig
-    state_dict_key: str = "model_state_dict"
 
     @endpoint
     async def setup(self):
@@ -407,41 +411,26 @@ async def setup(self):
     async def execute_model(self, schedule: SchedulerOutput):
         return self.worker.execute_model(schedule)
 
-    async def _load_tensor_parallel_state_dict(
-        self, current_state_dict: dict, version: int
-    ):
-        """
-        Load full state dict from torchstore into tensor parallel model with deterministic sharding.
-        """
-        sharding = VLLMSharding(
-            self.vllm_config.parallel_config.tensor_parallel_size, self.rank
-        )
-
-        for param_name in current_state_dict.keys():
-            current_tensor = current_state_dict[param_name]
-
-            # Load the full tensor from torchstore
-            # TODO: only get the part of the tensor that is needed
-            stored_tensor = await ts.get(
-                f"{self.state_dict_key}{DELIM}{version}{DELIM}{param_name}"
-            )
-            sharding.load_from_source_to_target(
-                param_name,
-                stored_tensor,
-                current_tensor,
-            )
-
     @endpoint
     async def update(self, version: int):
         """Update model weights by reading state dict from torchstore"""
-        key = f"{self.state_dict_key}{DELIM}{version}"
         model = self.worker.model_runner.model
-        current_state_dict = model.state_dict()
-        start = time.time()
-        await self._load_tensor_parallel_state_dict(current_state_dict, version)
-        self.logger.debug(
-            f"Loaded state dict from {key} in {time.time() - start} seconds"
-        )
+        prefix = get_param_prefix(version)
+        self.logger.debug(f"{prefix=}")
+        matching_keys = await ts.keys(prefix)
+        self.logger.debug(f"{matching_keys=}")
+        # TODO: find a way to save the original huggingface parameter names.
+        hf_names = [extract_param_name(key) for key in matching_keys]
+        self.logger.debug(f"{hf_names=}")
+        loaded_weights = set()
+        # We can't pass a generator since vllm load_weights is not async.
+        # Instead, we just call load_weights with one parameter at a time.
+        for name in hf_names:
+            param = await ts.get(get_param_key(version, name))
+            loaded = model.load_weights([(name, param)])
+            del param
+            loaded_weights.update(loaded)
+        self.logger.info(f"Updated {len(loaded_weights)} parameters")
 
     @endpoint
     async def setup_kv_cache(self):

src/forge/actors/torchstore_utils.py

@@ -0,0 +1,19 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+KEY_DELIM = "."
+
+
+def get_param_prefix(policy_version: int) -> str:
+    return f"policy_ver_{policy_version}"
+
+
+def get_param_key(policy_version: int, name: str) -> str:
+    return f"policy_ver_{policy_version}{KEY_DELIM}{name}"
+
+
+def extract_param_name(key: str) -> str:
+    return KEY_DELIM.join(key.split(KEY_DELIM)[1:])