API

trio.TrainingClient

class TrainingClient:
    def __init__(
        self,
        task_id: str,
        base_model: str,
        lora: LoraRunSpec,
    ):

TrainingClient 是用于 LoRA 模型训练的客户端，通过 ServiceClient.create_lora_training_client() 创建。

import pytrio as trio

client = trio.ServiceClient()
training_client = client.create_lora_training_client(base_model="Qwen/Qwen3.5-4B")

tokenizer = training_client.get_tokenizer()
tokens = tokenizer.encode("The meaning of life is")

data = [Datum(model_input=tokens, loss_fn_inputs={"target_tokens": tokens})]

# 前反向传播
future = training_client.forward_backward(data=data)
output = future.result()

# 梯度更新
training_client.optim_step(AdamParams(learning_rate=1e-4)).result()

属性

属性类型说明
task_idstr当前训练任务 ID
base_modelstr使用的基础模型名称
loraLoraRunSpecLoRA 初始化参数

方法

forward

def forward(
    self,
    data: list[Datum],
    loss_fn: str = "cross_entropy",
    loss_fn_config: dict[str, object] | None = None,
    auto_shift: bool = False,
) -> APIFuture[ForwardBackwardOutput]

仅执行前向传播，计算损失但不更新梯度。

参数

参数类型默认值说明
datalist[Datum]—样本列表，每个样本包含输入 token ids 和损失函数参数
loss_fnstr"cross_entropy"损失函数类型："cross_entropy" / "importance_sampling" / "ppo"
loss_fn_configdict | NoneNone损失函数的额外配置项
auto_shiftboolFalse为 True 时自动将 labels 偏移一位对齐预测目标

返回值

APIFuture[ForwardBackwardOutput] — 调用 .result() 获取输出。输出中包含以下两个值：

  • loss_fn_outputs: 损失函数输出，每个输入样本对应一项。
  • metrics: 前向传播指标。

示例

future = training_client.forward(data=data)
output = future.result()
print(output.metrics)

forward_backward

def forward_backward(
    self,
    data: list[Datum],
    loss_fn: str = "cross_entropy",
    loss_fn_config: dict[str, object] | None = None,
    auto_shift: bool = False,
) -> APIFuture[ForwardBackwardOutput]

执行前向 + 反向传播，计算并累积梯度。

参数

参数类型默认值说明
datalist[Datum]—样本列表，每个样本包含输入 token ids 和损失函数参数
loss_fnstr"cross_entropy"损失函数类型："cross_entropy" / "importance_sampling" / "ppo"
loss_fn_configdict | NoneNone损失函数的额外配置项
auto_shiftboolFalse为 True 时自动将 labels 偏移一位对齐预测目标

返回值

APIFuture[ForwardBackwardOutput] — 调用 .result() 获取输出。输出中包含以下两个值：

  • loss_fn_outputs: 损失函数输出，每个输入样本对应一项。
  • metrics: 前向传播指标。

示例

future = training_client.forward_backward(data=data, loss_fn="cross_entropy")
output = future.result()

forward_backward_custom

def forward_backward_custom(
    self,
    data: list[Datum],
    loss_fn: Callable[
        [list[Datum], list["torch.Tensor"]], tuple["torch.Tensor", dict[str, float]]
    ],
) -> APIFuture[ForwardBackwardOutput]

使用自定义 PyTorch 损失函数执行前向 + 反向传播。需要本地安装 torch。

参数

参数类型说明
datalist[Datum]样本列表
loss_fnCallable接收 (data, logprobs) 并返回 (loss_tensor, metrics_dict) 的函数

返回值

APIFuture[ForwardBackwardOutput] — 调用 .result() 获取输出。输出中包含以下两个值：

  • loss_fn_outputs: 损失函数输出，每个输入样本对应一项。
  • metrics: 前向传播指标。

示例

def my_loss(data, logprobs):
    loss = -sum(lp.mean() for lp in logprobs)
    return loss, {"my_loss": loss.item()}

future = training_client.forward_backward_custom(data=data, loss_fn=my_loss)
output = future.result()

optim_step

def optim_step(self, adam_params: AdamParams) -> APIFuture[OptimStepResponse]

根据当前累积的梯度执行一次 Adam 优化器更新，并清零梯度。

参数

参数类型说明
adam_paramsAdamParamsAdam 优化器参数，见 AdamParams

返回值

APIFuture[OptimStepResponse]， 调用 .result() 获取优化器指标。

示例

training_client.optim_step(AdamParams(learning_rate=1e-4)).result()

save_state

def save_state(
    self, name: str,
    ttl_seconds: int | None = None,
    overwrite: bool = False
) -> APIFuture[SaveWeightsResponse]

保存模型权重和优化器状态（完整 checkpoint），用于断点续训。

参数

参数类型说明
namestrcheckpoint 名称
ttl_seconds`intNone`
overwritebool为 True 时覆盖同名的现有存档

返回值

APIFuture[SaveWeightsResponse]， 调用 .result() 获取存档路径和模型名称。

  • path: 已保存 checkpoint 的 URL.
  • model: 已保存的模型名称。

示例

result = training_client.save_state(name="step-100").result()
print(result.path)

save_weights_for_sampler

def save_weights_for_sampler(
    self, name: str,
    ttl_seconds: int | None = None
) -> APIFuture[SaveWeightsForSamplerResponse]

仅保存模型权重（不含优化器状态），用于后续推理采样。

参数

参数类型说明
namestr权重保存名称
ttl_seconds`intNone`

返回值

APIFuture[SaveWeightsForSamplerResponse] ， 调用 .result() 获取存档路径、模型名称和权重大小。

  • path: 已保存权重的 Checkpoint Path。
  • model: 已保存的模型名称。
  • size: 权重大小，单位为字节。

