> ## Documentation Index > Fetch the complete documentation index at: https://docs.together.ai/llms.txt > Use this file to discover all available pages before exploring further. # Write a custom loss function > Train with any objective you compute in PyTorch from the model's log-probabilities, for SFT, DPO, or RL. A [custom loss function](/docs/post-training/training-methods#built-in-and-custom-loss-functions) is a training objective that you write and run in your own training loop, instead of one built into the service. The loss function decides what the model learns, so writing your own lets you train toward objectives the built-in ones can't express. Because you write it in ordinary PyTorch, you can also try out new objectives such as one from a recently published paper, without waiting for the service to add it. ## When to use a custom loss function Write a custom loss function when no built-in loss function matches your objective, for example: * **Preference objectives:** DPO, or another objective over chosen and rejected pairs. For a complete worked example, see [Direct preference optimization](/docs/post-training/dpo). * **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. A custom loss function sees only the log-probability that the model assigns to each target token, not the model's full output distribution. You can't use one for an objective that needs the probability of every token in the vocabulary, such as an exact entropy bonus. Before you write one, check whether a simpler option fits: * **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](/docs/post-training/grpo#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](/docs/post-training/grpo#variants). 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](/docs/post-training/reward-functions) instead. It's simpler and runs on-policy. ## How a custom loss function runs Each step of a custom loss function run splits the work between your process and the model resource: 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. The model weights don't change between `forward` and `custom_forward_backward`, so the parameter gradients match what the trainer would compute if it ran your loss function itself. Use `forward` for the first pass, not `forward_backward` with `return_loss_fn_outputs=True`. Even with that flag, `forward_backward` runs a backward pass through its built-in loss function, so those gradients add to yours before `optim_step` applies them. ## Write a loss function The cookbook's `post_training_cookbook.custom_loss` module runs steps 1 through 3 for you, so you write only the loss function. Your function receives a list of `float32` tensors, one per sample in batch order, each holding that sample's per-position log-probabilities. It must return a scalar tensor. Build each sample the same way as for SFT, with `target_tokens` and a `weights` row that marks the positions your loss function reads. Keep any other values it needs, such as reference log-probabilities or advantages, in your own code, and capture them in the function, as the example does with `ref`. This loss function is regularized SFT: cross-entropy plus a penalty for moving the model's log-probabilities away from a snapshot taken before training: ```python theme={null} 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) ``` `compute_logprobs` runs only the `forward` pass, so use it to snapshot a frozen reference before training. `compute_custom_forward_backward` runs steps 1 through 3 and returns the loss your function computed as `result.loss`. The service doesn't compute or report a loss for a custom step, so `result.loss` is the only record of it. Log it yourself to track training. Get these details right: * **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](/docs/post-training/model-resources#generator-replicas) with `num_generator_replicas=0`, which costs less. ## Use a custom loss function in an RL loop In an RL loop, a custom loss function replaces only the `forward_backward` call. Sample completions and score them as usual, then: 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. This loss function implements a plain policy-gradient objective. Replace it with your own: ```python theme={null} 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 The cookbook helpers wrap two SDK calls. To run them yourself, read the log-probabilities with `forward`, then send one list of gradients per sample, in batch order, to `custom_forward_backward`. This example reuses `regularized_sft_loss` from [Write a loss function](#write-a-loss-function): ```python theme={null} 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], ) ``` The `LossConfig` on `forward` only selects which positions return a log-probability, so every sample needs `target_tokens` and `weights`. Each gradient list must have one value per target position, the same length as the sample's `target_tokens`. ## Next steps Follow a complete recipe that uses a custom loss function. Check whether a built-in loss function already fits your objective. Run the RL loop that a custom policy-gradient loss function plugs into. Review how each training method uses a loss function. This documentation is built and hosted on [Mintlify](https://mintlify.com), a developer documentation platform.