RL diagnostics: documents-seen clock, model-version broadcast, reward + per-token model_version, data-wait

fast_llm/data/dataset/streaming.py

@@ -31,6 +31,10 @@ class RedisStreamingDocumentData(Config):
     rejected_span: tuple[int, int] | None = Field(default=None)
     advantage: float | None = Field(default=None)
     old_log_probabilities: torch.Tensor | None = Field(default=None)
+    # Raw (un-normalized) reward, a per-rollout scalar (broadcast per-token like `advantage`).
+    reward: float | None = Field(default=None)
+    # Model version each token was generated under (documents-seen units), one per token.
+    model_version: torch.Tensor | None = Field(default=None)
 
     def _validate(self):
         # Decode message
@@ -53,9 +57,15 @@ def _validate(self):
                 self.old_log_probabilities = torch.frombuffer(self.old_log_probabilities, dtype=torch.float32)
             elif isinstance(self.old_log_probabilities, (list, tuple)):
                 self.old_log_probabilities = torch.tensor(self.old_log_probabilities, dtype=torch.float32)
+            if isinstance(self.model_version, bytes):
+                self.model_version = torch.frombuffer(self.model_version, dtype=torch.int64)
+            elif isinstance(self.model_version, (list, tuple)):
+                self.model_version = torch.tensor(self.model_version, dtype=torch.int64)
         super()._validate()
         if self.old_log_probabilities is not None:
             Assert.eq(len(self.old_log_probabilities), self.num_tokens)
+        if self.model_version is not None:
+            Assert.eq(len(self.model_version), self.num_tokens)
 
     @functools.cached_property
     def num_tokens(self) -> int:
@@ -78,6 +88,8 @@ def to_message(self) -> dict[str, str | int | float | bytes]:
         message: dict[str, str | int | float | bytes] = {"tokens": self.tokens.numpy().tobytes()}
         if self.old_log_probabilities is not None:
             message["old_log_probabilities"] = self.old_log_probabilities.numpy().tobytes()
+        if self.model_version is not None:
+            message["model_version"] = self.model_version.numpy().tobytes()
         data = {}
         if self.loss_masking_spans is not None:
             data["loss_masking_spans"] = self.loss_masking_spans
@@ -87,6 +99,8 @@ def to_message(self) -> dict[str, str | int | float | bytes]:
             data["rejected_span"] = self.rejected_span
         if self.advantage is not None:
             data["advantage"] = self.advantage
+        if self.reward is not None:
+            data["reward"] = self.reward
         if data:
             message["data"] = json.dumps(data)
         return message
@@ -111,6 +125,12 @@ def to_document(self):
             old_log_probabilities=(
                 None if self.old_log_probabilities is None else TokenDataDocument(data=self.old_log_probabilities)
             ),
+            reward=(
+                None
+                if self.reward is None
+                else TokenDataDocument(data=torch.full([sample_size], self.reward, dtype=torch.float32))
+            ),
+            model_version=(None if self.model_version is None else TokenDataDocument(data=self.model_version)),
         )
 
 

fast_llm/data/document/language_model.py

@@ -26,6 +26,8 @@ class LanguageModelDocument(TokenDocument):
     image_patches: PatchDocument | None = None
     advantages: TokenDataDocument | None = None
     old_log_probabilities: TokenDataDocument | None = None
+    reward: TokenDataDocument | None = None
+    model_version: TokenDataDocument | None = None
 
 
 @dataclasses.dataclass(kw_only=True)
@@ -34,6 +36,8 @@ class LanguageModelTargetInput(ModelInput):
     mask: torch.Tensor | None = None
     advantages: torch.Tensor | None = None
     old_log_probabilities: torch.Tensor | None = None
+    reward: torch.Tensor | None = None
+    model_version: torch.Tensor | None = None
     label_counts: torch.Tensor | None = None
     num_labels: int | None = None
     num_labels_in_batch: int | None = None
@@ -83,6 +87,8 @@ def to_kwargs(self) -> dict[str, typing.Any]:
             LanguageModelKwargs.hidden_states: self.hidden_states,
             LanguageModelKwargs.advantages: [target.advantages for target in self.targets],
             LanguageModelKwargs.old_log_probabilities: [target.old_log_probabilities for target in self.targets],
