PyTRIO Skill

PyTRIO 是 TRIO 远程大模型后训练和推理服务的 Python SDK。本地代码负责准备数据和控制训练循环；前向传播、反向传播、优化器更新、采样和权重存储由 TRIO 服务执行。

安装与更新

安装这个 skill：执行 npx skills add SwanHubX/pytrio-skill -g -y，它会安装到当前用户级别的 .agents/skills 并供支持 Agent Skills 的 CLI 使用。
更新这个 skill：执行 npx skills update pytrio-skill -g -y，更新后重新启动或刷新 Agent 会话即可使用新版本。

使用 Skill 生成或调试代码前，建议先确认 PyTRIO SDK 为最新版：未锁定依赖时执行 python -m pip install --upgrade pytrio，再用 python -c 'from importlib.metadata import version; print(version("pytrio"))' 记录实际版本。若项目的 pyproject.toml、requirements 或 lockfile 已固定 PyTRIO 版本，不要擅自升级；应以项目锁定版本为准，并检查该版本的 API 签名。

使用定位

PyTRIO 代码通常由本地 Python 负责数据准备、训练循环和实验记录，由 TRIO 服务负责模型前向、反向、优化器更新和权重保存。

处理训练任务时，先按任务形态选择对应的 Datum 构造、loss 和 client 调用方式：

• SFT：监督微调、assistant-only loss mask、同步/异步训练、SwanLab 记录。
• 多模态：图文 ModelInput、ImageChunk.expected_tokens、多模态推理/SFT、Vision GRPO。
• GRPO：student rollout、reward、group-relative advantage、importance_sampling 更新。
• OPD：student rollout、teacher logprob、reverse-KL advantage、importance_sampling 更新。
• Search-R1：多轮搜索工具环境、结果 reward、group-relative advantage、observation mask。
• ReTool：多轮代码解释器环境、outcome reward、执行结果 mask、本地执行隔离。
• OPSD：同一初始模型的 Student / privileged Teacher、自轨迹逐 token reverse-KL。
• DAPO：Dynamic Sampling、Clip-Higher、token mean、Soft Overlong、自定义 PPO loss。
• GSPO：序列级重要性比率、序列级裁剪、forward_backward_custom。
• ALFWorld：同游戏独立环境、长轨迹工具交互、终局 reward、PPO 更新。
• Custom loss：本地 PyTorch objective、闭包元数据、forward_backward_custom / async 和 surrogate gradient。
• DPO：chosen/rejected 偏好数据、reference logprob、pairwise custom loss。

示例代码提供可改模板；关键字段和常见错误在 references/ 中说明。

使用顺序

1. 先按“安装与更新”确认项目使用的 PyTRIO SDK 版本；未锁定依赖时建议使用最新版，已锁定时尊重项目版本。
2. 判断任务类型：SFT、多模态推理/训练、Vision GRPO、GRPO、OPD、Search-R1、ReTool、OPSD、DAPO、GSPO、ALFWorld、custom loss、DPO、推理/保存权重，或 API 调试。
3. 按任务路由读取对应 reference 和 example；不要一开始展开全部官方文档。
4. 示例能覆盖时，按示例替换数据集、prompt、reward、teacher/student 模型、超参数和 SwanLab 配置。
5. 示例覆盖不了、API 行为不确定、报错难定位，或需要 checkpoint/OpenAI 兼容 API 等细节时，再读 references/doc-index.md，根据任务打开对应官方 Markdown 文档。
6. 若官方文档和本地 SDK 行为不一致，先检查已安装 SDK 签名或写最小复现，再给最终代码。

任务路由

核心规则

