Loss functions

Select a loss function with the loss_fn argument on forward_backward(). Each loss reads a different set of per-position fields from Datum.loss_fn_inputs, and some take a loss_fn_config dict to tune hyperparameters. Loops supports the following loss functions:

Supervised: cross_entropy
Preference: dpo Beta
Reinforcement learning: importance_sampling, ppo, cispo, dppo, dro

Some models restrict these losses; GLM 5.3’s long-context configuration, for example, doesn’t support cross_entropy or dpo. See Supported models for each model’s limits.

Gradient normalization

Every loss sums its per-position losses by default. To average them instead, set skip_token_normalization to false in loss_fn_config.

boolean

true sums per-position losses. false divides the sum by the number of active positions: valid targets with nonzero weight, or chosen and rejected pairs for dpo.Use the same value on every forward_backward() call before an optim_step(). The trainer rejects a step that mixes summed and averaged gradients. forward_backward_custom() always sends true.

`cross_entropy`

Standard supervised fine-tuning. Each Datum carries per-token targets under target_tokens, with -100 marking positions the loss should ignore (typically the prompt). forward_backward does not shift labels internally, so shift when you tokenize: the label at position i is the token at position i+1.

from baseten.loops import Datum, ModelInput, TensorData

datum = Datum(
    model_input=ModelInput.from_ints(tokens),
    loss_fn_inputs={
        "target_tokens": TensorData(
            data=targets, dtype="int64", shape=[len(targets)]
        ),
    },
)

fb = training_client.forward_backward(data=[datum]).result(timeout=600.0)

See Train on a chat dataset with the Loops SDK for the full tokenization step.

Weight individual positions

Pass an optional weights tensor to scale the per-position loss. Each position contributes -log(p(target)) * weight, so a weight of 0.0 drops a position and a larger weight increases its influence. A -100 target drops the position outright; use weights to scale what remains. Build weights alongside target_tokens so they align with the shifted labels:

from baseten.loops import Datum, ModelInput, TensorData

def build_weighted_datum(tokenizer, prompt, answer):
    p = tokenizer.encode(prompt, add_special_tokens=False)
    a = tokenizer.encode(answer, add_special_tokens=False)
    tokens = (p + a)[:-1]
    targets = [-100] * (len(p) - 1) + list(a)
    # One weight per target: ignore the prompt, train on the answer.
    weights = [0.0] * (len(p) - 1) + [1.0] * len(a)
    return Datum(
        model_input=ModelInput.from_ints(tokens),
        loss_fn_inputs={
            "target_tokens": TensorData(
                data=targets, dtype="int64", shape=[len(targets)]
            ),
            "weights": TensorData(
                data=weights, dtype="float32", shape=[len(weights)]
            ),
        },
    )

datum = build_weighted_datum(
    training_client.get_tokenizer(),
    prompt="What is the capital of France?\nAnswer:",
    answer=" Paris",
)
fb = training_client.forward_backward(data=[datum]).result(timeout=600.0)

Set weights on every Datum in the batch or on none of them.

Target shapes

target_tokens may have shape [N] for one target per token position or [N, K] for top-K soft targets. If provided, weights must match that shape. Loops sums the K weighted losses at each position and counts the position once. Every Datum in a batch must use the same rank and K.

Custom loss functions

Use forward_backward_custom() when the built-in losses do not express your objective. The callback receives the input Datum objects and one PyTorch log-probability tensor per datum. Each tensor matches the datum’s target_tokens shape, either [N] or [N, K]. Install the PyTorch dependency before using custom losses:

uv add 'baseten-loops[torch]'

For this method, each datum’s loss_fn_inputs may contain only target_tokens and optional weights. Capture other values that the loss needs, such as frozen reference or teacher logprobs, in the callback closure. The callback must return a scalar PyTorch tensor and a dictionary of numeric metrics:

import torch

def custom_loss(data, logprobs_list):
    per_datum = []
    for datum, logprobs in zip(data, logprobs_list, strict=True):
        targets = datum.loss_fn_inputs["target_tokens"].to_torch()
        active = targets.to(device=logprobs.device).ne(-100)
        per_datum.append(-logprobs[active].mean())
    loss = torch.stack(per_datum).mean()
    return loss, {"custom_loss": float(loss.detach())}

result = training_client.forward_backward_custom(
    data=[datum],
    loss_fn=custom_loss,
).result(timeout=600.0)

print(result.metrics["custom_loss"])

The loss_type_input parameter of forward_backward_custom() only accepts "logprobs". Inactive positions arrive as 0.0, so select active targets explicitly instead of averaging the entire tensor. The callback’s metrics are available in result.metrics, so include the custom loss there if you want to inspect it after the call. The loss must depend on every log-probability tensor in the batch.

`importance_sampling`, `ppo`, `dppo`, `cispo`, `dro`

On-policy reinforcement learning. Each Datum’s loss_fn_inputs must include per-position logprobs (from the behavior policy at rollout time) and advantages (from your reward model or return-to-go).

L = len(tokens)

datum = Datum(
    model_input=ModelInput.from_ints(tokens),
    loss_fn_inputs={
        "logprobs": TensorData(data=behavior_logprobs, dtype="float32", shape=[L]),
        "advantages": TensorData(data=advantages, dtype="float32", shape=[L]),
    },
)