+            LanguageModelKwargs.reward: [target.reward for target in self.targets],
+            LanguageModelKwargs.model_version: [target.model_version for target in self.targets],
             LanguageModelKwargs.label_counts: [target.label_counts for target in self.targets],
             LanguageModelKwargs.num_labels_in_batch: [target.num_labels_in_batch for target in self.targets],
         }
@@ -105,6 +111,8 @@ class LanguageModelBatch(TokenBatch):
     image_patches: PatchBatch | None = None
     advantages: TokenDataBatch | None = None
     old_log_probabilities: TokenDataBatch | None = None
+    reward: TokenDataBatch | None = None
+    model_version: TokenDataBatch | None = None
 
     @classmethod
     def from_documents(
@@ -123,6 +131,10 @@ def from_documents(
         batch.old_log_probabilities = TokenDataBatch.from_documents(
             [document.old_log_probabilities for document in documents], lengths, pad_to_size
         )
+        batch.reward = TokenDataBatch.from_documents([document.reward for document in documents], lengths, pad_to_size)
+        batch.model_version = TokenDataBatch.from_documents(
+            [document.model_version for document in documents], lengths, pad_to_size
+        )
         return batch
 
     def get_model_inputs(self, config: LanguageModelBatchPreprocessingConfig) -> list[LanguageModelInput]:
@@ -204,6 +216,11 @@ def _set_target_inputs(
                     target_input.old_log_probabilities = self.old_log_probabilities.get_cropped_data(
                         label_begin, label_end
                     )
+                    # Optional diagnostic data (present only when the producer sends it).
+                    if self.reward is not None:
+                        target_input.reward = self.reward.get_cropped_data(label_begin, label_end)
+                    if self.model_version is not None:
+                        target_input.model_version = self.model_version.get_cropped_data(label_begin, label_end)
 
                 model_input.targets.append(target_input)
 

fast_llm/engine/schedule/runner.py

@@ -66,6 +66,9 @@ def __repr__(self):
 
 class ScheduleRunner[ConfigType: ScheduleConfig](Configurable[ConfigType]):
     _is_setup: bool = False
+    # Whole-step document count (DP-summed) from the last `run_step`, or None when the data does not
+    # provide document counts (i.e. no loss requested `return_document_count`).
+    _num_documents_in_batch: int | None = None
     _compute_stream: torch.cuda.Stream | MockStream
     _data_stream: torch.cuda.Stream | MockStream
     _pipeline_stream: torch.cuda.Stream | MockStream
@@ -155,6 +158,9 @@ def run_step(
             assert self._support_training
 
         metrics = {} if return_metrics else None
+        if metrics is not None:
+            # Always present on logging steps so "no wait" shows as 0 rather than a gap.
+            metrics["data_wait_time_ms"] = 0.0
         # Set the context.
         context = BatchContext(
             iteration=iteration,
@@ -322,6 +328,8 @@ def _preprocess_data(
         model_inputs[0][0].share_batch_data(
             [model_input for model_inputs_ in model_inputs for model_input in model_inputs_], self._distributed
         )
+        # Whole-step DP-summed document count (the loss-normalization divisor), for `documents_seen`.
+        self._num_documents_in_batch = model_inputs[0][0].num_documents_in_batch
 
         for micro_batch, model_inputs_ in enumerate(model_inputs):
             Assert.eq(len(model_inputs_), self._config.micro_batch_splits)
@@ -426,6 +434,9 @@ def _get_forward_input(self, context: BatchContext, step: Step) -> torch.Tensor:
                 next(context.data_iterator)
 
             data_time = (time.perf_counter() - start_time) * 1000
+            if context.metrics is not None:
+                # Time the trainer spent blocked waiting for the data loader — how input-starved it is.
+                context.metrics["data_wait_time_ms"] = context.metrics.get("data_wait_time_ms", 0.0) + data_time
             if data_time > self._config.data_batch_warn_time_ms:
                 logger.warning(f"Data loading took {data_time:,.2f} ms")
         return context.inputs.pop(step.global_index).detach().requires_grad_(step.stage != 0)

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
 

fast_llm/engine/training/streaming.py

@@ -40,19 +40,20 @@ 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,
     ):
         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:

fast_llm/engine/training/trainer.py

@@ -39,6 +39,7 @@ class Trainer[ConfigType: TrainerConfig](Configurable[ConfigType], abc.ABC):
     _wandb: Wandb
     _optimizer: Optimizer | None
     _completed_steps: int
+    _documents_seen: int = 0
     _schedule: Schedule
 
     def __init__(self, config: TrainerConfig):
@@ -203,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()
@@ -227,6 +228,12 @@ def _train(self) -> tuple[bool, dict[PhaseType, dict[str, typing.Any]]]:
                     return_metrics=is_logging,
                 )
 
+                # Cumulative document count (the RL x-axis / model-version clock); `None` when the
+                # data provides no document counts (non-RL runs).
+                step_num_documents = self._runner._num_documents_in_batch
+                if step_num_documents is not None:
+                    self._documents_seen += step_num_documents
+
                 # Advanced, skipped, and Nan iterations.
                 if update_successful:
                     advanced_iters += 1
@@ -237,7 +244,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:
@@ -257,6 +270,11 @@ def _train(self) -> tuple[bool, dict[PhaseType, dict[str, typing.Any]]]:
                         metrics_key = PhaseType.training
                         metrics[metrics_key] = {
                             "batch_size": self._batch_size,
+                            **(
+                                {"num_documents": step_num_documents, "documents_seen": self._documents_seen}
+                                if step_num_documents is not None
+                                else {}
+                            ),
                             **{
                                 name: (value / advanced_iters if advanced_iters > 0 else float("nan"))
                                 for name, value in total_losses.items()
@@ -350,6 +368,7 @@ def _prepare_training_state(self) -> None:
             if self._do_train:
                 self._optimizer.reset_state()
             self._completed_steps = 0
+            self._documents_seen = 0
         else:
             log_main_rank(lambda: f"Loading checkpoint from iteration {last_iteration}...")
             self._load_checkpoint(self._config.training.checkpoint, last_iteration)
@@ -387,6 +406,7 @@ def _save_checkpoint(
         metadata = {
             "optimizer": self._optimizer.save(),
             "completed_steps": self._completed_steps,
+            "documents_seen": self._documents_seen,
         }
         if metrics is not None:
             metadata["metrics"] = metrics
@@ -430,6 +450,7 @@ def _load_checkpoint(self, config: TrainingCheckpointConfig, iteration: int) ->
             self._completed_steps = metadata["schedules"][PhaseType.training]["completed_steps"]
         else:
             self._completed_steps = metadata["completed_steps"]
+        self._documents_seen = metadata.get("documents_seen", 0)
         # TODO: Move barrier, ok file to FastLLMModel
         safe_barrier(
             self._distributed.world_group,

fast_llm/engine/training/wandb.py

@@ -5,6 +5,7 @@
 
 from fast_llm.config import Config
 from fast_llm.engine.config_utils.run import Run
+from fast_llm.engine.distributed.config import PhaseType
 from fast_llm.engine.training.config import WandbConfig
 
 
@@ -13,6 +14,9 @@ def __init__(self, config: WandbConfig, run: Run, experiment_config: Config):
         self._config = config
         self._is_setup = True
         self._run = run
+        # Whether the `documents_seen` custom x-axis has been declared (RL runs only, done lazily on
+        # the first log that carries it so non-RL runs are unaffected).
+        self._defined_documents_axis = False
         if self._config.entity_name is not None and self._run.is_main_rank:
             import wandb
             import wandb.sdk.lib.runid
@@ -51,6 +55,17 @@ def log_metrics(self, completed_steps: int, metrics: dict[str, dict[str, float |
         if self._wandb is not None:
             import wandb
 
+            if not self._defined_documents_axis and any(
+                isinstance(values, dict) and "documents_seen" in values for values in metrics.values()
+            ):
+                # Offer `documents_seen` as an alternative (cross-run) x-axis for the training metrics.
+                # `wandb.log(step=...)` still records the global step, so both axes remain available.
+                self._wandb.define_metric(f"{PhaseType.training}/documents_seen")
+                self._wandb.define_metric(
+                    f"{PhaseType.training}/*", step_metric=f"{PhaseType.training}/documents_seen"
+                )
+                self._defined_documents_axis = True
+
             wandb.log(metrics, step=completed_steps, commit=commit)  # noqa
 
     def alert(self, title, text, level="INFO", wait=0.001) -> None:

fast_llm/layers/language_model/loss/config.py

@@ -29,6 +29,8 @@ class LanguageModelLossKwargs(BlockKwargs):
     rejected_spans = "rejected_spans"
     advantages = "advantages"
     old_log_probabilities = "old_log_probabilities"
+    reward = "reward"
+    model_version = "model_version"
     label_counts = "label_counts"
     num_labels_in_batch = "num_labels_in_batch"