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 create or resume 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. The Session remains as the experiment record that connects Training Runs, metrics, sampling activity, and usage. See Core Concepts for the resource model.

from weaver import ServiceClient

with ServiceClient(
    organization="research",
    project="alignment",
    name="PPO baseline",
    labels={"dataset": "math", "environment": "staging"},
) 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	Resume an existing Session. Constructor metadata does not overwrite it.
name	Display name for a new Session (nex-weaver 1.11+).
labels	Searchable string labels for a new Session (nex-weaver 1.11+).
organization / project	Select scope by ID, slug, or display name.
organization_id / project_id	Select scope with canonical IDs.
default_tags	Legacy categorical tags attached to the Session.
user_metadata	Arbitrary Session metadata; unlike labels, it is not used for Console filtering.

Use name for a readable experiment title and labels for stable dimensions such as dataset, environment, or recipe. The Console can search names and IDs, combine exact label filters, and compare metrics from up to six Sessions.

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

print(training_client.training_run_id)

training_run_id is the canonical identifier shown in the Console. model_id remains an equivalent compatibility alias.

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	Arbitrary metadata attached to this Training Run.

Tokenizer

tokenizer = training_client.get_tokenizer()

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

get_tokenizer() returns the tokenizer that matches the selected 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 and accumulates gradients 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,
)

Operation Handles

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.

Async Clients

For an asyncio application, use the native async client stack instead of running the synchronous client in a thread:

from weaver import AsyncServiceClient

async with AsyncServiceClient(
    name="async-sft",
    labels={"pipeline": "online"},
) as service_client:
    training_client = await service_client.create_model(
        base_model="Qwen/Qwen3-8B",
    )
    handle = await training_client.forward_backward(
        datums, "cross_entropy", wait=False,
    )
    result = await handle.result()

AsyncServiceClient, AsyncTrainingClient, and AsyncSamplingClient mirror the synchronous APIs. Their network methods are awaited, and context exit terminates models created by that client.

User Metrics

Trainer loss and optimizer metrics are captured automatically. Use log_metrics() for application-side values such as reward or evaluation accuracy:

training_client.log_metrics(
    {"eval/accuracy": 0.82, "reward/mean": 1.35},
    step=100,
    labels={"split": "validation"},
)

These scalar series appear on Session details and on the nested Training Run detail reached from a Session. They can also be compared across Sessions.

Inspecting Training Runs

ServiceClient.list_training_runs() and get_training_run() expose the model-backed run records shown inside a Session, including checkpoint details. Training Runs are not a separate top-level Console workflow. list_models() and get_model() remain available as compatibility APIs.

Streaming Packed SFT

For large packed datasets, weaver.training_pipeline provides TokenBudgetBatcher for bounded-memory request construction and SubmitAheadQueue for limited asynchronous submit-ahead. Set dp_size and max_tokens_per_gpu to match your training configuration. See the SDK's streaming_sft.py example before using it.

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.

Pause and Resume Generation

Advanced rollout controllers can temporarily pause a sampling engine and resume it after draining requests or synchronizing weights:

sampling_client.pause_generation(mode="abort")
try:
    # Drain requests or synchronize weights.
    ...
finally:
    sampling_client.continue_generation()

Supported pause modes are abort, retract, and in_place. This freezes the sampling engine behind that Sampling Session, so always resume it in finally; ordinary evaluation scripts do not need these methods.

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 Catalog: discover models, prices, modes, and context limits.