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,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/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"),

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,46 @@ def test_gspo_metrics(
     )
 
 
+@pytest.mark.parametrize("include_model_version", (True, False))
+def test_policy_data_metrics(include_model_version):
+    """`_register_data_metrics` logs reward (and, when present, 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
+
+    kwargs = {
+        LanguageModelLossKwargs.labels: [labels],
+        LanguageModelLossKwargs.reward: [reward],
+        LanguageModelLossKwargs.model_version: [model_version] if include_model_version else [None],
+        LanguageModelLossKwargs.label_counts: [label_counts],
+        LanguageModelKwargs.num_documents_in_batch: num_documents,
+    }
+    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 (("reward", reward), ("model_version", model_version)):
+        if name == "model_version" and not include_model_version:
+            # Declared but not registered (data absent) -> reduces to 0 downstream, no entries here.
+            assert losses[f"grpo_{name}"] == []
+            continue
+        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))