Custom losses

Write your training objective in PyTorch and run its loss and backward pass on River's training backend. Custom losses let you experiment with objectives, auxiliary terms, and weighting schemes for RL, distillation, and supervised learning.

The client captures your function as a symbolic graph. River validates the graph and executes the loss and its backward pass on the training workers, alongside the model's logits.

For the design and application examples, read our blogpost: Custom Losses at Scale.

Define, prepare, and train

Install the custom-loss extra. The examples below use river-client 0.17.1 or later. Graph capture runs on the client's CPU.

pip install --upgrade 'river-client[custom-loss]>=0.17.1'

For an existing training model, such as the one in the SFT guide, define an objective and declare its input columns:

import torch.nn.functional as F
import river_client as river


def masked_cross_entropy(logits, inputs):
    ce = F.cross_entropy(logits, inputs["targets"], reduction="none").float()
    mask = inputs["mask"]
    loss = (ce * mask).sum() / mask.sum().clamp_min(1)
    return loss, {"token_count": mask.sum()}


loss = model.prepare_loss(river.CustomLoss(
    masked_cross_entropy,
    inputs={"targets": "int64", "mask": "float32"},
))

data = [{
    "input_ids": [1, 2, 3, 4],  # Illustrative token IDs.
    "loss_data": {
        "targets": [2, 3, 4, 0],
        "mask": [0.0, 1.0, 1.0, 0.0],
    },
}]
result = model.forward_backward(data, loss_fn=loss, zero_out=True)
model.optim_step(lr=1e-4)  # Runs only after backward succeeds.
print(result.metrics)  # Includes custom/token_count.

The function receives logits shaped [T, V] for one example: T token positions and V vocabulary entries. It returns a scalar tensor, optionally paired with a dictionary of scalar metrics. Supply the declared columns in loss_data; the schema specifies their names, dtypes, and shapes.

Keep the value returned by prepare_loss. Preparation captures the formula and its constants and checks support on that model before training. Prepare again when changing the formula, its constants, or the model/session. Put values that change between examples in input columns.

Use model.forward to evaluate the prepared loss. Custom objectives return scalar loss values and the metrics you declare, named custom/<name>.

How River executes the objective

A full sequence-by-vocabulary logits matrix can exceed GPU memory. River processes token rows in smaller groups and combines their results before finalizing each example's loss. It manages distributed execution and the loss workspace.

The backend accepts validated graphs of supported tensor operations. Each training request is checked for input and resource limits.

Example: choose among candidate decisions

A classifier can reuse the model's language-model head by mapping each choice to one token. At the decision position, normalize over those candidate tokens.

This example has one decision position per sequence. target is one-hot at that position and zero elsewhere. log_mask is zero there and a large finite negative value elsewhere; weight carries the example's weight at that position and is zero elsewhere.

import torch.nn.functional as F
import river_client as river


def decision_loss(logits, inputs):
    log_probs = F.log_softmax(logits, dim=-1)
    candidates = log_probs.gather(-1, inputs["candidates"]).float()
    normalizer = (candidates + inputs["log_mask"]).logsumexp((0, 1))
    target_log_prob = (inputs["target"] * candidates).sum()
    return (normalizer - target_log_prob) * inputs["weight"].sum()


K = 4  # Number of candidate decisions, each represented by a distinct token.
decision = model.prepare_loss(river.CustomLoss(
    decision_loss,
    inputs={
        "candidates": {"dtype": "int64", "shape": ("T", K)},
        "target": {"dtype": "float32", "shape": ("T", K)},
        "log_mask": {"dtype": "float32", "shape": ("T", 1)},
        "weight": "float32",
    },
))

Use valid candidate token IDs at every position, including masked positions. Set log_mask to a finite value such as -1e30 at non-decision positions. The same objective can use a vision-capable model's image-conditioned logits.

Use with AsyncTrainer

AsyncTrainer prepares the loss before starting rollouts. Use river.CustomLoss for the tensor objective and rl.CustomLoss for its input adapter, sampling requirements, and zero-advantage behavior. Keep reduction in the objective and omit normalize and built-in loss options from the trainer configuration.

By default, the trainer supplies any declared subset of:

To form importance ratios, compute current-policy logprobs from logits / engine.temperature inside the objective. Custom objectives require unfiltered stochastic sampling: top_p=1, top_k=-1 or 0, and a positive temperature.

For example, SAPO replaces hard clipping with a smooth surrogate. With the single-turn engine and rows from the RL recipe:

import torch.nn.functional as F
import river_client as river
from river_client import rl


temperature = engine.temperature


def sapo(logits, inputs):
    logprobs = -F.cross_entropy(
        logits / temperature, inputs["targets"], reduction="none"
    ).float()
    ratio = (logprobs - inputs["old_logprobs"]).exp()
    advantages = inputs["advantages"]
    tau_pos, tau_neg = 1.0, 1.05
    surrogate = (
        advantages.clamp_min(0) * (4.0 / tau_pos)
        / ((-tau_pos * (ratio - 1.0)).exp() + 1.0)
        + advantages.clamp_max(0) * (4.0 / tau_neg)
        / ((-tau_neg * (ratio - 1.0)).exp() + 1.0)
    )
    mask = inputs["mask"]
    return -(surrogate * mask).sum() / mask.sum().clamp_min(1)


loss = rl.CustomLoss(loss=river.CustomLoss(
    sapo,
    inputs={
        "targets": "int64",
        "old_logprobs": "float32",
        "advantages": "float32",
        "mask": "float32",
    },
))
trainer = rl.AsyncTrainer(
    engine=engine,
    optimizer=rl.Adam(lr=1e-5),
    completion=rl.GroupCompletion(mode="wait"),
    groups_per_step=2,
    group_size=8,
    max_staleness=2,
    advantage=rl.GroupCentered(standardize=True),
    loss=loss,
)
rl.run(trainer, rows, steps=100)

Each response contributes its own masked mean; the optimizer accumulates their sum. The tau constants control the surrogate's smoothness.

Request top-K logprobs

old_logprobs holds the sampled token's log probability at each position. For alternative-token probabilities, set requirements=rl.LossRequirements(top_logprobs=K) on rl.CustomLoss before starting rollouts. The sampling distribution stays unchanged.

A client-side input_adapter receives the requested arrays in datum["sampling_metadata"]:

These arrays are already aligned to prediction positions. Return the columns your objective declares, and use datum["loss_data"]["mask"] for eligibility. The following example uses retained sampling probabilities as soft targets:

import numpy as np
import torch.nn.functional as F
import river_client as river
from river_client import rl

K = 8


class TopKInputs:
    def __call__(self, datum):
        metadata = datum["sampling_metadata"]
        ids = np.asarray(metadata["sampler_topk_ids"], dtype=np.int64)
        probabilities = np.exp(metadata["sampler_topk_logprobs"])
        return {
            "candidate_ids": np.maximum(ids, 0),  # Replace unused -1 IDs.
            "probabilities": np.where(ids >= 0, probabilities, 0).astype(np.float32),
            "mask": datum["loss_data"]["mask"],
        }

    def contract(self):
        return {"version": 1, "top_logprobs": K}


def topk_soft_targets(temperature):
    def objective(logits, inputs):
        logp = F.log_softmax(logits / temperature, dim=-1)
        selected = logp.gather(-1, inputs["candidate_ids"]).float()
        ce = -(selected * inputs["probabilities"]).sum(-1)
        mask = inputs["mask"]
        return (ce * mask).sum() / mask.sum().clamp_min(1)

    return rl.CustomLoss(
        loss=river.CustomLoss(objective, inputs={
            "candidate_ids": {"dtype": "int64", "shape": ("T", K)},
            "probabilities": {"dtype": "float32", "shape": ("T", K)},
            "mask": "float32",
        }),
        input_adapter=TopKInputs(),
        requirements=rl.LossRequirements(top_logprobs=K),
        train_zero_advantage_batches=True,
    )


# Use this loss in the AsyncTrainer construction above.
loss = topk_soft_targets(engine.temperature)

The sampler's original top-K probabilities serve as cross-entropy weights.

CISPO supports score centering directly through rl.CISPO(normalize="token", score_centering=rl.ScoreCentering(top_k=K)), which requests the sampling metadata automatically.

Supported operations and validation

Custom losses use a supported subset of PyTorch:

Model-dependent rows stay independent until reduction. Keep reduced model-dependent values in scalar computations, including detached values. Input-only scalars, such as mask counts, can normalize token values. Use declared columns for input data and Python numeric constants for fixed coefficients. The graph must have the same structure across tensor values, sequence lengths, devices, dtypes, and autograd states.

Preparation checks the graph. Submitted batches are checked for shape, input size, and memory use. The loss workspace is one part of total training memory. Use validation errors to identify unsupported expressions, fix input data, or reduce a request's size. Contact River support for model or endpoint availability errors.