Loop building blocks

Every cookbook on this site links back to these building blocks instead of a repository file. Copy what you need — none of it depends on a provider SDK, only the standard library.

Datum helpers

A “datum” is one Whitney training example: right-shifted input/target tokens plus whatever loss-specific tensors your loss_function needs. The completion’s last token only ever appears as a target, never as an input, so input and target_tokens are one token shorter than prompt_tokens + completion_tokens.

Tensor = dict[str, object]
Datum = dict[str, object]


def tensor(dtype: str, values: list[int] | list[float]) -> Tensor:
    """Build one Whitney dense Tensor: {dtype, shape, values}."""
    return {"dtype": dtype, "shape": [len(values)], "values": list(values)}


def _rightshift(prompt_tokens, completion_tokens):
    full = list(prompt_tokens) + list(completion_tokens)
    if len(full) < 2:
        raise ValueError("prompt_tokens + completion_tokens must contain at least two tokens")
    return full[:-1], full[1:], len(prompt_tokens) - 1


def _aligned(pad_len: int, values: list[float]) -> list[float]:
    return [0.0] * pad_len + list(values)


def sft_datum(prompt_tokens: list[int], completion_tokens: list[int]) -> Datum:
    """cross_entropy datum: supervision only over completion_tokens."""
    input_tokens, target_tokens, pad_len = _rightshift(prompt_tokens, completion_tokens)
    weights = _aligned(pad_len, [1.0] * len(completion_tokens))
    return {
        "input": {"token_ids": input_tokens},
        "loss_inputs": {
            "target_tokens": tensor("int64", target_tokens),
            "weights": tensor("float32", weights),
        },
    }


def rl_datum(
    prompt_tokens: list[int],
    sampled_tokens: list[int],
    sampled_logprobs: list[float],
    advantage: float | list[float],
    *,
    clip_low: float | None = None,
    clip_high: float | None = None,
) -> Datum:
    """importance_sampling-family datum: target_tokens/weights/logprobs/advantages,
    all aligned to the same right-shifted, prompt-masked length.

    advantage is either one scalar broadcast across every completion token
    (standard GRPO-family use) or a pre-computed per-token list of the same
    length as sampled_tokens (see kl_adjust_advantages below). clip_low/
    clip_high attach clip_low_threshold/clip_high_threshold tensors — a
    Tinker-Cloud-only mechanism. Pass None for both on Modal; use scalar
    loss_config keys there instead (see each algorithm's section).
    """
    if len(sampled_logprobs) != len(sampled_tokens):
        raise ValueError("sampled_logprobs must have one entry per sampled_tokens element")
    input_tokens, target_tokens, pad_len = _rightshift(prompt_tokens, sampled_tokens)
    completion_len = len(sampled_tokens)
    if isinstance(advantage, (int, float)):
        per_token_advantage = [float(advantage)] * completion_len
    else:
        per_token_advantage = [float(value) for value in advantage]
    loss_inputs: dict[str, Tensor] = {
        "target_tokens": tensor("int64", target_tokens),
        "weights": tensor("float32", _aligned(pad_len, [1.0] * completion_len)),
        "logprobs": tensor("float32", _aligned(pad_len, sampled_logprobs)),
        "advantages": tensor("float32", _aligned(pad_len, per_token_advantage)),
    }
    if clip_low is not None:
        loss_inputs["clip_low_threshold"] = tensor("float32", _aligned(pad_len, [clip_low] * completion_len))
    if clip_high is not None:
        loss_inputs["clip_high_threshold"] = tensor("float32", _aligned(pad_len, [clip_high] * completion_len))
    return {"input": {"token_ids": input_tokens}, "loss_inputs": loss_inputs}

See JSON primitives for the wire shape these produce.

Advantage helpers

def group_mean_advantages(rewards: list[float], *, normalize_std: bool = False) -> list[float]:
    """Group-relative advantage: reward minus the group mean.

    Standard GRPO additionally divides by the group standard deviation
    (normalize_std=True); the no-std variant is the raw centered value.
    """
    if len(rewards) < 2:
        raise ValueError("group_mean_advantages needs at least two rewards to form a group")
    mean_reward = sum(rewards) / len(rewards)
    centered = [reward - mean_reward for reward in rewards]
    if not normalize_std:
        return centered
    variance = sum(value * value for value in centered) / len(centered)
    std = variance**0.5
    if std < 1e-6:
        return centered
    return [value / std for value in centered]


def all_zero(advantages: list[float]) -> bool:
    """True when a group produced zero gradient signal (uniform reward)."""
    return all(value == 0 for value in advantages)


def leave_one_out_baseline(rewards: list[float]) -> list[float]:
    """RLOO: each trajectory's advantage is its reward minus the mean of the
    other trajectories in its group."""
    count = len(rewards)
    if count < 2:
        raise ValueError("leave_one_out_baseline needs at least two rewards to form a group")
    total = sum(rewards)
    return [reward - (total - reward) / (count - 1) for reward in rewards]


def length_normalize(advantage: float, completion_length: int) -> float:
    """Scale a scalar advantage by 1 / len(completion)."""
    if completion_length <= 0:
        raise ValueError("completion_length must be positive")
    return advantage / completion_length


