Write a custom loss function

​

When to use a custom loss function

• Preference objectives: DPO, or another objective over chosen and rejected pairs. For a complete worked example, see Direct preference optimization.
• Regularized SFT: Cross-entropy plus an extra term, such as a penalty for drifting from the starting model.
• New RL objectives: A policy-gradient variant that the built-in loss functions don’t cover.
• Research prototypes: Any objective you want to try without a change to the service.

• Built-in loss function: The service ships cross-entropy for SFT and several policy-gradient loss functions for RL, listed in Built-in loss functions. The policy-gradient ones already cover clipped objectives such as PPO and CISPO in a single forward_backward call, and the cookbook’s GRPO loop can switch to any of them. Use one if it fits your objective. With a custom loss function, each step runs the model’s forward pass twice, once in forward and again in custom_forward_backward, so it takes longer than it would with a built-in loss function on the same batch.
• Reward function: If your objective is a score per completion rather than a function of the model’s log-probabilities, use GRPO with a reward function instead. It’s simpler and runs on-policy.

​

How a custom loss function runs

1. forward scores the batch on the model resource’s trainer without accumulating gradients and returns the model’s log-probability of each target token.
2. Your local process runs your loss function on those log-probabilities in PyTorch, then calls .backward() on the loss to get its gradient with respect to each log-probability.
3. custom_forward_backward sends those gradients to the model resource’s trainer, which runs the forward pass again and backpropagates them through the model to accumulate parameter gradients.
4. optim_step applies the accumulated gradients, as in every other loop.

​

Write a loss function

import torch

from post_training_cookbook.custom_loss import (
    compute_custom_forward_backward,
    compute_logprobs,
)
from post_training_cookbook.samples import extract_loss_mask
from together.post_training import AdamParams

# 1. Snapshot the reference logprobs once, with no gradient.
ref = [torch.tensor(r) for r in compute_logprobs(session, batch)]
masks = [
    torch.tensor(extract_loss_mask(s), dtype=torch.float32) for s in batch
]


# 2. Define the loss function over the live (autograd) logprobs.
beta = 0.1


def regularized_sft_loss(logprobs: list[torch.Tensor]) -> torch.Tensor:
    return torch.stack(
        [
            (-lp * m).sum() / m.sum()
            + beta * ((lp - r) ** 2 * m).sum() / m.sum()
            for lp, r, m in zip(logprobs, ref, masks)
        ]
    ).mean()


# 3. One training step.
result = compute_custom_forward_backward(session, batch, regularized_sft_loss)
session.trainer.optim_step(
    adam_params=AdamParams(learning_rate=1e-6, weight_decay=0.1),
)
print(result.loss)

• Reduce over the trained positions only: forward returns a real log-probability only where the sample’s weights entry is non-zero, and 0.0 at every other position. The service ignores gradients at those zero-weight positions, but the zeros still count in your loss if you average over the whole row. Divide by the mask’s sum, as the example does with m.sum().
• Return a scalar tensor: compute_custom_forward_backward calls .backward() on what your function returns and raises a ValueError for anything other than a 0-dimensional tensor.
• Snapshot a frozen reference once: Call compute_logprobs for it before training starts. Recomputing the reference every step repeats the same forward passes.
• Chunk large forward passes: forward runs the whole batch at once. If it runs out of memory or returns no log-probabilities, call compute_logprobs on smaller slices of the batch. The DPO recipe does this with precompute_ref_batch_size.
• Skip generators when the loop never samples: Run a loop that never calls sample, such as SFT or DPO, on a trainer-only model resource with num_generator_replicas=0, which costs less.

​

Use a custom loss function in an RL loop

1. Build one sample per completion with the cookbook’s build_cross_entropy_sample. Pass the prompt and completion as one token sequence, with a loss mask of the same length that’s 0 over the prompt and 1 over the completion. Keep each completion’s advantage in your own code, not in the sample.
2. Run compute_custom_forward_backward with a loss function that reads the advantages, then call optim_step.
3. Call weights_sync so the next sample call uses the updated policy.

import torch

from post_training_cookbook.custom_loss import compute_custom_forward_backward
from post_training_cookbook.samples import extract_loss_mask
from together.post_training import AdamParams

# batch has one sample per completion, and advantages one float per sample.
masks = [
    torch.tensor(extract_loss_mask(s), dtype=torch.float32) for s in batch
]


def policy_gradient_loss(logprobs: list[torch.Tensor]) -> torch.Tensor:
    return torch.stack(
        [
            -(adv * lp * m).sum() / m.sum()
            for lp, m, adv in zip(logprobs, masks, advantages)
        ]
    ).mean()


compute_custom_forward_backward(session, batch, policy_gradient_loss)
session.trainer.optim_step(adam_params=AdamParams(learning_rate=1e-6))
session.trainer.weights_sync()

​

Call the SDK without the cookbook

import torch

from together.post_training import LossConfig

# 1. Read the logprobs without accumulating gradients.
scored = session.trainer.forward(
    samples=batch,
    loss=LossConfig(type="LOSS_TYPE_CROSS_ENTROPY"),
)
logprobs = [
    torch.tensor(out.tensors["logprobs"].data, requires_grad=True)
    for out in scored.loss_fn_outputs
]

# 2. Compute the loss and its gradient for every logprob.
loss = regularized_sft_loss(logprobs)
loss.backward()

# 3. Send the gradients back.
session.trainer.custom_forward_backward(
    samples=batch,
    gradients=[{"data": lp.grad.tolist()} for lp in logprobs],
)

​

Next steps