Custom loss functions

The built-in policy losses rLLM ships, and how to register your own — one function that runs unchanged on verl, Tinker, and Fireworks.

Module: rllm.trainer.algorithms.loss — exposed at top level as rllm.register_loss and rllm.LossContext.

This page covers:

1. How the loss hook works (one function, all three backends)
2. Built-in losses and how to select one
3. The LossContext contract and its reducers
4. Loss-aggregation modes (loss_agg_mode)
5. Writing and registering a custom loss (worked example: IcePop)

Core concept

A policy loss is a single function that returns the complete scalar objective and does its own masking and aggregation — the same convention as verl’s POLICY_LOSS_REGISTRY. You select it by name; there is no list of losses and no separate auxiliary-loss framework (an additive term is composed inside a loss body).

The same function runs on every backend. It reads a normalized LossContext and calls the reducer the backend injects, so a loss body never needs to know whether it is running in verl’s in-process kernel or a Tinker/Fireworks forward_backward_custom pass.

import rllm

@rllm.register_loss("my_loss")                 # same decorator style as @rllm.rollout
def my_loss(ctx: rllm.LossContext):
    ratio = (ctx.logp_curr - ctx.logp_old).exp()
    pg = -ctx.advantages * ratio.clamp(max=20).detach() * ctx.logp_curr
    return ctx.aggregate(pg, ctx.action_mask), {"ratio_mean": ratio.mean().item()}

algorithm:
  loss_fn: my_loss                    # single selector
  loss_params: {clip: 0.2}            # loss-specific params → ctx.params
  loss_agg_mode: seq-mean-token-mean  # one knob, same semantics on every backend
  loss_plugins: ["my_pkg.losses"]     # imported at startup so @register_loss fires

A loss returns (scalar_loss, metrics) where metrics is a dict[str, float] logged under actor/… (verl) or train/… (managed backends).

Built-in losses

Ported from verl 0.8 and verified against its kernels. Select any of these with algorithm.loss_fn — no plugin needed.

loss_fn	Description
ppo_clip	Standard PPO/GRPO clipped surrogate (verl vanilla)
dppo_tv / dppo_kl	DPPO (arXiv:2602.04879) — replaces ratio-clipping with a per-token divergence mask (total-variation / KL). Uses the untruncated ratio (C=∞) per the paper §5.4.
cispo	CISPO (arXiv:2506.13585) — clips the importance-sampling weight (stop-grad) but keeps every token’s gradient
gspo	GSPO (arXiv:2507.18071) — sequence-level importance ratio via seq_reduce; pins seq-mean-token-mean
icepop	IcePop (arXiv:2510.18855) — double-sided train/inference discrepancy mask: zeros a token’s gradient when exp(logp_curr − logp_rollout) leaves [icepop_alpha, icepop_beta] (default 0.5/5.0). Reports icepop/ratio_{mean,min,max} + mask_frac.
reinforce	Clip-free REINFORCE-style policy gradient (verl gpg)
echo	ECHO (arXiv:2605.24517) — PPO term plus an observation-token cross-entropy term, composed in one body

Common hyperparameters (eps_clip, eps_clip_high, kl_beta) come from the top-level algorithm config; anything else goes under algorithm.loss_params and lands in ctx.params.

The LossContext contract

Every loss receives a LossContext with these fields (all tensors share a shape — (batch, response_len) on verl, (num_tokens,) per datum on the managed path):

Field	Meaning	verl analog
logp_curr	Current policy log-probs — the only gradient-carrying input	log_prob
logp_old	The importance-ratio denominator (old/behavior policy)	old_log_probs
advantages	Per-token advantage estimates	—
action_mask	1.0 on assistant/action tokens (where the PG applies)	—
obs_mask	1.0 on environment-observation tokens (for ECHO-style terms)	—
logp_ref	Reference-policy log-probs for a KL term, or None	ref_log_prob
logp_rollout	Raw inference/sampling log-probs from the rollout (behavior policy), or None	rollout_log_probs
params	Loss hyperparameters from loss_params	—

logp_rollout is the true sampling distribution and equals logp_old under bypass_mode; they differ only when a proximal old-policy was recomputed (non-bypass). Reach for it when a loss must reference the rollout distribution independently of the ratio denominator — e.g. train/inference-mismatch corrections like TIS or IcePop.

Only logp_curr carries gradient; everything else is a detached constant. That is what lets the managed backends reproduce dL/dθ from your loss with a forward_backward_custom surrogate.

Two reducers are injected on the context:

