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

apps/toy_rl/sumdigits-tp.yaml

@@ -0,0 +1,68 @@
+# Toy app Training Configuration
+
+# Global configuration
+group_size: 16
+batch_size: 64
+max_req_tokens: 64
+max_res_tokens: 64
+model: "Qwen/Qwen2.5-0.5B-Instruct"
+
+# Dataset configuration
+dataset:
+  model: ${model}
+
+# Policy configuration
+policy:
+  engine_config:
+    model: ${model}
+    tensor_parallel_size: 2
+    pipeline_parallel_size: 1
+    enforce_eager: false
+  sampling_config:
+    n: ${group_size}
+    max_tokens: ${max_res_tokens}
+    temperature: 1.0
+    top_p: 1.0
+  use_vllm_builtin_load: true
+
+# Trainer configuration
+trainer:
+  model_name: ${model}
+  learning_rate: 1e-5
+  use_vllm_builtin_load: true
+
+# Reference model configuration
+ref_model:
+  model_name: ${model}
+
+# Replay buffer configuration
+replay_buffer:
+  batch_size: ${batch_size}
+  max_policy_age: 1 # Async by 1
+  dp_size: 1
+
+services:
+  dataset:
+    procs: 1
+    num_replicas: 1
+    with_gpus: false
+  policy:
+    procs: 1
+    num_replicas: 1
+    with_gpus: true
+  trainer:
+    procs: 1
+    num_replicas: 1
+    with_gpus: true
+  replay_buffer:
+    procs: 1
+    num_replicas: 1
+    with_gpus: false
+  reward_actor:
+    procs: 1
+    num_replicas: 1
+    with_gpus: false
+  ref_model:
+    procs: 1
+    num_replicas: 1
+    with_gpus: true

apps/toy_rl/sumdigits.py

@@ -4,30 +4,32 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
-# Usage: python -m apps.grpo.main --config apps/grpo/qwen3_1_7b.yaml
+# Usage: python -m apps.toy_rl.sumdigits --config apps/toy_rl/sumdigits.yaml
 
 import asyncio
 import random
 import time
 import uuid
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 from typing import Any
 
 import torch
 import torch.nn.functional as F
 import torchstore as ts
 from forge.actors.policy import Policy
 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
 from forge.controller.provisioner import shutdown
+
 from forge.losses.grpo_loss import SimpleGRPOLoss
 from forge.util.metric_logging import get_metric_logger
-
 from forge.util.ops import selective_log_softmax
 from monarch.actor import endpoint
 from omegaconf import DictConfig
+from torch.functional import _return_counts
 
 from torchstore.state_dict_utils import DELIM
 from transformers import AutoModelForCausalLM
@@ -250,10 +252,14 @@ async def forward(self, episode: Episode) -> torch.Tensor:
 class Trainer(ForgeActor):
     """Reinforce Loss Trainer implementation for policy optimization."""
 
-    model_name: str
+    model_name: str = ""
     learning_rate: float = 1e-5
     device: torch.device | None = None
     state_dict_key: str = "model_state_dict"
+    use_vllm_builtin_load: bool = False
+
+    def __post_init__(self):
+        super().__init__()
 
     @endpoint
     async def setup(self):
@@ -340,6 +346,9 @@ def train_step(self, episodes: list[Episode]) -> float:
     @endpoint
     async def push_weights(self, version: int):
         """Update policy model weights with trainer's current weights."""
+        if self.use_vllm_builtin_load:
+            await self._push_weights_hf_nonsharded(version)
+            return None
         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=self.model.config.num_hidden_layers
@@ -351,6 +360,13 @@ async def push_weights(self, version: int):
             f"Pushed weights to {key} in {end_time - start_time:.2f} seconds"
         )
 
+    async def _push_weights_hf_nonsharded(self, policy_version: int) -> None:
+        """Push weights to torchstore in HF format, non-sharded."""
+        hf_state_dict = self.model.state_dict()
+        for name, param in hf_state_dict.items():
+            key = get_param_key(policy_version, name)
+            await ts.put(key, param)
+
 
 @dataclass
 class RewardActor(ForgeActor):
@@ -440,6 +456,8 @@ async def main(cfg: DictConfig):
     )
 
     # ---- Setup services ---- #
