Custom Loss Functions

For use cases beyond the built-in loss functions, you can choose a more flexible custom loss function: pass a manually implemented loss function to the forward_backward_custom method to compute the loss and other metrics.

Custom loss functions always run in the local Python process. PyTRIO obtains the current model's per-token logprobs from the server, calls the user-defined function locally, and sends the loss gradients with respect to those logprobs back to the server to complete backpropagation through the model parameters.

Custom loss functions are usually slower, so we recommend using the async method forward_backward_custom_async whenever possible to avoid blocking the training loop. See Async for an introduction to the async methods.

Usage

Define the Loss Function

First, define a loss function with the following signature:

def logprob_squared_loss(
    data: list[trio.Datum],
    logprobs: list[torch.Tensor],
) -> tuple[torch.Tensor, dict[str, float]]:
    ...

Where:

• data is the list of input data. Each element is a Datum object, in the same order as the data passed to forward_backward_custom;
• logprobs is the list of per-token log probabilities from the current model's forward pass. It corresponds one-to-one with data, and each tensor has the same length as its corresponding Datum.model_input;

Prepare Datum and Auxiliary Data

A custom loss may use three types of data:

loss_fn_inputs contains strongly typed tensor inputs for server-side loss functions and only stores fields defined by the corresponding loss schema. With forward_backward_custom, users only need to provide target_tokens with the same length as model_input. Custom keys, scalars, dataclasses, and other Python objects cannot be stored there.

When a custom loss requires additional local data, define a factory function and use a closure to bind that data to the final two-argument loss function. The following example uses a simplified sequence-level objective:

from collections.abc import Callable
from dataclasses import dataclass


@dataclass(frozen=True)
class SequenceLossMeta:
    sampling_logprobs: list[float]
    advantage: float
    completion_tokens: int


def make_sequence_loss_fn(
    metas: list[SequenceLossMeta],
) -> Callable[
    [list[trio.Datum], list[torch.Tensor]],
    tuple[torch.Tensor, dict[str, float]],
]:
    if not metas:
        raise ValueError("metas must not be empty")

    # Preserve the ordering of the current batch.
    batch_metas = tuple(metas)

    def sequence_loss_fn(
        data: list[trio.Datum],
        logprobs: list[torch.Tensor],
    ) -> tuple[torch.Tensor, dict[str, float]]:
        if not (len(data) == len(logprobs) == len(batch_metas)):
            raise ValueError("data, logprobs and metas must have the same length")

        objectives = []
        for meta, current in zip(batch_metas, logprobs, strict=True):
            if meta.completion_tokens <= 0:
                raise ValueError("completion_tokens must be positive")
            if current.numel() < meta.completion_tokens:
                raise ValueError("logprob sequence is shorter than the completion")

            current_completion = current[-meta.completion_tokens :].float()
            sampling = torch.as_tensor(
                meta.sampling_logprobs,
                dtype=current_completion.dtype,
                device=current_completion.device,
            )
            if current_completion.numel() != sampling.numel():
                raise ValueError("sampling logprobs must match completion tokens")

            sequence_ratio = torch.exp((current_completion - sampling).mean())
            objectives.append(sequence_ratio * meta.advantage)

        loss = -torch.stack(objectives).mean()
        return loss, {"sequence_loss": float(loss.detach().item())}

    return sequence_loss_fn

The returned sequence_loss_fn still matches the fixed signature required by PyTRIO, while batch_metas remains in the local closure:

loss_fn = make_sequence_loss_fn(batch_metas)
future = training_client.forward_backward_custom(
    data=batch_data,
    loss_fn=loss_fn,
)
result = future.result()

batch_data[i], batch_metas[i], and the logprobs[i] returned by PyTRIO must all describe the same sample. Create a separate closure for each batch, and compute reference logprobs and similar data before calling forward_backward_custom; do not make network requests from inside the loss function.

For a complete GSPO sequence-level clipping implementation, see 07-gspo/loss.py. Its GSPOMeta and make_gspo_loss_fn use the same data-separation pattern.

Call the Loss Function

Once the loss function is defined, pass it to the forward_backward_custom method:

future = training_client.forward_backward_custom(
    data=data,
    loss_fn=your_custom_loss_fn,
)
result = future.result()

The async version is recommended:

future = await training_client.forward_backward_custom_async(
    data=data,
    loss_fn=your_custom_loss_fn,
)
result = await future

loss = result.get("metrics", {}).get("loss:sum", 0.0)
print(loss)

Example

The following is a simple example: it defines a logprob_squared_loss function that computes the sum of squared per-token log probabilities and optimizes it as the loss.

import asyncio

import swanlab
import torch
from datasets import load_dataset

import pytrio as trio

EPOCHS = 10
BATCH_SIZE = 2


# Custom loss function: sum of squared per-token log probabilities
def logprob_squared_loss(
    _: list[trio.Datum], logprobs: list[torch.Tensor]
) -> tuple[torch.Tensor, dict[str, float]]:  
    flat = torch.cat([x.reshape(-1) for x in logprobs])
    loss = (flat**2).sum()
    return loss, {"logprob_squared_loss": float(loss.detach().item())}


