TrainingClient

改编自 Tinker SDK（Apache-2.0）

参考文档改编自 Tinker SDK，已为 Mint Recursive 做了调整。详见归属说明与 Apache 2.0 许可。

class mint.TrainingClient ( holder , model_seq_id , model_id )¶

用于通过前向/反向传播和优化步骤训练模型的客户端。

TrainingClient 对应一个可训练、可采样的微调模型。通常通过 service_client.create_lora_training_client() 获取实例。 主要方法：

forward_backward() —— 计算训练梯度
optim_step() —— 用 Adam 优化器更新模型参数
save_weights_and_get_sampling_client() —— 导出训练好的模型用于推理

training_client = service_client.create_lora_training_client(base_model="Qwen/Qwen3-8B")
fwdbwd_future = training_client.forward_backward(training_data, "cross_entropy")
optim_future = training_client.optim_step(types.AdamParams(learning_rate=1e-4))
fwdbwd_result = fwdbwd_future.result() # Wait for gradients
optim_result = optim_future.result() # Wait for parameter update
sampling_client = training_client.save_weights_and_get_sampling_client("my-model")

参数：

holder ( InternalClientHolder ) —— 管理 HTTP 连接和异步操作的内部客户端
model_seq_id ( int )
model_id ( types. ModelID ) —— 要训练的模型的唯一标识。训练操作必须提供。

forward ( data , loss_fn , loss_fn_config=None )¶

计算前向传播，不计算梯度。

参数：

data ( Listtypes. [Datum ] ) —— 训练数据样本列表
loss_fn ( types. LossFnType ) —— 损失函数类型（如 "cross_entropy"）
loss_fn_config ( Dict[str, float] | None ，默认： None ) —— 损失函数的可选配置

返回： 包含前向传播输出和 loss 的 APIFuture

data = [types.Datum(
model_input=types.ModelInput.from_ints(tokenizer.encode("Hello")),
loss_fn_inputs={"target_tokens": types.ModelInput.from_ints(tokenizer.encode("world"))}
)]
future = training_client.forward(data, "cross_entropy")
result = await future
print(f"Loss: {result.loss}")

异步变体： forward_async()

forward_backward ( data , loss_fn , loss_fn_config=None )¶

计算前向和反向传播，得到梯度。

参数：

data ( Listtypes. [Datum ] ) —— 训练数据样本列表
loss_fn ( types. LossFnType ) —— 损失函数类型（如 "cross_entropy"）
loss_fn_config ( Dict[str, float] | None ，默认： None ) —— 损失函数的可选配置

返回： 包含前向/反向传播输出、loss 和梯度的 APIFuture

data = [types.Datum(
model_input=types.ModelInput.from_ints(tokenizer.encode("Hello")),
loss_fn_inputs={"target_tokens": types.ModelInput.from_ints(tokenizer.encode("world"))}
)]
# Compute gradients
fwdbwd_future = training_client.forward_backward(data, "cross_entropy")
# Update parameters
optim_future = training_client.optim_step(
types.AdamParams(learning_rate=1e-4)
)
fwdbwd_result = await fwdbwd_future
print(f"Loss: {fwdbwd_result.loss}")

异步变体： forward_backward_async()

forward_backward_custom ( data , loss_fn , loss_type_input='logprobs' )¶

用自定义损失函数计算前向/反向传播。

允许你定义在对数概率上操作的自定义损失函数。自定义函数接收 logprobs，计算 loss 和梯度。

参数：

data ( Listtypes. [Datum ] ) —— 训练数据样本列表
loss_fn ( CustomLossFnV1 ) —— 自定义损失函数，接收 (data, logprobs)，返回 (loss, metrics)
loss_type_input ( Literal['logprobs'] ，默认： 'logprobs' ) —— loss_fn 的输入空间。目前唯一支持的值是 "logprobs"。

返回： 包含自定义 loss 前向/反向传播输出的 APIFuture

def custom_loss(data, logprobs_list):
# Custom loss computation
loss = torch.mean(torch.stack([torch.mean(lp) for lp in logprobs_list]))
metrics = {"custom_metric": loss.item()}
return loss, metrics
future = training_client.forward_backward_custom(data, custom_loss)
result = future.result()
print(f"Custom loss: {result.loss}")
print(f"Metrics: {result.metrics}")

异步变体： forward_backward_custom_async()

optim_step ( adam_params )¶

用 Adam 优化器更新模型参数。

mindlab-toolkit 使用的 Adam 优化器与 torch.optim.AdamW 相同。注意，与 PyTorch 不同，Mint Recursive 的默认 weight decay 值为 0.0（无 weight decay）。

参数：

adam_params ( types. AdamParams ) —— Adam 优化器参数（learning_rate、betas、eps、weight_decay）

返回： 包含优化步骤响应的 APIFuture

# First compute gradients
fwdbwd_future = training_client.forward_backward(data, "cross_entropy")
# Then update parameters
optim_future = training_client.optim_step(
types.AdamParams(
learning_rate=1e-4,
weight_decay=0.01
)
)
# Wait for both to complete
fwdbwd_result = await fwdbwd_future
optim_result = await optim_future

