diff --git a/fast_llm/data/dataset/streaming.py b/fast_llm/data/dataset/streaming.py index ec8fe7bd1..35d61b3f2 100644 --- a/fast_llm/data/dataset/streaming.py +++ b/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)), ) diff --git a/fast_llm/data/document/language_model.py b/fast_llm/data/document/language_model.py index 16114cb80..592f9e3cf 100644 --- a/fast_llm/data/document/language_model.py +++ b/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) diff --git a/fast_llm/layers/language_model/loss/config.py b/fast_llm/layers/language_model/loss/config.py index 4f3cfae8d..70e0ad62b 100644 --- a/fast_llm/layers/language_model/loss/config.py +++ b/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" diff --git a/fast_llm/layers/language_model/loss/policy_gradient.py b/fast_llm/layers/language_model/loss/policy_gradient.py index c07c81e05..b2a3f533b 100644 --- a/fast_llm/layers/language_model/loss/policy_gradient.py +++ b/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 diff --git a/tests/data/test_streaming.py b/tests/data/test_streaming.py index e938de7d0..6388e8bee 100644 --- a/tests/data/test_streaming.py +++ b/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"), diff --git a/tests/layers/test_lm_losses.py b/tests/layers/test_lm_losses.py index 78d588cb1..60ac845ca 100644 --- a/tests/layers/test_lm_losses.py +++ b/tests/layers/test_lm_losses.py @@ -8,7 +8,7 @@ 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 @@ -16,6 +16,8 @@ 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))