grpo.py

from collections.abc import Callable
from typing import Literal

import torch
from transformers import AutoTokenizer, PreTrainedModel, PreTrainedTokenizerBase

def tokenize_prompt_and_output(prompt: list[str], output: list[str], tokenizer: AutoTokenizer) -> dict[str, torch.Tensor]:
    assert len(prompt) == len(output), "Prompt and output must have the same length"
    squence = []
    squence_mask = []
    for prompt_token, output_token in zip(prompt, output):
        prompt_token_ids = tokenizer.encode(prompt_token, add_special_tokens=False)
        output_token_ids = tokenizer.encode(output_token, add_special_tokens=False)
        full_token_ids = prompt_token_ids + output_token_ids
        full_token_mask = [0] * len(prompt_token_ids) + [1] * len(output_token_ids)
        squence.append(full_token_ids)
        squence_mask.append(full_token_mask)
    max_length = max([len(i) for i in squence])

    padded_squence =[]
    padded_squence_mask =[]
    for sequence, sequence_mask in zip(squence, squence_mask):
        padded_sequence = sequence + [tokenizer.pad_token_id] * (max_length - len(sequence))
        padded_sequence_mask = sequence_mask + [0] * (max_length - len(sequence_mask))
        padded_squence.append(padded_sequence)
        padded_squence_mask.append(padded_sequence_mask)

    padded_token_ids = torch.tensor(padded_squence, dtype = torch.long)
    padded_token_mask = torch.tensor(padded_squence_mask, dtype = torch.long)

    input_ids = padded_token_ids[:, :-1]
    labels = padded_token_ids[:, 1:]
    response_mask = padded_token_mask[:, 1:]
    return {
        "input_ids": input_ids,
        "labels": labels,
        "response_mask": response_mask,
    }



def _response_log_probs_from_logits(
    logits: torch.Tensor,
    labels: torch.Tensor,
    return_token_entropy: bool,
) -> dict[str, torch.Tensor]:
    # Gathering the target logits before normalizing avoids materializing a
    # [batch, sequence, vocab] log-probability tensor.  This is exactly the
    # selected entry of log_softmax and keeps its gradient with respect to the
    # model parameters.
    # Keep probability arithmetic in float32 for half-precision models. In
    # particular, bfloat16 cannot represent GSPO's 1 +/- 3e-4 clip boundaries.
    # We still avoid constructing the full vocabulary log-softmax tensor.
    probability_logits = (
        logits.float() if logits.dtype in {torch.float16, torch.bfloat16} else logits
    )
    log_normalizer = torch.logsumexp(probability_logits, dim=-1)
    selected_logits = probability_logits.gather(
        dim=-1, index=labels.unsqueeze(-1)
    ).squeeze(-1)
    selected_log_probs = selected_logits - log_normalizer

    result = {
        "log_probs": selected_log_probs
    }
    if return_token_entropy:
        # Entropy is only a reported metric, so it must not retain an autograd
        # graph.  Work on short sequence chunks and reuse the softmax buffer
        # in-place; peak temporary storage is therefore bounded instead of
        # being another full [batch, sequence, vocab] tensor.
        max_workspace_bytes = 16 * 1024 * 1024
        bytes_per_sequence_position = (
            logits.shape[0] * logits.shape[-1] * probability_logits.element_size()
        )
        sequence_chunk_size = max(
            1,
            max_workspace_bytes // max(1, bytes_per_sequence_position),
        )

        with torch.no_grad():
            detached_logits = probability_logits.detach()
            detached_log_normalizer = log_normalizer.detach()
            token_entropy = torch.empty_like(detached_log_normalizer)
            for start in range(0, logits.shape[1], sequence_chunk_size):
                end = min(start + sequence_chunk_size, logits.shape[1])
                chunk_logits = detached_logits[:, start:end, :]
                probabilities = torch.softmax(chunk_logits, dim=-1)
                expected_logit = probabilities.mul_(chunk_logits).sum(dim=-1)
                token_entropy[:, start:end] = (
                    detached_log_normalizer[:, start:end] - expected_logit
                )
        result["token_entropy"] = token_entropy
    return result