示例

result = training_client.save_weights_for_sampler(name="step-100").result()
print(result.path)

create_sampling_client

def create_sampling_client(
    self,
    base_model: str,
    model_path: str | None = None,
) -> SamplingClient

基于指定的 LoRA 权重路径创建 SamplingClient，可在训练过程中随时启动推理。

参数

参数类型说明
base_modelstr基础模型名称
model_pathstr | NoneNone

返回值

SamplingClient

示例

sampling_client = training_client.create_sampling_client(model_path="/path/to/weights")

save_weights_and_get_sampling_client

def save_weights_and_get_sampling_client(self) -> SamplingClient

保存当前模型权重到一个临时（匿名）存档，并立即返回一个已加载该权重的 SamplingClient。该方法无需指定存档名称，所保存的权重仅供训练过程中的临时推理采样使用，适用于 Agent-RL 等需要在训练循环中持续从最新策略采样的场景。

返回值

SamplingClient

示例

sampling_client = training_client.save_weights_and_get_sampling_client()

get_tokenizer

def get_tokenizer(self)

获取与当前基础模型匹配的 tokenizer，基于 transformers / modelscope 提供的 AutoTokenizer。

示例

tokenizer = training_client.get_tokenizer()
tokens = tokenizer.encode("The meaning of life is")

异步方法

forward_async

async def forward_async(
    self,
    data: list[Datum],
    loss_fn: str = "cross_entropy",
    loss_fn_config: dict[str, object] | None = None,
    auto_shift: bool = False,
) -> APIFuture[ForwardBackwardOutput]

forward 的异步版本，参数相同。

返回值

APIFuture[ForwardBackwardOutput] — 调用 .result() 或 await future 获取输出。输出中包含以下两个值：

  • loss_fn_outputs: 损失函数输出，每个输入样本对应一项。
  • metrics: 前向传播指标。
future = await training_client.forward_async(data=data)
output = future.result()

forward_backward_async

async def forward_backward_async(
    self,
    data: list[Datum],
    loss_fn: str = "cross_entropy",
    loss_fn_config: dict[str, object] | None = None,
    auto_shift: bool = False,
) -> APIFuture[ForwardBackwardOutput]

forward_backward 的异步版本，参数相同。

返回值

APIFuture[ForwardBackwardOutput] — 调用 .result() 或 await future 获取输出。输出中包含以下两个值：

  • loss_fn_outputs: 损失函数输出，每个输入样本对应一项。
  • metrics: 前向传播指标。
future = await training_client.forward_backward_async(data=data)
output = future.result()

forward_backward_custom_async

async def forward_backward_custom_async(
    self,
    data: list[Datum],
    loss_fn: Callable[
        [list[Datum], list["torch.Tensor"]],
        tuple["torch.Tensor", dict[str, float]]
    ],
) -> APIFuture[ForwardBackwardOutput]

forward_backward_custom 的异步版本，参数相同。

返回值

APIFuture[ForwardBackwardOutput] — 调用 .result() 或 await future 获取输出。输出中包含以下两个值：

  • loss_fn_outputs: 损失函数输出，每个输入样本对应一项。
  • metrics: 前向传播指标。
future = await training_client.forward_backward_custom_async(data=data, loss_fn=my_loss)
output = future.result()

optim_step_async

async def optim_step_async(self, adam_params: AdamParams) -> APIFuture[OptimStepResponse]

optim_step 的异步版本，参数相同。

返回值

APIFuture[OptimStepResponse]，调用 .result() 或 await future 获取优化器指标。

future = await training_client.optim_step_async(AdamParams(learning_rate=1e-4))
await future

save_state_async

async def save_state_async(
    self, name: str,
    ttl_seconds: int | None = None,
    overwrite: bool = False
) -> APIFuture[SaveWeightsResponse]

save_state 的异步版本，参数相同。

返回值

APIFuture[SaveWeightsResponse]，调用 .result() 或 await future 获取存档路径和模型名称。

  • path: 已保存 checkpoint 的 URL.
  • model: 已保存的模型名称。
future = await training_client.save_state_async(name="step-100")
result = await future

save_weights_for_sampler_async

async def save_weights_for_sampler_async(
    self, name: str,
    ttl_seconds: int | None = None
) -> APIFuture[SaveWeightsForSamplerResponse]

save_weights_for_sampler 的异步版本，参数相同。

返回值

APIFuture[SaveWeightsForSamplerResponse]，调用 .result() 或 await future 获取存档路径、模型名称和权重大小。

  • path: 已保存权重的 Checkpoint Path。
  • model: 已保存的模型名称。
  • size: 权重大小，单位为字节。
future = await training_client.save_weights_for_sampler_async(name="step-100")
result = await future

create_sampling_client_async

async def create_sampling_client_async(
    self,
    base_model: str,
    model_path: str | None = None,
) -> SamplingClient

create_sampling_client 的异步版本，参数相同。

返回值

SamplingClient

sampling_client = await training_client.create_sampling_client_async(model_path="/path/to/weights")

save_weights_and_get_sampling_client_async

async def save_weights_and_get_sampling_client_async(self) -> SamplingClient

save_weights_and_get_sampling_client 的异步版本，参数相同。

返回值

SamplingClient

sampling_client = await training_client.save_weights_and_get_sampling_client_async()

trio.ServiceClient

上一页

trio.SamplingClient

下一页

本页目录

属性
方法
forward
forward_backward
forward_backward_custom
optim_step
save_state
save_weights_for_sampler
create_sampling_client
save_weights_and_get_sampling_client
get_tokenizer
异步方法
forward_async
forward_backward_async
forward_backward_custom_async
optim_step_async
save_state_async
save_weights_for_sampler_async
create_sampling_client_async
save_weights_and_get_sampling_client_async