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
Reinforcement learning: importance_sampling, ppo, cispo, dppo, dro

Loss availability depends on the base model and its configuration.

`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)

cross_entropy accepts a loss_fn_config key:

bool

Divide the summed per-position losses by the count of active positions (valid targets, plus nonzero weight when you set weights). Set true to skip the division and keep the summed gradient scale; forward_backward_custom() sets it automatically. The trainer rejects an optimizer step that combines normalized and unnormalized gradients.

See the quickstart 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"])

loss_type_input 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",  # or "importance_sampling"
).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. Takes no loss_fn_config.

`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.

float

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

float

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.

float

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

float

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.

float

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

float

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

float

Scale of advantages inside the policy gradient.

float

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.

float

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

`dpo`

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

# 1. Snapshot the reference policy. This blocks until the snapshot is
#    published and returns a SamplingClient directly.
reference = training_client.save_weights_and_get_sampling_client(name="dpo-ref")

# 2. Build interleaved chosen/rejected datums with target_tokens and
#    weights set as above, then score them once.
datums = attach_reference_logprobs(datums, reference)

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

loss_fn_config keys:

float

KL-penalty strength against the reference.

float

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

Development

Deployment

Engines

Model Labs

Organization

Hosting options

Observability

Troubleshooting

`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

​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

`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