Training and Sampling

This page covers the core Weaver SDK workflow: creating clients, preparing Datum objects, training, exporting weights, sampling, and computing logprobs.

Workflow

A typical workflow looks like this:

1. Create a ServiceClient to connect to Weaver and manage a session.
2. Call create_model() to create a TrainingClient.
3. Convert examples into types.Datum.
4. Call forward_backward() to accumulate gradients.
5. Call optim_step() to update model parameters.
6. Export sampler weights and use SamplingClient.sample() for evaluation or rollout.

Creating Clients

ServiceClient

ServiceClient is the SDK entry point. It creates or reuses a session and cleans up models created in the context manager when it exits.

from weaver import ServiceClient

with ServiceClient() as service_client:
    print(service_client.session_id)

Common parameters:

Parameter	Description
api_key	Weaver API key. Defaults to WEAVER_API_KEY.
base_url	Weaver service URL. Defaults to the SDK endpoint.
session_id	Reuse an existing session.
default_tags	Tags written to sessions for experiment tracking.

TrainingClient

Create a training model:

from weaver import types

training_client = service_client.create_model(
    base_model="Qwen/Qwen3-8B",
    training_mode="lora",
    lora_config=types.LoraConfig(rank=32, seed=42),
)

Common parameters:

Parameter	Description
base_model	Base model name, such as Qwen/Qwen3-8B. Long-context variants can use :<max_seq_len>, for example Qwen/Qwen3-8B:262144.
training_mode	lora or full_ft. If omitted, the server defaults to LoRA.
lora_config	LoRA configuration. Default rank is 32 with attention, MLP, and unembedding enabled.
performance_tier	Optional throughput tier, such as normal, fast, or flash. Higher tiers usually mean higher throughput and higher cost.
user_metadata	Experiment metadata passed to the server.

Tokenizer

tokenizer = training_client.get_tokenizer()

tokens = tokenizer.encode("Hello, world!", add_special_tokens=True)
text = tokenizer.decode(tokens)

If the server returns a tokenizer_path, the SDK uses it first. Otherwise it uses base_model.

Preparing Training Data

Weaver represents a training example as types.Datum. Each Datum contains:

• model_input: model input tokens.
• loss_fn_inputs: tensors required by the loss function.
• metadata: optional metadata, such as router replay information.

SFT Format

Cross-entropy training usually needs target_tokens and weights:

import torch
from weaver import types


def process_example(prompt, completion, tokenizer):
    prompt_tokens = tokenizer.encode(prompt, add_special_tokens=True)
    completion_tokens = tokenizer.encode(completion, add_special_tokens=False)
    tokens = prompt_tokens + completion_tokens

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

    return 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),
        },
    )

weights controls token-level loss contribution:

• 0.0: ignore the token, such as prompt tokens.
• 1.0: include normally, such as completion tokens.
• Other non-negative values: scale the token contribution.

Training APIs

forward()

Runs a forward pass without accumulating gradients. Use it for evaluation, logprob collection, or the first step of custom losses.

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

forward_backward()

Runs forward and backward, accumulating gradients on the training backend until the next optim_step().

result = training_client.forward_backward(
    datums,
    "cross_entropy",
    wait=True,
)

Parameters:

Parameter	Description
data	Sequence[types.Datum].
loss_fn	Loss name, such as cross_entropy or importance_sampling.
loss_fn_config	Optional loss configuration.
metadata	Optional request-level metadata.
wait	True blocks for a result. False returns an OperationHandle.

Results usually contain per-example outputs and aggregate metrics:

{
    "result": {
        "loss_fn_outputs": [...],
        "metrics": {
            "loss": 0.5
        }
    }
}

optim_step()

Apply accumulated gradients with Adam:

training_client.optim_step(
    types.AdamParams(learning_rate=1e-4),
    wait=True,
)

AdamParams defaults:

types.AdamParams(
    learning_rate=1e-4,
    beta1=0.9,
    beta2=0.95,
    eps=1e-8,
    weight_decay=1e-2,
    grad_clip_norm=1.0,
)

Async Operations

Most training calls support wait=False:

handle = training_client.forward_backward(
    datums,
    "cross_entropy",
    wait=False,
)

result = handle.result()

This is useful when data preparation, rollout, or evaluation can overlap with remote training.

Custom Loss

forward_backward_custom() is useful for research losses. The SDK runs forward(..., "forward_logprob"), converts returned logprobs into differentiable PyTorch tensors, lets your function compute a scalar loss, then sends the logprob gradients back through Weaver's surrogate backward path.

def my_loss_fn(data, logprob_tensors):
    loss = -sum(t.mean() for t in logprob_tensors)
    metrics = {"custom_loss": float(loss.detach())}
    return loss, metrics


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

Sampling

After training, export sampler weights and create a SamplingClient:

sampling_client = training_client.save_weights_and_get_sampling_client(
    name="my-model-step-100",
)

Then call sample():

prompt_tokens = tokenizer.encode("Hello, ", add_special_tokens=True)

result = sampling_client.sample(
    prompt=types.ModelInput.from_ints(prompt_tokens),
    sampling_params=types.SamplingParams(
        max_tokens=50,
        temperature=0.8,
        top_p=0.95,
        stop=["\n"],
    ),
    num_samples=1,
)

print(result["sequences"][0]["text"])

SamplingParams

types.SamplingParams(
    max_tokens=100,
    temperature=1.0,
    top_p=1.0,
    top_k=-1,
    stop=["\n", 151645],
    seed=42,
)

sample() also supports include_prompt_logprobs, topk_prompt_logprobs, return_sampling_mask, return_old_logprob, and return_moe_topk_indices. These are useful for RL rollout, PPO/GRPO, and MoE router replay.

Computing Logprobs

tokens = tokenizer.encode("Hello, world!", add_special_tokens=True)

logprobs = sampling_client.compute_logprobs(
    prompt=types.ModelInput.from_ints(tokens),
)

print(logprobs)

The returned list is aligned to prompt tokens. Its length equals the number of prompt tokens, and the first item is None because the first token has no previous context.

Next Steps

• Loss Functions: learn SFT, RL, and custom losses.
• Saving and Loading: save sampler weights and durable checkpoints.
• Model Lineup: choose training models, modes, and context length.