def get_response_log_probs(
    model: PreTrainedModel,
    input_ids: torch.Tensor,
    labels: torch.Tensor,
    return_token_entropy: bool = False,
    logits_chunk_size: int | None = None,
) -> dict[str, torch.Tensor]:
    """Score tokens, optionally recomputing OLMo2 vocabulary chunks in backward.

    ``None`` uses ``model.config.grpo_logits_chunk_size`` when configured;
    zero uses the generic model forward. Positive sizes limit the sequence
    positions passed to OLMo2's output head in each call. Checkpointing the
    entire head and probability calculation avoids retaining all vocabulary
    chunks in the autograd graph. Decoder checkpointing is configured on the
    model independently.
    """
    if logits_chunk_size is None:
        logits_chunk_size = getattr(
            getattr(model, "config", None), "grpo_logits_chunk_size", 0,
        )
    if (
        not isinstance(logits_chunk_size, int)
        or isinstance(logits_chunk_size, bool)
        or logits_chunk_size < 0
    ):
        raise ValueError("logits_chunk_size must be a nonnegative integer")
    if logits_chunk_size == 0:
        return _response_log_probs_from_logits(
            model(input_ids).logits, labels, return_token_entropy,
        )

    # This bypasses the causal-LM forward, so support only an architecture
    # whose complete logit computation is the decoder followed by this head.
    # Other architectures can add scaling or soft-capping after the head.
    decoder = getattr(model, "model", None)
    get_output_embeddings = getattr(model, "get_output_embeddings", None)
    if (
        getattr(getattr(model, "config", None), "model_type", None) != "olmo2"
        or not isinstance(decoder, torch.nn.Module)
        or not callable(get_output_embeddings)
    ):
        raise ValueError(
            "Chunked log probabilities require OLMo2 with model.model and an output head"
        )
    output_head = get_output_embeddings()
    if not isinstance(output_head, torch.nn.Module):
        raise ValueError("Chunked log probabilities require an OLMo2 output head")

    from torch.utils.checkpoint import checkpoint

    hidden_states = decoder(input_ids=input_ids, use_cache=False).last_hidden_state

    def score_chunk(hidden_chunk, label_chunk):
        return _response_log_probs_from_logits(
            output_head(hidden_chunk), label_chunk, return_token_entropy,
        )

    chunks = []
    for start in range(0, hidden_states.shape[1], logits_chunk_size):
        hidden_chunk = hidden_states[:, start:start + logits_chunk_size, :]
        label_chunk = labels[:, start:start + logits_chunk_size]
        if torch.is_grad_enabled():
            chunk = checkpoint(
                score_chunk, hidden_chunk, label_chunk, use_reentrant=False,
            )
        else:
            # Cached old-policy scores need neither an autograd graph nor
            # recomputation, while still keeping head allocations bounded.
            chunk = score_chunk(hidden_chunk, label_chunk)
        chunks.append(chunk)

    result = {"log_probs": torch.cat([chunk["log_probs"] for chunk in chunks], dim=1)}
    if return_token_entropy:
        result["token_entropy"] = torch.cat(
            [chunk["token_entropy"] for chunk in chunks], dim=1,
        )
    return result


def compute_rollout_rewards(
    reward_fn: Callable[[str, str], dict[str, float]],
    rollout_responses: list[str],
    repeated_ground_truths: list[str],
) -> tuple[torch.Tensor, dict[str, float]]:
    if len(rollout_responses) != len(repeated_ground_truths):
        raise ValueError("Responses and ground truths must have the same length")

    reward_results = [
        reward_fn(response, ground_truth)
        for response, ground_truth in zip(rollout_responses, repeated_ground_truths)
    ]

    raw_rewards = torch.tensor(
        [result["reward"] for result in reward_results],
        dtype=torch.float32,
    )
    format_rewards = torch.tensor(
        [result["format_reward"] for result in reward_results],
        dtype=torch.float32,
    )
    answer_rewards = torch.tensor(
        [result["answer_reward"] for result in reward_results],
        dtype=torch.float32,
    )

    metadata = {
        "reward_mean": raw_rewards.mean().item(),
        "format_reward_mean": format_rewards.mean().item(),
        "answer_reward_mean": answer_rewards.mean().item(),
    }
    return raw_rewards, metadata


