exp(weight-transfer): add sparse filesystem updates

packages/prime-rl-configs/src/prime_rl/configs/inference.py

@@ -209,7 +209,7 @@ def auto_setup_api_server_count(self):
 
 
 class WeightBroadcastConfig(BaseConfig):
-    type: Literal["nccl", "filesystem", "nixl"] = "filesystem"
+    type: Literal["nccl", "filesystem", "sparse_filesystem", "nixl"] = "filesystem"
     """Weight broadcast transport."""
 
 

packages/prime-rl-configs/src/prime_rl/configs/orchestrator.py

@@ -378,6 +378,10 @@ class FileSystemWeightBroadcastConfig(BaseWeightBroadcastConfig):
     type: Literal["filesystem"] = "filesystem"
 
 
+class SparseFileSystemWeightBroadcastConfig(BaseWeightBroadcastConfig):
+    type: Literal["sparse_filesystem"] = "sparse_filesystem"
+
+
 class InMemoryWeightBroadcastConfig(BaseWeightBroadcastConfig):
     host: str = "localhost"
     """Weight transfer host."""
@@ -410,7 +414,10 @@ class NIXLWeightBroadcastConfig(InMemoryWeightBroadcastConfig):
 
 
 WeightBroadcastConfig: TypeAlias = Annotated[
-    FileSystemWeightBroadcastConfig | NCCLWeightBroadcastConfig | NIXLWeightBroadcastConfig,
+    FileSystemWeightBroadcastConfig
+    | SparseFileSystemWeightBroadcastConfig
+    | NCCLWeightBroadcastConfig
+    | NIXLWeightBroadcastConfig,
     Field(discriminator="type"),
 ]
 

packages/prime-rl-configs/src/prime_rl/configs/rl.py

@@ -20,6 +20,9 @@
 from prime_rl.configs.orchestrator import (
     OrchestratorConfig,
 )
