Skip to main content
The Loops trainer accepts several loss functions through 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.

​
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)
All five 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. All are typically driven by 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. There are three additional pieces compared to SFT:
  1. Snapshot the reference before any optim_step. The initial weights are the reference for offline DPO.
  2. 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.
  3. Score reference logprobs once with attach_reference_logprobs and reuse across every step.
Each Datum needs three loss_fn_inputs fields. You build the first two yourself when tokenizing; the third 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