• 使用 import pytrio as trio 导入。
• 创建 client 前先用 trio login 或 trio login -k <API_KEY> 完成认证。CLI 命令是 trio，不是 pytrio。
• 使用 trio.ServiceClient() 作为主入口。
• 纯文本 Datum.model_input 和 SamplingClient.sample(prompt=...) 使用 trio.ModelInput.from_ints(...)；多模态输入显式构造包含文本与图片 chunk 的 trio.ModelInput(chunks=[...])。
• 采样参数传 trio.SamplingParams(...) 对象，不要传普通 dict。
• 同步远程调用通常返回 future，需要调用 .result() 取得结果。
• 异步 API 有两种返回边界：sample_async()、compute_logprobs_async()、save_weights_and_get_sampling_client_async() 通常一次 await 直接得到结果；forward_backward_async()、optim_step_async() 的第一次 await 返回 APIFuture，还要再次 await 该 future 才算远程任务完成。始终以当前 SDK 签名为准。
• ServiceClient、TrainingClient、SamplingClient 都没有 close() / close_async() 等显式关闭方法；训练结束、保存权重后直接退出即可，不要自行添加 client 清理调用。
• SFT 默认手动做自回归右移，除非明确使用 auto_shift=True：model_input=tokens[:-1]、target_tokens=tokens[1:]、weights=weights[1:]。
• SFT 用 weights 屏蔽 prompt token，不要使用 HuggingFace 风格的 -100 labels。
• 多模态能力要求 pytrio>=0.2.7。先确认基模支持图片输入，再按模型 chat template 的真实顺序组织 EncodedTextChunk 与 ImageChunk；不要把 Qwen 的特殊 token 或 <|image_pad|> 规则直接套到其他模型。
• ImageChunk.data 传原始图片字节，format 与真实编码保持一致。多模态训练和任何依赖本地序列长度的对齐都要用匹配远端模型的 image processor 计算 expected_tokens，不要写死视觉 token 数。
• chat template 已写入特殊 token 时，拆分图片占位符后的文本使用 add_special_tokens=False 编码。采样后校验 response.input_tokens == len(prompt)，长度不一致时停止训练并修正 image processor 或 chunk 顺序。
• 多模态 SFT 与 Vision GRPO 都保留完整图文 prompt，并只拼接 completion[:-1]；prompt/image 区间从 prompt_length - 1 开始右移对齐。SFT 的 target_tokens/weights 在该区间填零，Vision GRPO 的 target_tokens/old logprobs/advantages 在该区间填零。
• GRPO、OPD、Search-R1 和 OPSD 通常使用 loss_fn="importance_sampling"，Datum.loss_fn_inputs 必须包含右移对齐后的 target_tokens、旧策略 logprobs 和 advantages。
• importance_sampling 的 prompt/observation 区间不训练时，用 target_tokens=0、logprobs=0.0、advantages=0.0 占位，保证长度和 model_input 一致。
• GRPO 的 advantage 来自同一 prompt 的 group 内 reward 相对均值；如果整组 reward 完全相同，advantage 全为 0，通常跳过该组。
• OPD 的 teacher 只对 student 实际采样出来的 completion 计算 logprob；不要用 teacher 自己生成的 completion 替代 student 轨迹。
• Search-R1 的 tool observation 进入后续上下文但不进入 loss；先在完整同题 group 内计算 advantage，再拆 micro-batch 累积梯度，搜索后端本身不参与训练。
• ReTool 的代码执行结果进入后续上下文但不进入 loss；保持真实 token 前缀连续。本地 subprocess 只能限制资源，处理不可信模型代码时使用容器或专用沙箱，并移除环境凭证。
• OPSD 的 Student 和 Teacher 来自同一初始模型；Teacher 固定在 step 0，通过包含参考解答的 privileged prompt 对 Student completion 计算 logprob，参考解答不是 SFT label。
• DAPO 使用 forward_backward_custom 实现非对称 PPO clip 和 token-level reduction；Dynamic Sampling 被丢弃的候选也要计入 token 与 wall-clock 成本。
• GSPO 的 sequence ratio 是 completion token 概率比的几何平均；只裁剪一次完整序列比率，prompt token 不参与 ratio。
• ALFWorld 的同组轨迹必须来自同一个游戏和相同初始 observation，同时持有独立环境状态；环境 observation token 的 advantage 为 0。
• Custom loss 的函数签名是 loss_fn(data, logprobs) -> (loss, metrics)；函数在本地用 torch 执行，返回的 loss 必须保留 autograd 计算图。
• forward_backward_custom 的 Datum.loss_fn_inputs 只放与 model_input 等长的 target_tokens；sampling/reference logprob、mask、advantage、分组和归一化参数使用每个 batch 独立的闭包传入。
• Custom loss 回调中不要采样、访问网络或调用其他 PyTRIO client。异步版第一次 await 返回 APIFuture，第二次 await 才等待远程任务完成。
• DPO 使用 forward_backward_custom(data, loss_fn) 完成 pairwise objective；不要用 cross_entropy 或 importance_sampling 代替。
• DPO batch 保持 [chosen_0, rejected_0, chosen_1, rejected_1, ...] 顺序，reference logprobs 要提前计算并通过闭包传给 custom loss。
• 推理用权重保存使用 save_weights_for_sampler()；完整断点续训使用 save_state()。
• 模型名优先使用当前官方文档或 client.get_supported_models() 返回值，不要硬编码旧模型名。
• 训练脚本建议默认接入 SwanLab。SFT 记录 loss；Vision GRPO 记录 reward、format rate、degenerate group 比例、completion 长度和 trainer metrics；GRPO 记录 reward、degenerate group 比例和 trainer metrics；OPD/OPSD 记录 reverse KL、completion token 数和 trainer metrics；Search-R1 记录 format、search success/error、turns 和 search calls；ReTool 记录 tool-call、sandbox 与轨迹指标；DAPO 记录候选/有效组、token 利用率和 refill 时间；GSPO 记录 sequence ratio、clip 与原始分母；ALFWorld 记录成功率、非法动作、环境步数和 prefill/训练成本；DPO 记录 loss、accuracy、margin、chosen/rejected reward。

