Broadcast documents_seen as the model version in weights_ready

fast_llm/engine/training/config.py

@@ -388,7 +388,7 @@ def new_setup():
 
 class TrainerCallback[ConfigType: TrainerCallbackConfig](Configurable[ConfigType]):
     # TODO: Make a more exhaustive set of events and arguments.
-    def run_begin(self, step: int):
+    def run_begin(self, step: int, documents_seen: int):
         pass
 
     def step_end(
@@ -397,6 +397,7 @@ def step_end(
         reduced_losses: dict[str, float | int],
         update_successful: bool,
         train_metrics: dict[str, typing.Any] | None,
+        documents_seen: int,
     ) -> dict[str, typing.Any] | None:
         """Optionally return a dict of scalar metrics to merge into the step's training logs."""
         return None

fast_llm/engine/training/streaming.py

@@ -54,23 +54,24 @@ def __init__(self, config: ConfigType, model: "FastLLMModel"):
             self._process_group = self._pool.get_process_group(range(world_size), 0)
             logger.info(f"Weights broadcast rendezvous at {init_method} connected")
 
-    def run_begin(self, step: int):
+    def run_begin(self, step: int, documents_seen: int):
         # TODO: ====== Send a train / run begin signal? ======
-        self._broadcast_weights(step)
+        self._broadcast_weights(step, documents_seen)
 
     def step_end(
         self,
         step: int,
         reduced_losses: dict[str, float | int],
         update_successful: bool,
         train_metrics: dict[str, typing.Any] | None,
+        documents_seen: int,
     ) -> dict[str, typing.Any] | None:
         # Harvest the previous broadcast's timing before re-recording the events on this step's sync.
         if self._weight_sync_pending and self._weight_sync_end.query():
             self._weight_sync_time_ms = self._weight_sync_start.elapsed_time(self._weight_sync_end)
             self._weight_sync_pending = False
         if update_successful:
-            self._broadcast_weights(step)
+            self._broadcast_weights(step, documents_seen)
         # Report the last measured sync (not reset between steps): a broadcast runs on ~every step so
         # this stays current, and holding the value keeps a skipped or not-yet-harvested step from
         # dropping the metric.
@@ -93,10 +94,15 @@ def _clear(self):
             del self._pool
             del self._process_group
 
-    def _broadcast_weights(self, step: int):
+    def _broadcast_weights(self, step: int, documents_seen: int):
         if self._do_broadcast:
             self._client.xadd(
-                REDIS_TRAINING_STREAM, {REDIS_TRAINING_FIELD: json.dumps({"type": "weights_ready", "step": step})}
+                REDIS_TRAINING_STREAM,
+                {
+                    REDIS_TRAINING_FIELD: json.dumps(
+                        {"type": "weights_ready", "step": step, "documents_seen": documents_seen}
+                    )
+                },
             )
         weight_sync_start_time = None
         if self._use_cuda_timing:

fast_llm/engine/training/trainer.py

@@ -204,7 +204,7 @@ def _train(self) -> tuple[bool, dict[PhaseType, dict[str, typing.Any]]]:
         safe_barrier(self._distributed.world_group, "train begin")
 
         for callback in self._callbacks.values():
-            callback.run_begin(self._completed_steps)
+            callback.run_begin(self._completed_steps, self._documents_seen)
 
         if torch.cuda.is_available():
             torch.cuda.synchronize()
@@ -242,7 +242,11 @@ def _train(self) -> tuple[bool, dict[PhaseType, dict[str, typing.Any]]]:
                 callback_metrics = {}
                 for callback in self._callbacks.values():
                     returned_metrics = callback.step_end(
-                        self._completed_steps, reduced_losses, update_successful, train_metrics
+                        self._completed_steps,
+                        reduced_losses,
+                        update_successful,
+                        train_metrics,
+                        self._documents_seen,
                     )
                     if returned_metrics:
                         callback_metrics.update(returned_metrics)

tests/models/test_streaming.py

@@ -93,6 +93,7 @@ def _run_event_consumer(
         process_group = pool.get_process_group(range(world_size), consumer_rank)
         timeout_ms = int(streaming_config.timeout * 1000)
         last_id = "0-0"
+        last_documents_seen = 0
         while True:
             result = client.xread(
                 streams={REDIS_TRAINING_STREAM: last_id},
@@ -112,6 +113,9 @@ def _run_event_consumer(
                 if message["type"] == "training_finished":
                     return
                 elif message["type"] == "weights_ready":
+                    Assert.incl("documents_seen", message)
+                    Assert.geq(message["documents_seen"], last_documents_seen)
+                    last_documents_seen = message["documents_seen"]
                     weights = {}
                     while True:
                         meta = _broadcast_object(None, process_group, src=0)