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.

`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 takes no loss_fn_config. 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.

Keep loss inputs one-dimensional

Keep target_tokens and weights one-dimensional, one value per token position, with shape [N] where N is the number of tokens in model_input. Loops reads the flat TensorData.data buffer and ignores any second dimension, so it does not support the two-dimensional (N, K) form Tinker uses for top-K soft targets.

`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: importance_sampling doesn’t clip the ratio; ppo clips it; dppo, cispo, and dro add further variants. Drive them from a rollout loop that samples from a paired SamplingClient and scores each trajectory before calling forward_backward.

The trainer accepts dppo and dpo, but they aren’t in the SDK’s LossFnType literal yet, so type-checked code needs # type: ignore[arg-type] on those calls.

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

beta (float, default 0.1): KL-penalty strength against the reference.
label_smoothing (float, default 0.0, range [0, 0.5]): enables conservative DPO (cDPO).

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.

Overview

Get started

Model APIs

Inference

Development

Deployment

Engines

Model Labs

Organization

Observability

Troubleshooting

`cross_entropy`

Weight individual positions

Keep loss inputs one-dimensional

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

`dpo`

Next steps

​cross_entropy

​Weight individual positions

​Keep loss inputs one-dimensional

​importance_sampling, ppo, dppo, cispo, dro

​dpo

​Next steps

`cross_entropy`

Weight individual positions

Keep loss inputs one-dimensional

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

`dpo`

Next steps