Carry reward and per-token model_version, log reward/version 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,10 @@ def _set_target_inputs(
                     target_input.old_log_probabilities = self.old_log_probabilities.get_cropped_data(
                         label_begin, label_end
                     )
+                    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

@@ -29,6 +29,7 @@
 class BatchContext:
     iteration: int
     schedule: Schedule
+    documents_seen: int = 0
     # Index and data: (iteration, data_index, input, kwargs)
     data_iterator: typing.Iterator[tuple[int, torch.Tensor, dict]] = None
     inputs: dict[int, torch.Tensor] = dataclasses.field(default_factory=dict)
@@ -149,6 +150,7 @@ def run_step(
         schedule: Schedule,
         *,
         iteration: int = 1,
+        documents_seen: int = 0,
         return_metrics: bool = False,
     ) -> tuple[dict[str, float | int], bool, dict[str, typing.Any] | None, int]:
         assert self._is_setup
@@ -161,6 +163,7 @@ def run_step(
         context = BatchContext(
             iteration=iteration,
             schedule=schedule,
+            documents_seen=documents_seen,
             losses={loss_def: [] for loss_def in self._loss_definitions},
             metrics=metrics,
         )
@@ -336,6 +339,7 @@ def _preprocess_data(
                     metrics=context.metrics,
                     extra_kwargs={
                         "grad_output": grad_output,
+                        "documents_seen": context.documents_seen,
                         "micro_batch": micro_batch,
                         "num_micro_batches": self._config.sequential_micro_batches,
                         "micro_batch_splits": self._config.micro_batch_splits,

fast_llm/engine/training/trainer.py

@@ -225,6 +225,7 @@ def _train(self) -> tuple[bool, dict[PhaseType, dict[str, typing.Any]]]:
                     train_iterator,
                     self._schedule,
                     iteration=self._completed_steps,
+                    documents_seen=self._documents_seen,
                     return_metrics=is_logging,
                 )
 

fast_llm/layers/language_model/config.py

@@ -25,6 +25,8 @@ class LanguageModelKwargs(LanguageModelLossKwargs):
     sample_map = "sample_map"
     embedding_map = "embedding_map"
     num_documents_in_batch = "num_documents_in_batch"
+    # Cumulative document count at the start of the step.
+    documents_seen = "documents_seen"
     # TODO: These are generic
     phase = "phase"
     loss_mask = "loss_mask"

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

@@ -48,6 +48,18 @@ class LanguageModelPolicyGradientLoss[ConfigType: LanguageModelPolicyGradientLos
 ):
     """Shared scaffolding for policy-gradient losses (GRPO, GSPO)."""
 
+    # Per-token diagnostics supplied by the rollout producer, logged (mean/max/min) under the given
+    # metric name. Reward is logged as `train_samples_reward`: averaged over the sample-filtered
+    # training batch it is biased, so it is a diagnostic, not a valid policy-performance metric.
+    # Model version is logged as `staleness` (`documents_seen - model_version`, documents-seen units):
+    # how many documents were trained since each token's generating policy was synced. When a field's
+    # reference key is set, its per-token value is subtracted from that whole-batch scalar.
+    _DATA_METRIC_FIELDS = (
+        # (metric_name, data_key, reference_key)
+        ("train_samples_reward", LanguageModelLossKwargs.reward, None),
+        ("staleness", LanguageModelLossKwargs.model_version, LanguageModelKwargs.documents_seen),
+    )
+
     def __init__(
         self,
         config: ConfigType,
@@ -109,6 +121,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 +143,46 @@ 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 _register_data_metrics(self, kwargs: dict[str, typing.Any], losses: dict, split_index: int) -> None:
+        # The values are constant / near-constant within a document, so the per-document mean and the
+        # token extrema are the natural summaries.
+        num_documents = kwargs[LanguageModelKwargs.num_documents_in_batch]
+        loss_mask = self._get_labels(kwargs, split_index) >= 0
+        label_counts = self._prepare_target(kwargs[LanguageModelLossKwargs.label_counts], split_index)
+        document_weight = loss_mask.float() / label_counts.float().clamp(min=1)
+        negative_infinity = document_weight.new_full((), float("-inf"))
+        positive_infinity = document_weight.new_full((), float("inf"))
+        for metric_name, data_key, reference_key in self._DATA_METRIC_FIELDS:
+            values = self._prepare_target(kwargs[data_key], split_index)
+            if reference_key is not None:
+                # Subtract before casting: both are large document counts, so casting first would
+                # lose the small difference to float32 rounding.
+                values = kwargs[reference_key] - values
+            values = values.float()
+            self._register_loss(
+                f"{self._name}_{metric_name}", (values * document_weight).sum() / num_documents, losses
+            )
+            self._register_loss(
+                f"{self._name}_max_{metric_name}",
+                torch.where(loss_mask, values, negative_infinity).max(),
+                losses,
+                reduce_op=torch.distributed.ReduceOp.MAX,
+            )
+            self._register_loss(
+                f"{self._name}_min_{metric_name}",
+                torch.where(loss_mask, values, positive_infinity).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"),

tests/layers/test_lm_losses.py

@@ -8,14 +8,16 @@
 from fast_llm.core.ops import split_op
 from fast_llm.engine.config_utils import data_type
 from fast_llm.engine.config_utils.data_type import DataType
-from fast_llm.engine.distributed.config import DistributedBackend
+from fast_llm.engine.distributed.config import DistributedBackend, DistributedConfig
 from fast_llm.functional.config import EntropyLossType, TargetFormat
 from fast_llm.functional.entropy_loss import fused_entropy_loss_forward_backward, torch_entropy_loss_forward_backward
 from fast_llm.functional.triton import triton_available
 from fast_llm.functional.triton.entropy_loss import triton_entropy_loss_forward_backward
 from fast_llm.functional.triton.grpo_loss import triton_grpo_loss_forward_backward
 from fast_llm.functional.triton.gspo_loss import triton_gspo_loss_forward_backward
 from fast_llm.functional.triton.z_loss import triton_z_loss_forward_backward
+from fast_llm.layers.language_model.config import LanguageModelKwargs
+from fast_llm.layers.language_model.loss.config import LanguageModelGRPOLossConfig, LanguageModelLossKwargs
 from fast_llm.layers.language_model.loss.dpo import dpo_loss
 from fast_llm.layers.language_model.loss.loss import loss_forward_backward
 from fast_llm.layers.language_model.loss.policy_gradient import (
@@ -840,6 +842,43 @@ def test_gspo_metrics(
     )
 
 
+def test_policy_data_metrics():
+    """`_register_data_metrics` logs reward and staleness (documents_seen - model_version) mean/max/min."""
+    config = LanguageModelGRPOLossConfig.from_dict({"metrics": "basic"})
+    loss = config.get_layer(DistributedConfig.from_dict({}), name="grpo", prediction_distance=1, prediction_heads=1)
+
+    # 6 tokens, 2 documents (3 each), token 2 masked. reward is constant per document.
+    labels = torch.tensor([1, 2, -100, 3, 4, 5])
+    loss_mask = labels >= 0
+    reward = torch.tensor([1.0, 1.0, 1.0, 0.0, 0.0, 0.0])
+    model_version = torch.tensor([5, 5, 5, 7, 8, 9], dtype=torch.int64)
+    label_counts = torch.tensor([2, 2, 2, 3, 3, 3], dtype=torch.int32)
+    num_documents = 2
+    documents_seen = 100
+
+    kwargs = {
+        LanguageModelLossKwargs.labels: [labels],
+        LanguageModelLossKwargs.reward: [reward],
+        LanguageModelLossKwargs.model_version: [model_version],
+        LanguageModelLossKwargs.label_counts: [label_counts],
+        LanguageModelKwargs.num_documents_in_batch: num_documents,
+        LanguageModelKwargs.documents_seen: documents_seen,
+    }
+    losses = {loss_def.name: [] for loss_def in loss.get_loss_definitions()}
+    loss._register_data_metrics(kwargs, losses, 0)
+
+    def reference(values: torch.Tensor) -> tuple[float, float, float]:
+        masked = loss_mask.float() / label_counts.float().clamp(min=1)
+        values = values.float()
+        return (values * masked).sum() / num_documents, values[loss_mask].max(), values[loss_mask].min()
+
+    for name, values in (("train_samples_reward", reward), ("staleness", documents_seen - model_version)):
+        mean, maximum, minimum = reference(values)
+        Assert.rms_close_relative(losses[f"grpo_{name}"][0], mean, 1e-6)
+        Assert.rms_close_relative(losses[f"grpo_max_{name}"][0], maximum, 1e-6)
+        Assert.rms_close_relative(losses[f"grpo_min_{name}"][0], minimum, 1e-6)
+
+
 @pytest.mark.skip(reason="DPO loss is broken")
 def test_dpo_loss():
     logits = torch.normal(0, 1, (200, 100))