+    print(f"{cfg.policy=}")
+    print(f"{cfg.services.policy=}")
     await ts.initialize()
     (
         dataloader,

src/forge/actors/policy.py

@@ -43,12 +43,16 @@
 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.data_models.completion import Completion
 from forge.data_models.prompt import to_prompt
-
 from forge.interfaces import Policy as PolicyInterface
 from forge.types import ProcessConfig
 
@@ -127,6 +131,8 @@ def create_vllm_config(self) -> VllmConfig:
 class Policy(PolicyInterface):
     engine_config: EngineConfig | Mapping = field(default_factory=EngineConfig)
     sampling_config: SamplingConfig | Mapping = field(default_factory=SamplingConfig)
+    use_vllm_builtin_load: bool = False
+    test_blah_blah: int = 0
     available_devices: str | None = None
     # Gets set up by setup
     sampling_params: SamplingParams | None = None
@@ -145,6 +151,7 @@ def __post_init__(self):
             self.engine_config = EngineConfig.from_dict(self.engine_config)
         if isinstance(self.sampling_config, Mapping):
             self.sampling_config = SamplingConfig.from_dict(self.sampling_config)
+        # No conversion needed for boolean flag
 
     @classmethod
     async def launch(  # pyright: ignore[reportIncompatibleMethodOverride]
@@ -153,6 +160,7 @@ async def launch(  # pyright: ignore[reportIncompatibleMethodOverride]
         process_config: ProcessConfig,
         engine_config: EngineConfig | Mapping = EngineConfig(),
         sampling_config: SamplingConfig | Mapping = SamplingConfig(),
+        use_vllm_builtin_load: bool = False,
         available_devices: str | None = None,
         **kwargs,
     ) -> "Policy":
@@ -191,6 +199,7 @@ async def launch(  # pyright: ignore[reportIncompatibleMethodOverride]
             cls,
             engine_config=engine_config,
             sampling_config=sampling_config,
+            use_vllm_builtin_load=use_vllm_builtin_load,
             available_devices=available_devices,
             policy_worker=workers,
         )
@@ -384,7 +393,10 @@ async def update_weights(self, policy_version: int):
             await asyncio.gather(*curr_requests)
 
         logger.debug(f"Starting weight update on {self.__class__.__name__}")
-        await self.policy_worker.update.call(version=policy_version)
+        if self.use_vllm_builtin_load:
+            await self.policy_worker._update_hf_nonsharded.call(version=policy_version)
+        else:
+            await self.policy_worker._update_sharded.call(version=policy_version)
         self.policy_version = policy_version
         logger.info(f"Weight update completed (now v{self.policy_version})")
 
@@ -496,7 +508,7 @@ async def _load_tensor_parallel_state_dict(
             )
 
     @endpoint
-    async def update(self, version: int):
+    async def _update_sharded(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
@@ -505,6 +517,27 @@ async def update(self, version: int):
         await self._load_tensor_parallel_state_dict(current_state_dict, version)
         logger.debug(f"Loaded state dict from {key} in {time.time() - start} seconds")
 
+    @endpoint
+    async def _update_hf_nonsharded(self, version: int):
+        """Update model weights by reading state dict from torchstore"""
+        model = self.worker.model_runner.model
+        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):
         """Based on vllm/v1/engine/core.py:EngineCore._initialize_kv_caches

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:])

src/forge/actors/trainer.py

@@ -36,6 +36,8 @@
 from torchtitan.experiments.forge.engine import ForgeEngine
 from torchtitan.experiments.forge.job_config import ForgeJobConfig
 
+from forge.actors.torchstore_utils import get_param_key
+
 from forge.controller import ForgeActor
 from forge.data.utils import batch_to_device
 
@@ -93,6 +95,7 @@ class RLTrainer(ForgeActor):
     activation_checkpoint: ActivationCheckpoint = field(
         default_factory=ActivationCheckpoint
     )
+    use_vllm_builtin_load: bool = False
     compile: Compile = field(default_factory=Compile)
     float8: Float8 = field(default_factory=Float8)
     comm: Comm = field(default_factory=Comm)
@@ -142,7 +145,7 @@ def __post_init__(self):
     async def setup(self):
         # TODO: update ForgeEngine to not use ForgeJobConfig
         engine_config = {f.name: getattr(self, f.name) for f in fields(self)}
-        for key in {"loss", "state_dict_key", "use_dcp"}:
+        for key in {"loss", "state_dict_key", "use_dcp", "use_vllm_builtin_load"}:
             engine_config.pop(key)  # Not part of job config
         self.engine = ForgeEngine(ForgeJobConfig(**engine_config))
         self.engine.checkpointer.load(step=self.step)
@@ -248,6 +251,12 @@ def train_step(
 
     @endpoint
     async def push_weights(self, policy_version: int) -> None:
+        if self.use_vllm_builtin_load:
+            await self._push_weights_hf_nonsharded(policy_version)
+        else:
+            await self._push_weights_sharded(policy_version)
+
+    async def _push_weights_sharded(self, policy_version: int) -> None:
         # Save to torchstore. Hacking in to the Checkpointer's prepped state-dict for now.
         # TODO:
         # 1. Checkpoint invokes state-dict flattening during dcp_save for [MODEL].
@@ -290,6 +299,22 @@ async def push_weights(self, policy_version: int) -> None:
 
         logger.debug(f"Pushed weights to {key} in {end_time - start_time:.2f} seconds")
 
+    async def _push_weights_hf_nonsharded(self, policy_version: int) -> None:
+        """Push weights to torchstore in HF format, non-sharded."""
+        if "model" not in self.engine.checkpointer.states:
+            raise RuntimeError("Model state not found in checkpointer state")
+
+        sd = self.engine.checkpointer.states["model"].state_dict()
+        flattened_state_dict, _ = flatten_state_dict(sd)
+        if self.engine.checkpointer.sd_adapter is None:
+            raise RuntimeError(
+                "Trying to save checkpoint in HF safetensors format, but sd_adapter is not provided."
+            )
+        hf_state_dict = self.engine.checkpointer.sd_adapter.to_hf(flattened_state_dict)
+        for name, param in hf_state_dict.items():
+            key = get_param_key(policy_version, name)
+            await ts.put(key, param)
+
     @endpoint
     async def cleanup(self) -> None:
         if self.engine.checkpointer: