From a0c251521c2760ab2a7a954ff01279910b95c5ed Mon Sep 17 00:00:00 2001 From: Joel Lamy-Poirier Date: Tue, 7 Jul 2026 14:17:01 -0400 Subject: [PATCH 1/3] Broadcast documents_seen as the model version in weights_ready The `weights_ready` event now carries `document_count` (= documents_seen) alongside the existing `step` (= completed step). Consumers stamp `document_count` onto rollouts so staleness is measured in documents, aligning with DeepSpeed's document clock, while `step` remains available for logging. Keeping `step` makes the change backward-compatible. Threads `documents_seen` through the `TrainerCallback.run_begin` / `step_end` hooks (a training-progress counter alongside `step`). Co-Authored-By: Claude Opus 4.8 --- fast_llm/engine/training/config.py | 3 ++- fast_llm/engine/training/streaming.py | 18 +++++++++++++----- fast_llm/engine/training/trainer.py | 10 ++++++++-- 3 files changed, 23 insertions(+), 8 deletions(-) diff --git a/fast_llm/engine/training/config.py b/fast_llm/engine/training/config.py index d3a3ce566..6081fd4ad 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, ): pass diff --git a/fast_llm/engine/training/streaming.py b/fast_llm/engine/training/streaming.py index aec14530f..f36ebbcaf 100644 --- a/fast_llm/engine/training/streaming.py +++ b/fast_llm/engine/training/streaming.py @@ -40,9 +40,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, @@ -50,9 +50,10 @@ def step_end( reduced_losses: dict[str, float | int], update_successful: bool, train_metrics: dict[str, typing.Any] | None, + documents_seen: int, ): if update_successful: - self._broadcast_weights(step) + self._broadcast_weights(step, documents_seen) def train_end(self, step: int): # TODO: ====== Send something on unsuccessful ends? ====== @@ -69,10 +70,17 @@ 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: + # `document_count` is the model version consumers stamp onto rollouts (aligning staleness + # with DeepSpeed's document clock); `step` is kept so consumers can also log the raw step. 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, "document_count": documents_seen} + ) + }, ) for shard_name, layer_name, tensor in self._model.iter_checkpoint(self._config.export, {}): if self._do_broadcast: diff --git a/fast_llm/engine/training/trainer.py b/fast_llm/engine/training/trainer.py index 1ed18c449..c9d2fd1fe 100644 --- a/fast_llm/engine/training/trainer.py +++ b/fast_llm/engine/training/trainer.py @@ -203,7 +203,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() @@ -237,7 +237,13 @@ def _train(self) -> tuple[bool, dict[PhaseType, dict[str, typing.Any]]]: nan_iters += not all(math.isfinite(loss) for loss in reduced_losses.values()) for callback in self._callbacks.values(): - callback.step_end(self._completed_steps, reduced_losses, update_successful, train_metrics) + callback.step_end( + self._completed_steps, + reduced_losses, + update_successful, + train_metrics, + self._documents_seen, + ) # Logging. metrics = {} if is_logging: From bb9162784adcbb54cb64863c61859b1c361e7342 Mon Sep 17 00:00:00 2001 From: Joel Lamy-Poirier Date: Wed, 8 Jul 2026 17:36:36 -0400 Subject: [PATCH 2/3] Address review: rename weights_ready field to documents_seen, assert it in test - Rename the wire field document_count -> documents_seen for consistency with the internal counter and the rest of the codebase (no shipped consumer yet). - Reword the broadcast comment to describe the fields locally, without naming downstream consumers. - Assert the new documents_seen field is present and non-negative in the streaming consumer test. Co-Authored-By: Claude Opus 4.8 --- fast_llm/engine/training/streaming.py | 6 +++--- tests/models/test_streaming.py | 2 ++ 2 files changed, 5 insertions(+), 3 deletions(-) diff --git a/fast_llm/engine/training/streaming.py b/fast_llm/engine/training/streaming.py index f36ebbcaf..02df68343 100644 --- a/fast_llm/engine/training/streaming.py +++ b/fast_llm/engine/training/streaming.py @@ -72,13 +72,13 @@ def _clear(self): def _broadcast_weights(self, step: int, documents_seen: int): if self._do_broadcast: - # `document_count` is the model version consumers stamp onto rollouts (aligning staleness - # with DeepSpeed's document clock); `step` is kept so consumers can also log the raw step. + # `documents_seen` is the cumulative document count, which doubles as the model version; + # `step` is the raw training step. self._client.xadd( REDIS_TRAINING_STREAM, { REDIS_TRAINING_FIELD: json.dumps( - {"type": "weights_ready", "step": step, "document_count": documents_seen} + {"type": "weights_ready", "step": step, "documents_seen": documents_seen} ) }, ) diff --git a/tests/models/test_streaming.py b/tests/models/test_streaming.py index 3a7e27f4e..716a64410 100644 --- a/tests/models/test_streaming.py +++ b/tests/models/test_streaming.py @@ -112,6 +112,8 @@ 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"], 0) weights = {} while True: meta = _broadcast_object(None, process_group, src=0) From 3aa0aeb3f7d9ec90dd6b24a9d360d98adb3db5a5 Mon Sep 17 00:00:00 2001 From: Joel Lamy-Poirier Date: Fri, 10 Jul 2026 07:40:46 -0400 Subject: [PATCH 3/3] Address re-review: assert non-decreasing documents_seen in test, drop restating comment Co-Authored-By: Claude Opus 4.8 --- fast_llm/engine/training/streaming.py | 2 -- tests/models/test_streaming.py | 4 +++- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/fast_llm/engine/training/streaming.py b/fast_llm/engine/training/streaming.py index 388419abe..823843850 100644 --- a/fast_llm/engine/training/streaming.py +++ b/fast_llm/engine/training/streaming.py @@ -96,8 +96,6 @@ def _clear(self): def _broadcast_weights(self, step: int, documents_seen: int): if self._do_broadcast: - # `documents_seen` is the cumulative document count, which doubles as the model version; - # `step` is the raw training step. self._client.xadd( REDIS_TRAINING_STREAM, { diff --git a/tests/models/test_streaming.py b/tests/models/test_streaming.py index 716a64410..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}, @@ -113,7 +114,8 @@ def _run_event_consumer( return elif message["type"] == "weights_ready": Assert.incl("documents_seen", message) - Assert.geq(message["documents_seen"], 0) + 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)