name: pytrio description: Guide for using the PyTRIO SDK (v0.1.13b0) — remote LLM training and inference. Trigger when code imports pytrio, or user mentions PyTRIO / TRIO training / TRIO SDK / remote LLM training.
PyTRIO SDK Skill
PyTRIO 是什么
PyTRIO 是远程 LLM 训练与推理 SDK。本地代码只做数据准备和调度,实际的前向传播、反向传播、优化器步骤和推理采样都在远端 GPU 集群上执行。不需要本地 GPU。
- 安装:
pip install pytrio transformers torch(或uv add pytrio transformers torch) - 导入约定:
import pytrio as trio,所有 API 通过trio.xxx访问 - Tokenizer: 使用
train.get_tokenizer()或sampler.get_tokenizer()(需安装 modelscope) - 环境搭建详见
references/00-overview.md
认证(重要)
使用 PyTRIO 前必须先完成认证,否则所有 API 调用都会失败。
检查是否已登录: 查看 ~/.pytrio/config.toml 是否存在且包含 api_key 字段。
如果未登录,按以下流程操作:
- 提示用户提供 API Key(可从 https://pytrio.cn/dashboard 获取)
- 用户提供 key 后,执行非交互式登录命令:
trio login -k <用户提供的API_KEY> - 登录成功后凭证保存在
~/.pytrio/config.toml,后续ServiceClient()无需传 api_key
在代码中也可以直接传入(适合临时使用):
client = ServiceClient(api_key="YOUR_API_KEY")
注意:CLI 命令是 trio login,不是 pytrio login。
PyTRIO ↔ PyTorch 概念速查
| PyTorch | PyTRIO | 说明 |
|---|---|---|
model = AutoModel(...) |
client = ServiceClient(api_key=...) |
无本地模型对象 |
model(input_ids) |
train.forward(data) |
远端前向 |
loss.backward() |
train.forward_backward(data) |
前向+反向合一 |
optimizer.step() |
train.optim_step(AdamParams(...)) |
远端 Adam |
optimizer.zero_grad() |
不需要 | 远端自动管理 |
torch.save(...) |
train.save_state(name) |
保存到云端 |
model.generate(...) |
sampler.sample(prompt, ...) |
远端推理 |
torch.Tensor |
TensorData / plain list |
自动互转 |
核心 API 速查
入口
import pytrio as trio
client = trio.ServiceClient(api_key="...")
models = client.get_supported_models()
模型选择规则:必须使用 Qwen/Qwen3-4B-Instruct-2507,不要使用 Qwen/Qwen3-4B(Base 模型训练端不可用)。
训练
import pytrio as trio
train = client.create_lora_training_client(
base_model="Qwen/Qwen3-4B-Instruct-2507",
rank=32, # [4, 64]
train_mlp=True,
train_attn=True,
train_unembed=False,
)
tokenizer = train.get_tokenizer()
# tokenizer 注意事项(官方示例的标准模式):
# 序列起始段(prompt)用 add_special_tokens=True —— 保留 BOS
# 拼接延续段(completion)用 add_special_tokens=False —— 避免重复 BOS
# 推理时 encode 默认 add_special_tokens=True,保留 BOS 作为生成起始
# apply_chat_template() 已自行管理 special tokens,返回文本再 encode 时也用 False
# 官方标准 SFT 数据构造模式:
prompt_tokens = tokenizer.encode(prompt_text, add_special_tokens=True)
prompt_weights = [0] * len(prompt_tokens)
completion_tokens = tokenizer.encode(completion_text, add_special_tokens=False)
completion_weights = [1] * len(completion_tokens)
tokens = prompt_tokens + completion_tokens
weights = prompt_weights + completion_weights
# 手动偏移对齐: model_input[i] 预测 target_tokens[i]
input_tokens = tokens[:-1]
target_tokens = tokens[1:]
weights = weights[1:] # weights 必须同步偏移
sample = trio.Datum(
model_input=trio.ModelInput.from_ints(input_tokens), # 必须用 ModelInput.from_ints() 包装
loss_fn_inputs={
"target_tokens": target_tokens, # 值可直接传 list,会自动转 TensorData
"weights": weights, # 可选,默认全 1;SFT 场景强烈建议传
},
)
# 训练步骤
fb = train.forward_backward(data=[sample], loss_fn="cross_entropy")
result = fb.result() # APIFuture -> ForwardBackwardOutput
# result.metrics 的 key 是 "loss:sum"(注意冒号),不是 "loss"
# per-token 归一化 loss(两种等价写法,需 import numpy as np):
# 方式 A: loss_sum / total_weight_sum
# total_weight_sum = sum(sum(ex.loss_fn_inputs["weights"].tolist()) for ex in data)
# 方式 B: -np.dot(logprobs, weights) / weights.sum()
# logprobs = np.concatenate([o["logprobs"].tolist() for o in result.loss_fn_outputs])
# weights = np.concatenate([ex.loss_fn_inputs["weights"].tolist() for ex in data])
loss_sum = result.metrics.get("loss:sum", 0.0)
train.optim_step(trio.AdamParams(learning_rate=1e-4)).result()
损失函数
| loss_fn | 必需 loss_fn_inputs | 用途 |
|---|---|---|
"cross_entropy" |
target_tokens, 可选 weights |
SFT |
"importance_sampling" |
target_tokens, logprobs, advantages,强烈建议 weights |
GRPO / 离线 RL |
"ppo" |
target_tokens, logprobs, advantages,强烈建议 weights |
在线 RL |
RL 场景的 weights 用于对 prompt 部分做 mask(0=ignore,1=计入 loss),不传则 prompt 也会参与策略梯度,通常会跑偏。官方所有 RL 示例都带 weights。
ppo 可通过 loss_fn_config 自定义裁剪阈值(默认 ε=0.2,即 ratio 裁到 [0.8, 1.2]):
train.forward_backward(
data=data, loss_fn="ppo",
loss_fn_config={"clip_low_threshold": 0.9, "clip_high_threshold": 1.1},
)
自定义 loss(需要本地 torch):
def my_loss(samples, logprobs_list):
loss = ... # torch.Tensor
return loss, {"metric": loss.item()}
train.forward_backward_custom(data=samples, loss_fn=my_loss)
保存与恢复
# 保存权重(仅用于推理)
save = train.save_weights_for_sampler("name").result()
# save.path -> 远端路径
# 保存完整 checkpoint(含优化器,用于断点续训)
state = train.save_state("name").result()
# 断点续训
train2 = client.create_training_client_from_state_with_optimizer(state.path)
# 仅加载权重续训(不含优化器动量)
train3 = client.create_training_client_from_state(state.path)
推理
import pytrio as trio
# 从训练权重推理
sampler = train.create_sampling_client(model_path=save.path)
# 或独立推理
sampler = client.create_sampling_client(base_model="Qwen/Qwen3-4B-Instruct-2507")
tokenizer = sampler.get_tokenizer()
input_ids = tokenizer.encode("Hello")
result = sampler.sample(
prompt=trio.ModelInput.from_ints(input_ids), # 必须用 trio.ModelInput 包装
num_samples=1,
sampling_params=trio.SamplingParams(max_tokens=100, temperature=0.7),
).result()
text = result.sequences[0].text
logprobs = result.sequences[0].logprobs
管理
rest = client.create_rest_client()
rest.list_weights() # 已保存权重
rest.list_training_runs() # 训练记录
rest.list_checkpoints(run_id) # 某次训练的 checkpoint
rest.list_user_checkpoints() # 所有 checkpoint
rest.get_archive_url(id) # 下载链接
rest.delete_checkpoint(run_id, cp_id)
AdamParams SDK 默认值
注意与 PyTorch 默认值不同:beta2=0.95(PyTorch 默认 0.999)、eps=1e-12(PyTorch 默认 1e-8)。如需与 PyTorch 行为一致,请显式传入 AdamParams(beta2=0.999, eps=1e-8)。
错误处理
APIFuture.result(timeout=None)支持 timeout 参数(秒),超时抛TimeoutError- 服务端异常抛
pydantic.ValidationError或ValueError - WebSocket 断连不自动重连,需重建
ServiceClient
陷阱清单
- 所有方法返回
APIFuture— 必须调用.result()获取实际值,否则拿到的是 Future 对象 - model_input 和 target_tokens 长度必须一致 — 在
auto_shift=False(默认)时,用户负责偏移对齐 importance_sampling/ppo需要四个等长字段 —model_input,target_tokens,logprobs,advantages必须等长trainable_token_indices和train_unembed互斥 — 不能同时启用ignore_eos=True时必须设置max_tokens— 否则抛异常- checkpoint 路径是远端路径 —
save_state()/save_weights_for_sampler()返回的 path 是云端路径,不是本地文件 forward_backward_custom需要本地 torch — 自定义 loss 函数在本地计算梯度后传回远端- 导入约定
import pytrio as trio— 所有 API 通过trio.xxx访问(如trio.ServiceClient、trio.Datum),不要用from pytrio import ... - 登录命令是
trio login— CLI 命令仍然是trio,不是pytrio - 所有 loss_fn_inputs 的值会自动转为 TensorData — 可以直接传 list、numpy array 或 torch tensor
- model_input 必须用
trio.ModelInput.from_ints(list)包装 — 不能直接传list[int],否则 Pydantic 验证报错 - prompt 参数同理必须用
trio.ModelInput.from_ints(list)包装 —sampler.sample(prompt=trio.ModelInput.from_ints(ids), ...) - tokenizer 使用
get_tokenizer()—train.get_tokenizer()或sampler.get_tokenizer()(需安装 modelscope) - 运行时需绕过本地代理 — 如有系统代理,需设置
no_proxy=pytrio.cn https_proxy="" http_proxy=""环境变量 - 必须使用 Instruct 模型 —
Qwen/Qwen3-4B(Base) 训练端不可用,必须用Qwen/Qwen3-4B-Instruct-2507 - 与 HuggingFace API 命名不同 — PyTRIO 的参数名和 HF transformers 有差异,不要混用:
- 采样参数是
max_tokens(不是max_new_tokens) - 推理结果是
result.sequences[0].text(不是result.samples) - 前向传播必须用关键字参数
train.forward_backward(data=[...])(不是位置参数) - 优化器参数是
learning_rate(不是lr) - loss_fn_inputs 的 key 是
target_tokens(不是labels) - 使用
import pytrio as trio而非from pytrio import ... - 不要导入内部模块(如
pytrio.types.forward.encoded_text_chunk)
- 采样参数是
- model_input 和 target_tokens 必须偏移一位 —
model_input=tokens[:-1],target_tokens=tokens[1:],不能传相同的 tokens - loss key 是
"loss:sum"(带冒号) —fb_result.metrics.get("loss:sum"),不是"loss"、"train_loss"或.loss属性 - 不要用 -100 做 token masking — HF 用
-100标记忽略的 token,但 PyTRIO 不支持这个机制。PyTRIO 用weights字段控制哪些 token 参与 loss 计算(0.0=忽略,1.0=计算),见best-practices/sft.md sampling_params必须传trio.SamplingParams对象 — 不能传 dict,sampler.sample(sampling_params=trio.SamplingParams(max_tokens=100))add_special_tokens按段区分 — prompt 起始段用True(保留 BOS),completion 拼接段用False(避免重复 BOS)。不要全 True(会多 BOS),也不要全 False(模型失去起始标记,会训偏)weights必须和target_tokens同步偏移 — 官方标准模式是先构造full_weights,再weights = weights[1:],和target_tokens = tokens[1:]严格对齐- RL 场景(IS/PPO)也要传
weights— 对 prompt 部分 mask 掉,否则 prompt token 也会参与策略梯度 - 异步版本以
_async结尾 — 高并发 / RL rollout 推荐用异步版,见best-practices/async.md
详细文档
完整 API 文档在本 Skill 的 references/ 目录下:
references/00-overview.md— 架构总览与概念映射references/01-service-client.md— ServiceClient APIreferences/02-training-client.md— TrainingClient APIreferences/03-sampling-client.md— SamplingClient APIreferences/04-rest-client.md— RestClient APIreferences/05-data-types.md— 数据类型参考references/06-debug.md— 调试指南(分诊流程、错误解码、自检清单)
示例代码
可运行的最小闭环示例在本 Skill 的 examples/ 目录下:
examples/01_train_sft.py— SFT 训练examples/02_inference.py— 推理examples/03_checkpoint_resume.py— 断点续训examples/04_model_management.py— 模型管理examples/05_importance_sampling.py— GRPO 风格训练
场景最佳实践
针对特定训练场景的推荐代码模板和注意事项,在 best-practices/ 目录下:
best-practices/sft.md— SFT 监督微调:prompt masking(weights 用法)、EOS 追加、tokenizer 选项best-practices/grpo.md— GRPO/PPO 强化学习:真实 rollout→reward→forward_backward 闭环best-practices/async.md— 异步训练:_asyncAPI、asyncio.gather、并发调度
部署
references/07-openai-compat.md— 训完模型通过 OpenAI 兼容 API 对外提供服务(chat.completions/completions)