From 7b8ca77eb11826807545d140ca89760636500b89 Mon Sep 17 00:00:00 2001 From: Joel Lamy-Poirier Date: Tue, 7 Jul 2026 14:04:52 -0400 Subject: [PATCH 1/4] Unify policy-loss metrics config and add GSPO metrics Rename `GRPOMetricsLevel` to `PolicyLossMetrics` and lift the `metrics` field from the GRPO config to the shared `LanguageModelPolicyGradientLossConfig`, so GSPO gains the same diagnostic levels (none/basic/with_entropy). The `pipeline_parallel == 1` guard moves to the shared `LanguageModelPolicyGradientLoss.__init__`. Add `compute_gspo_metrics` (segment-level analog of `compute_grpo_metrics`) and GSPO's `_register_extra_metrics`/`get_loss_definitions`. GSPO clips per segment, so its statistics (ratio, KL, clip fraction) are per-segment; they are accumulated as token-weighted sums (weight `mask / num_labels_in_seq`, which sums to 1 per document across SDP/SP ranks) so they partition and reduce exactly like the per-token loss. Metric names match GRPO where the meaning aligns (`_num_segments` replaces `_num_tokens`). Tested via `compute_gspo_metrics` against an independent plain-Python segment-loop reference (`reference_gspo_metrics` + `test_gspo_metrics`). Co-Authored-By: Claude Opus 4.8 --- fast_llm/layers/language_model/loss/config.py | 22 +- .../language_model/loss/policy_gradient.py | 237 +++++++++++++++--- tests/layers/test_lm_losses.py | 150 +++++++++++ 3 files changed, 364 insertions(+), 45 deletions(-) diff --git a/fast_llm/layers/language_model/loss/config.py b/fast_llm/layers/language_model/loss/config.py index 9a220aacf..13fc24b0a 100644 --- a/fast_llm/layers/language_model/loss/config.py +++ b/fast_llm/layers/language_model/loss/config.py @@ -205,7 +205,7 @@ def loss_class(self) -> "type[LanguageModelZLoss]": return LanguageModelZLoss -class GRPOMetricsLevel(enum.StrEnum): +class PolicyLossMetrics(enum.StrEnum): none = "none" basic = "basic" with_entropy = "with_entropy" @@ -219,6 +219,16 @@ class LanguageModelPolicyGradientLossConfig(LanguageModelLossConfig): epsilon_low: float = Field(default=0.2, desc="Lower clip parameter for ratio of log probs") epsilon_high: float = Field(default=0.2, desc="Upper clip parameter for ratio of log probs") + metrics: PolicyLossMetrics = Field( + default=PolicyLossMetrics.none, + desc=( + "Additional diagnostic metrics to log. " + "`basic`: importance-ratio, KL and advantage statistics. " + "`with_entropy`: also log the policy entropy. " + "Not supported with pipeline_parallel > 1." + ), + hint=FieldHint.feature, + ) @property def loss_class(self) -> "type[LanguageModelPolicyGradientLoss]": @@ -236,16 +246,6 @@ class LanguageModelGRPOLossConfig(LanguageModelPolicyGradientLossConfig): desc="Enable triton implementation. Default: use if available.", hint=FieldHint.expert, ) - metrics: GRPOMetricsLevel = Field( - default=GRPOMetricsLevel.none, - desc=( - "Additional GRPO metrics to log. " - "`basic`: per-token ratio, KL, and advantage statistics. " - "`with_entropy`: also log per-token entropy. " - "Not supported with pipeline_parallel > 1." - ), - hint=FieldHint.feature, - ) @property def loss_class(self) -> "type[LanguageModelGRPOLoss]": diff --git a/fast_llm/layers/language_model/loss/policy_gradient.py b/fast_llm/layers/language_model/loss/policy_gradient.py index a024d4232..ba004944d 100644 --- a/fast_llm/layers/language_model/loss/policy_gradient.py +++ b/fast_llm/layers/language_model/loss/policy_gradient.py @@ -12,11 +12,11 @@ from fast_llm.layers.block.config import BlockKwargs from fast_llm.layers.language_model.config import LanguageModelKwargs from fast_llm.layers.language_model.loss.config import ( - GRPOMetricsLevel, LanguageModelGRPOLossConfig, LanguageModelGSPOLossConfig, LanguageModelLossKwargs, LanguageModelPolicyGradientLossConfig, + PolicyLossMetrics, ) from fast_llm.layers.language_model.loss.loss import LanguageModelLoss from fast_llm.utils import Assert @@ -36,39 +36,26 @@ class GRPOMetrics(typing.NamedTuple): entropy: torch.Tensor | None +class GSPOMetrics(typing.NamedTuple): + # Statistics over segments (documents); GSPO clips per segment. `*_sum` fields are raw sums over + # segments (divided by the document count at registration time to form means / fractions). + old_logprobs: torch.Tensor + ratio_sum: torch.Tensor + ratio_squared_sum: torch.Tensor + kl_sum: torch.Tensor + clipped_sum: torch.Tensor + advantage_sum: torch.Tensor + max_advantage: torch.Tensor + min_advantage: torch.Tensor + num_segments: torch.Tensor + entropy: torch.Tensor | None + + class LanguageModelPolicyGradientLoss[ConfigType: LanguageModelPolicyGradientLossConfig]( LanguageModelLoss[ConfigType] ): """Shared scaffolding for policy-gradient losses (GRPO, GSPO).""" - def _register_new_logprobs( - self, - new_logprobs_mean: torch.Tensor | None, - kwargs: dict[str, typing.Any], - losses: dict | None, - ) -> None: - if new_logprobs_mean is not None: - new_logprobs_mean = new_logprobs_mean / kwargs[LanguageModelKwargs.num_documents_in_batch] - self._register_loss( - self._logprob_metric_name, new_logprobs_mean, losses, reduce_op=torch.distributed.ReduceOp.SUM - ) - - def get_loss_definitions(self) -> list[LossDef]: - defs = super().get_loss_definitions() - defs.append(LossDef(self._logprob_metric_name)) - return defs - - def get_preprocessing_config(self) -> dict[str, typing.Any]: - return {"use_grpo_data": True, "return_label_counts": True, "return_document_count": True} - - @functools.cached_property - def _logprob_metric_name(self) -> str: - return f"{self._name}_new_logprobs" - - -class LanguageModelGRPOLoss[ConfigType: LanguageModelGRPOLossConfig](LanguageModelPolicyGradientLoss[ConfigType]): - """GRPO: per-token IS-ratio clipping.""" - def __init__( self, config: ConfigType, @@ -95,12 +82,41 @@ def __init__( weight=weight, register_loss=register_loss, ) + # The extra metrics need a second softmax over the full logits, which pipeline parallelism splits. Assert.custom( - lambda metrics, pipeline_parallel: metrics == GRPOMetricsLevel.none or pipeline_parallel == 1, + lambda metrics, pipeline_parallel: metrics == PolicyLossMetrics.none or pipeline_parallel == 1, config.metrics, distributed_config.pipeline_parallel, ) + def _register_new_logprobs( + self, + new_logprobs_mean: torch.Tensor | None, + kwargs: dict[str, typing.Any], + losses: dict | None, + ) -> None: + if new_logprobs_mean is not None: + new_logprobs_mean = new_logprobs_mean / kwargs[LanguageModelKwargs.num_documents_in_batch] + self._register_loss( + self._logprob_metric_name, new_logprobs_mean, losses, reduce_op=torch.distributed.ReduceOp.SUM + ) + + def get_loss_definitions(self) -> list[LossDef]: + defs = super().get_loss_definitions() + defs.append(LossDef(self._logprob_metric_name)) + return defs + + def get_preprocessing_config(self) -> dict[str, typing.Any]: + return {"use_grpo_data": True, "return_label_counts": True, "return_document_count": True} + + @functools.cached_property + def _logprob_metric_name(self) -> str: + return f"{self._name}_new_logprobs" + + +class LanguageModelGRPOLoss[ConfigType: LanguageModelGRPOLossConfig](LanguageModelPolicyGradientLoss[ConfigType]): + """GRPO: per-token IS-ratio clipping.""" + def _forward_backward( self, logits: "torch.Tensor", @@ -137,7 +153,7 @@ def _forward_backward( self._register_new_logprobs(new_logprobs_mean, kwargs, losses) # Skip the extra softmax pass when there is nothing to register. - if losses is not None and self._config.metrics != GRPOMetricsLevel.none: + if losses is not None and self._config.metrics != PolicyLossMetrics.none: self._register_extra_metrics(logits, kwargs, losses, split_index) return loss, grad @@ -159,7 +175,7 @@ def _register_extra_metrics( self._config.epsilon_high, self._logits_scale_factor, group=self._parallel_dim.group if self._vocab_parallel else None, - compute_entropy=self._config.metrics == GRPOMetricsLevel.with_entropy, + compute_entropy=self._config.metrics == PolicyLossMetrics.with_entropy, ) num_documents = kwargs[LanguageModelKwargs.num_documents_in_batch] @@ -198,7 +214,7 @@ def _register_extra_metrics( def get_loss_definitions(self) -> list[LossDef]: defs = super().get_loss_definitions() - if self._config.metrics != GRPOMetricsLevel.none: + if self._config.metrics != PolicyLossMetrics.none: defs.extend( [ LossDef(f"{self._name}_old_logprobs"), @@ -213,7 +229,7 @@ def get_loss_definitions(self) -> list[LossDef]: LossDef(f"{self._name}_num_tokens"), ] ) - if self._config.metrics == GRPOMetricsLevel.with_entropy: + if self._config.metrics == PolicyLossMetrics.with_entropy: defs.append(LossDef(f"{self._name}_entropy")) return defs @@ -298,8 +314,79 @@ def _forward_backward( ) self._register_new_logprobs(new_logprobs_mean, kwargs, losses) + + # Skip the extra softmax pass when there is nothing to register. + if losses is not None and self._config.metrics != PolicyLossMetrics.none: + self._register_extra_metrics(logits, kwargs, losses, split_index, document_index_zero_based, num_segments) + return loss, grad + def _register_extra_metrics( + self, + logits: torch.Tensor, + kwargs: dict[str, typing.Any], + losses: dict | None, + split_index: int, + document_index_zero_based: torch.Tensor, + num_segments: int, + ) -> None: + metrics = compute_gspo_metrics( + logits, + self._get_labels(kwargs, split_index), + self._prepare_target(kwargs[LanguageModelLossKwargs.advantages], split_index), + self._prepare_target(kwargs[LanguageModelLossKwargs.old_log_probabilities], split_index), + document_index_zero_based, + num_segments, + self._prepare_target(kwargs[LanguageModelLossKwargs.label_counts], split_index), + self._config.epsilon_low, + self._config.epsilon_high, + self._logits_scale_factor, + group=self._parallel_dim.group if self._vocab_parallel else None, + sdp_group=self._sequence_data_dim.group if self._sequence_data_active else None, + sp_group=self._parallel_dim.group if self._sequence_parallel else None, + compute_entropy=self._config.metrics == PolicyLossMetrics.with_entropy, + ) + + num_documents = kwargs[LanguageModelKwargs.num_documents_in_batch] + + self._register_loss(f"{self._name}_old_logprobs", metrics.old_logprobs / num_documents, losses) + self._register_loss(f"{self._name}_ratio_new_old", metrics.ratio_sum / num_documents, losses) + self._register_loss(f"{self._name}_ratio_new_old_sum", metrics.ratio_sum, losses) + self._register_loss(f"{self._name}_ratio_new_old_squared_sum", metrics.ratio_squared_sum, losses) + self._register_loss(f"{self._name}_kl_new_old", metrics.kl_sum / num_documents, losses) + self._register_loss(f"{self._name}_clipped_ratio_fraction", metrics.clipped_sum / num_documents, losses) + self._register_loss(f"{self._name}_advantage", metrics.advantage_sum / num_documents, losses) + self._register_loss( + f"{self._name}_max_advantage", metrics.max_advantage, losses, reduce_op=torch.distributed.ReduceOp.MAX + ) + self._register_loss( + f"{self._name}_min_advantage", metrics.min_advantage, losses, reduce_op=torch.distributed.ReduceOp.MIN + ) + self._register_loss(f"{self._name}_num_segments", metrics.num_segments, losses) + if metrics.entropy is not None: + self._register_loss(f"{self._name}_entropy", metrics.entropy / num_documents, losses) + + def get_loss_definitions(self) -> list[LossDef]: + defs = super().get_loss_definitions() + if self._config.metrics != PolicyLossMetrics.none: + defs.extend( + [ + LossDef(f"{self._name}_old_logprobs"), + LossDef(f"{self._name}_ratio_new_old"), + LossDef(f"{self._name}_ratio_new_old_sum"), + LossDef(f"{self._name}_ratio_new_old_squared_sum"), + LossDef(f"{self._name}_kl_new_old"), + LossDef(f"{self._name}_clipped_ratio_fraction"), + LossDef(f"{self._name}_advantage"), + LossDef(f"{self._name}_max_advantage", reduction=ReductionType.maximum), + LossDef(f"{self._name}_min_advantage", reduction=ReductionType.minimum), + LossDef(f"{self._name}_num_segments"), + ] + ) + if self._config.metrics == PolicyLossMetrics.with_entropy: + defs.append(LossDef(f"{self._name}_entropy")) + return defs + def get_preprocessing_config(self) -> dict[str, typing.Any]: return super().get_preprocessing_config() | {"return_document_index": True} @@ -359,6 +446,88 @@ def compute_grpo_metrics( ) +# Not @torch.compile for the same reason as `fused_gspo_loss_forward_backward`: the Python-int +# `num_segments` argument trips dynamo. Metrics run only on logging steps, so eager is fine. +def compute_gspo_metrics( + logits: torch.Tensor, # (*batch, vocab_local) + target: torch.Tensor, # (*batch,) + advantages: torch.Tensor, # (*batch,) + old_log_probabilities: torch.Tensor, # (*batch,) + document_index_zero_based: torch.Tensor, # (*batch,) int — segment ID per token, 0-based + num_segments: int, + num_labels_in_seq: torch.Tensor, # (*batch,) — per-document labeled-token count broadcast per token + epsilon_low: float = 0.2, + epsilon_high: float = 0.2, + logits_scale_factor: float = 1.0, + group: torch.distributed.ProcessGroup | None = None, + sdp_group: torch.distributed.ProcessGroup | None = None, + sp_group: torch.distributed.ProcessGroup | None = None, + compute_entropy: bool = False, +) -> GSPOMetrics: + """Segment-level GSPO diagnostics (GSPO clips per document/segment). + + Reuses the loss's segment aggregation to build the per-segment geometric-mean ratio and + advantage. Per-segment sums are then accumulated as token-weighted sums with weight + `mask / num_labels_in_seq` — which sums to 1 per document across SDP/SP ranks — so they partition + correctly across ranks and reduce (SUM) exactly like the per-token loss, rather than summing the + SDP/SP-duplicated segment buffers. Advantage extrema use MAX/MIN reduction, which is idempotent + under that duplication. + """ + loss_mask = target >= 0 + + logits_norm, exp_logits, sum_exp_logits, _ = fused_softmax_base(logits, logits_scale_factor, group) + predicted_logits, _, _ = fused_predicted_logits_from_labels(logits_norm, target, loss_mask, group) + new_log_probs = predicted_logits - sum_exp_logits.log() + log_ratio = new_log_probs - old_log_probabilities + + flat_document_index = document_index_zero_based.reshape(-1).long() + flat_mask = loss_mask.reshape(-1).to(log_ratio.dtype) + mean_token_weight = flat_mask / num_labels_in_seq.reshape(-1).to(log_ratio.dtype).clamp(min=1) + + mean_log_ratio_per_segment = log_ratio.new_zeros(num_segments).index_add_( + 0, flat_document_index, log_ratio.reshape(-1) * mean_token_weight + ) + mean_advantage_per_segment = log_ratio.new_zeros(num_segments).index_add_( + 0, flat_document_index, advantages.reshape(-1).to(log_ratio.dtype) * mean_token_weight + ) + for reduce_group in (sdp_group, sp_group): + if reduce_group is not None: + all_reduce(mean_log_ratio_per_segment, op=ReduceOp.SUM, group=reduce_group) + all_reduce(mean_advantage_per_segment, op=ReduceOp.SUM, group=reduce_group) + + segment_ratio = mean_log_ratio_per_segment.exp() # geometric-mean IS ratio per segment + clipped_segment = (segment_ratio < 1.0 - epsilon_low) | (segment_ratio > 1.0 + epsilon_high) + kl_segment = segment_ratio - mean_log_ratio_per_segment - 1.0 # k3 estimator, per segment + + ratio_per_token = segment_ratio[flat_document_index] + + # Advantage extrema over the masked tokens (idempotent under SDP/SP duplication, and robust to + # empty segments). Advantage is constant within a document, so this equals the per-segment extrema. + neg_inf = advantages.new_full((), float("-inf")) + pos_inf = advantages.new_full((), float("inf")) + + entropy: torch.Tensor | None = None + if compute_entropy: + weighted_logits_sum = (exp_logits * logits_norm).sum(-1) + if group is not None: + all_reduce(weighted_logits_sum, op=ReduceOp.SUM, group=group) + entropy_per_token = sum_exp_logits.log() - weighted_logits_sum / sum_exp_logits + entropy = (entropy_per_token.reshape(-1) * mean_token_weight).sum() + + return GSPOMetrics( + old_logprobs=(old_log_probabilities.reshape(-1) * mean_token_weight).sum(), + ratio_sum=(ratio_per_token * mean_token_weight).sum(), + ratio_squared_sum=(ratio_per_token * ratio_per_token * mean_token_weight).sum(), + kl_sum=(kl_segment[flat_document_index] * mean_token_weight).sum(), + clipped_sum=(clipped_segment[flat_document_index].to(log_ratio.dtype) * mean_token_weight).sum(), + advantage_sum=(mean_advantage_per_segment[flat_document_index] * mean_token_weight).sum(), + max_advantage=torch.where(loss_mask, advantages, neg_inf).max(), + min_advantage=torch.where(loss_mask, advantages, pos_inf).min(), + num_segments=mean_token_weight.sum(), + entropy=entropy, + ) + + @torch.compile def fused_grpo_loss_forward_backward( logits: torch.Tensor, # (*batch, vocab) diff --git a/tests/layers/test_lm_losses.py b/tests/layers/test_lm_losses.py index d7d14ad3e..799a74502 100644 --- a/tests/layers/test_lm_losses.py +++ b/tests/layers/test_lm_losses.py @@ -20,7 +20,9 @@ from fast_llm.layers.language_model.loss.loss import loss_forward_backward from fast_llm.layers.language_model.loss.policy_gradient import ( GRPOMetrics, + GSPOMetrics, compute_grpo_metrics, + compute_gspo_metrics, fused_grpo_loss_forward_backward, fused_gspo_loss_forward_backward, ) @@ -169,6 +171,73 @@ def reference_grpo_metrics( ) +def reference_gspo_metrics( + logits: torch.Tensor, + target: torch.Tensor, + advantages: torch.Tensor, + old_log_probabilities: torch.Tensor, + document_index: torch.Tensor, + num_segments: int, + num_labels_in_seq: torch.Tensor, + epsilon_low: float, + epsilon_high: float, + logits_scale_factor: float, + compute_entropy: bool, +) -> GSPOMetrics: + log_softmax = torch.nn.functional.log_softmax(logits.float() * logits_scale_factor, dim=-1) + loss_mask = target >= 0 + new_log_probs = log_softmax.gather(-1, (target * loss_mask).unsqueeze(-1)).squeeze(-1) + log_ratio = new_log_probs - old_log_probabilities.float() + + flat_doc = document_index.reshape(-1).long() + flat_mask = loss_mask.reshape(-1) + flat_log_ratio = log_ratio.reshape(-1) + flat_advantages = advantages.reshape(-1).float() + flat_old = old_log_probabilities.reshape(-1).float() + + old_logprobs = log_ratio.new_zeros(()) + ratio_sum = log_ratio.new_zeros(()) + ratio_squared_sum = log_ratio.new_zeros(()) + kl_sum = log_ratio.new_zeros(()) + clipped_sum = log_ratio.new_zeros(()) + advantage_sum = log_ratio.new_zeros(()) + num_segments_count = log_ratio.new_zeros(()) + for segment in range(num_segments): + in_segment = (flat_doc == segment) & flat_mask + count = in_segment.sum() + if int(count) == 0: + continue + count_float = count.float() + mean_log_ratio = flat_log_ratio[in_segment].sum() / count_float + ratio = mean_log_ratio.exp() + ratio_sum = ratio_sum + ratio + ratio_squared_sum = ratio_squared_sum + ratio * ratio + kl_sum = kl_sum + (ratio - mean_log_ratio - 1.0) + clipped_sum = clipped_sum + ((ratio < 1 - epsilon_low) | (ratio > 1 + epsilon_high)).float() + advantage_sum = advantage_sum + flat_advantages[in_segment].sum() / count_float + old_logprobs = old_logprobs + flat_old[in_segment].sum() / count_float + num_segments_count = num_segments_count + 1.0 + + entropy = None + if compute_entropy: + entropy_per_token = -(log_softmax.exp() * log_softmax).sum(-1).reshape(-1) + masked = flat_mask.float() / num_labels_in_seq.reshape(-1).float().clamp(min=1) + entropy = (entropy_per_token * masked).sum() + + return GSPOMetrics( + old_logprobs=old_logprobs, + ratio_sum=ratio_sum, + ratio_squared_sum=ratio_squared_sum, + kl_sum=kl_sum, + clipped_sum=clipped_sum, + advantage_sum=advantage_sum, + max_advantage=advantages[loss_mask].max(), + min_advantage=advantages[loss_mask].min(), + num_segments=num_segments_count, + entropy=entropy, + ) + + def reference_gspo_loss( logits: torch.Tensor, labels: torch.Tensor, @@ -531,6 +600,64 @@ def _test_grpo_metrics( _check_grpo_metrics(ref, got, threshold=5e-5 if dtype == DataType.float32 else 1e-4) +def _check_gspo_metrics(ref: GSPOMetrics, got: GSPOMetrics, threshold: float) -> None: + for name in GSPOMetrics._fields: + ref_value = getattr(ref, name) + got_value = getattr(got, name) + if ref_value is None: + assert got_value is None, name + else: + Assert.rms_close_relative(got_value, ref_value, threshold) + + +def _test_gspo_metrics( + batch_shape, num_columns, logits_scale_factor, loss_masking, dtype, compute_entropy, num_segments, group=None +): + logits, target, advantages, old_log_probabilities = _get_grpo_loss_inputs( + num_columns, loss_masking, batch_shape, dtype + ) + # Per-token segment IDs + per-document labeled-token counts, mirroring `_test_gspo_loss`. + seq_len = batch_shape[-1] if len(batch_shape) > 1 else batch_shape[0] + span = max(seq_len // num_segments, 1) + base = torch.arange(seq_len, device=target.device) // span + document_index = base.clamp(max=num_segments - 1).expand(batch_shape).contiguous() + flat_doc = document_index.reshape(-1).long() + flat_target = target.reshape(-1) + labels_per_document = torch.zeros(num_segments, dtype=torch.int32, device=target.device).scatter_add( + 0, flat_doc, (flat_target >= 0).to(torch.int32) + ) + num_labels_in_seq = labels_per_document[flat_doc].reshape(target.shape) + + ref = reference_gspo_metrics( + logits, + target, + advantages, + old_log_probabilities, + document_index, + num_segments, + num_labels_in_seq, + epsilon_low=0.2, + epsilon_high=0.2, + logits_scale_factor=logits_scale_factor, + compute_entropy=compute_entropy, + ) + got = compute_gspo_metrics( + split_op(logits, group, -1).contiguous(), + target, + advantages, + old_log_probabilities, + document_index, + num_segments, + num_labels_in_seq, + epsilon_low=0.2, + epsilon_high=0.2, + logits_scale_factor=logits_scale_factor, + group=group, + compute_entropy=compute_entropy, + ) + _check_gspo_metrics(ref, got, threshold=5e-5 if dtype == DataType.float32 else 1e-4) + + def _test_z_loss( batch_shape, num_columns, grad_output, logits_scale_factor, loss_masking, dtype, block_size, accumulate, group=None ): @@ -698,6 +825,29 @@ def test_grpo_metrics( _test_grpo_metrics(batch_shape, num_columns, logits_scale_factor, loss_masking, dtype, compute_entropy) +@pytest.mark.slow +@pytest.mark.parametrize("batch_shape", _BATCH_SHAPES) +@pytest.mark.parametrize( + ("num_columns", "grad_output", "logits_scale_factor", "loss_masking", "dtype", "num_segments", "accumulate"), + _GSPO_PARAMETERS, +) +@pytest.mark.parametrize("compute_entropy", (False, True)) +def test_gspo_metrics( + batch_shape, + num_columns, + grad_output, + logits_scale_factor, + loss_masking, + dtype, + num_segments, + accumulate, + compute_entropy, +): + _test_gspo_metrics( + batch_shape, num_columns, logits_scale_factor, loss_masking, dtype, compute_entropy, num_segments + ) + + @pytest.mark.skip(reason="DPO loss is broken") def test_dpo_loss(): logits = torch.normal(0, 1, (200, 100)) From e4973fb89538f503fd094ef76401cb2f00988e3e Mon Sep 17 00:00:00 2001 From: Joel Lamy-Poirier Date: Wed, 8 Jul 2026 11:06:30 -0400 Subject: [PATCH 2/4] GSPO metrics: log num_tokens, dedupe loss defs, add distributed test - Log GSPO num_tokens (labeled-token count) so mean response length and ratio variance are derivable, matching GRPO. - Extract the shared GRPO/GSPO metric LossDef list into LanguageModelPolicyGradientLoss._policy_metric_definitions(*extra); GSPO passes its num_segments def. Order-preserving. - Add a GSPO-metrics subtest to _run_lm_loss_distributed (and to the test_lm_loss_distributed loss_type list) so compute_gspo_metrics' vocab-parallel group all-reduce path is exercised and reported. Co-Authored-By: Claude Opus 4.8 --- .../language_model/loss/policy_gradient.py | 63 ++++++++----------- tests/layers/test_lm_losses.py | 18 ++++++ 2 files changed, 43 insertions(+), 38 deletions(-) diff --git a/fast_llm/layers/language_model/loss/policy_gradient.py b/fast_llm/layers/language_model/loss/policy_gradient.py index ba004944d..f5602541f 100644 --- a/fast_llm/layers/language_model/loss/policy_gradient.py +++ b/fast_llm/layers/language_model/loss/policy_gradient.py @@ -48,6 +48,7 @@ class GSPOMetrics(typing.NamedTuple): max_advantage: torch.Tensor min_advantage: torch.Tensor num_segments: torch.Tensor + num_tokens: torch.Tensor entropy: torch.Tensor | None @@ -101,6 +102,26 @@ def _register_new_logprobs( self._logprob_metric_name, new_logprobs_mean, losses, reduce_op=torch.distributed.ReduceOp.SUM ) + def _policy_metric_definitions(self, *extra: LossDef) -> list[LossDef]: + if self._config.metrics == PolicyLossMetrics.none: + return [] + defs = [ + LossDef(f"{self._name}_old_logprobs"), + LossDef(f"{self._name}_ratio_new_old"), + LossDef(f"{self._name}_ratio_new_old_sum"), + LossDef(f"{self._name}_ratio_new_old_squared_sum"), + LossDef(f"{self._name}_kl_new_old"), + LossDef(f"{self._name}_clipped_ratio_fraction"), + LossDef(f"{self._name}_advantage"), + LossDef(f"{self._name}_max_advantage", reduction=ReductionType.maximum), + LossDef(f"{self._name}_min_advantage", reduction=ReductionType.minimum), + *extra, + LossDef(f"{self._name}_num_tokens"), + ] + if self._config.metrics == PolicyLossMetrics.with_entropy: + defs.append(LossDef(f"{self._name}_entropy")) + return defs + def get_loss_definitions(self) -> list[LossDef]: defs = super().get_loss_definitions() defs.append(LossDef(self._logprob_metric_name)) @@ -213,25 +234,7 @@ def _register_extra_metrics( self._register_loss(f"{self._name}_entropy", metrics.entropy / num_documents, losses) def get_loss_definitions(self) -> list[LossDef]: - defs = super().get_loss_definitions() - if self._config.metrics != PolicyLossMetrics.none: - defs.extend( - [ - LossDef(f"{self._name}_old_logprobs"), - LossDef(f"{self._name}_ratio_new_old"), - LossDef(f"{self._name}_ratio_new_old_sum"), - LossDef(f"{self._name}_ratio_new_old_squared_sum"), - LossDef(f"{self._name}_kl_new_old"), - LossDef(f"{self._name}_clipped_ratio_fraction"), - LossDef(f"{self._name}_advantage"), - LossDef(f"{self._name}_max_advantage", reduction=ReductionType.maximum), - LossDef(f"{self._name}_min_advantage", reduction=ReductionType.minimum), - LossDef(f"{self._name}_num_tokens"), - ] - ) - if self._config.metrics == PolicyLossMetrics.with_entropy: - defs.append(LossDef(f"{self._name}_entropy")) - return defs + return super().get_loss_definitions() + self._policy_metric_definitions() class LanguageModelGSPOLoss[ConfigType: LanguageModelGSPOLossConfig](LanguageModelPolicyGradientLoss[ConfigType]): @@ -363,29 +366,12 @@ def _register_extra_metrics( f"{self._name}_min_advantage", metrics.min_advantage, losses, reduce_op=torch.distributed.ReduceOp.MIN ) self._register_loss(f"{self._name}_num_segments", metrics.num_segments, losses) + self._register_loss(f"{self._name}_num_tokens", metrics.num_tokens, losses) if metrics.entropy is not None: self._register_loss(f"{self._name}_entropy", metrics.entropy / num_documents, losses) def get_loss_definitions(self) -> list[LossDef]: - defs = super().get_loss_definitions() - if self._config.metrics != PolicyLossMetrics.none: - defs.extend( - [ - LossDef(f"{self._name}_old_logprobs"), - LossDef(f"{self._name}_ratio_new_old"), - LossDef(f"{self._name}_ratio_new_old_sum"), - LossDef(f"{self._name}_ratio_new_old_squared_sum"), - LossDef(f"{self._name}_kl_new_old"), - LossDef(f"{self._name}_clipped_ratio_fraction"), - LossDef(f"{self._name}_advantage"), - LossDef(f"{self._name}_max_advantage", reduction=ReductionType.maximum), - LossDef(f"{self._name}_min_advantage", reduction=ReductionType.minimum), - LossDef(f"{self._name}_num_segments"), - ] - ) - if self._config.metrics == PolicyLossMetrics.with_entropy: - defs.append(LossDef(f"{self._name}_entropy")) - return defs + return super().get_loss_definitions() + self._policy_metric_definitions(LossDef(f"{self._name}_num_segments")) def get_preprocessing_config(self) -> dict[str, typing.Any]: return super().get_preprocessing_config() | {"return_document_index": True} @@ -524,6 +510,7 @@ def compute_gspo_metrics( max_advantage=torch.where(loss_mask, advantages, neg_inf).max(), min_advantage=torch.where(loss_mask, advantages, pos_inf).min(), num_segments=mean_token_weight.sum(), + num_tokens=flat_mask.sum(), entropy=entropy, ) diff --git a/tests/layers/test_lm_losses.py b/tests/layers/test_lm_losses.py index 799a74502..6bd630de3 100644 --- a/tests/layers/test_lm_losses.py +++ b/tests/layers/test_lm_losses.py @@ -234,6 +234,7 @@ def reference_gspo_metrics( max_advantage=advantages[loss_mask].max(), min_advantage=advantages[loss_mask].min(), num_segments=num_segments_count, + num_tokens=loss_mask.sum(), entropy=entropy, ) @@ -939,6 +940,21 @@ def _run_lm_loss_distributed(test_context: DistributedTestContext, base_path: pa compute_entropy, test_context.group, ) + # GSPO metrics + for compute_entropy in (False, True): + with test_context.subtest(base_path, f"gspo_metrics-{compute_entropy}-{suffix}", 2) as subtest: + if subtest.do_run: + torch.manual_seed((seed + hash(subtest.name)) % 2**32) + _test_gspo_metrics( + batch_shape, + num_columns, + logits_scale_factor, + loss_masking, + dtype, + compute_entropy, + 4, + test_context.group, + ) @pytest.mark.slow @@ -981,6 +997,8 @@ def test_run_lm_loss_distributed(run_parallel_script, result_path): "grpo", "grpo_metrics-False", "grpo_metrics-True", + "gspo_metrics-False", + "gspo_metrics-True", ), ) def test_lm_loss_distributed( From fed92da1f157fa5d41acde287213994cedcc9a8e Mon Sep 17 00:00:00 2001 From: Joel Lamy-Poirier Date: Wed, 8 Jul 2026 11:58:54 -0400 Subject: [PATCH 3/4] Dedupe GRPO/GSPO metrics into a shared core The two metric paths differed only in how the importance ratio is formed (per-token for GRPO, per-segment geometric mean for GSPO). Collapse the duplication: - One `PolicyMetrics` tuple replaces `GRPOMetrics`/`GSPOMetrics`. - `_policy_metrics_log_ratio` holds the shared softmax -> log-ratio -> entropy pass; `_policy_metrics_reduce` holds the shared weighted-sum reduction, parameterized by the ratio-variance weight (token mask vs document weight). - One `_register_policy_metrics` replaces the two near-identical registrations. `compute_grpo_metrics` / `compute_gspo_metrics` now only derive the per-token `(ratio, effective_log_ratio)` -- their actual difference. GSPO's token-level `exp`-after-gather is bit-identical to the previous per-segment-then-broadcast. Behavior-preserving: the independent test references (repackaged into `PolicyMetrics`, math unchanged) still pass. Co-Authored-By: Claude Opus 4.8 --- .../language_model/loss/policy_gradient.py | 301 ++++++++---------- tests/layers/test_lm_losses.py | 43 +-- 2 files changed, 156 insertions(+), 188 deletions(-) diff --git a/fast_llm/layers/language_model/loss/policy_gradient.py b/fast_llm/layers/language_model/loss/policy_gradient.py index f5602541f..ef7ba8b76 100644 --- a/fast_llm/layers/language_model/loss/policy_gradient.py +++ b/fast_llm/layers/language_model/loss/policy_gradient.py @@ -22,7 +22,13 @@ from fast_llm.utils import Assert -class GRPOMetrics(typing.NamedTuple): +class PolicyMetrics(typing.NamedTuple): + # Weighted sums for the policy-gradient diagnostics, shared by the per-token (GRPO) and per-segment + # (GSPO) paths. `ratio_new_old`, `kl_new_old`, `clipped_ratio_fraction`, `advantage` and `entropy` + # are document-weighted and divided by the document count at registration to form means / fractions; + # `ratio_new_old_sum` / `_squared_sum` stay raw (variance is derived downstream, over tokens for the + # per-token path and over segments for the per-segment one). `num_segments` is populated for the + # per-segment path only. old_logprobs: torch.Tensor ratio_new_old: torch.Tensor ratio_new_old_sum: torch.Tensor @@ -33,22 +39,7 @@ class GRPOMetrics(typing.NamedTuple): max_advantage: torch.Tensor min_advantage: torch.Tensor num_tokens: torch.Tensor - entropy: torch.Tensor | None - - -class GSPOMetrics(typing.NamedTuple): - # Statistics over segments (documents); GSPO clips per segment. `*_sum` fields are raw sums over - # segments (divided by the document count at registration time to form means / fractions). - old_logprobs: torch.Tensor - ratio_sum: torch.Tensor - ratio_squared_sum: torch.Tensor - kl_sum: torch.Tensor - clipped_sum: torch.Tensor - advantage_sum: torch.Tensor - max_advantage: torch.Tensor - min_advantage: torch.Tensor - num_segments: torch.Tensor - num_tokens: torch.Tensor + num_segments: torch.Tensor | None entropy: torch.Tensor | None @@ -122,6 +113,23 @@ def _policy_metric_definitions(self, *extra: LossDef) -> list[LossDef]: defs.append(LossDef(f"{self._name}_entropy")) return defs + def _register_policy_metrics(self, metrics: PolicyMetrics, kwargs: dict[str, typing.Any], losses: dict) -> None: + num_documents = kwargs[LanguageModelKwargs.num_documents_in_batch] + for name in ("old_logprobs", "ratio_new_old", "kl_new_old", "clipped_ratio_fraction", "advantage"): + self._register_loss(f"{self._name}_{name}", getattr(metrics, name) / num_documents, losses) + for name in ("ratio_new_old_sum", "ratio_new_old_squared_sum", "num_tokens"): + self._register_loss(f"{self._name}_{name}", getattr(metrics, name), losses) + self._register_loss( + f"{self._name}_max_advantage", metrics.max_advantage, losses, reduce_op=torch.distributed.ReduceOp.MAX + ) + self._register_loss( + f"{self._name}_min_advantage", metrics.min_advantage, losses, reduce_op=torch.distributed.ReduceOp.MIN + ) + if metrics.num_segments is not None: + self._register_loss(f"{self._name}_num_segments", metrics.num_segments, losses) + if metrics.entropy is not None: + self._register_loss(f"{self._name}_entropy", metrics.entropy / num_documents, losses) + def get_loss_definitions(self) -> list[LossDef]: defs = super().get_loss_definitions() defs.append(LossDef(self._logprob_metric_name)) @@ -198,40 +206,7 @@ def _register_extra_metrics( group=self._parallel_dim.group if self._vocab_parallel else None, compute_entropy=self._config.metrics == PolicyLossMetrics.with_entropy, ) - - num_documents = kwargs[LanguageModelKwargs.num_documents_in_batch] - - for attr in ( - "old_logprobs", - "ratio_new_old", - "kl_new_old", - "clipped_ratio_fraction", - "advantage", - ): - self._register_loss(f"{self._name}_{attr}", getattr(metrics, attr) / num_documents, losses) - - for attr in ( - "ratio_new_old_sum", - "ratio_new_old_squared_sum", - "num_tokens", - ): - self._register_loss(f"{self._name}_{attr}", getattr(metrics, attr), losses) - - self._register_loss( - f"{self._name}_max_advantage", - metrics.max_advantage, - losses, - reduce_op=torch.distributed.ReduceOp.MAX, - ) - self._register_loss( - f"{self._name}_min_advantage", - metrics.min_advantage, - losses, - reduce_op=torch.distributed.ReduceOp.MIN, - ) - - if metrics.entropy is not None: - self._register_loss(f"{self._name}_entropy", metrics.entropy / num_documents, losses) + self._register_policy_metrics(metrics, kwargs, losses) def get_loss_definitions(self) -> list[LossDef]: return super().get_loss_definitions() + self._policy_metric_definitions() @@ -349,26 +324,7 @@ def _register_extra_metrics( sp_group=self._parallel_dim.group if self._sequence_parallel else None, compute_entropy=self._config.metrics == PolicyLossMetrics.with_entropy, ) - - num_documents = kwargs[LanguageModelKwargs.num_documents_in_batch] - - self._register_loss(f"{self._name}_old_logprobs", metrics.old_logprobs / num_documents, losses) - self._register_loss(f"{self._name}_ratio_new_old", metrics.ratio_sum / num_documents, losses) - self._register_loss(f"{self._name}_ratio_new_old_sum", metrics.ratio_sum, losses) - self._register_loss(f"{self._name}_ratio_new_old_squared_sum", metrics.ratio_squared_sum, losses) - self._register_loss(f"{self._name}_kl_new_old", metrics.kl_sum / num_documents, losses) - self._register_loss(f"{self._name}_clipped_ratio_fraction", metrics.clipped_sum / num_documents, losses) - self._register_loss(f"{self._name}_advantage", metrics.advantage_sum / num_documents, losses) - self._register_loss( - f"{self._name}_max_advantage", metrics.max_advantage, losses, reduce_op=torch.distributed.ReduceOp.MAX - ) - self._register_loss( - f"{self._name}_min_advantage", metrics.min_advantage, losses, reduce_op=torch.distributed.ReduceOp.MIN - ) - self._register_loss(f"{self._name}_num_segments", metrics.num_segments, losses) - self._register_loss(f"{self._name}_num_tokens", metrics.num_tokens, losses) - if metrics.entropy is not None: - self._register_loss(f"{self._name}_entropy", metrics.entropy / num_documents, losses) + self._register_policy_metrics(metrics, kwargs, losses) def get_loss_definitions(self) -> list[LossDef]: return super().get_loss_definitions() + self._policy_metric_definitions(LossDef(f"{self._name}_num_segments")) @@ -377,6 +333,68 @@ def get_preprocessing_config(self) -> dict[str, typing.Any]: return super().get_preprocessing_config() | {"return_document_index": True} +def _policy_metrics_log_ratio( + logits: torch.Tensor, + target: torch.Tensor, + old_log_probabilities: torch.Tensor, + logits_scale_factor: float, + group: torch.distributed.ProcessGroup | None, + compute_entropy: bool, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]: + """Shared softmax pass for both policy-gradient metric paths. Returns the loss mask, the per-token + new/old log-ratio, and (when requested) the per-token policy entropy. `exp_logits` / `logits_norm` + are local vocab slices, so the entropy's weighted-logit sum is all-reduced across the tensor-parallel + group to recover the global expectation before dividing by the already-global `sum_exp_logits`.""" + loss_mask = target >= 0 + logits_norm, exp_logits, sum_exp_logits, _ = fused_softmax_base(logits, logits_scale_factor, group) + predicted_logits, _, _ = fused_predicted_logits_from_labels(logits_norm, target, loss_mask, group) + log_ratio = predicted_logits - sum_exp_logits.log() - old_log_probabilities + + entropy_per_token: torch.Tensor | None = None + if compute_entropy: + weighted_logits_sum = (exp_logits * logits_norm).sum(-1) + if group is not None: + all_reduce(weighted_logits_sum, op=ReduceOp.SUM, group=group) + entropy_per_token = sum_exp_logits.log() - weighted_logits_sum / sum_exp_logits + return loss_mask, log_ratio, entropy_per_token + + +def _policy_metrics_reduce( + ratio: torch.Tensor, # per token + effective_log_ratio: torch.Tensor, # per token; the log-ratio whose exp gives `ratio` + advantage_for_sum: torch.Tensor, # per token; advantage entering the document-weighted mean + advantages: torch.Tensor, # per token; raw advantage for the extrema + old_log_probabilities: torch.Tensor, # per token + loss_mask: torch.Tensor, # per token + document_weight: torch.Tensor, # per token: mask / labeled-token count (sums to 1 per document) + variance_weight: torch.Tensor, # per token: token mask (per-token path) or `document_weight` (per-segment) + epsilon_low: float, + epsilon_high: float, + entropy_per_token: torch.Tensor | None, + num_segments: torch.Tensor | None, +) -> PolicyMetrics: + """Shared metric reduction. Document-weighted sums divided by the document count give per-document + means; `variance_weight` selects the ratio-variance granularity (token vs segment).""" + kl = ratio - effective_log_ratio - 1.0 + clipped = ((ratio < 1.0 - epsilon_low) | (ratio > 1.0 + epsilon_high)).to(document_weight.dtype) + neg_inf = advantages.new_full((), float("-inf")) + pos_inf = advantages.new_full((), float("inf")) + return PolicyMetrics( + old_logprobs=(old_log_probabilities * document_weight).sum(), + ratio_new_old=(ratio * document_weight).sum(), + ratio_new_old_sum=(ratio * variance_weight).sum(), + ratio_new_old_squared_sum=(ratio * ratio * variance_weight).sum(), + kl_new_old=(kl * document_weight).sum(), + clipped_ratio_fraction=(clipped * document_weight).sum(), + advantage=(advantage_for_sum * document_weight).sum(), + max_advantage=torch.where(loss_mask, advantages, neg_inf).max(), + min_advantage=torch.where(loss_mask, advantages, pos_inf).min(), + num_tokens=loss_mask.to(document_weight.dtype).sum(), + num_segments=num_segments, + entropy=None if entropy_per_token is None else (entropy_per_token * document_weight).sum(), + ) + + @torch.compile def compute_grpo_metrics( logits: torch.Tensor, # (*batch, vocab_local) @@ -389,46 +407,26 @@ def compute_grpo_metrics( logits_scale_factor: float = 1.0, group: torch.distributed.ProcessGroup | None = None, compute_entropy: bool = False, -) -> GRPOMetrics: - loss_mask = target >= 0 - mask = loss_mask.float() - masked = mask / label_counts.float().clamp(min=1) - - logits_norm, exp_logits, sum_exp_logits, _ = fused_softmax_base(logits, logits_scale_factor, group) - predicted_logits, _, _ = fused_predicted_logits_from_labels(logits_norm, target, loss_mask, group) - new_log_probs = predicted_logits - sum_exp_logits.log() - - log_ratio = new_log_probs - old_log_probabilities - ratio = log_ratio.exp() - clipped = (ratio < 1.0 - epsilon_low) | (ratio > 1.0 + epsilon_high) - kl = ratio - log_ratio - 1.0 - - neg_inf = advantages.new_full((), float("-inf")) - pos_inf = advantages.new_full((), float("inf")) - - entropy: torch.Tensor | None = None - if compute_entropy: - # exp_logits and logits_norm are local vocab slices — sum over the local slice, then all-reduce - # across the tensor-parallel group to recover the global E_p[logit_norm] before dividing by the - # already-global sum_exp_logits. - weighted_logits_sum = (exp_logits * logits_norm).sum(-1) - if group is not None: - all_reduce(weighted_logits_sum, op=ReduceOp.SUM, group=group) - entropy_per_token = sum_exp_logits.log() - weighted_logits_sum / sum_exp_logits - entropy = (entropy_per_token * masked).sum() - - return GRPOMetrics( - old_logprobs=(old_log_probabilities * masked).sum(), - ratio_new_old=(ratio * masked).sum(), - ratio_new_old_sum=(ratio * mask).sum(), - ratio_new_old_squared_sum=(ratio * ratio * mask).sum(), - kl_new_old=(kl * masked).sum(), - clipped_ratio_fraction=(clipped.float() * masked).sum(), - advantage=(advantages * masked).sum(), - max_advantage=torch.where(loss_mask, advantages, neg_inf).max(), - min_advantage=torch.where(loss_mask, advantages, pos_inf).min(), - num_tokens=mask.sum(), - entropy=entropy, +) -> PolicyMetrics: + """Per-token diagnostics: the importance ratio and its clip / KL are token-level, and the ratio + variance is over tokens (`variance_weight` = the token mask).""" + loss_mask, log_ratio, entropy_per_token = _policy_metrics_log_ratio( + logits, target, old_log_probabilities, logits_scale_factor, group, compute_entropy + ) + mask = loss_mask.to(log_ratio.dtype) + return _policy_metrics_reduce( + ratio=log_ratio.exp(), + effective_log_ratio=log_ratio, + advantage_for_sum=advantages, + advantages=advantages, + old_log_probabilities=old_log_probabilities, + loss_mask=loss_mask, + document_weight=mask / label_counts.to(log_ratio.dtype).clamp(min=1), + variance_weight=mask, + epsilon_low=epsilon_low, + epsilon_high=epsilon_high, + entropy_per_token=entropy_per_token, + num_segments=None, ) @@ -449,69 +447,48 @@ def compute_gspo_metrics( sdp_group: torch.distributed.ProcessGroup | None = None, sp_group: torch.distributed.ProcessGroup | None = None, compute_entropy: bool = False, -) -> GSPOMetrics: - """Segment-level GSPO diagnostics (GSPO clips per document/segment). - - Reuses the loss's segment aggregation to build the per-segment geometric-mean ratio and - advantage. Per-segment sums are then accumulated as token-weighted sums with weight - `mask / num_labels_in_seq` — which sums to 1 per document across SDP/SP ranks — so they partition - correctly across ranks and reduce (SUM) exactly like the per-token loss, rather than summing the - SDP/SP-duplicated segment buffers. Advantage extrema use MAX/MIN reduction, which is idempotent - under that duplication. - """ - loss_mask = target >= 0 - - logits_norm, exp_logits, sum_exp_logits, _ = fused_softmax_base(logits, logits_scale_factor, group) - predicted_logits, _, _ = fused_predicted_logits_from_labels(logits_norm, target, loss_mask, group) - new_log_probs = predicted_logits - sum_exp_logits.log() - log_ratio = new_log_probs - old_log_probabilities - +) -> PolicyMetrics: + """Segment-level diagnostics (clipping is per document/segment): the ratio is the per-segment + geometric mean, broadcast back to tokens, and the ratio variance is over segments + (`variance_weight` = `document_weight`). The per-segment log-ratio / advantage are token-weighted + by `document_weight` (which sums to 1 per document across SDP/SP ranks) then all-reduced, so they + partition correctly across ranks and the token-level reduction matches the per-token loss.""" + loss_mask, log_ratio, entropy_per_token = _policy_metrics_log_ratio( + logits, target, old_log_probabilities, logits_scale_factor, group, compute_entropy + ) flat_document_index = document_index_zero_based.reshape(-1).long() flat_mask = loss_mask.reshape(-1).to(log_ratio.dtype) - mean_token_weight = flat_mask / num_labels_in_seq.reshape(-1).to(log_ratio.dtype).clamp(min=1) + document_weight = flat_mask / num_labels_in_seq.reshape(-1).to(log_ratio.dtype).clamp(min=1) mean_log_ratio_per_segment = log_ratio.new_zeros(num_segments).index_add_( - 0, flat_document_index, log_ratio.reshape(-1) * mean_token_weight + 0, flat_document_index, log_ratio.reshape(-1) * document_weight ) mean_advantage_per_segment = log_ratio.new_zeros(num_segments).index_add_( - 0, flat_document_index, advantages.reshape(-1).to(log_ratio.dtype) * mean_token_weight + 0, flat_document_index, advantages.reshape(-1).to(log_ratio.dtype) * document_weight ) for reduce_group in (sdp_group, sp_group): if reduce_group is not None: all_reduce(mean_log_ratio_per_segment, op=ReduceOp.SUM, group=reduce_group) all_reduce(mean_advantage_per_segment, op=ReduceOp.SUM, group=reduce_group) - segment_ratio = mean_log_ratio_per_segment.exp() # geometric-mean IS ratio per segment - clipped_segment = (segment_ratio < 1.0 - epsilon_low) | (segment_ratio > 1.0 + epsilon_high) - kl_segment = segment_ratio - mean_log_ratio_per_segment - 1.0 # k3 estimator, per segment - - ratio_per_token = segment_ratio[flat_document_index] - - # Advantage extrema over the masked tokens (idempotent under SDP/SP duplication, and robust to - # empty segments). Advantage is constant within a document, so this equals the per-segment extrema. - neg_inf = advantages.new_full((), float("-inf")) - pos_inf = advantages.new_full((), float("inf")) - - entropy: torch.Tensor | None = None - if compute_entropy: - weighted_logits_sum = (exp_logits * logits_norm).sum(-1) - if group is not None: - all_reduce(weighted_logits_sum, op=ReduceOp.SUM, group=group) - entropy_per_token = sum_exp_logits.log() - weighted_logits_sum / sum_exp_logits - entropy = (entropy_per_token.reshape(-1) * mean_token_weight).sum() - - return GSPOMetrics( - old_logprobs=(old_log_probabilities.reshape(-1) * mean_token_weight).sum(), - ratio_sum=(ratio_per_token * mean_token_weight).sum(), - ratio_squared_sum=(ratio_per_token * ratio_per_token * mean_token_weight).sum(), - kl_sum=(kl_segment[flat_document_index] * mean_token_weight).sum(), - clipped_sum=(clipped_segment[flat_document_index].to(log_ratio.dtype) * mean_token_weight).sum(), - advantage_sum=(mean_advantage_per_segment[flat_document_index] * mean_token_weight).sum(), - max_advantage=torch.where(loss_mask, advantages, neg_inf).max(), - min_advantage=torch.where(loss_mask, advantages, pos_inf).min(), - num_segments=mean_token_weight.sum(), - num_tokens=flat_mask.sum(), - entropy=entropy, + # Broadcast the per-segment geometric-mean log-ratio back to tokens; taking `exp` after the gather is + # identical to gathering the per-segment ratio, and likewise carries through the clip / KL. Advantage + # is constant within a document, so the per-segment mean advantage matches the raw advantage on masked + # tokens, but the reduced buffer is used for the sum so it stays correct across SDP/SP ranks. + effective_log_ratio = mean_log_ratio_per_segment[flat_document_index] + return _policy_metrics_reduce( + ratio=effective_log_ratio.exp(), + effective_log_ratio=effective_log_ratio, + advantage_for_sum=mean_advantage_per_segment[flat_document_index], + advantages=advantages.reshape(-1), + old_log_probabilities=old_log_probabilities.reshape(-1), + loss_mask=loss_mask.reshape(-1), + document_weight=document_weight, + variance_weight=document_weight, + epsilon_low=epsilon_low, + epsilon_high=epsilon_high, + entropy_per_token=None if entropy_per_token is None else entropy_per_token.reshape(-1), + num_segments=document_weight.sum(), ) diff --git a/tests/layers/test_lm_losses.py b/tests/layers/test_lm_losses.py index 6bd630de3..14e144ad3 100644 --- a/tests/layers/test_lm_losses.py +++ b/tests/layers/test_lm_losses.py @@ -19,8 +19,7 @@ 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 ( - GRPOMetrics, - GSPOMetrics, + PolicyMetrics, compute_grpo_metrics, compute_gspo_metrics, fused_grpo_loss_forward_backward, @@ -139,7 +138,7 @@ def reference_grpo_metrics( epsilon_high: float, logits_scale_factor: float, compute_entropy: bool, -) -> GRPOMetrics: +) -> PolicyMetrics: log_softmax = torch.nn.functional.log_softmax(logits.float() * logits_scale_factor, dim=-1) loss_mask = target >= 0 mask = loss_mask.float() @@ -156,7 +155,7 @@ def reference_grpo_metrics( entropy_per_token = -(log_softmax.exp() * log_softmax).sum(-1) entropy = (entropy_per_token * masked).sum() - return GRPOMetrics( + return PolicyMetrics( old_logprobs=(old_log_probabilities.float() * masked).sum(), ratio_new_old=(ratio * masked).sum(), ratio_new_old_sum=(ratio * mask).sum(), @@ -167,6 +166,7 @@ def reference_grpo_metrics( max_advantage=advantages[loss_mask].max(), min_advantage=advantages[loss_mask].min(), num_tokens=mask.sum(), + num_segments=None, entropy=entropy, ) @@ -183,7 +183,7 @@ def reference_gspo_metrics( epsilon_high: float, logits_scale_factor: float, compute_entropy: bool, -) -> GSPOMetrics: +) -> PolicyMetrics: log_softmax = torch.nn.functional.log_softmax(logits.float() * logits_scale_factor, dim=-1) loss_mask = target >= 0 new_log_probs = log_softmax.gather(-1, (target * loss_mask).unsqueeze(-1)).squeeze(-1) @@ -224,17 +224,18 @@ def reference_gspo_metrics( masked = flat_mask.float() / num_labels_in_seq.reshape(-1).float().clamp(min=1) entropy = (entropy_per_token * masked).sum() - return GSPOMetrics( + return PolicyMetrics( old_logprobs=old_logprobs, - ratio_sum=ratio_sum, - ratio_squared_sum=ratio_squared_sum, - kl_sum=kl_sum, - clipped_sum=clipped_sum, - advantage_sum=advantage_sum, + ratio_new_old=ratio_sum, + ratio_new_old_sum=ratio_sum, + ratio_new_old_squared_sum=ratio_squared_sum, + kl_new_old=kl_sum, + clipped_ratio_fraction=clipped_sum, + advantage=advantage_sum, max_advantage=advantages[loss_mask].max(), min_advantage=advantages[loss_mask].min(), - num_segments=num_segments_count, num_tokens=loss_mask.sum(), + num_segments=num_segments_count, entropy=entropy, ) @@ -554,8 +555,8 @@ def _test_gspo_loss( Assert.rms_close_relative(new_logprobs_triton, new_logprobs_fused, 1e-5, 1e-6) -def _check_grpo_metrics(ref: GRPOMetrics, got: GRPOMetrics, threshold: float) -> None: - for name in GRPOMetrics._fields: +def _check_policy_metrics(ref: PolicyMetrics, got: PolicyMetrics, threshold: float) -> None: + for name in PolicyMetrics._fields: ref_value = getattr(ref, name) got_value = getattr(got, name) if ref_value is None: @@ -598,17 +599,7 @@ def _test_grpo_metrics( group=group, compute_entropy=compute_entropy, ) - _check_grpo_metrics(ref, got, threshold=5e-5 if dtype == DataType.float32 else 1e-4) - - -def _check_gspo_metrics(ref: GSPOMetrics, got: GSPOMetrics, threshold: float) -> None: - for name in GSPOMetrics._fields: - ref_value = getattr(ref, name) - got_value = getattr(got, name) - if ref_value is None: - assert got_value is None, name - else: - Assert.rms_close_relative(got_value, ref_value, threshold) + _check_policy_metrics(ref, got, threshold=5e-5 if dtype == DataType.float32 else 1e-4) def _test_gspo_metrics( @@ -656,7 +647,7 @@ def _test_gspo_metrics( group=group, compute_entropy=compute_entropy, ) - _check_gspo_metrics(ref, got, threshold=5e-5 if dtype == DataType.float32 else 1e-4) + _check_policy_metrics(ref, got, threshold=5e-5 if dtype == DataType.float32 else 1e-4) def _test_z_loss( From 1c2005fae0ea101e84c67c72bfd9ffede5c3c1d1 Mon Sep 17 00:00:00 2001 From: Joel Lamy-Poirier Date: Wed, 8 Jul 2026 12:55:10 -0400 Subject: [PATCH 4/4] Fine-review fixes: naming and annotations - Rename enum `PolicyLossMetrics` -> `PolicyMetricsLevel` to disambiguate from the `PolicyMetrics` value tuple (restores the `Level` suffix). - `compute_gspo_metrics` label-count parameter `num_labels_in_seq` -> `label_counts` to match the sibling `compute_grpo_metrics` and the kwarg source. - Tighten `losses: dict | None` -> `dict` in both `_register_extra_metrics` (only reached under `losses is not None`). - Name the `num_segments = 4` local in the distributed GSPO-metrics subtest. Co-Authored-By: Claude Opus 4.8 --- fast_llm/layers/language_model/loss/config.py | 6 ++--- .../language_model/loss/policy_gradient.py | 24 +++++++++---------- tests/layers/test_lm_losses.py | 3 ++- 3 files changed, 17 insertions(+), 16 deletions(-) diff --git a/fast_llm/layers/language_model/loss/config.py b/fast_llm/layers/language_model/loss/config.py index 13fc24b0a..4f3cfae8d 100644 --- a/fast_llm/layers/language_model/loss/config.py +++ b/fast_llm/layers/language_model/loss/config.py @@ -205,7 +205,7 @@ def loss_class(self) -> "type[LanguageModelZLoss]": return LanguageModelZLoss -class PolicyLossMetrics(enum.StrEnum): +class PolicyMetricsLevel(enum.StrEnum): none = "none" basic = "basic" with_entropy = "with_entropy" @@ -219,8 +219,8 @@ class LanguageModelPolicyGradientLossConfig(LanguageModelLossConfig): epsilon_low: float = Field(default=0.2, desc="Lower clip parameter for ratio of log probs") epsilon_high: float = Field(default=0.2, desc="Upper clip parameter for ratio of log probs") - metrics: PolicyLossMetrics = Field( - default=PolicyLossMetrics.none, + metrics: PolicyMetricsLevel = Field( + default=PolicyMetricsLevel.none, desc=( "Additional diagnostic metrics to log. " "`basic`: importance-ratio, KL and advantage statistics. " diff --git a/fast_llm/layers/language_model/loss/policy_gradient.py b/fast_llm/layers/language_model/loss/policy_gradient.py index ef7ba8b76..c07c81e05 100644 --- a/fast_llm/layers/language_model/loss/policy_gradient.py +++ b/fast_llm/layers/language_model/loss/policy_gradient.py @@ -16,7 +16,7 @@ LanguageModelGSPOLossConfig, LanguageModelLossKwargs, LanguageModelPolicyGradientLossConfig, - PolicyLossMetrics, + PolicyMetricsLevel, ) from fast_llm.layers.language_model.loss.loss import LanguageModelLoss from fast_llm.utils import Assert @@ -76,7 +76,7 @@ def __init__( ) # The extra metrics need a second softmax over the full logits, which pipeline parallelism splits. Assert.custom( - lambda metrics, pipeline_parallel: metrics == PolicyLossMetrics.none or pipeline_parallel == 1, + lambda metrics, pipeline_parallel: metrics == PolicyMetricsLevel.none or pipeline_parallel == 1, config.metrics, distributed_config.pipeline_parallel, ) @@ -94,7 +94,7 @@ def _register_new_logprobs( ) def _policy_metric_definitions(self, *extra: LossDef) -> list[LossDef]: - if self._config.metrics == PolicyLossMetrics.none: + if self._config.metrics == PolicyMetricsLevel.none: return [] defs = [ LossDef(f"{self._name}_old_logprobs"), @@ -109,7 +109,7 @@ def _policy_metric_definitions(self, *extra: LossDef) -> list[LossDef]: *extra, LossDef(f"{self._name}_num_tokens"), ] - if self._config.metrics == PolicyLossMetrics.with_entropy: + if self._config.metrics == PolicyMetricsLevel.with_entropy: defs.append(LossDef(f"{self._name}_entropy")) return defs @@ -182,7 +182,7 @@ def _forward_backward( self._register_new_logprobs(new_logprobs_mean, kwargs, losses) # Skip the extra softmax pass when there is nothing to register. - if losses is not None and self._config.metrics != PolicyLossMetrics.none: + if losses is not None and self._config.metrics != PolicyMetricsLevel.none: self._register_extra_metrics(logits, kwargs, losses, split_index) return loss, grad @@ -191,7 +191,7 @@ def _register_extra_metrics( self, logits: torch.Tensor, kwargs: dict[str, typing.Any], - losses: dict | None, + losses: dict, split_index: int, ) -> None: metrics = compute_grpo_metrics( @@ -204,7 +204,7 @@ def _register_extra_metrics( self._config.epsilon_high, self._logits_scale_factor, group=self._parallel_dim.group if self._vocab_parallel else None, - compute_entropy=self._config.metrics == PolicyLossMetrics.with_entropy, + compute_entropy=self._config.metrics == PolicyMetricsLevel.with_entropy, ) self._register_policy_metrics(metrics, kwargs, losses) @@ -294,7 +294,7 @@ def _forward_backward( self._register_new_logprobs(new_logprobs_mean, kwargs, losses) # Skip the extra softmax pass when there is nothing to register. - if losses is not None and self._config.metrics != PolicyLossMetrics.none: + 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) return loss, grad @@ -303,7 +303,7 @@ def _register_extra_metrics( self, logits: torch.Tensor, kwargs: dict[str, typing.Any], - losses: dict | None, + losses: dict, split_index: int, document_index_zero_based: torch.Tensor, num_segments: int, @@ -322,7 +322,7 @@ def _register_extra_metrics( group=self._parallel_dim.group if self._vocab_parallel else None, sdp_group=self._sequence_data_dim.group if self._sequence_data_active else None, sp_group=self._parallel_dim.group if self._sequence_parallel else None, - compute_entropy=self._config.metrics == PolicyLossMetrics.with_entropy, + compute_entropy=self._config.metrics == PolicyMetricsLevel.with_entropy, ) self._register_policy_metrics(metrics, kwargs, losses) @@ -439,7 +439,7 @@ def compute_gspo_metrics( old_log_probabilities: torch.Tensor, # (*batch,) document_index_zero_based: torch.Tensor, # (*batch,) int — segment ID per token, 0-based num_segments: int, - num_labels_in_seq: torch.Tensor, # (*batch,) — per-document labeled-token count broadcast per token + label_counts: torch.Tensor, # (*batch,) — per-document labeled-token count broadcast per token epsilon_low: float = 0.2, epsilon_high: float = 0.2, logits_scale_factor: float = 1.0, @@ -458,7 +458,7 @@ def compute_gspo_metrics( ) flat_document_index = document_index_zero_based.reshape(-1).long() flat_mask = loss_mask.reshape(-1).to(log_ratio.dtype) - document_weight = flat_mask / num_labels_in_seq.reshape(-1).to(log_ratio.dtype).clamp(min=1) + document_weight = flat_mask / label_counts.reshape(-1).to(log_ratio.dtype).clamp(min=1) mean_log_ratio_per_segment = log_ratio.new_zeros(num_segments).index_add_( 0, flat_document_index, log_ratio.reshape(-1) * document_weight diff --git a/tests/layers/test_lm_losses.py b/tests/layers/test_lm_losses.py index 14e144ad3..78d588cb1 100644 --- a/tests/layers/test_lm_losses.py +++ b/tests/layers/test_lm_losses.py @@ -932,6 +932,7 @@ def _run_lm_loss_distributed(test_context: DistributedTestContext, base_path: pa test_context.group, ) # GSPO metrics + num_segments = 4 for compute_entropy in (False, True): with test_context.subtest(base_path, f"gspo_metrics-{compute_entropy}-{suffix}", 2) as subtest: if subtest.do_run: @@ -943,7 +944,7 @@ def _run_lm_loss_distributed(test_context: DistributedTestContext, base_path: pa loss_masking, dtype, compute_entropy, - 4, + num_segments, test_context.group, )