def compute_group_normalized_rewards(
    raw_rewards: torch.Tensor,
    group_size: int,
    baseline: Literal["mean", "none"] = "mean",
    advantage_eps: float = 1e-6,
    advantage_normalizer: Literal["std", "none", "mean"] = "std",
) -> tuple[torch.Tensor, dict[str, float]]:
    if raw_rewards.ndim != 1:
        raise ValueError("raw_rewards must be a one-dimensional tensor")
    if group_size <= 0:
        raise ValueError("group_size must be positive")
    if raw_rewards.numel() % group_size != 0:
        raise ValueError("The number of rewards must be divisible by group_size")

    grouped_rewards = raw_rewards.reshape(-1, group_size)

    group_mean = grouped_rewards.mean(dim=1, keepdim=True)
    # The assignment specifies torch.std's default (sample-standard-deviation)
    # behavior for standard GRPO.
    group_std = grouped_rewards.std(dim=1, keepdim=True)

    if baseline == "mean":
        advantages = grouped_rewards - group_mean
    elif baseline == "none":
        # RFT uses the raw binary reward as its weight.
        advantages = grouped_rewards
    else:
        raise ValueError(f"Unsupported baseline: {baseline}")

    if advantage_normalizer == "std":
        advantages = advantages / (group_std + advantage_eps)
    elif advantage_normalizer == "mean":
        # MaxRL upweights prompts whose sampled group has a lower mean reward.
        advantages = advantages / (group_mean + advantage_eps)
    elif advantage_normalizer != "none":
        raise ValueError(
            f"Unsupported advantage normalizer: {advantage_normalizer}"
        )

    advantages = advantages.flatten()

    metadata = {
        "advantage_mean": advantages.mean().item(),
        "advantage_std": advantages.std().item(),
    }
    return advantages, metadata


