diff --git a/fast_llm/engine/training/config.py b/fast_llm/engine/training/config.py index 2ba27dbfd..953300b2e 100644 --- a/fast_llm/engine/training/config.py +++ b/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 diff --git a/fast_llm/engine/training/streaming.py b/fast_llm/engine/training/streaming.py index 2387ce0cd..823843850 100644 --- a/fast_llm/engine/training/streaming.py +++ b/fast_llm/engine/training/streaming.py @@ -54,9 +54,9 @@ 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, @@ -64,13 +64,14 @@ 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: # 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: diff --git a/fast_llm/engine/training/trainer.py b/fast_llm/engine/training/trainer.py index fa4a25262..fb57c6a66 100644 --- a/fast_llm/engine/training/trainer.py +++ b/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) diff --git a/tests/models/test_streaming.py b/tests/models/test_streaming.py index 3a7e27f4e..dee8d7ad1 100644 --- a/tests/models/test_streaming.py +++ b/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)