ping-pong weight sync (#763)

Author: amirafzali

apps/grpo/main.py

@@ -27,7 +27,6 @@
 from forge.rl import collate, ComputeAdvantages, Episode, RewardActor
 from forge.rl.loss import DAPOLoss, GRPOLoss
 from forge.types import LauncherConfig, ProvisionerConfig
-from forge.util.checkpoint import drop_weights
 from forge.util.config import parse
 from forge.util.logging import get_logger
 from omegaconf import DictConfig, OmegaConf
@@ -335,10 +334,6 @@ async def continuous_training():
                 await generator.update_weights.fanout(training_step)
                 t.step("update_weights")
 
-                if training_step >= 2:
-                    await drop_weights(training_step - 1)
-                    t.step("drop_weights")
-
                 t.stop()
                 restart_tracer = True
 

src/forge/actors/_torchstore_utils.py

@@ -10,13 +10,25 @@
 
 KEY_DELIM = "."
 
+# Alternate between two storage version IDs
+# This reuses allocations instead of incrementing versions and deleting old ones
+VERSION_A = 0
+VERSION_B = 1
+
+
+def get_storage_version(step: int) -> int:
+    """Map incrementing step to ping-pong storage version (0 or 1)."""
+    return VERSION_A if step % 2 == 0 else VERSION_B
+
 
 def get_param_prefix(policy_version: int) -> str:
-    return f"policy_ver_{policy_version:010d}"
+    storage_version = get_storage_version(policy_version)
+    return f"policy_ver_{storage_version:010d}"
 
 
 def get_param_key(policy_version: int, name: str) -> str:
-    return f"policy_ver_{policy_version:010d}{KEY_DELIM}{name}"
+    storage_version = get_storage_version(policy_version)
+    return f"policy_ver_{storage_version:010d}{KEY_DELIM}{name}"
 
 
 def extract_param_name(key: str) -> str:

src/forge/actors/trainer/titan.py

@@ -327,9 +327,12 @@ async def push_weights(self, policy_version: int) -> None:
                 "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)
+
+        entries = [
+            (get_param_key(policy_version, name), param)
+            for name, param in hf_state_dict.items()
+        ]
+        await ts.put_batch(entries)
         end_time = time.perf_counter()
         logger.info("Completed weights push in %.2f seconds", end_time - start_time)