训练与采样

本页介绍 Weaver 的核心 SDK 工作流：创建 client、准备 Datum、执行训练、导出权重、采样与计算 logprobs。

工作流概览

典型流程如下：

创建 ServiceClient，连接 Weaver 服务并创建或复用训练会话。
调用 create_model() 创建 TrainingClient。
把样本转换为 types.Datum。
调用 forward_backward() 累积梯度。
调用 optim_step() 更新模型参数。
导出 sampler 权重，并通过 SamplingClient.sample() 做评估或 rollout。

创建 Client

ServiceClient

ServiceClient 是 SDK 的入口。它会创建或复用训练会话，并在上下文管理器退出时清理本次创建的模型实例。训练会话会作为实验记录继续保留，将训练运行、指标、采样活动和用量关联起来。资源关系见核心概念。

python

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)

常用参数：

参数	说明
`api_key`	Weaver API 密钥；不传时读取 `WEAVER_API_KEY`。
`base_url`	Weaver 服务地址；不传时使用 SDK 默认地址。
`session_id`	复用已有训练会话；constructor 元信息不会覆盖原记录。
`name`	新训练会话的展示名称（`nex-weaver` 1.11+）。
`labels`	新训练会话的可搜索字符串 labels（`nex-weaver` 1.11+）。
`organization` / `project`	按 ID、slug 或展示名称选择范围。
`organization_id` / `project_id`	使用规范 ID 选择范围。
`default_tags`	附加到训练会话的旧版分类 tags。
`user_metadata`	任意训练会话元信息；与 `labels` 不同，不用于 Console 筛选。

建议用 name 表示易读的实验名称，用 labels 记录 dataset、environment、recipe 等稳定维度。Console 支持搜索名称和 ID、组合多个精确 label 筛选，并比较最多六个训练会话的指标。

TrainingClient

创建训练模型：

python

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 是 Console 中使用的规范标识；model_id 仍作为等价的兼容别名保留。

常用参数：

参数	说明
`base_model`	基座模型名称，例如 `Qwen/Qwen3-8B`。长上下文变体可使用 `:<max_seq_len>` 后缀，例如 `Qwen/Qwen3-8B:262144`。
`training_mode`	训练模式；可传 `lora` 或 `full_ft`。不传时服务端默认使用 LoRA。
`lora_config`	LoRA 配置；默认 rank 为 32，并训练 attention、MLP 和 unembedding。
`performance_tier`	可选吞吐档位，例如 `normal`、`fast`、`flash`；更高档位通常意味着更高吞吐和更高成本。
`user_metadata`	附加到此训练运行的任意元信息。

Tokenizer

python

tokenizer = training_client.get_tokenizer()

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

get_tokenizer() 会返回与所选模型匹配的 tokenizer。

准备训练数据

Weaver 使用 types.Datum 表示单条训练样本。一个 Datum 包含：

model_input：模型输入 tokens。
loss_fn_inputs：loss 函数需要的额外张量。
metadata：可选元信息，例如 router replay 信息。

SFT 数据格式

交叉熵训练通常需要 target_tokens 和 weights：

python

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 用来控制哪些 token 参与 loss：

0.0：忽略，例如 prompt token。
1.0：正常参与，例如 completion token。
其他非负值：按比例放大或缩小该 token 的贡献。

训练 API

forward()

只执行前向计算，不累积梯度。它常用于评估 loss、获取 logprobs，或作为自定义 loss 的第一步。

python

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

forward_backward()

执行前向和反向，并累积梯度，等待后续 optim_step() 更新参数。

python

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

参数说明：

参数	说明
`data`	`Sequence[types.Datum]`。
`loss_fn`	loss 名称，例如 `cross_entropy`、`importance_sampling`。
`loss_fn_config`	可选 loss 配置。
`metadata`	可选请求级元信息。
`wait`	`True` 时阻塞等待结果；`False` 时返回 `OperationHandle`。

返回结果通常包含每条样本输出和聚合指标：

python

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

optim_step()

用累积梯度执行优化器更新：

python

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

AdamParams 默认值：

python

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

异步操作（OperationHandle）

大多数训练操作支持 wait=False：