+from prime_rl.configs.orchestrator import (
+    SparseFileSystemWeightBroadcastConfig as OrchestratorSparseFileSystemWeightBroadcastConfig,
+)
 from prime_rl.configs.shared import (
     EnvVars,
     ResumeConfig,
@@ -37,6 +40,9 @@
 from prime_rl.configs.trainer import (
     NIXLWeightBroadcastConfig as TrainerNIXLWeightBroadcastConfig,
 )
+from prime_rl.configs.trainer import (
+    SparseFileSystemWeightBroadcastConfig as TrainerSparseFileSystemWeightBroadcastConfig,
+)
 from prime_rl.configs.trainer import (
     TokenizerConfig,
     TrainerConfig,
@@ -166,8 +172,18 @@ class SharedFileSystemWeightBroadcastConfig(BaseConfig):
     """Timeout in seconds for the broadcast handshake and transfer."""
 
 
+class SharedSparseFileSystemWeightBroadcastConfig(BaseConfig):
+    type: Literal["sparse_filesystem"] = "sparse_filesystem"
+
+    timeout: int = 1200
+    """Timeout in seconds for the broadcast handshake and transfer."""
+
+
 SharedWeightBroadcastConfig: TypeAlias = Annotated[
-    SharedFileSystemWeightBroadcastConfig | SharedNCCLWeightBroadcastConfig | SharedNIXLWeightBroadcastConfig,
+    SharedFileSystemWeightBroadcastConfig
+    | SharedSparseFileSystemWeightBroadcastConfig
+    | SharedNCCLWeightBroadcastConfig
+    | SharedNIXLWeightBroadcastConfig,
     Field(discriminator="type"),
 ]
 
@@ -506,6 +522,13 @@ def auto_setup_weight_broadcast(self):
             self.orchestrator.weight_broadcast = OrchestratorFileSystemWeightBroadcastConfig(
                 timeout=self.weight_broadcast.timeout, broadcast_final=broadcast_final
             )
+        elif self.weight_broadcast.type == "sparse_filesystem":
+            self.trainer.weight_broadcast = TrainerSparseFileSystemWeightBroadcastConfig(
+                timeout=self.weight_broadcast.timeout, broadcast_final=broadcast_final
+            )
+            self.orchestrator.weight_broadcast = OrchestratorSparseFileSystemWeightBroadcastConfig(
+                timeout=self.weight_broadcast.timeout, broadcast_final=broadcast_final
+            )
         if self.inference is not None:
             self.inference.weight_broadcast = InferenceWeightBroadcastConfig(type=self.weight_broadcast.type)
 

packages/prime-rl-configs/src/prime_rl/configs/trainer.py

@@ -550,6 +550,11 @@ class FileSystemWeightBroadcastConfig(BaseWeightBroadcastConfig):
     type: Literal["filesystem"] = "filesystem"
 
 
+class SparseFileSystemWeightBroadcastConfig(BaseWeightBroadcastConfig):
+    type: Literal["sparse_filesystem"] = "sparse_filesystem"
+    """Experimental per-rank sorted-index/absolute-value filesystem updates."""
+
+
 class InMemoryWeightBroadcastConfig(BaseWeightBroadcastConfig):
     host: str = "localhost"
     """Weight transfer host."""
@@ -583,7 +588,10 @@ class NIXLWeightBroadcastConfig(InMemoryWeightBroadcastConfig):
 
 
 WeightBroadcastConfig: TypeAlias = Annotated[
-    FileSystemWeightBroadcastConfig | NCCLWeightBroadcastConfig | NIXLWeightBroadcastConfig,
+    FileSystemWeightBroadcastConfig
+    | SparseFileSystemWeightBroadcastConfig
+    | NCCLWeightBroadcastConfig
+    | NIXLWeightBroadcastConfig,
     Field(discriminator="type"),
 ]
 

src/prime_rl/inference/vllm/server.py

@@ -63,6 +63,7 @@ def models(request: Request) -> OpenAIServingModels:
 WORKER_EXTENSION_CLS = {
     "nccl": "prime_rl.inference.vllm.worker.nccl.NCCLWeightUpdateWorker",
     "filesystem": "prime_rl.inference.vllm.worker.filesystem.FileSystemWeightUpdateWorker",
+    "sparse_filesystem": "prime_rl.inference.vllm.worker.filesystem.FileSystemWeightUpdateWorker",
     "nixl": "prime_rl.inference.vllm.worker.nixl.NIXLWeightUpdateWorker",
 }
 

src/prime_rl/inference/vllm/worker/filesystem.py

@@ -1,9 +1,12 @@
+from pathlib import Path
 from typing import TYPE_CHECKING
 
+import torch
 from torch.nn import Module
 from vllm.model_executor.model_loader import DefaultModelLoader, get_model_loader
 
 from prime_rl.inference.vllm.worker.weight_transfer import load_weights_checkpoint_layerwise
+from prime_rl.trainer.sparse_update import apply_sparse_update
 
 # This is to get type hints for the Worker class but not actually extend it at runtime as this is required by vLLM worker extension
 if TYPE_CHECKING:
@@ -18,38 +21,82 @@ class FileSystemWeightUpdateWorker(Worker):
     """vLLM worker extension for updating weights in-place using shared filesystem."""
 
     def init_broadcaster(self) -> None:
-        """Initialize the broadcaster."""
-        ...
+        """Initialize sparse-update receiver state."""
+        self._sparse_state_dict: dict[str, torch.Tensor] | None = None
+        self._sparse_step = 0
+        self._sparse_base_path: str | None = None
 
     def liveness_probe(self) -> None:
         """No-op RPC used by the API server liveness endpoint."""
         return None
 
     def update_weights_from_path(self, weight_path: str) -> None:
-        """Update weights from a specified path in shared filesystem containing a HF-compatible checkpoint."""
-        # Get vLLM model runner and model
-        # When enforce_eager=True, model isn't wrapped by torch.compile so no .runnable attr
-        model_runner = self.model_runner
-        if hasattr(model_runner.model, "runnable"):
-            model = model_runner.model.runnable
-        else:
-            model = model_runner.model
+        """Load a full checkpoint or apply a sparse HF patch chain."""
+        if not hasattr(self, "_sparse_step"):
+            self.init_broadcaster()
+        path = Path(weight_path)
+        if (path / "sparse_manifest.json").exists():
+            self._ensure_sparse_cache()
+            self._sparse_step = apply_sparse_update(self._sparse_state_dict, path, expected_base_step=self._sparse_step)
+            model = (
+                self.model_runner.model.runnable
+                if hasattr(self.model_runner.model, "runnable")
+                else self.model_runner.model
+            )
+            load_weights_checkpoint_layerwise(
+                model,
+                self._sparse_state_dict.items(),
+                self.model_runner.model_config,
+                self.vllm_config,
+            )
+            return
+
+        model = (
+            self.model_runner.model.runnable
+            if hasattr(self.model_runner.model, "runnable")
+            else self.model_runner.model
+        )
         assert isinstance(model, Module)
+        weights_iterator = self._weights_iterator(model, path)
+        load_weights_checkpoint_layerwise(model, weights_iterator, self.model_runner.model_config, self.vllm_config)
+        self._sparse_state_dict = None
+        self._sparse_base_path = weight_path
+        self._sparse_step = self._extract_step(path) or 0
 
-        # Get vLLM model loader
+    def _weights_iterator(self, model: Module, path: Path):
         model_loader = get_model_loader(self.load_config)
         assert isinstance(model_loader, DefaultModelLoader)
-        local_source = DefaultModelLoader.Source(
-            weight_path,
-            revision=None,  # TODO: Check that this is correct or if we should use the default (model_config.revision)
+        source = DefaultModelLoader.Source(
+            str(path),
+            revision=None,
             prefix="",
             fall_back_to_pt=getattr(model, "fall_back_to_pt_during_load", True),
             allow_patterns_overrides=getattr(model, "allow_patterns_overrides", None),
         )
-        weights_iterator = model_loader._get_weights_iterator(local_source)
-        load_weights_checkpoint_layerwise(
-            model,
-            weights_iterator,
-            self.model_runner.model_config,
-            self.vllm_config,
+        return model_loader._get_weights_iterator(source)
+
+    def _ensure_sparse_cache(self) -> None:
+        if self._sparse_state_dict is not None:
+            return
+        model = (
+            self.model_runner.model.runnable
+            if hasattr(self.model_runner.model, "runnable")
+            else self.model_runner.model
         )
+        source = Path(self._sparse_base_path or self.model_runner.model_config.model)
+        self._sparse_state_dict = {
+            name: tensor.detach().to("cpu", dtype=torch.bfloat16).contiguous()
+            for name, tensor in self._weights_iterator(model, source)
+        }
+        if "lm_head.weight" not in self._sparse_state_dict and "model.embed_tokens.weight" in self._sparse_state_dict:
+            self._sparse_state_dict["lm_head.weight"] = self._sparse_state_dict["model.embed_tokens.weight"].clone()
+
+    @staticmethod
+    def _extract_step(path: Path) -> int | None:
+        for parent in (path, *path.parents):
+            if parent.name.startswith("step_"):
+                try:
+                    return int(parent.name.removeprefix("step_"))
+                except ValueError:
+                    return None
+        return None

src/prime_rl/trainer/rl/train.py

@@ -243,6 +243,30 @@ def train(config: TrainerConfig):
     else:
         logger.info("Starting from scratch")
 
+    if config.weight_broadcast.type == "sparse_filesystem" and weight_sender is not None:
+        from prime_rl.trainer.models import PreTrainedModelPrimeRL
+        from prime_rl.trainer.sparse_update import OptimizerSparseUpdateHook, SparseUpdateWriter
+
+        if checkpoint_step is not None:
+            raise ValueError("Sparse filesystem weight updates do not support checkpoint resume in this POC")
+        state_dict = model.state_dict()
+        if isinstance(model, PreTrainedModelPrimeRL) and model.is_prime_state_dict(state_dict):
+            raise ValueError(
+                "Sparse filesystem updates currently require trainer state-dict names/shapes to match HF format."
+            )
+        sparse_writer = SparseUpdateWriter(
+            config.output_dir / ".sparse_weight_updates",
+            rank=world.rank,
+        )
+        sparse_writer.initialize(state_dict, base_step=0)
+        sparse_writer.write(state_dict, target_step=0)
+        OptimizerSparseUpdateHook(
+            optimizer,
+            sparse_writer,
+            model.state_dict,
+            lambda: progress.step,
+        )
+
     # Set up the data loader (Optionally, use a fake data loader for debugging)
     logger.info(f"Initializing data loader ({config.data})")
     t0 = time.perf_counter()