[Feature] Add CI test for weight version update

slime/backends/fsdp_utils/actor.py

@@ -1,5 +1,6 @@
 import logging
 import os
+import random
 from argparse import Namespace
 from itertools import accumulate
 
@@ -791,6 +792,15 @@ def update_weights(self) -> None:  # type: ignore[override]
             dist.barrier(group=get_gloo_group())
 
         self.weight_updater.update_weights()
+
+        if self.args.ci_test and len(rollout_engines) > 0:
+            engine = random.choice(rollout_engines)
+            engine_version = ray.get(engine.get_weight_version.remote())
+            if str(engine_version) != str(self.weight_updater.weight_version):
+                raise RuntimeError(
+                    f"Weight version mismatch! Engine: {engine_version}, Updater: {self.weight_updater.weight_version}"
+                )
+
         clear_memory()
 
     def _create_ref_model(self, ref_load_path: str | None):

slime/backends/fsdp_utils/update_weight_utils.py

@@ -33,6 +33,7 @@ class UpdateWeight(abc.ABC):
     def __init__(self, args: Namespace, model: torch.nn.Module) -> None:
         self.args = args
         self.model = model
+        self.weight_version = 0
 
     @abc.abstractmethod
     def connect_rollout_engines(
@@ -43,6 +44,7 @@ def connect_rollout_engines(
         pass
 
     def update_weights(self) -> None:
+        self.weight_version += 1
         bucket = []
         bucket_size = 0
         for name, param in self.model.state_dict().items():
@@ -71,10 +73,10 @@ def update_weights(self) -> None:
 
     def wait_and_update_bucket_weights(self, bucket):
         bucket = [(name, param.wait()) if hasattr(param, "wait") else (name, param) for name, param in bucket]
-        self.update_bucket_weights(bucket)
+        self.update_bucket_weights(bucket, weight_version=self.weight_version)
 
     @abc.abstractmethod
-    def update_bucket_weights(self, named_tensors) -> None:
+    def update_bucket_weights(self, named_tensors, weight_version=None) -> None:
         pass
 
 
@@ -114,7 +116,7 @@ def connect_rollout_engines(
                 # Calculate TP rank within this SGLang engine group
                 self.tp_rank = dist.get_rank() - start_rank
 
-    def update_bucket_weights(self, named_tensors) -> None:
+    def update_bucket_weights(self, named_tensors, weight_version=None) -> None:
         monkey_patch_torch_reductions()
         # Use flattened bucket approach similar to Megatron
         logger.info("Using flattened tensor bucket")
@@ -162,6 +164,7 @@ def update_bucket_weights(self, named_tensors) -> None:
                     "serialized_named_tensors": [tensors[i] for tensors in gathered_serialized_batches],
                     "load_format": "flattened_bucket",
                     "flush_cache": False,
+                    "weight_version": str(weight_version),
                 }
                 ref = self._ipc_engine.update_weights_from_tensor.remote(**kwargs)
                 ray.get(ref)
@@ -174,10 +177,6 @@ def update_bucket_weights(self, named_tensors) -> None:
 class UpdateWeightFromDistributed(UpdateWeight):
     """Broadcast weights via a temporary NCCL group to rollout engines."""
 
-    def __init__(self, args: Namespace, model: torch.nn.Module) -> None:
-        self.args = args
-        self.model = model
-
     def connect_rollout_engines(
         self,
         rollout_engines: Sequence[ActorHandle],
@@ -220,7 +219,7 @@ def connect_rollout_engines(
             )
             ray.get(refs)
 
-    def update_bucket_weights(self, named_tensors) -> None:
+    def update_bucket_weights(self, named_tensors, weight_version=None) -> None:
         """Send names/dtypes/shapes metadata to engines, then broadcast tensors.
 
         Ensures tensors are contiguous; when `world_size == 1`, converts DTensors
@@ -235,6 +234,7 @@ def update_bucket_weights(self, named_tensors) -> None:
                 dtypes=[param.dtype for _, param in named_tensors],
                 shapes=[param.shape for _, param in named_tensors],
                 group_name=self._group_name,
+                weight_version=str(weight_version),
             )
             for engine in self.rollout_engines
         ]

slime/backends/megatron_utils/actor.py

@@ -1,5 +1,6 @@
 import logging
 import os
+import random
 import socket
 from argparse import Namespace
 from contextlib import nullcontext
@@ -474,6 +475,14 @@ def update_weights(self) -> None:
             self.weight_updater.update_weights()
             print_memory("after update_weights")
 
+            if self.args.ci_test and len(rollout_engines) > 0:
+                engine = random.choice(rollout_engines)
+                engine_version = ray.get(engine.get_weight_version.remote())
+                if str(engine_version) != str(self.weight_updater.weight_version):
+                    raise RuntimeError(
+                        f"Weight version mismatch! Engine: {engine_version}, Updater: {self.weight_updater.weight_version}"
+                    )
+
             if getattr(self.args, "keep_old_actor", False):
                 if self.args.update_weights_interval == 1:
                     logger.info("updating model queue: rollout_actor -> old_actor, actor -> rollout_actor")

slime/ray/rollout.py

@@ -330,7 +330,7 @@ def _split_train_data_by_dp(self, data, dp_size):
 
 def init_rollout_engines(args, pg, all_rollout_engines):
     if args.debug_train_only:
-        return 0, None
+        return 0
 
     num_gpu_per_engine = min(args.rollout_num_gpus_per_engine, args.num_gpus_per_node)
     num_engines = args.rollout_num_gpus // num_gpu_per_engine
@@ -391,7 +391,7 @@ def init_rollout_engines(args, pg, all_rollout_engines):
     num_new_engines = len(rollout_engines)
 
     if num_new_engines == 0:
-        return num_new_engines, None
+        return num_new_engines
 
     if args.rollout_external:
         addr_and_ports = _allocate_rollout_engine_addr_and_ports_external(args=args, rollout_engines=rollout_engines)