pytrio-skill

name	pytrio-skill
description	使用 PyTRIO/TRIO 编写、调试或解释远程大模型训练与推理代码。用户提到 pytrio、TRIO、PyTRIO SDK、ServiceClient、TrainingClient、SamplingClient、LoRA 训练、SFT、异步 SFT、多模态推理、多模态训练、Vision GRPO、视觉 RLVR、ImageChunk、RL、GRPO、OPD、Search-R1、ReTool、DAPO、GSPO、ALFWorld、TextWorld、Agentic RL、OPSD、On-Policy Self-Distillation、DPO、custom loss、自定义损失函数、forward_backward_custom、forward_backward_custom_async、工具调用强化学习、HuggingFace datasets、TRIO OpenAI 兼容 API、SwanLab 训练记录、权重保存、checkpoint 或远程 LLM 后训练时使用。
metadata	{"version":"0.2.1"}

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 签名或写最小复现，再给最终代码。

任务路由

用户任务	读取内容
安装、登录、第一次训练或推理	references/doc-index.md -> 快速开始
编写简单 SFT 训练代码	references/sft.md；examples/quickstart_sft.py
编写角色 SFT 或 Chat-甄嬛类微调	references/sft.md；references/chat-huanhuan.md；examples/chat-huanhuan.py
编写异步 SFT 或异步记录 SwanLab	references/sft.md；examples/chat-huanhuan-async.py
编写多轮对话 SFT 蒸馏	references/sft.md；examples/sft-distill-conversation.py 或 examples/sft-distill-conversation-async.py
编写多模态推理 / 图像输入 / 多模态 SFT	references/vision.md；API 细节再读 references/doc-index.md -> 多模态
编写 Vision GRPO / GeoQA / 视觉 RLVR	references/vision.md；examples/vision-grpo.py；需要固定评测时读取其中的完整源码目录
编写 GRPO / GSM8K / reward-based RLVR	references/grpo.md；examples/grpo-gsm8k.py 或 examples/grpo-gsm8k-async.py
编写 OPD / Medical OPD / multi-teacher / teacher-KL 蒸馏	references/opd.md；examples/opd-deepmath.py 或 examples/opd-deepmath-async.py
编写 Search-R1 / Agentic RL / 多轮搜索工具训练	references/search-r1.md；完整项目读取其中的源码目录
编写 ReTool / 代码解释器 Agent / code-interlaced RL	references/retool.md；完整项目读取其中的源码目录
编写 OPSD / On-Policy Self-Distillation / privileged teacher	references/opsd.md；完整项目读取其中的源码目录
编写 DAPO / Dynamic Sampling / Clip-Higher / Soft Overlong	references/dapo.md；完整项目读取其中的源码目录
编写 GSPO / sequence-level ratio / sequence clipping	references/gspo.md；完整项目读取其中的源码目录
编写 ALFWorld / TextWorld / 长轨迹环境 Agentic RL	references/alfworld.md；完整项目读取其中的源码目录
编写通用自定义损失函数 / / async

核心规则

• 使用 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/ 在该区间填零，Vision GRPO 的 /old / 在该区间填零。

代码骨架

先用本地 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()。