• ctx.aggregate(per_token, mask, mode=None) — collapses a per-token tensor to the scalar objective, realizing the configured loss_agg_mode with global normalization spanning the whole optimizer step. Pass mode= to force a specific reduction regardless of config (GSPO does this).
• ctx.seq_reduce(values, mask, reduction) — reduces per sequence and broadcasts back to per-token (a row on verl’s (B, T); the whole datum on the managed path). This is what enables sequence-level losses like GSPO.

Write the loss body shape-agnostically: elementwise math on the fields, then ctx.aggregate.

Loss-aggregation modes

loss_agg_mode is a single, backend-agnostic knob using verl’s names. The loss body just calls ctx.aggregate; each backend realizes the mode with global normalization over the entire optimizer step (all micro-batches, grad-accumulation passes, and DP ranks).

Mode	Semantics
seq-mean-token-mean (default)	Mean within each sequence, then mean over sequences — every sequence weighted equally
token-mean	Every token weighted equally across the batch
seq-mean-token-sum	Sum within each sequence, then mean over sequences

A loss whose math requires a specific reduction pins it at registration and ignores the config:

@rllm.register_loss("gspo", agg_mode="seq-mean-token-mean")
def gspo(ctx): ...

How each backend runs it

loss_fn resolution is native-first: if the backend has a fused kernel for that name, rLLM uses it; only losses without a native kernel take the custom path.

Backend	Mechanism
verl	Runs in-process over a LossContext; aggregate wraps verl’s agg_loss (global-batch norm); logp_old = old_log_probs. Native kernels still handle non-rLLM loss_fn.
Tinker	A forward_backward_custom closure, single pass; logp_old = the sampling log-probs; normalized client-side.
Fireworks	The same forward_backward_custom; the client returns a raw sum and optim_step normalizes server-side from the resolved loss_agg_mode (accumulation-invariant).

Writing a custom loss

A loss is any function (ctx) -> (scalar, metrics). The example below is IcePop (arXiv:2510.18855) — it zeros a token’s gradient when its train/inference ratio exp(logp_curr − logp_rollout) leaves a band [alpha, beta] (masking out-of-range tokens rather than clipping them). rLLM already ships this as the built-in icepop (see the table above); it’s reproduced here under a different name as a template you can adapt.

# my_losses.py
import rllm
import torch

@rllm.register_loss("my_icepop")
def my_icepop(ctx: rllm.LossContext):
    """Mask tokens whose train/inference ratio exp(logp_curr - logp_rollout) leaves [alpha, beta]."""
    if ctx.logp_rollout is None:
        raise ValueError("my_icepop needs rollout log-probs, but ctx.logp_rollout is None")
    alpha = float(ctx.params.get("alpha", 0.5))
    beta = float(ctx.params.get("beta", 5.0))
    ratio = torch.exp(ctx.logp_curr.detach() - ctx.logp_rollout)   # detached weight, no grad
    keep = ((ratio >= alpha) & (ratio <= beta)).to(ctx.logp_curr.dtype)
    pg = -ctx.advantages * ratio * ctx.logp_curr * keep            # grad flows via logp_curr
    am = ctx.action_mask
    mask_frac = ((1.0 - keep) * am).sum() / am.sum().clamp(min=1.0)
    return ctx.aggregate(pg, am), {"my_icepop/mask_frac": mask_frac.item()}

Three conventions worth internalizing:

• logp_curr is the only differentiable term. Anything used as a weight (the ratio, a mask) should be .detach()ed so gradient flows only through the score logp_curr.
• Use logp_rollout for the true sampling distribution and guard on None rather than falling back to logp_old (which is the proximal, not the inference log-prob). It’s always populated on the managed backends.
• Never reduce by hand. Return ctx.aggregate(per_token, mask) so your loss inherits the configured aggregation and correct global normalization.

Selecting and registering it

Point loss_fn at the registered name and list the defining module in loss_plugins so rLLM imports it at startup (which runs the @register_loss decorator on every worker):

algorithm:
  loss_fn: icepop
  loss_params: {alpha: 0.5, beta: 5.0}
  loss_plugins: ["my_losses"]        # importable module path

Alternatively, expose the module through an rllm.losses entry point in your package, or simply import it before constructing the trainer — any of these makes the name resolvable.

Notes

• bypass_mode (async). With one optimizer step per batch, a recomputed proximal equals the current policy, so logp_old defaults to the rollout log-probs (bypass_mode=True on Fireworks) — making the ratio a real trust region against the behavior policy. Set false only for multi-update-per-batch PPO. This is training infrastructure, not part of a loss body.
• Sequence-level losses need ctx.seq_reduce; see gspo for the pattern (length-normalized sequence ratio, then applied per token via the stop-gradient identity).
• Tinker normalization is exact only when fwd_bwd_group_size == mini_batch_size; under grad accumulation the divisor is per-pass. Fireworks and verl normalize globally and are accumulation-invariant.