fb = training_client.forward_backward(
    data=[datum],
    loss_fn="ppo",
).result(timeout=600.0)

These losses take the same loss_fn_inputs shape and differ in clipping and regularization. Drive them from a rollout loop that samples from a paired SamplingClient and scores each trajectory before calling forward_backward.

`importance_sampling`

REINFORCE with an importance ratio: each token contributes -advantage * exp(new_logprobs - old_logprobs), with no clipping.

`ppo`

Clipped surrogate: the probability ratio is clamped to [clip_low_threshold, clip_high_threshold] and the surrogate takes the pessimistic minimum of the raw and clamped advantage-weighted terms.

number

Lower bound of the clamp on the per-token probability ratio exp(new_logprobs - old_logprobs).

number

Upper bound of the clamp. Both bounds must be nonnegative, and the high at least the low.

`cispo`

Clipped importance weight: the ratio is clamped to the same bounds as ppo, but the clamp is detached, so it scales the raw new-policy gradient without backpropagating through the clip.

number

Lower bound of the clamp on the importance weight exp(new_logprobs - old_logprobs).

number

Upper bound of the clamp. Both bounds must be nonnegative, and the high at least the low.

`dppo`

Policy gradient with total-variation masking and a squared log-ratio KL trust region: tokens whose probability moved too far in the direction their advantage rewards are dropped from the gradient, and the loss adds a squared log-ratio penalty.

number

Drops a token from the policy gradient when its probability rose by more than this and its advantage is positive.

number

Drops a token from the policy gradient when its probability fell by more than this and its advantage is negative.

number

Scale of advantages inside the policy gradient.

number

Weight of the squared log-ratio trust-region penalty.

`dro`

Direct reward optimization: the objective is the advantage-weighted log-probability minus a squared log-ratio penalty.

number

Strength of the squared log-ratio penalty on new_logprobs - old_logprobs. Must be nonnegative.

`dpo`

DPO is in beta. Its inputs, loss_fn_config keys, and behavior can change between releases.

Direct Preference Optimization. Trains on chosen/rejected preference pairs against a frozen reference policy. DPO adds the following steps to SFT:

Snapshot the reference before any optim_step. The initial weights are the reference for offline DPO.
Interleave the batch as [chosen_0, rejected_0, chosen_1, rejected_1, ...]. Pairing is by position, and the SDK rejects odd-length batches before submitting anything.
Score reference logprobs once with attach_reference_logprobs and reuse across every step.

Each Datum needs the following loss_fn_inputs fields. Build target_tokens and weights yourself when tokenizing; ref_logprobs comes from attach_reference_logprobs():

target_tokens: labels shifted by one, -100 over the prompt positions.
weights: 0.0 over the prompt, 1.0 over the response. The DPO log-ratio sums logprobs weighted by this mask.
ref_logprobs: frozen reference logprobs, one per position. Populated automatically by attach_reference_logprobs.

from baseten.loops import AdamParams, Datum
from baseten.loops.helpers.datum import attach_reference_logprobs

# Blocks until the snapshot is published, then returns a SamplingClient.
reference = training_client.save_weights_and_get_sampling_client(name="dpo-ref")

# Interleaved chosen/rejected datums, with target_tokens and weights as above.
datums = attach_reference_logprobs(datums, reference)

fb = training_client.forward_backward(
    data=datums,
    loss_fn="dpo",
    loss_fn_config={"beta": 0.1},
).result(timeout=600.0)
training_client.optim_step(AdamParams(learning_rate=1e-5)).result(timeout=600.0)

loss_fn_config keys:

number

KL-penalty strength against the reference.

number

Conservative DPO (cDPO) smoothing, in [0, 0.5]. 0.0 is standard DPO.

Metrics on fb.metrics: the DPO pass reports loss, the mean DPO loss over the batch. A fresh policy starts near ln 2 (about 0.693) and decreases as the policy learns to prefer chosen responses.

Next steps

TrainingClient reference for the full forward_backward() signature.
Helpers reference for attach_reference_logprobs.
Loops concepts explains the paired trainer + sampler model that makes reference-policy snapshots possible.

Get started

Model APIs

Inference

Training

Baseten Grid

Frontier Gateway

Sandboxes

Development

Deployment

Engines

Model Labs

Organization

Hosting options

Observability

Troubleshooting

Gradient normalization

`cross_entropy`

Weight individual positions

Target shapes

Custom loss functions

`importance_sampling`, `ppo`, `dppo`, `cispo`, `dro`

`importance_sampling`

`ppo`

`cispo`

`dppo`

`dro`

`dpo`

Next steps

​Gradient normalization

​cross_entropy

​Weight individual positions

​Target shapes

​Custom loss functions

​importance_sampling, ppo, dppo, cispo, dro

​importance_sampling

​ppo

​cispo

​dppo

​dro

​dpo

​Next steps

Gradient normalization

`cross_entropy`

Weight individual positions

Target shapes

Custom loss functions

`importance_sampling`, `ppo`, `dppo`, `cispo`, `dro`

`importance_sampling`

`ppo`

`cispo`

`dppo`

`dro`

`dpo`

Next steps