Loss Functions

Weaver selects server-side loss functions with the loss_fn string and passes token-aligned tensors through Datum.loss_fn_inputs. Each loss has its own input contract, but the core pattern is consistent: model_input provides context, target_tokens provides next-token targets, and other tensors control token-level weights or policy gradients.

Common Losses

Loss	Typical use	Main inputs
cross_entropy	SFT, distillation, masked next-token training	target_tokens, weights
forward_logprob	Forward-only logprob collection	target_tokens, optional sampling_mask
importance_sampling	RL policy gradient and sampling correction	target_tokens, logprobs, advantages, loss_mask
grpo	GRPO-style training	Similar to importance_sampling, with group advantages produced by the trainer
ppo_clip	PPO clipped surrogate	target_tokens, logprobs, advantages, loss_mask
truncated_importance_sampling	IS with truncation	Similar to importance_sampling
opd_kl_importance_sampling	OPD/IS with teacher KL	old_logprobs, teacher_logprobs, and related fields
surrogate	Internal backward path for SDK custom losses	surrogate_weights, loss_mask

TIP

Most users start with cross_entropy and importance_sampling. Advanced PPO/GRPO/OPD losses are usually produced by NexRL recipes or custom trainers.

Cross Entropy

cross_entropy is the standard supervised fine-tuning loss.

Inputs

Each Datum needs:

• target_tokens: Tensor[int64], next-token targets aligned to model_input.
• weights: Tensor[float32], token-level loss weights.

datum = types.Datum(
    model_input=types.ModelInput.from_ints(tokens[:-1]),
    loss_fn_inputs={
        "target_tokens": torch.tensor(tokens[1:], dtype=torch.int64),
        "weights": torch.tensor(weights[1:], dtype=torch.float32),
    },
)

Formula

Where is the target token, is the token weight, and is the model probability.

Masking

The common SFT mask ignores prompts and trains completions:

weights = [0.0] * len(prompt_tokens) + [1.0] * len(completion_tokens)

Make sure weights[1:] and target_tokens have the same length.

forward_logprob

forward_logprob computes token logprobs without backward propagation. It is useful for likelihood evaluation, perplexity, custom losses, and storing old policy logprobs for RL.

result = training_client.forward(
    datums,
    "forward_logprob",
    wait=True,
)

Unlike SamplingClient.compute_logprobs(), trainer-side forward_logprob is aligned to explicit target_tokens and does not include a leading None placeholder.

Importance Sampling

importance_sampling is designed for policy optimization. It uses rollout-time logprobs and advantages to increase probability for high-advantage tokens and decrease it for low-advantage tokens.

Inputs

Each Datum usually needs:

• target_tokens: target tokens.
• logprobs: rollout or old-policy token logprobs.
• advantages: advantages or reward signals aligned to target_tokens.
• loss_mask: 0/1 mask for tokens participating in RL loss.

Optional fields include ref_logprobs for KL regularization and sampling_mask for structured sampling or router replay.

datum = types.Datum(
    model_input=types.ModelInput.from_ints(input_tokens),
    loss_fn_inputs={
        "target_tokens": torch.tensor(target_tokens, dtype=torch.int64),
        "logprobs": torch.tensor(old_logprobs, dtype=torch.float32),
        "advantages": torch.tensor(advantages, dtype=torch.float32),
        "loss_mask": torch.tensor(loss_mask, dtype=torch.int64),
    },
)

Intuition

Simplified:

where is loss_mask, is advantage, and is the policy ratio correction. The implementation can also apply ratio transforms, clipping, KL regularization, and aggregation options from loss_fn_config.

PPO / GRPO Related Losses

Weaver trainer exposes grpo, ppo_clip, truncated_importance_sampling, and opd_kl_importance_sampling. These are mostly intended for NexRL or custom trainers because they require group rewards, old policy logprobs, reference policy logprobs, teacher logprobs, or extra KL configuration.

Recommendations:

• With NexRL, let the recipe/trainer generate the required Datum fields.
• With a custom trainer, ensure every per-token tensor has the same length as target_tokens.
• Start with a tiny batch and wait=True to inspect loss outputs and metrics.

Custom Losses

Use forward_backward_custom() when built-in losses are not enough:

def ranking_loss(data, logprob_tensors):
    chosen, rejected = logprob_tensors
    loss = -(chosen.sum() - rejected.sum()).sigmoid().log()
    metrics = {"ranking_loss": float(loss.detach())}
    return loss, metrics


result = training_client.forward_backward_custom(datums, ranking_loss)
training_client.optim_step(types.AdamParams(), wait=True)

The SDK runs a forward logprob pass, calls your local PyTorch loss, backpropagates through returned logprob tensors, and sends the gradients back to Weaver through surrogate.

Debugging Tips

• Check that target_tokens, weights, advantages, and loss_mask lengths match.
• Use zero masks or weights for prompt tokens unless you intentionally train on prompts.
• RL logprobs should come from rollout or an old policy, not the already-updated model.
• For a new loss, first run a tiny forward() batch and inspect logprobs and metrics.

Next Steps

• Training and Sampling: plug losses into the training loop.
• Saving and Loading: save durable checkpoints.
• Finetune Models: orchestrate RL data flow with NexRL.