From 5975c8d98af7a3f349db0f7e3198f45c2432f5cf Mon Sep 17 00:00:00 2001 From: Joel Lamy-Poirier Date: Tue, 7 Jul 2026 14:39:24 -0400 Subject: [PATCH 1/7] Carry reward and per-token model_version, log reward/version metrics MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add two optional per-token fields to the RL streaming schema and thread them through the data pipeline alongside `advantages` / `old_log_probabilities` (reusing `TokenDataDocument`/`TokenDataBatch`): - `reward`: the raw (un-normalized) reward, a per-rollout scalar broadcast per-token — distinct from the group-relative `advantage`. - `model_version`: the model version each token was generated under (documents-seen units), one per token, for staleness diagnostics. Both are optional (absent when the producer does not send them), so the batch/target plumbing guards on presence. The shared policy-gradient loss logs mean/max/min of each supplied field when `metrics != none` (GRPO and GSPO), registered only when the data is present. Staleness is `documents_seen - model_version`, derivable from the logged version stats and `documents_seen`. Co-Authored-By: Claude Opus 4.8 --- fast_llm/data/dataset/streaming.py | 20 +++++++ fast_llm/data/document/language_model.py | 17 ++++++ fast_llm/layers/language_model/loss/config.py | 2 + .../language_model/loss/policy_gradient.py | 55 +++++++++++++++++++ tests/data/test_streaming.py | 28 ++++++++++ tests/layers/test_lm_losses.py | 44 ++++++++++++++- 6 files changed, 165 insertions(+), 1 deletion(-) 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)) From e08726e5759ae6165855b121a9e9598f1e57e083 Mon Sep 17 00:00:00 2001 From: Joel Lamy-Poirier Date: Fri, 10 Jul 2026 09:18:13 -0400 Subject: [PATCH 2/7] Review fixes: inline single-use target helper, hoist inf sentinels, trim comments Inline the one-caller `_get_optional_target` into `_register_data_metrics`, dropping its dead `targets is None` disjunct (`to_kwargs` always sets the key); build the max/min sentinels once in the first-iteration init block; narrow the `losses` annotation to `dict`; move `_DATA_METRIC_FIELDS` above `__init__`; drop two restating/downstream-referencing comments. Co-Authored-By: Claude Opus 4.8 --- fast_llm/data/document/language_model.py | 1 - .../language_model/loss/policy_gradient.py | 34 ++++++++----------- 2 files changed, 14 insertions(+), 21 deletions(-) diff --git a/fast_llm/data/document/language_model.py b/fast_llm/data/document/language_model.py index 592f9e3cf..b70814c88 100644 --- a/fast_llm/data/document/language_model.py +++ b/fast_llm/data/document/language_model.py @@ -216,7 +216,6 @@ 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: diff --git a/fast_llm/layers/language_model/loss/policy_gradient.py b/fast_llm/layers/language_model/loss/policy_gradient.py index b6260e724..36c59727f 100644 --- a/fast_llm/layers/language_model/loss/policy_gradient.py +++ b/fast_llm/layers/language_model/loss/policy_gradient.py @@ -48,6 +48,13 @@ class LanguageModelPolicyGradientLoss[ConfigType: LanguageModelPolicyGradientLos ): """Shared scaffolding for policy-gradient losses (GRPO, GSPO).""" + # 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 __init__( self, config: ConfigType, @@ -81,13 +88,6 @@ 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, @@ -138,29 +138,23 @@ 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: + def _register_data_metrics(self, kwargs: dict[str, typing.Any], losses: dict, 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`. + # the token extrema are the natural summaries. 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: + targets = kwargs[key] + if targets[self._prediction_distance - 1] is None: continue + values = self._prepare_target(targets, split_index).float() 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")) + 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}", From 28b3fd62d844bcdbd3aace239dcc89e4d7a60944 Mon Sep 17 00:00:00 2001 From: Joel Lamy-Poirier Date: Fri, 10 Jul 2026 09:30:51 -0400 Subject: [PATCH 3/7] Log reward metric as train_samples_reward to flag filtering bias The reward is averaged over the sample-filtered training batch, so its mean is biased and not a valid policy-performance metric. Rename only the logged metric label (the data field and kwargs key stay `reward`), and note the caveat. Co-Authored-By: Claude Opus 4.8 --- fast_llm/layers/language_model/loss/policy_gradient.py | 8 +++++--- tests/layers/test_lm_losses.py | 2 +- 2 files changed, 6 insertions(+), 4 deletions(-) diff --git a/fast_llm/layers/language_model/loss/policy_gradient.py b/fast_llm/layers/language_model/loss/policy_gradient.py index 36c59727f..ddc972db8 100644 --- a/fast_llm/layers/language_model/loss/policy_gradient.py +++ b/fast_llm/layers/language_model/loss/policy_gradient.py @@ -48,10 +48,12 @@ class LanguageModelPolicyGradientLoss[ConfigType: LanguageModelPolicyGradientLos ): """Shared scaffolding for policy-gradient losses (GRPO, GSPO).""" - # 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. + # Per-token diagnostics supplied by the rollout producer, logged (mean/max/min) under the given + # 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 the version each token was generated under. _DATA_METRIC_FIELDS = ( - ("reward", LanguageModelLossKwargs.reward), + ("train_samples_reward", LanguageModelLossKwargs.reward), ("model_version", LanguageModelLossKwargs.model_version), ) diff --git a/tests/layers/test_lm_losses.py b/tests/layers/test_lm_losses.py index 60ac845ca..59dad6cb9 100644 --- a/tests/layers/test_lm_losses.py +++ b/tests/layers/test_lm_losses.py @@ -871,7 +871,7 @@ def reference(values: torch.Tensor) -> tuple[float, float, float]: 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)): + for name, values in (("train_samples_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}"] == [] From b225da1ae3929d56d414a47e90a95cfbceaf6333 Mon Sep 17 00:00:00 2001 From: Joel Lamy-Poirier Date: Fri, 10 Jul 2026 09:56:45 -0400 Subject: [PATCH 4/7] Log staleness (documents_seen - model_version) instead of raw model version Staleness must be computed per-step (model-version metrics are summed across steps and averaged only at log time, so it can't be derived from a single documents_seen afterwards). Thread the step's documents_seen into the loss kwargs through the same channel iteration uses (run_step -> BatchContext -> extra_kwargs), and subtract it from the per-token model version so the loss logs staleness directly. Other run_step callers take the default 0. Co-Authored-By: Claude Opus 4.8 --- fast_llm/engine/schedule/runner.py | 4 +++ fast_llm/engine/training/trainer.py | 1 + fast_llm/layers/language_model/config.py | 2 ++ .../language_model/loss/policy_gradient.py | 31 +++++++++++-------- tests/layers/test_lm_losses.py | 8 +++-- 5 files changed, 30 insertions(+), 16 deletions(-) diff --git a/fast_llm/engine/schedule/runner.py b/fast_llm/engine/schedule/runner.py index fb3ef538d..a8cc4f6d1 100644 --- a/fast_llm/engine/schedule/runner.py +++ b/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, diff --git a/fast_llm/engine/training/trainer.py b/fast_llm/engine/training/trainer.py index fb57c6a66..e11307e1e 100644 --- a/fast_llm/engine/training/trainer.py +++ b/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, ) diff --git a/fast_llm/layers/language_model/config.py b/fast_llm/layers/language_model/config.py index bde33f297..842f36001 100644 --- a/fast_llm/layers/language_model/config.py +++ b/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; the staleness reference for `model_version`. + documents_seen = "documents_seen" # TODO: These are generic phase = "phase" loss_mask = "loss_mask" diff --git a/fast_llm/layers/language_model/loss/policy_gradient.py b/fast_llm/layers/language_model/loss/policy_gradient.py index ddc972db8..ce3925cde 100644 --- a/fast_llm/layers/language_model/loss/policy_gradient.py +++ b/fast_llm/layers/language_model/loss/policy_gradient.py @@ -49,12 +49,15 @@ 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 - # 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 the version each token was generated under. + # 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 = ( - ("train_samples_reward", LanguageModelLossKwargs.reward), - ("model_version", LanguageModelLossKwargs.model_version), + # (metric_name, data_key, reference_key) + ("train_samples_reward", LanguageModelLossKwargs.reward, None), + ("staleness", LanguageModelLossKwargs.model_version, LanguageModelKwargs.documents_seen), ) def __init__( @@ -141,31 +144,33 @@ def _register_policy_metrics(self, metrics: PolicyMetrics, kwargs: dict[str, typ 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: - # 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. + # Mean (per document), max and min of each supplied per-token diagnostic. 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 = None - for name, key in self._DATA_METRIC_FIELDS: + for metric_name, key, reference_key in self._DATA_METRIC_FIELDS: targets = kwargs[key] if targets[self._prediction_distance - 1] is None: continue values = self._prepare_target(targets, split_index).float() + if reference_key is not None: + values = kwargs[reference_key] - values 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) 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}_{metric_name}", (values * masked).sum() / num_documents, losses) self._register_loss( - f"{self._name}_max_{name}", + f"{self._name}_max_{metric_name}", torch.where(loss_mask, values, neg_inf).max(), losses, reduce_op=torch.distributed.ReduceOp.MAX, ) self._register_loss( - f"{self._name}_min_{name}", + f"{self._name}_min_{metric_name}", torch.where(loss_mask, values, pos_inf).min(), losses, reduce_op=torch.distributed.ReduceOp.MIN, @@ -173,7 +178,7 @@ def _register_data_metrics(self, kwargs: dict[str, typing.Any], losses: dict, sp def _data_metric_definitions(self) -> list[LossDef]: defs = [] - for name, _ in self._DATA_METRIC_FIELDS: + 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)) diff --git a/tests/layers/test_lm_losses.py b/tests/layers/test_lm_losses.py index 59dad6cb9..8802082e9 100644 --- a/tests/layers/test_lm_losses.py +++ b/tests/layers/test_lm_losses.py @@ -844,7 +844,7 @@ 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.""" + """`_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) @@ -855,6 +855,7 @@ def test_policy_data_metrics(include_model_version): 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], @@ -862,6 +863,7 @@ def test_policy_data_metrics(include_model_version): LanguageModelLossKwargs.model_version: [model_version] if include_model_version else [None], 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) @@ -871,8 +873,8 @@ def reference(values: torch.Tensor) -> tuple[float, float, float]: 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), ("model_version", model_version)): - if name == "model_version" and not include_model_version: + for name, values in (("train_samples_reward", reward), ("staleness", documents_seen - model_version)): + if name == "staleness" and not include_model_version: # Declared but not registered (data absent) -> reduces to 0 downstream, no entries here. assert losses[f"grpo_{name}"] == [] continue From 9daf611536f248a926a9cd337d09fe818b2d4489 Mon Sep 17 00:00:00 2001 From: Joel Lamy-Poirier Date: Fri, 10 Jul 2026 11:04:20 -0400 Subject: [PATCH 5/7] Review fixes: rename masked->document_weight, loop var key->data_key Match the file's established name for mask/labeled-token-count weight, and align the loop variable with the tuple's `data_key` comment. Co-Authored-By: Claude Opus 4.8 --- fast_llm/layers/language_model/loss/policy_gradient.py | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/fast_llm/layers/language_model/loss/policy_gradient.py b/fast_llm/layers/language_model/loss/policy_gradient.py index ce3925cde..1d2167b71 100644 --- a/fast_llm/layers/language_model/loss/policy_gradient.py +++ b/fast_llm/layers/language_model/loss/policy_gradient.py @@ -149,8 +149,8 @@ def _register_data_metrics(self, kwargs: dict[str, typing.Any], losses: dict, sp # are the natural summaries. num_documents = kwargs[LanguageModelKwargs.num_documents_in_batch] loss_mask = None - for metric_name, key, reference_key in self._DATA_METRIC_FIELDS: - targets = kwargs[key] + for metric_name, data_key, reference_key in self._DATA_METRIC_FIELDS: + targets = kwargs[data_key] if targets[self._prediction_distance - 1] is None: continue values = self._prepare_target(targets, split_index).float() @@ -159,10 +159,12 @@ def _register_data_metrics(self, kwargs: dict[str, typing.Any], losses: dict, sp 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) + document_weight = loss_mask.float() / label_counts.float().clamp(min=1) neg_inf = values.new_full((), float("-inf")) pos_inf = values.new_full((), float("inf")) - self._register_loss(f"{self._name}_{metric_name}", (values * masked).sum() / num_documents, losses) + 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, neg_inf).max(), From f64ce7c74d93b3e6bca824b50346b3abaeac025e Mon Sep 17 00:00:00 2001 From: Joel Lamy-Poirier Date: Fri, 10 Jul 2026 11:20:31 -0400 Subject: [PATCH 6/7] Register reward/staleness metrics unconditionally (fields are required) Registering them conditionally on per-batch data presence made the metric's cross-DP all-reduce data-dependent: with metrics on (pipeline_parallel forced to 1), a rank that registered the metric all-reduces while a rank that skipped it does not, so mixed presence across DP ranks would hang. reward/model_version are always sent when these metrics are enabled, so register unconditionally like the other policy metrics; absence now fails uniformly instead of hanging. Hoist the shared per-document weight and extrema sentinels out of the loop. Co-Authored-By: Claude Opus 4.8 --- .../language_model/loss/policy_gradient.py | 17 ++++++----------- tests/layers/test_lm_losses.py | 9 ++------- 2 files changed, 8 insertions(+), 18 deletions(-) diff --git a/fast_llm/layers/language_model/loss/policy_gradient.py b/fast_llm/layers/language_model/loss/policy_gradient.py index 1d2167b71..435de39e9 100644 --- a/fast_llm/layers/language_model/loss/policy_gradient.py +++ b/fast_llm/layers/language_model/loss/policy_gradient.py @@ -148,20 +148,15 @@ def _register_data_metrics(self, kwargs: dict[str, typing.Any], losses: dict, sp # 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 = None + 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) + neg_inf = document_weight.new_full((), float("-inf")) + pos_inf = document_weight.new_full((), float("inf")) for metric_name, data_key, reference_key in self._DATA_METRIC_FIELDS: - targets = kwargs[data_key] - if targets[self._prediction_distance - 1] is None: - continue - values = self._prepare_target(targets, split_index).float() + values = self._prepare_target(kwargs[data_key], split_index).float() if reference_key is not None: values = kwargs[reference_key] - values - 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) - document_weight = loss_mask.float() / label_counts.float().clamp(min=1) - neg_inf = values.new_full((), float("-inf")) - pos_inf = values.new_full((), float("inf")) self._register_loss( f"{self._name}_{metric_name}", (values * document_weight).sum() / num_documents, losses ) diff --git a/tests/layers/test_lm_losses.py b/tests/layers/test_lm_losses.py index 8802082e9..dc6ca3bf3 100644 --- a/tests/layers/test_lm_losses.py +++ b/tests/layers/test_lm_losses.py @@ -842,8 +842,7 @@ def test_gspo_metrics( ) -@pytest.mark.parametrize("include_model_version", (True, False)) -def test_policy_data_metrics(include_model_version): +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) @@ -860,7 +859,7 @@ def test_policy_data_metrics(include_model_version): kwargs = { LanguageModelLossKwargs.labels: [labels], LanguageModelLossKwargs.reward: [reward], - LanguageModelLossKwargs.model_version: [model_version] if include_model_version else [None], + LanguageModelLossKwargs.model_version: [model_version], LanguageModelLossKwargs.label_counts: [label_counts], LanguageModelKwargs.num_documents_in_batch: num_documents, LanguageModelKwargs.documents_seen: documents_seen, @@ -874,10 +873,6 @@ def reference(values: torch.Tensor) -> tuple[float, float, 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)): - if name == "staleness" 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) From 91f674d2552d6f5743ab191e147c1919d801346f Mon Sep 17 00:00:00 2001 From: Joel Lamy-Poirier Date: Fri, 10 Jul 2026 11:43:50 -0400 Subject: [PATCH 7/7] Review fixes: int-domain staleness subtraction, naming/comment nits Compute staleness (documents_seen - model_version) in the integer domain before casting to float, so the small difference isn't lost to float32 rounding when both counts exceed ~2^24. Also spell out neg/pos infinity sentinels, and drop a code-restating comment and a downstream-consumer reference. Co-Authored-By: Claude Opus 4.8 --- fast_llm/layers/language_model/config.py | 2 +- .../language_model/loss/policy_gradient.py | 18 ++++++++++-------- 2 files changed, 11 insertions(+), 9 deletions(-) diff --git a/fast_llm/layers/language_model/config.py b/fast_llm/layers/language_model/config.py index 842f36001..298e691b9 100644 --- a/fast_llm/layers/language_model/config.py +++ b/fast_llm/layers/language_model/config.py @@ -25,7 +25,7 @@ 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; the staleness reference for `model_version`. + # Cumulative document count at the start of the step. documents_seen = "documents_seen" # TODO: These are generic phase = "phase" diff --git a/fast_llm/layers/language_model/loss/policy_gradient.py b/fast_llm/layers/language_model/loss/policy_gradient.py index 435de39e9..25c483f6c 100644 --- a/fast_llm/layers/language_model/loss/policy_gradient.py +++ b/fast_llm/layers/language_model/loss/policy_gradient.py @@ -144,31 +144,33 @@ def _register_policy_metrics(self, metrics: PolicyMetrics, kwargs: dict[str, typ 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: - # Mean (per document), max and min of each supplied per-token diagnostic. The values are - # constant / near-constant within a document, so the per-document mean and the token extrema - # are the natural summaries. + # 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) - neg_inf = document_weight.new_full((), float("-inf")) - pos_inf = document_weight.new_full((), float("inf")) + 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).float() + 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, neg_inf).max(), + 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, pos_inf).min(), + torch.where(loss_mask, values, positive_infinity).min(), losses, reduce_op=torch.distributed.ReduceOp.MIN, )