def kl_adjust_advantages(
    advantages: list[float],
    sampled_logprobs: list[float],
    reference_logprobs: list[float],
    *,
    kl_coef: float,
) -> list[float]:
    """Subtract a per-token KL-vs-reference penalty from per-token advantages.
    Used by REINFORCE++'s K2/reference-KL term — see reference_logprobs from
    a frozen second run in [Distillation](/distillation)."""
    return [
        advantage - kl_coef * (sampled_logprob - reference_logprob)
        for advantage, sampled_logprob, reference_logprob in zip(
            advantages, sampled_logprobs, reference_logprobs, strict=True
        )
    ]

Building a forward payload

def forward_backward_payload(
    data: list[Datum],
    loss_function: str,
    loss_config: dict[str, float] | None = None,
) -> dict:
    """Assemble a Whitney ForwardPayload: {data, loss_function, loss_config}."""
    if not data:
        raise ValueError("forward_backward_payload needs at least one datum")
    return {"data": data, "loss_function": loss_function, "loss_config": dict(loss_config or {})}

The build_loss contract

Most RL loops in the cookbooks are a single callable plugged into the same loop skeleton below:

build_loss(sample, cycle, request, ctx) -> dict

sample — the raw sampler result: {"sequences": [{"tokens", "logprobs", "stop_reason"}, ...], ...}.
cycle — the zero-based cycle index.
request — your original sample request object, needed for request["prompt"]["token_ids"] (the sampler result alone carries no prompt tokens).
ctx — a context object bundling a live client, run_id, sampler_id, and provider, so your callable can issue further ordered operations before returning its payload — e.g. an extra logprobs call for a reference-policy KL term, or another sample call for dynamic resampling. ctx.client.sampler_operation(ctx.run_id, ctx.sampler_id, kind, payload) is the same method the loop itself uses for sample. Any such call still advances the run’s ordered seq_id — it is a real, billed Whitney operation, not a side channel.

Check ctx.provider against the provider loss support table before using a provider-specific loss name.

Ordered RL loop skeleton

for cycle in range(cycles):
    sampler = client.run_operation(run_id, "save_weights_for_sampler", {})
    sampler_id = result_sampler_id(sampler)
    sample = client.sampler_operation(run_id, sampler_id, "sample", sample_request)

    loss_payload = build_loss(sample, cycle, sample_request, ctx)
    backward = client.run_operation(run_id, "forward_backward", loss_payload)
    optimized = client.run_operation(
        run_id, "optim_step", {"optimizer": adamw_optimizer(learning_rate=1e-6)}
    )

resume = client.run_operation(run_id, "save_state", {})
artifact = client.run_operation(run_id, "export_lora", {})
client.finish(run_id)
client.close_session(session_id)

On failure or interruption: cancel the run, poll until terminal, then close the session. See SFT for the equivalent single-pass loop without sampling, and Distillation for the two-run variant where a frozen teacher run answers logprobs instead of training.

Provider loss support

GET /v1/training/capabilities does not return a loss-function list — it cannot, because Whitney forwards loss_function straight through to the provider. The table below is a caller-side hint that saves a wasted allocation, not a live capability; call require_advertised_loss before you build a payload with a provider-specific loss name.

Loss	Modal	Tinker
`cross_entropy`	✅	✅
`importance_sampling`	✅	✅
`ppo`	✅	✅
`cispo`	✅	✅
`dro`		✅
`gspo`	✅
`ppo_critic`	✅
`dppo`	✅

Scalar loss_config keys Modal actually validates, per loss (an empty row means that loss accepts no loss_config keys at all):

Loss	Keys
`cross_entropy`	—
`importance_sampling`	—
`ppo`	`clip_low_threshold`, `clip_high_threshold`, `value_clip`
`gspo`	`clip_low_threshold`, `clip_high_threshold`
`cispo`	`clip_low_threshold`, `clip_high_threshold`
`ppo_critic`	`value_clip`
`dppo`	`delta_low`, `delta_high`

On Tinker Cloud, the equivalent clip range is set per-token on the datum itself (rl_datum(..., clip_low=, clip_high=)), not through loss_config.

Requiring an advertised loss

def require_advertised_loss(provider: str, loss_function: str, support: dict[str, frozenset[str]]) -> None:
    """Fail closed before allocation when a loss name isn't one Whitney forwards for this provider."""
    supported = support.get(provider)
    if supported is None:
        raise ValueError(f"Unknown provider {provider!r}")
    if loss_function not in supported:
        raise ValueError(
            f"{provider} does not advertise loss_function={loss_function!r} "
            f"(known losses: {', '.join(sorted(supported))})"
        )

Get started

How it works

Cookbooks

Distillation

Coding agents

Loop building blocks

Datum helpers

Advantage helpers

Building a forward payload

The build_loss contract

Ordered RL loop skeleton

Provider loss support

Requiring an advertised loss

​Datum helpers

​Advantage helpers

​Building a forward payload

​The build_loss contract

​Ordered RL loop skeleton

​Provider loss support

​Requiring an advertised loss

Datum helpers

Advantage helpers

Building a forward payload

The build_loss contract

Ordered RL loop skeleton

Provider loss support

Requiring an advertised loss