python

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

result = handle.result()

这适合把数据准备、rollout 或评估与远端训练并行起来。

原生异步 Client

python

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 和 AsyncSamplingClient 与同步 API 对应。网络方法需要 await，退出 context 时也会终止该 client 创建的模型。

用户指标

Trainer 产生的 loss 和 optimizer 指标会自动记录。reward、评估准确率等应用侧指标使用 log_metrics()：

python

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

这些标量会显示在训练会话详情，以及从训练会话进入的训练运行详情中，也可以跨训练会话比较。

查询训练运行

ServiceClient.list_training_runs() 和 get_training_run() 返回训练会话内的模型记录及检查点详情。训练运行不是独立的 Console 顶层工作流。list_models() 和 get_model() 作为兼容 API 继续保留。

流式 Packed SFT

处理大型 packed 数据集时，weaver.training_pipeline 提供 TokenBudgetBatcher，用于有界内存的请求构造；SubmitAheadQueue 用于受控的异步 submit-ahead。dp_size 和 max_tokens_per_gpu 需要与训练配置一致。使用前请参考 SDK 的 streaming_sft.py 示例。

自定义 Loss

forward_backward_custom() 适合研究型 loss。它会先调用 forward(..., "forward_logprob") 取回 logprobs，然后让你的 Python 函数在本地计算标量 loss 并反传，最后把 logprob 梯度作为 surrogate backward 发送给 Weaver。

python

def my_loss_fn(data, logprob_tensors):
    # logprob_tensors 是带 requires_grad=True 的 torch.Tensor 列表。
    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)

采样

训练后，先导出 sampler 权重并创建 SamplingClient：

python

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

随后调用 sample()：

python

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

python

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

参数	说明
`max_tokens`	最多生成 token 数；不传时使用模型默认值。
`temperature`	温度；`0.0` 更确定，较高值更发散。
`top_p`	nucleus sampling 参数。
`top_k`	top-k sampling；`-1` 表示不额外覆盖默认行为。
`stop`	停止条件，可混合字符串和 token id。
`seed` / `sampling_seed`	采样随机种子。

sample() 还支持：

include_prompt_logprobs
topk_prompt_logprobs
return_sampling_mask
return_old_logprob
return_moe_topk_indices

这些选项常用于 RL rollout、PPO/GRPO 训练或 MoE router replay。

暂停与恢复生成

高级 rollout controller 可以暂时暂停 sampling engine，在排空请求或同步权重后恢复：

python

sampling_client.pause_generation(mode="abort")
try:
    # 排空请求或同步权重。
    ...
finally:
    sampling_client.continue_generation()

暂停模式包括 abort、retract 和 in_place。该操作会冻结此采样会话背后的 sampling engine，因此必须在 finally 中恢复；普通评估脚本不需要使用这些方法。

计算 Logprobs

python

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

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

print(logprobs)

返回值与 prompt token 对齐：长度等于 prompt token 数量，第一项为 None，因为第一个 token 没有前文条件概率。

下一步

损失函数：了解 SFT、RL 和自定义 loss。
保存与加载：保存 sampler 权重和持久检查点。
支持的模型：查询模型、价格、训练模式和上下文限制。

训练与采样 ​

工作流概览 ​

创建 Client ​

ServiceClient ​

TrainingClient ​

Tokenizer ​

准备训练数据 ​

SFT 数据格式 ​

训练 API ​

forward() ​

forward_backward() ​

optim_step() ​

异步操作（OperationHandle） ​

原生异步 Client ​

用户指标 ​

查询训练运行 ​

流式 Packed SFT ​

自定义 Loss ​

采样 ​

SamplingParams ​

暂停与恢复生成 ​

计算 Logprobs ​

下一步 ​

训练与采样

工作流概览

创建 Client

ServiceClient

TrainingClient

Tokenizer

准备训练数据

SFT 数据格式

训练 API

forward()

forward_backward()

optim_step()

异步操作（OperationHandle）

原生异步 Client

用户指标

查询训练运行

流式 Packed SFT

自定义 Loss

采样

SamplingParams

暂停与恢复生成

计算 Logprobs

下一步