async def main():
    # 1. Connect to TRIO
    service_client = trio.ServiceClient()

    # 2. Create a LoRA training client
    training_client = await service_client.create_lora_training_client_async(
        base_model="Qwen/Qwen3-4B-Instruct-2507",
        seed=42,
        train_mlp=True,
        train_attn=True,
        train_unembed=False,
    )
    tokenizer = training_client.get_tokenizer()

    # 3. Load the dataset and convert it to the training format
    dataset = load_dataset(
        "HuggingFaceTB/smoltalk",
        "everyday-conversations",
        split="train[:10]",
    )

    all_samples: list[trio.Datum] = []

    for example in dataset:
        text = tokenizer.apply_chat_template(example["messages"], tokenize=False)
        tokens = tokenizer.encode(text, add_special_tokens=False)

        input_ids = tokens[:-1]
        target_tokens = tokens[1:]

        all_samples.append(
            trio.Datum(
                model_input=trio.ModelInput.from_ints(input_ids),
                loss_fn_inputs={
                    "target_tokens": target_tokens,
                },
            )
        )

    batches = [
        all_samples[i : i + BATCH_SIZE] for i in range(0, len(all_samples), BATCH_SIZE)
    ]

    # 4. Initialize SwanLab to record training metrics
    swanlab.init(
        project="trio-custom-loss",
        experiment_name="logprob-squared-loss",
    )

    # 5. Training
    for epoch in range(EPOCHS):
        futures = []

        for batch in batches:
            future = await training_client.forward_backward_custom_async(  
                data=batch,
                loss_fn=logprob_squared_loss,  
            )
            tokens_in_batch = sum(len(d.loss_fn_inputs["target_tokens"]) for d in batch)
            futures.append((future, tokens_in_batch))

            await training_client.optim_step_async(
                trio.AdamParams(
                    learning_rate=1e-4,
                    beta1=0.9,
                    beta2=0.999,
                    eps=1e-8,
                    weight_decay=0,
                )
            )

        for future, tokens_in_batch in futures:
            result = await future
            loss_sum = float(result.get("metrics", {}).get("loss:sum", 0.0))
            loss_mean = loss_sum / tokens_in_batch
            swanlab.log({"train/loss": loss_mean})

        print(f"Epoch {epoch} completed")


if __name__ == "__main__":
    asyncio.run(main())

Note: logprob_squared_loss is only an example loss function and does not work well in practice. Do not use it for your own training.

How It Works

The model's forward computation graph lives on the server, while the user-defined loss graph is created in the local Python process. There is no single continuous computation graph across the network, so PyTRIO applies the chain rule in two stages to obtain the derivative of the loss with respect to the model parameters:

$\frac{\partial \text{loss}}{\partial \theta}$

The trainer's default Cross Entropy Loss is computed as follows:

loss_elementwise = -logprobs * weights
loss = loss_elementwise.sum()

In typical learning frameworks, $w_i \in [0, 1]$ represents the loss weight of each position:

• when $w_i = 0$, the loss at that position is not computed;
• when $w_i = 1$, the loss at that position is computed normally;
• when $w_i \in (0, 1)$, the loss at that position is scaled.

The gradient of the loss with respect to the parameters $\theta$ is:

$\frac{\partial \text{loss}}{\partial \theta} = \sum_i \frac{\partial (-\log p_i \cdot w_i)}{\partial \theta}$

If $w_i$ is treated as a constant (i.e., $w_i$ does not depend on $\theta$), then:

$\frac{\partial \text{loss}}{\partial \theta} = \sum_i -w_i \frac{\partial \log p_i}{\partial \theta}$

In the custom-loss flow, the SDK generates surrogate weights from the logprob gradients computed locally. Denote these weights by $\tilde{w}_i$:

$\tilde{w}_i = -\frac{\partial L}{\partial \log p_i} \in \mathbb{R}$

Substituting them into the Cross Entropy form gives:

$\frac{\partial \tilde{L}}{\partial \theta} = \sum_i \frac{\partial L}{\partial \log p_i} \frac{\partial \log p_i}{\partial \theta}$

This is exactly the chain rule:

$\frac{\partial L}{\partial \theta} = \sum_i \frac{\partial L}{\partial \log p_i} \frac{\partial \log p_i}{\partial \theta}$

Here $\frac{\partial \log p_i}{\partial \theta}$ is computed on the server, while $\frac{\partial L}{\partial \log p_i}$ is computed by local PyTorch autograd and sent back to the server.

In other words, this is equivalent to constructing a surrogate objective that is linear in the logprobs:

$\tilde{L}(\theta) = \sum_i \log p_i \cdot \frac{\partial L}{\partial \log p_i}$

The server treats $\frac{\partial L}{\partial \log p_i}$ as a constant. Although this surrogate objective differs in form from the original loss, it produces exactly equivalent model-parameter gradients for this forward-backward pass.

The surrogate weights $\tilde{w}_i$ generated internally by the SDK are not restricted to $[0, 1]$ and may be negative. They only carry the locally computed logprob gradients back to the server and cannot be used to store additional user metadata.

Execution Flow

forward_backward_custom completes gradient computation in two phases between the client and the server:

1. Prepare data: the client constructs a list of Datum objects and prepares the target tokens. Other local data required by the algorithm can be bound to the loss-function closure.
2. Forward computation: the server performs one forward pass and computes the target-token logprobs.
3. Compute the custom loss on the client: the client reconstructs the returned values as differentiable PyTorch tensors and calls the user-defined loss_fn(data, logprobs). Additional data captured by the closure also participates in this step.
4. Backpropagate to logprobs on the client: the client backpropagates through this loss to obtain $\frac{\partial L}{\partial \text{logprobs}}$, the gradient of each logprob with respect to the final loss.
5. Run surrogate forward-backward on the server: the server uses these gradients as weights to construct the surrogate loss and runs forward-backward on it, producing parameter gradients exactly equivalent to those from the original custom loss.

Why the Custom Function Does Not Need to Be Uploaded

In this design, the server only needs to:

• compute the target-token logprobs;
• receive $\frac{\partial L}{\partial \text{logprobs}}$ from the client;
• run standard gradient computation on the surrogate objective.

Therefore, the user-defined Python function always stays on the client. PyTRIO does not pickle it or send it to the server.