异步变体： optim_step_async()

save_state ( name , ttl_seconds=None , overwrite=False )¶

将模型权重保存到持久化存储。

参数：

name ( str ) —— 保存的 checkpoint 名称
ttl_seconds ( int | None ，默认： None ) —— checkpoint 的可选 TTL，单位秒（None = 永不过期）
overwrite ( bool ，默认： False ) —— 为 True 时，覆盖同名的已有 checkpoint

返回： 包含保存响应的 APIFuture，其中带有 checkpoint 路径

# Save after training
save_future = training_client.save_state("checkpoint-001")
result = await save_future
print(f"Saved to: {result.path}")

异步变体： save_state_async()

load_state ( path , weights_access_token=None )¶

从已保存的 checkpoint 加载模型权重。

此方法只加载模型权重，不加载优化器状态（如 Adam 动量）。如需同时恢复优化器状态，请使用 load_state_with_optimizer。

参数：

path ( str ) —— 已保存权重的 Mint Recursive 路径（如 "mindlab-toolkit://run-id/weights/checkpoint-001"）
weights_access_token ( str | None ，默认： None ) —— 用于加载其他账户下 checkpoint 的可选访问令牌。

返回： 包含加载响应的 APIFuture

# Load checkpoint to continue training (weights only, optimizer resets)
load_future = training_client.load_state("mindlab-toolkit://run-id/weights/checkpoint-001")
await load_future
# Continue training from loaded state

异步变体： load_state_async()

load_state_with_optimizer ( path , weights_access_token=None )¶

从 checkpoint 加载模型权重和优化器状态。

参数：

path ( str ) —— 已保存权重的 Mint Recursive 路径（如 "mindlab-toolkit://run-id/weights/checkpoint-001"）
weights_access_token ( str | None ，默认： None ) —— 用于加载其他账户下 checkpoint 的可选访问令牌。

返回： 包含加载响应的 APIFuture

# Resume training with optimizer state
load_future = training_client.load_state_with_optimizer(
"mindlab-toolkit://run-id/weights/checkpoint-001"
)
await load_future
# Continue training with restored optimizer momentum

异步变体： load_state_with_optimizer_async()

save_weights_for_sampler ( name , ttl_seconds=None )¶

保存模型权重，供 SamplingClient 使用。

参数：

name ( str ) —— 保存的采样器权重名称
ttl_seconds ( int | None ，默认： None ) —— checkpoint 的可选 TTL，单位秒（None = 永不过期）

返回： 包含保存响应的 APIFuture，其中带有采样器路径

# Save weights for inference
save_future = training_client.save_weights_for_sampler("sampler-001")
result = await save_future
print(f"Sampler weights saved to: {result.path}")
# Use the path to create a sampling client
sampling_client = service_client.create_sampling_client(
model_path=result.path
)

异步变体： save_weights_for_sampler_async()

get_info ()¶

获取当前模型的信息。

返回： GetInfoResponse，包含模型配置和元数据

info = training_client.get_info()
print(f"Model ID: {info.model_data.model_id}")
print(f"Base model: {info.model_data.model_name}")
print(f"LoRA rank: {info.model_data.lora_rank}")

异步变体： get_info_async()

get_tokenizer ()¶

获取当前模型对应的 tokenizer。

返回： 与模型兼容的 PreTrainedTokenizer

tokenizer = training_client.get_tokenizer()
tokens = tokenizer.encode("Hello world")
text = tokenizer.decode(tokens)

create_sampling_client ( model_path , retry_config=None )¶

从已保存的权重创建 SamplingClient。

参数：

model_path ( str ) —— 已保存权重的 Mint Recursive 路径
retry_config ( RetryConfig | None ，默认： None ) —— 失败请求的可选重试配置

返回： 配置了指定权重的 SamplingClient

sampling_client = training_client.create_sampling_client(
"mindlab-toolkit://run-id/weights/checkpoint-001"
)
# Use sampling_client for inference

异步变体： create_sampling_client_async()

save_weights_and_get_sampling_client ( name=None , retry_config=None )¶

保存当前权重并创建用于推理的 SamplingClient。

参数：

name ( str | None ，默认： None ) —— 已废弃，无实际作用。将在未来版本中移除。
retry_config ( RetryConfig | None ，默认： None ) —— 失败请求的可选重试配置

返回： 配置了当前模型权重的 SamplingClient

# After training, create a sampling client directly
sampling_client = training_client.save_weights_and_get_sampling_client()
# Now use it for inference
prompt = types.ModelInput.from_ints(tokenizer.encode("Hello"))
params = types.SamplingParams(max_tokens=20)
result = sampling_client.sample(prompt, 1, params).result()

异步变体： save_weights_and_get_sampling_client_async()

被引用¶

ServiceClient.create_lora_training_client
ServiceClient.create_training_client_from_state
ServiceClient.create_training_client_from_state_with_optimizer