def compute_policy_gradient_loss(
    raw_rewards_or_advantages: torch.Tensor,
    policy_log_probs: torch.Tensor,
    importance_reweighting_method: Literal["none", "noclip", "grpo", "gspo"] = "none",
    old_log_probs: torch.Tensor | None = None,
    cliprange: float | None = None,
    response_mask: torch.Tensor | None = None,
) -> tuple[torch.Tensor, dict[str, torch.Tensor]]:
    if policy_log_probs.ndim != 2:
        raise ValueError(
            "policy_log_probs must have shape (batch_size, sequence_length)"
        )

    if policy_log_probs.dtype in {torch.float16, torch.bfloat16}:
        policy_log_probs = policy_log_probs.float()
    if response_mask is not None and response_mask.shape != policy_log_probs.shape:
        raise ValueError("response_mask and policy_log_probs must have the same shape")
    token_mask = (
        response_mask.bool() if response_mask is not None
        else torch.ones_like(policy_log_probs, dtype=torch.bool)
    )
    token_count = token_mask.sum().detach()

    def token_mean(values: torch.Tensor) -> torch.Tensor:
        return values.detach().masked_fill(~token_mask, 0.0).sum() / token_count.clamp_min(1)

    advantages = raw_rewards_or_advantages
    if advantages.ndim == 1:
        advantages = advantages.unsqueeze(-1)
    elif advantages.ndim != 2 or advantages.shape[1] != 1:
        raise ValueError("Advantages must have shape (batch_size,) or (batch_size, 1)")

    if advantages.shape[0] != policy_log_probs.shape[0]:
        raise ValueError("Advantages and policy_log_probs must have the same batch size")

    if importance_reweighting_method == "none":
        per_token_loss = -advantages * policy_log_probs
        metadata = {
            "on_policy_loss_mean": token_mean(per_token_loss),
        }

    elif importance_reweighting_method == "noclip":
        if old_log_probs is None:
            raise ValueError(
                "old_log_probs must be provided for noclip importance reweighting"
            )
        if old_log_probs.shape != policy_log_probs.shape:
            raise ValueError("old_log_probs and policy_log_probs must have the same shape")

        log_ratio = policy_log_probs - old_log_probs.detach()
        importance_weights = torch.exp(log_ratio)
        per_token_loss = -advantages * importance_weights

        metadata = {
            "importance_ratio_mean": token_mean(importance_weights),
            "importance_ratio_count": token_count,
            "off_policy_loss_mean": token_mean(per_token_loss),
        }

    elif importance_reweighting_method == "grpo":
        if old_log_probs is None:
            raise ValueError("old_log_probs must be provided for GRPO")
        if old_log_probs.shape != policy_log_probs.shape:
            raise ValueError("old_log_probs and policy_log_probs must have the same shape")
        if cliprange is None or cliprange <= 0:
            raise ValueError("cliprange must be positive")
        if response_mask is not None and response_mask.shape != policy_log_probs.shape:
            raise ValueError("response_mask and policy_log_probs must have the same shape")

        W_t = torch.exp(policy_log_probs - old_log_probs.detach())
        clipped_W_t = torch.clamp(W_t, min=1 - cliprange, max=1 + cliprange)

        # Take the smaller objective, then negate for gradient descent.
        unclipped_objective = advantages * W_t
        clipped_objective = advantages * clipped_W_t
        per_token_loss = -torch.minimum(unclipped_objective, clipped_objective)

        # Only count clipping that actually removes a nonzero policy gradient.
        clip_mask = (
            ((advantages > 0) & (W_t > 1 + cliprange))
            | ((advantages < 0) & (W_t < 1 - cliprange))
        )
        clip_count = (token_mask & clip_mask).sum()
        clip_fraction = clip_count / token_count.clamp_min(1)

        metadata = {
            "clip_fraction": clip_fraction.detach(),
            "clip_count": clip_count.detach(),
            "clip_total": token_count,
            "importance_ratio_mean": token_mean(W_t),
            "importance_ratio_count": token_count,
            "off_policy_loss_mean": token_mean(per_token_loss),
        }

    elif importance_reweighting_method == "gspo":
        if old_log_probs is None:
            raise ValueError("old_log_probs must be provided for GSPO")
        if old_log_probs.shape != policy_log_probs.shape:
            raise ValueError("old_log_probs and policy_log_probs must have the same shape")
        if cliprange is None or cliprange <= 0:
            raise ValueError("cliprange must be positive")
        if response_mask is None:
            raise ValueError("response_mask must be provided for GSPO")
        if response_mask.shape != policy_log_probs.shape:
            raise ValueError("response_mask and policy_log_probs must have the same shape")

        # Average log-ratios over response tokens, then exponentiate to obtain
        # one geometric-mean importance weight per response: [batch, 1].
        log_ratio = policy_log_probs - old_log_probs.detach()
        valid_mask = response_mask.bool()
        length = valid_mask.sum(dim=1, keepdim=True)
        mean_log_ratio = (
            log_ratio.masked_fill(~valid_mask, 0.0).sum(dim=1, keepdim=True)
            / length.clamp_min(1)
        )
        S_i = torch.exp(mean_log_ratio)
        clipped_S_i = torch.clamp(S_i, min=1 - cliprange, max=1 + cliprange)

        unclipped_objective = advantages * S_i
        clipped_objective = advantages * clipped_S_i
        sequence_loss = -torch.minimum(unclipped_objective, clipped_objective)
        # The existing sequence aggregation averages these repeated values.
        # Do not divide by length again here.
        per_token_loss = sequence_loss.expand_as(policy_log_probs)

        clip_mask = (
            ((advantages > 0) & (S_i > 1 + cliprange))
            | ((advantages < 0) & (S_i < 1 - cliprange))
        )
        valid_sequences = length > 0
        clip_count = (clip_mask & valid_sequences).sum()
        sequence_count = valid_sequences.sum()
        clip_fraction = clip_count / sequence_count.clamp_min(1)
        metadata = {
            "importance_ratio_mean": (
                S_i.detach().masked_fill(~valid_sequences, 0.0).sum()
                / sequence_count.clamp_min(1)
            ),
            "clipped_importance_ratio_mean": (
                clipped_S_i.detach().masked_fill(~valid_sequences, 0.0).sum()
                / sequence_count.clamp_min(1)
            ),
            "importance_ratio_count": sequence_count.detach(),
            "clip_fraction": clip_fraction.detach(),
            "clip_count": clip_count.detach(),
            "clip_total": sequence_count.detach(),
            "off_policy_loss_mean": token_mean(per_token_loss),
        }
    else:
        raise NotImplementedError(
            f"{importance_reweighting_method} importance reweighting is not implemented"
        )

    return per_token_loss, metadata


