RL diagnostics: documents_seen, weights_ready document-count version, reward/model_version staleness metrics

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/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

@@ -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()
@@ -244,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:

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"
 

fast_llm/layers/language_model/loss/policy_gradient.py

@@ -81,6 +81,13 @@ def __init__(
             distributed_config.pipeline_parallel,
         )
 
+    # Per-token diagnostic data supplied by the rollout producer (mean/max/min logged when present).
+    # `reward` is the raw reward; `model_version` the version each token was generated under.
+    _DATA_METRIC_FIELDS = (
+        ("reward", LanguageModelLossKwargs.reward),
+        ("model_version", LanguageModelLossKwargs.model_version),
+    )
+
     def _register_new_logprobs(
         self,
         new_logprobs_mean: torch.Tensor | None,
@@ -109,6 +116,7 @@ def _policy_metric_definitions(self, *extra: LossDef) -> list[LossDef]:
             *extra,
             LossDef(f"{self._name}_num_tokens"),
         ]
+        defs.extend(self._data_metric_definitions())
         if self._config.metrics == PolicyMetricsLevel.with_entropy:
             defs.append(LossDef(f"{self._name}_entropy"))
         return defs
@@ -130,6 +138,51 @@ def _register_policy_metrics(self, metrics: PolicyMetrics, kwargs: dict[str, typ
         if metrics.entropy is not None:
             self._register_loss(f"{self._name}_entropy", metrics.entropy / num_documents, losses)
 
+    def _get_optional_target(self, kwargs: dict[str, typing.Any], key: str, split_index: int) -> torch.Tensor | None:
+        targets = kwargs.get(key)
+        if targets is None or targets[self._prediction_distance - 1] is None:
+            return None
+        return self._prepare_target(targets, split_index)
+
+    def _register_data_metrics(self, kwargs: dict[str, typing.Any], losses: dict | None, split_index: int) -> None:
+        # Mean (per document), max and min of each supplied per-token diagnostic. `reward` and
+        # `model_version` are constant / near-constant within a document, so the per-document mean and
+        # the token extrema are the natural summaries; staleness is `documents_seen - model_version`.
+        num_documents = kwargs[LanguageModelKwargs.num_documents_in_batch]
+        loss_mask = None
+        for name, key in self._DATA_METRIC_FIELDS:
+            values = self._get_optional_target(kwargs, key, split_index)
+            if values is None:
+                continue
+            if loss_mask is None:
+                loss_mask = self._get_labels(kwargs, split_index) >= 0
+                label_counts = self._prepare_target(kwargs[LanguageModelLossKwargs.label_counts], split_index)
+                masked = loss_mask.float() / label_counts.float().clamp(min=1)
+            values = values.float()
+            neg_inf = values.new_full((), float("-inf"))
+            pos_inf = values.new_full((), float("inf"))
+            self._register_loss(f"{self._name}_{name}", (values * masked).sum() / num_documents, losses)
+            self._register_loss(
+                f"{self._name}_max_{name}",
+                torch.where(loss_mask, values, neg_inf).max(),
+                losses,
+                reduce_op=torch.distributed.ReduceOp.MAX,
+            )
+            self._register_loss(
+                f"{self._name}_min_{name}",
+                torch.where(loss_mask, values, pos_inf).min(),
+                losses,
+                reduce_op=torch.distributed.ReduceOp.MIN,
+            )
+
+    def _data_metric_definitions(self) -> list[LossDef]:
+        defs = []
+        for name, _ in self._DATA_METRIC_FIELDS:
+            defs.append(LossDef(f"{self._name}_{name}"))
+            defs.append(LossDef(f"{self._name}_max_{name}", reduction=ReductionType.maximum))
+            defs.append(LossDef(f"{self._name}_min_{name}", reduction=ReductionType.minimum))
+        return defs
+
     def get_loss_definitions(self) -> list[LossDef]:
         defs = super().get_loss_definitions()
         defs.append(LossDef(self._logprob_metric_name))
@@ -184,6 +237,7 @@ def _forward_backward(
         # Skip the extra softmax pass when there is nothing to register.
         if losses is not None and self._config.metrics != PolicyMetricsLevel.none:
             self._register_extra_metrics(logits, kwargs, losses, split_index)
+            self._register_data_metrics(kwargs, losses, split_index)
 
         return loss, grad
 
@@ -296,6 +350,7 @@ def _forward_backward(
         # Skip the extra softmax pass when there is nothing to register.
         if losses is not None and self._config.metrics != PolicyMetricsLevel.none:
             self._register_extra_metrics(logits, kwargs, losses, split_index, document_index_zero_based, num_segments)
+            self._register_data_metrics(kwargs, losses, split_index)
 
         return loss, grad
 

tests/data/test_streaming.py

@@ -48,6 +48,22 @@ def fake_redis(monkeypatch):
             {"tokens": list(range(3)), "advantage": 0.33, "old_log_probabilities": [0.25, -0.52, 0.99]},
             {"tokens": list(range(4)), "advantage": 0.7, "old_log_probabilities": [1, 2, 3, 4]},
         ),
+        (
+            {
+                "tokens": list(range(3)),
+                "advantage": 0.33,
+                "old_log_probabilities": [0.25, -0.52, 0.99],
+                "reward": 1.0,
+                "model_version": [5, 5, 5],
+            },
+            {
+                "tokens": list(range(4)),
+                "advantage": 0.7,
+                "old_log_probabilities": [1, 2, 3, 4],
+                "reward": 0.0,
+                "model_version": [7, 8, 8, 9],
+            },
+        ),
     ],
 )
 def test_streaming_dataset(
@@ -97,6 +113,18 @@ def test_streaming_dataset(
         else:
             assert sampled_document.old_log_probabilities is None
 
+        if "reward" in document:
+            Assert.rms_close(
+                sampled_document.reward.data, torch.full([len(document["tokens"])], document["reward"]), 1e-8
+            )
+        else:
+            assert sampled_document.reward is None
+
+        if "model_version" in document:
+            Assert.eq(sampled_document.model_version.data.tolist(), document["model_version"])
+        else:
+            assert sampled_document.model_version is None
+
 
 @pytest.mark.parametrize(
     ("messages", "expected_samples", "expected_lengths"),