代码骨架

先用本地 reference + example 解决问题；示例不够时，再读官方文档补足 API 细节。若官方文档和本地 SDK 行为不一致，先检查已安装 SDK 签名或写最小复现。

SFT 的核心是 assistant-only weights 和自回归右移：

tokens = prompt_tokens + completion_tokens
weights = [0.0] * len(prompt_tokens) + [1.0] * len(completion_tokens)
datum = trio.Datum(
    model_input=trio.ModelInput.from_ints(tokens[:-1]),
    loss_fn_inputs={
        "target_tokens": np.asarray(tokens[1:], dtype=np.int32),
        "weights": np.asarray(weights[1:], dtype=np.float32),
    },
)
training_client.forward_backward([datum], loss_fn="cross_entropy").result()

GRPO/OPD 都用 importance_sampling，区别只在 advantage 来源：

obs_len = len(prompt_tokens) - 1
datum = trio.Datum(
    model_input=trio.ModelInput.from_ints(prompt_tokens + completion_tokens[:-1]),
    loss_fn_inputs={
        "target_tokens": np.asarray([0] * obs_len + completion_tokens, dtype=np.int64),
        "logprobs": np.asarray([0.0] * obs_len + old_logprobs, dtype=np.float32),
        "advantages": np.asarray([0.0] * obs_len + advantages, dtype=np.float32),
    },
)
training_client.forward_backward([datum], loss_fn="importance_sampling").result()

• GRPO：advantages = reward - mean(group_rewards)，old logprobs 来自同组 rollout 的 student sampler。
• Vision GRPO：沿用 GRPO 的 reward 与 advantage；把纯文本 prompt 换成有正确 expected_tokens 的图文 chunks，并按 references/vision.md 对齐 prompt/image 零占位。
• OPD：teacher 对 prompt + student_completion 调 compute_logprobs，advantages = -kl_coef * (student_logprobs - teacher_logprobs)。
• OPSD：沿用 OPD 的逐 token advantage，但 teacher 是同一初始模型，并通过 privileged prompt 看到参考解答。
• Search-R1：沿用 GRPO 的轨迹级 advantage；assistant token 使用 rollout old logprob 和 trajectory advantage，tool observation token 的 logprob/advantage 都填 0。
• ReTool：沿用多轮 observation mask，最终答案给 outcome reward；本地代码执行器属于环境，不参与训练。
• DAPO：用 custom PPO loss 同时实现 Clip-Higher 和 token mean；Dynamic Sampling 与 Soft Overlong 在 rollout/reward 层完成。
• GSPO：用 custom loss 对 completion 计算一个 sequence ratio，再做序列级裁剪。
• ALFWorld：按同一游戏计算轨迹级 advantage；assistant action token 参与 PPO，环境 observation token 全部 mask。

通用 custom loss 的 Datum 只放服务端 forward 所需的 target_tokens，算法元数据通过闭包传入；完整契约见 references/custom-loss.md。DPO 的 pairwise loss 形态如下：

def loss_fn(data, logprobs_list):
    # data、reference_logprobs、response_masks 保持 chosen/rejected 同序。
    loss = compute_dpo_loss(logprobs_list, reference_logprobs, response_masks)
    return loss, {"dpo/loss": float(loss.detach().item())}

training_client.forward_backward_custom([chosen, rejected], loss_fn).result()

每次 forward_backward* 后调用 training_client.optim_step(trio.AdamParams(...))；保存推理权重使用 save_weights_for_sampler(...).result()。