def aggregate_loss_across_microbatch(
    per_token_policy_gradient_loss: torch.Tensor,
    mask: torch.Tensor,
    loss_normalization: Literal["sequence", "constant"] = "sequence",
    normalization_constant: int | None = None,
) -> torch.Tensor:
    if per_token_policy_gradient_loss.ndim != 2:
        raise ValueError("per-token loss must have shape (batch_size, sequence_length)")
    if mask.shape != per_token_policy_gradient_loss.shape:
        raise ValueError("mask and per-token loss must have the same shape")

    if loss_normalization == "sequence":
        masked_loss = per_token_policy_gradient_loss * mask
        response_lengths = mask.sum(dim=1).clamp_min(1)
        response_loss = masked_loss.sum(dim=1) / response_lengths
        microbatch_loss = response_loss.mean()
        return microbatch_loss

    if loss_normalization == "constant":
        if normalization_constant is None or normalization_constant <= 0:
            raise ValueError(
                "normalization_constant must be positive for constant normalization"
            )
        masked_loss = per_token_policy_gradient_loss * mask
        loss = masked_loss.sum() / normalization_constant
        return loss

    raise ValueError(f"Unsupported loss normalization: {loss_normalization}")


def grpo_train_step(
    model: torch.nn.Module,
    tokenizer: PreTrainedTokenizerBase,
    optimizer: torch.optim.Optimizer,
    gradient_accumulation_steps: int,
    max_grad_norm: float | None,
    reward_fn: Callable[[str, str], dict[str, float]],
    repeated_prompts: list[str],
    rollout_responses: list[str],
    repeated_ground_truths: list[str],
    group_size: int,
    baseline: Literal["mean", "none"] = "mean",
    advantage_eps: float = 1e-6,
    advantage_normalizer: Literal["std", "none", "mean"] = "std",
    importance_reweighting_method: Literal["none", "noclip", "grpo", "gspo"] = "none",
    old_log_probs: torch.Tensor | None = None,
    cliprange: float | None = None,
    loss_normalization: Literal["sequence", "constant"] = "sequence",
    normalization_constant: int | None = None,
) -> tuple[torch.Tensor, dict[str, torch.Tensor | float]]:
    batch_size = len(rollout_responses)
    if batch_size == 0:
        raise ValueError("rollout_responses must not be empty")
    if len(repeated_prompts) != batch_size or len(repeated_ground_truths) != batch_size:
        raise ValueError("Prompts, responses, and ground truths must have the same length")
    if group_size <= 0 or batch_size % group_size != 0:
        raise ValueError("group_size must be positive and divide the rollout batch size")
    if gradient_accumulation_steps <= 0 or gradient_accumulation_steps > batch_size:
        raise ValueError("gradient_accumulation_steps must be between 1 and the batch size")
    if loss_normalization not in {"sequence", "constant"}:
        raise ValueError(f"Unsupported loss normalization: {loss_normalization}")
    if loss_normalization == "constant" and (
        normalization_constant is None or normalization_constant <= 0
    ):
        raise ValueError(
            "normalization_constant must be positive for constant normalization"
        )

    device = next(model.parameters()).device

    raw_rewards, reward_metadata = compute_rollout_rewards(
        reward_fn=reward_fn,
        rollout_responses=rollout_responses,
        repeated_ground_truths=repeated_ground_truths,
    )
    advantages, advantage_metadata = compute_group_normalized_rewards(
        raw_rewards=raw_rewards,
        group_size=group_size,
        baseline=baseline,
        advantage_eps=advantage_eps,
        advantage_normalizer=advantage_normalizer,
    )

    optimizer.zero_grad(set_to_none=True)

    # Zero-advantage responses contribute exactly zero to the policy gradient,
    # so exclude them from token scoring while retaining full-batch reward and
    # advantage statistics above.
    active_indices = torch.nonzero(advantages != 0, as_tuple=False).flatten()
    if active_indices.numel() == 0:
        zero_loss = torch.zeros((), dtype=torch.float32, device=device)
        metadata: dict[str, torch.Tensor | float] = {
            **reward_metadata,
            **advantage_metadata,
            ("on_policy_loss_mean" if importance_reweighting_method == "none"
             else "off_policy_loss_mean"): zero_loss.clone(),
            "loss": zero_loss,
            "gradient_norm": zero_loss.clone(),
            "token_entropy": zero_loss.clone(),
            "response_token_count": torch.zeros(
                (), dtype=torch.long, device=device
            ),
            "entropy_count": zero_loss.clone(),
            "optimizer_step": 0.0,
            "active_response_count": 0.0,
        }
        if importance_reweighting_method != "none":
            metadata.update(importance_ratio_mean=zero_loss.clone(), importance_ratio_count=zero_loss.clone())
        if importance_reweighting_method in {"grpo", "gspo"}:
            metadata.update(clip_count=zero_loss.clone(), clip_total=zero_loss.clone(), clip_fraction=zero_loss.clone())
        return zero_loss, metadata

    tokenized_rollouts = tokenize_prompt_and_output(
        prompt=repeated_prompts,
        output=rollout_responses,
        tokenizer=tokenizer,
    )

    input_ids = tokenized_rollouts["input_ids"].index_select(0, active_indices)
    labels = tokenized_rollouts["labels"].index_select(0, active_indices)
    response_mask = tokenized_rollouts["response_mask"].index_select(
        0, active_indices
    )
    active_advantages = advantages.index_select(0, active_indices)
    active_old_log_probs = None
    if old_log_probs is not None:
        if old_log_probs.ndim != 2 or old_log_probs.shape[0] != batch_size:
            raise ValueError(
                "old_log_probs must have shape (batch_size, sequence_length)"
            )
        active_old_log_probs = old_log_probs.index_select(
            0, active_indices.to(old_log_probs.device)
        )

    total_loss = torch.zeros((), dtype=torch.float32, device=device)
    entropy_sum = torch.zeros((), dtype=torch.float32, device=device)
    response_token_count = torch.zeros((), dtype=torch.long, device=device)

    # Preserve the configured microbatch size while reducing the number of
    # accumulation steps in proportion to the number of retained responses.
    microbatch_size = (
        batch_size + gradient_accumulation_steps - 1
    ) // gradient_accumulation_steps
    microbatch_indices = torch.arange(active_indices.numel()).split(
        microbatch_size
    )
    # These sums cover every scored microbatch, not just the last one. Zero-
    # advantage rows are excluded, so all these diagnostics are active-only.
    diagnostic_sums: dict[str, torch.Tensor] = {}
    diagnostic_weights: dict[str, torch.Tensor] = {}
    diagnostic_counts: dict[str, torch.Tensor] = {}

    for indices in microbatch_indices:
        current_microbatch_size = indices.numel()
        microbatch_input_ids = input_ids[indices].to(device)
        microbatch_labels = labels[indices].to(device)
        microbatch_response_mask = response_mask[indices].to(device)
        microbatch_advantages = active_advantages[indices].to(device)
        microbatch_old_log_probs = (
            active_old_log_probs.index_select(
                0, indices.to(active_old_log_probs.device)
            ).to(device)
            if active_old_log_probs is not None
            else None
        )

        log_prob_results = get_response_log_probs(
            model=model,
            input_ids=microbatch_input_ids,
            labels=microbatch_labels,
            return_token_entropy=True,
        )
        microbatch_policy_log_probs = log_prob_results["log_probs"]
        microbatch_token_entropy = log_prob_results["token_entropy"]

        per_token_loss, loss_metadata = compute_policy_gradient_loss(
            raw_rewards_or_advantages=microbatch_advantages,
            policy_log_probs=microbatch_policy_log_probs,
            importance_reweighting_method=importance_reweighting_method,
            old_log_probs=microbatch_old_log_probs,
            cliprange=cliprange,
            response_mask=microbatch_response_mask,
        )
        microbatch_loss = aggregate_loss_across_microbatch(
            per_token_policy_gradient_loss=per_token_loss,
            mask=microbatch_response_mask,
            loss_normalization=loss_normalization,
            normalization_constant=normalization_constant,
        )
        if loss_normalization == "sequence":
            scale = current_microbatch_size / batch_size
        else:
            scale = 1.0
        scaled_microbatch_loss = microbatch_loss * scale
        if not torch.isfinite(scaled_microbatch_loss):
            raise RuntimeError("Non-finite GRPO loss; optimizer step aborted")
        microbatch_entropy_sum = (
            microbatch_token_entropy * microbatch_response_mask
        ).sum().detach()
        if not torch.isfinite(microbatch_entropy_sum):
            raise RuntimeError("Non-finite token entropy; optimizer step aborted")
        # Clipping can hide an overflowing importance ratio in the loss while
        # backward still produces NaNs. Check the already-reduced diagnostics.
        for key, value in loss_metadata.items():
            if key.endswith("_mean") and not torch.isfinite(value):
                raise RuntimeError(f"Non-finite {key}; optimizer step aborted")
        scaled_microbatch_loss.backward()
        total_loss = total_loss + scaled_microbatch_loss.detach()

        for key, value in loss_metadata.items():
            if key in {"clip_count", "clip_total", "importance_ratio_count"}:
                diagnostic_counts[key] = diagnostic_counts.get(key, torch.zeros_like(value)) + value
            elif key.endswith("_mean"):
                weight = (
                    loss_metadata["importance_ratio_count"]
                    if "importance_ratio" in key else microbatch_response_mask.sum().detach()
                )
                diagnostic_sums[key] = diagnostic_sums.get(key, torch.zeros_like(value)) + value * weight
                diagnostic_weights[key] = diagnostic_weights.get(key, torch.zeros_like(weight)) + weight

        entropy_sum = entropy_sum + microbatch_entropy_sum
        response_token_count = response_token_count + microbatch_response_mask.sum().detach()

    if max_grad_norm is not None:
        gradient_norm = torch.nn.utils.clip_grad_norm_(
            model.parameters(), max_grad_norm, error_if_nonfinite=True,
        )
    else:
        parameter_gradient_norms = [
            parameter.grad.detach().norm(2)
            for parameter in model.parameters()
            if parameter.grad is not None
        ]
        gradient_norm = (
            torch.stack(parameter_gradient_norms).norm(2)
            if parameter_gradient_norms
            else torch.zeros((), device=device)
        )
        if not torch.isfinite(gradient_norm):
            raise RuntimeError("Non-finite gradient norm; optimizer step aborted")

    optimizer.step()
    optimizer.zero_grad(set_to_none=True)

    mean_token_entropy = entropy_sum / response_token_count.clamp_min(1)
    loss_diagnostics = {
        key: value / diagnostic_weights[key].clamp_min(1)
        for key, value in diagnostic_sums.items()
    }
    loss_diagnostics.update(diagnostic_counts)
    if "clip_total" in diagnostic_counts:
        loss_diagnostics["clip_fraction"] = (
            diagnostic_counts["clip_count"] / diagnostic_counts["clip_total"].clamp_min(1)
        )
    metadata: dict[str, torch.Tensor | float] = {
        **reward_metadata,
        **advantage_metadata,
        **loss_diagnostics,
        "loss": total_loss.detach(),
        "gradient_norm": gradient_norm.detach(),
        "token_entropy": mean_token_entropy.detach(),
        "response_token_count": response_token_count.detach(),
        "entropy_count": response_token_count.detach(),
        "optimizer_step": 1.0,
        "active_response_count": float(active_indices.numel()),
    }
    return total_loss.detach(), metadata