人工智能
在 Amazon SageMaker AI 上通过知识蒸馏构建高效智能体你在前沿模型上构建了一个购物助手,而且它运行良好。购物者提出“一家至少 4.5 星店铺里二十美元以下的黑色腰包”,智能体便会搜索商品目录、检查候选商品、核对约束条件,并最终给出推荐。但这样一次交互的实际成本是多少?智能体助手并不是一次模型调用,而是一个循环:搜索、读取结果、细化、验证、决策。一个完成的任务可能需要对那个前沿模型进行多次调用,而每一次调用、每一位购物者,你都要付费。在生产规模下,这种按任务计算的成本会不断累积,侵蚀项目的投资回报率,并阻碍项目上线。
让智能体应用更高效的一种方法是使用更小的模型,它响应更快,托管成本也更低。问题在于质量:较小的模型通常无法胜任该任务,而这正是当初选择较大模型的原因。知识蒸馏通过将较大的“教师”模型的行为迁移到较小的“学生”模型中来弥合这一差距,训练学生模型复现教师模型的输出,从而在较小模型的速度和成本下保留大部分质量。
在本文中,你将学习如何使用知识蒸馏,特别是序列级知识蒸馏(SeqKD),在 Amazon SageMaker AI上微调购物智能体。我们将介绍如何用前沿教师模型生成高质量的多步轨迹并保留那些成功解决任务的轨迹,如何对它们进行编码以用于监督微调,如何在无服务器基础设施上运行训练,以及如何针对留出基准评估学生模型。
图 1:一个智能体任务是一个循环,而不是单次调用:智能体搜索、读取结果、细化、验证并决策,每一步都是对前沿模型的一次单独计费调用。一个完成的任务可能包含多次调用,且针对每一位购物者
什么是序列级知识蒸馏(SeqKD)?
最简便适用的蒸馏形式是序列级知识蒸馏(SeqKD)。它使用教师模型生成高质量响应,然后通过监督微调(SFT)在这些响应上训练学生模型。SeqKD 只需要教师模型的 文本, 无需访问教师 logits,也不需要强化学习基础设施,因此即使面对闭源或仅提供 API 的教师模型也能奏效。与基于 logit 或 on-policy 等其他方法相比,它处于蒸馏谱系中最具成本效益且最易移植的一端。教师写出答案,学生便在这些答案上进行微调。由于这些答案会作为数据集持久保存,你可以复用它们来训练或测试不同的学生模型,而无需重新生成数据。
原则上,你会希望学生模型匹配教师模型在可能输出上的完整分布。但问题在于,你做不到。可能序列的空间极其庞大,因此精确匹配并不可行。SeqKD 通过用教师实际产生的样本进行近似,从而绕开了这一问题。
对于智能体而言,教师产生的序列不是一段散文,而是一条轨迹,即推理、工具调用及其返回观察结果的有序运行,其中每一步都以之前的一切为条件。在整个序列上训练学生模型,能教会它通向每个动作的推理以及多步行动路线本身,而不仅仅是一段答案文本。这使得智能体任务成为序列级 KD 最清晰的例证之一,因为被复制的序列就是行为本身。
图 2:SeqKD 仅通过教师的文本输出迁移行为:教师生成轨迹,学生在其上训练,结果在较小模型的速度和成本下保留了教师的大部分质量。训练期间无需教师权重、logits 或梯度,因此它作为普通的监督微调运行
何时 SeqKD 是正确选择
SeqKD 适用于一组特定情形:
你有输入但没有高质量的目标输出
这是 SeqKD 的核心用例,它本质上是数据生成。你提供输入,教师写出目标,因此你可以在没有现成高质量标签的任务或领域上进行微调。目标可以简单到单个答案,也可以复杂到多轮工具使用轨迹。本文选取了该谱系中要求最高的一端——智能体轨迹,这类数据尤其难以获取。单个任务会分支出许多部分正确的运行,而区分已解决的运行与失败的运行需要可靠的评分器。带有程序化评分器的基准恰好提供了这一点,因此你可以让教师生成轨迹,并只保留那些成功解决任务的轨迹。
你的教师是你只能通过 API 访问的前沿模型
这是前沿模型的常见情况。你无法访问权重来自行托管,无法读取其 logits,也无法对其进行微调。这就排除了基于 logit 的蒸馏和 on-policy 方法,它们都需要这种内部访问权限。SeqKD 只需要教师的文本响应,因此它是唯一仍然可行的方法,并且它让你在学生模型接手后免于为每次请求支付前沿模型的价格。
你想要一个强大且具成本效益的第一步
一些最强的小型推理模型在很大程度上是用 SeqKD(在教师推理轨迹上进行监督微调)构建的,因此在条件合适时,它是推理迁移的有效起点。 DeepSeek-R1 蒸馏模型 就是一个著名例子:大型教师的推理通过在教师生成的推理轨迹上进行微调,被蒸馏到较小的 Qwen 和 Llama 学生模型中。
先决条件
要运行此示例,你需要:
- 一个 AWS 账户,并具有用于 SageMaker AI 的 AWS Identity and Access Management (IAM) 角色 该角色对 Amazon Simple Storage Service (Amazon S3) 存储桶具有读写访问权限,并拥有调用教师模型的权限。
- 用于生成轨迹的前沿教师模型(本演练使用 Amazon Bedrock 上的一个模型)。 Amazon Bedrock).
- 访问 SageMaker AI 模型定制 位于受支持的 AWS 区域。
- Amazon SageMaker Studio (或本地 IDE,并配备 已配置 AWS CLI),运行 Python 3.12 或更高版本的内核,并使用 SageMaker Python SDK v3 以及
boto3(固定你的版本)。使用一个 ml.m5.4xlarge 实例 至少对于笔记本环境而言,因为数据准备和本地评分步骤都在其 CPU 上运行。 - 具备 Python、监督微调以及使用工具的智能体的实用知识。
- 配套的代码仓库,以及一个带有程序化评分器的任务(本演练使用 ShoppingBench 测试集。请替换为你自己的任务和评分器)。解决方案概述本节将端到端地介绍整个流程。一个前沿教师模型,即 Amazon Bedrock 上的 GLM-5.3,充当 ShoppingBench 智能体来生成轨迹。基准自带的评分器只保留那些解出任务的轨迹,而这些已解出的轨迹会通过 SageMaker 上的无服务器监督微调来微调一个 Qwen3.5-4B 学生模型。 配套代码仓库 以及一个带有程序化评分器的任务(本演练使用 ShoppingBench 测试集。请替换为你自己的任务和评分器)。
解决方案概述
本节将端到端地介绍整个流程。一个前沿教师模型,即 Amazon Bedrock 上的 GLM-5.3,充当 ShoppingBench 智能体来生成轨迹。基准自带的评分器只保留那些解出任务的轨迹,而这些已解出的轨迹会通过 SageMaker AI 模型定制上的无服务器监督微调来微调一个 Qwen3.5-4B 学生模型。随后,蒸馏后的学生模型会在留出的 ShoppingBench 划分上进行评估,并部署在运行工具循环的 SageMaker AI 推理实时端点之后。
图 3:在 Amazon SageMaker AI 上蒸馏购物智能体:Amazon Bedrock 上的前沿教师模型在训练问题上充当 ShoppingBench 智能体,ShoppingBench 自带的评分器只保留解出任务的轨迹(拒绝采样),每条已解出的轨迹按步骤编码,学生模型通过 Amazon SageMaker AI 模型定制进行无服务器微调,最后在留出的 ShoppingBench 测试划分上评估结果
准备数据集
寻找值得模仿的轨迹才是难点所在,而 ShoppingBench 评分器正是区分成功运行与失败运行的关键。ShoppingBench 发布了 750 道合成的购物问题,其中商品、店铺和优惠券意图各 250 道,每道题都配有隐藏的标准答案以及定义成功与否的约束条件。我们一次性将它们划分为 426 道训练题和 324 道留出评估题(商品 103 道、店铺 105 道、优惠券 116 道),并以固定索引划分,使两个集合互不重叠。
对于每道训练题,一个前沿教师模型(Amazon Bedrock 上的 GLM-5.3,开启推理且温度设为 1.0)充当 ShoppingBench 智能体。它会搜索约 270 万件商品的目录,检查候选商品,并推荐满足约束条件的商品。我们对每道题采样教师模型两次,共得到 850 条轨迹,然后让 ShoppingBench 自带的评分器只保留解出任务的轨迹。温度设为 1.0 是有意为之。采用采样而非贪心解码,能让两次尝试探索不同的查询和推理路径,从而使更多题目至少有一条解出的轨迹,并得到更多样化的学习集合。这一拒绝采样步骤将一个有能力但不完美的教师模型转化为干净的训练集,因为教师模型只需足够频繁地答对,使解出的运行值得模仿即可。
每条解出的轨迹都是一个完整回合:购物请求、教师在每一步的推理与工具调用(搜索、检查、推荐、终止),以及环境返回的观察结果。我们将推理过程与动作一并保留,结果部分表明,正是这一点让学生模型在很大程度上达到了教师模型的水平。
为训练编码轨迹
序列级蒸馏是对已存在的轨迹进行监督学习,因此训练过程中不会调用任何工具,应用也只会看到文本。由于这些轨迹是多步工具使用对话,我们将每条轨迹编码为一个 messages 聊天轮次列表,并让学员自己的聊天模板来渲染它,这样训练器看到的 token 就与智能体在推理时生成的 token 一致。
每条轨迹的每一步都会成为一行训练数据。 user 轮次携带系统提示词以及到目前为止的对话。 assistant 轮次承载该步骤的推理及其引出的工具调用,即一个 … 块,后跟一个以 harness 原生形式表示的 [{…}]。损失只落在 assistant 轮次上,因此模型只学习生成下一步的推理与动作,而不是预测环境返回的观察结果。一条包含多个步骤的轨迹会展开成多行。经过过滤后,这里解出的轨迹扩展为约 2,600 行训练数据和 130 行验证数据。
单行数据将到目前为止的对话(用户)与下一步要学习的内容(助手)配对:
{
"messages": [
{"role": "user", "content": "You are a shopping assistant. ...\n# Task\nFind a black waist bag under $20 from a shop rated 4.5 or higher.\n# Observations so far\n(none yet)"},
{"role": "assistant", "content": "<think>Start with a broad search and let the price filter work before I narrow by shop rating.</think>\n<tool_call>[{\"name\": \"find_product\", \"parameters\": {\"q\": \"black waist bag\", \"price\": \"0-20\"}}]</tool_call>"}
]
}
后续每一步都会成为单独一行,新的观察结果追加到用户内容中,而下一步的推理与工具调用则作为 assistant 目标,一直到最后的 recommend_product 和 terminate 调用。用你模型的确切聊天模板渲染每一行。
按题目而非按行拆分数据。每道题对教师模型采样两次,而每条轨迹会展开成多个步骤,因此许多行共享同一道题,按行拆分会让训练集和验证集之间泄漏近乎相同的上下文。留出整道题(约百分之五)能让验证集保持诚实。
每步编码的成本在于上下文被重复,因为每一行都会重述不断增长的对话,而无服务器微调会截断超过该限制的任何内容 dataset_max_len,因此超出限制的行会丢失它所需的部分内容。与其让限制决定截断什么,不如让各行保持在预算之内。最有效的解决办法是压缩工具观察结果,因为它们在 token 数量中占主导地位。对于商品搜索结果,只保留下一步决策所依赖的字段(id、价格、评分以及正在检查的约束条件),丢弃原始列表文本,并限制每次搜索保留的结果数量。如果轨迹仍然超出预算,就限制每条轨迹的步数。这里的 dataset_max_len 为 16,384 个 token,在裁剪观察结果后能容纳大多数轨迹。
在每个切分对应的 JSONL 文件都进入 S3 之后,将每个文件注册为一个定制数据集资产,并使用 SageMaker Python SDK v3 中的 获取训练作业所消费的 和 。 数据集资产 , DataSet.create ,该函数返回 train_dataset 和 val_dataset 训练作业所消耗的。
微调学生模型
训练使用 SageMaker AI 模型定制 使用无服务器 SFT。你从模型中心选择一个基础模型,让作业指向已编码的数据集,托管服务会配置计算资源、运行作业并返回定制后的模型,指标则在 Amazon SageMaker with MLflow 中跟踪。由于数据已经是受支持的格式,因此无需管理训练脚本、容器镜像或实例。作业会将其输出写入一个模型包组,该组会为其产出的每个模型进行版本管理,因此请先用 创建该组(如果名称已存在则复用),并将其传递给 th 模型中心,让作业指向已编码的数据集,托管服务会配置计算资源、运行作业并返回定制后的模型,指标则在 Amazon SageMaker with MLflow由于数据已经是受支持的格式,因此无需管理训练脚本、容器镜像或实例。
作业会将其输出写入一个 模型包组 该组会为其产出的每个模型进行版本管理,因此请先用 ModelPackageGroup.create (如果名称已存在则复用),并将其传递给训练器。然后配置 定制作业选择基础模型和训练类型,将其指向训练数据集和验证数据集以及模型包组,然后命名一个 MLflow 实验,以便在作业运行时将指标流式传输到 MLflow。
from sagemaker.train.sft_trainer import SFTTrainer
from sagemaker.train.common import TrainingType
sft_trainer = SFTTrainer(
model="huggingface-vlm-qwen3-5-4b", # Qwen3.5-4B from the JumpStart model hub
training_type=TrainingType.LORA, # TrainingType.FULL for full fine-tuning
training_dataset=train_dataset.arn,
validation_dataset=val_dataset.arn,
model_package_group=model_package_group,
mlflow_experiment_name="seqkd-shopping-exp",
s3_output_path="s3://<amzn-s3-demo-bucket>/seqkd/output/",
accept_eula=True,
)
训练器从一组有文档记录的默认超参数开始,你可以使用 sft_trainer.hyperparameters.to_dict() 并在启动前覆盖。本次运行中,我们提高序列长度以适配多步轨迹,使用 64 的 Low-Rank Adaptation (LoRA) 秩,并在小数据集上训练三个 epoch:
sft_trainer.hyperparameters.dataset_max_len = 16384
sft_trainer.hyperparameters.max_epochs = 3
sft_trainer.hyperparameters.global_batch_size = 64
sft_trainer.hyperparameters.learning_rate = 0.0001
sft_trainer.hyperparameters.lora_rank = 64
sft_trainer.hyperparameters.lora_alpha = 256 # alpha / rank = 4
sft_trainer.hyperparameters.lr_warmup_steps_ratio = 0.1
sft_trainer.hyperparameters.train_val_split_ratio = 0.95
先用试运行验证配置,然后启动托管作业,并在作业完成后读取已注册的 Model Package ARN。该 ARN 就是可供部署的产物,供评估和部署服务步骤使用:
from sagemaker.core.resources import TrainingJob
sft_trainer.train(dry_run=True) # checks the config without launching
training_job = sft_trainer.train(wait=True) # launch and wait
model_package_arn = TrainingJob.get(
training_job_name=training_job.training_job_name
).output_model_package_arn
评估学生模型
分两个阶段进行评估,先采用成本最低的方式。完整的端到端评估既耗时又耗费资源。它需要针对数百万个商品启动搜索服务器、部署端点,并对每个留出问题运行多步智能体循环。在运行完整评估之前,先使用 SageMaker AI 模型评估执行一次低成本的托管预检查,该检查会针对教师模型的真实动作,对每一步生成的步骤进行评分。
预检查需要一个评分函数。模型评估支持你提供的自定义评分器奖励函数,该函数对每个样本运行一次,将学生模型生成的动作与教师模型的动作进行比较。对于工具调用蒸馏,应评估工具选择(模型是否按正确顺序调用了正确的工具,忽略参数格式),因为这是蒸馏最需要迁移的行为。为了将重点放在选择而非语法上,奖励函数会从学生模型的动作和教师模型的参考中读取工具名称,而不要求精确匹配。 自定义评分器 你提供的奖励函数对每个样本运行一次,将学生模型生成的动作与教师模型的动作进行比较。对于工具调用蒸馏,应评估工具选择(模型是否按正确顺序调用了正确的工具,忽略参数格式),因为这是蒸馏最需要迁移的行为。为了将重点放在选择而非语法上,奖励函数会从学生模型的动作和教师模型的参考中读取工具名称,而不要求精确匹配。
def reward_function(sample, index):
pred = sample.get("model_response", "")
ref = sample.get("reference_answer", "") or sample.get("response", "")
ref_names = tool_names(ref) # tool names from the teacher's <tool_call>[...]</tool_call>
pred_names = tool_names(pred) # tool names the student emitted, parsed leniently
toolset = 1.0 if (pred_names and pred_names == ref_names) else 0.0
return {
"id": str(sample.get("id", f"sample-{index:03d}")),
"aggregate_reward_score": toolset, # the gate metric: tool selection
"metrics_list": [
{"name": "toolset_match", "value": toolset, "type": "Reward"},
],
}
将该函数注册为带版本的 Evaluator ,然后调用 为每个任务运行一次 Custom Scorer 作业,在评分学生模型的同时也评分基础模型,以便每个结果都有对比。 Evaluator.create(),然后调用 evaluate() 为每个任务运行一次 Custom Scorer 作业,在评分学生模型的同时也评分基础模型,以便每个结果都有对比。
决策规则很简单。学生模型必须在这个低成本预检查中胜过基础模型。如果做到了,说明蒸馏已经生效,值得运行更全面的端到端评估。如果没有做到,就先修复训练。
部署蒸馏后的智能体
蒸馏后的模型不会自行运行工具。微调教会它在正确的时机以文本形式发出正确的工具调用,但仍然需要有东西来执行该调用并反馈结果。在推理时,你需要用一个小的确定性执行框架包裹模型,该框架会提示模型、解析它发出的工具调用、运行工具、追加观察结果,然后再次提示,如此循环,直到模型给出最终答案。这是普通的应用代码,或者现成的服务框架,而不是强化学习,也不是同策略训练循环。
学生模型是一个基于开放权重基础模型的低秩适配(LoRA)适配器,这决定了你的部署方式。与其合并适配器并另外提供一份完整的权重副本,不如只提供一次基础模型,并将学生模型作为适配器附加在其上,作为同一个实时端点上的两个推理组件。在调用时指定推理组件即可选择由哪一个来应答,因此同一个端点可以同时服务基础模型和蒸馏后的学生模型,便于并排比较。
我们使用 Deep Java Library 大模型推理(DJL LMI)容器来服务这个开放权重模型,该容器捆绑了用于高吞吐生成的 vLLM 引擎,并固定到特定版本,而不是 latest,因此重新部署是可复现的。配套的 notebook 将推理组件调用封装在一个小的 kd_serve 辅助函数,因此启用 LoRA 服务部署基础组件、再附加学生适配器,就简化为:
from sagemaker.core import image_uris
import kd_serve
# DJL LMI container (vLLM engine bundled in); pin the version for reproducible redeploys
image_uri = image_uris.retrieve("djl-lmi", region=region, version="0.36.0")
# 1. Deploy the base model as an inference component, with LoRA serving enabled.
endpoint_name, base_ic = kd_serve.deploy_base_endpoint(
endpoint_name="seqkd-shopping-agent",
role=role,
region=region,
image_uri=image_uri,
base_model_id="Qwen/Qwen3.5-4B",
instance_types=["ml.g5.2xlarge", "ml.g6.2xlarge", "ml.g6e.2xlarge"],
max_lora_rank=64,
max_loras=4,
max_model_len=131072,
)
# 2. Resolve the student's adapter artifacts from its Model Package and attach them
# as a second inference component on the same base.
adapter_uri = kd_serve.resolve_adapter_uri(model_package_arn, region)
student_ic = kd_serve.attach_adapter(
endpoint_name=endpoint_name,
base_ic_name=base_ic,
adapter_component_name=f"{endpoint_name}-student",
adapter_s3_uri=adapter_uri,
region=region,
)
该 instance_types 列表是按顺序排列的偏好。SageMaker AI Inference 会部署在第一个有可用容量的类型上,并沿列表向下回退,因此学生模型会部署在 ml.g5.2xlarge, ml.g6.2xlarge,或 ml.g6e.2xlarge 具体取决于可用容量。基础权重只加载一次,两个推理组件共享这些权重,因此添加适配器所消耗的内存仅为适配器本身的大小,而无需再复制一份模型。
结果
我们在留出的 ShoppingBench 测试集上评估蒸馏后的学生模型,并将其与一组参考基准进行对比,这些基准界定了何为良好表现:
- 基础 Qwen3.5-4B: 未经微调的现成学生模型。这是下限。
- 蒸馏学生模型(SeqKD): 在教师轨迹上微调的同一个 4B 模型,也就是本文所描述的方法。这才是真正重要的结果。
- GLM-5.3 教师模型: 生成训练数据的前沿模型,运行在 Bedrock 上。我们同时报告了采样(温度 1.0,即生成数据时所用的设置)和贪心(温度 0)两种结果。
- 中等规模前沿模型: 一个作为智能体运行的强大专有模型,作为实际性能上限纳入对比。
每个模型都在三种购物意图(商品、店铺和优惠券)上运行同一个未经修改的 ShoppingBench 智能体循环。我们报告绝对成功率(ASR),即智能体解决的留出问题所占比例,在三种意图上取平均,缺失的轨迹计为失败。当推荐的一个或多个产品满足所有陈述的约束时,该问题即计为已解决。以下数字来自使用 LoRA 学生模型的一次运行,因此预计会有几个百分点的运行间波动。
任务成功率
主要结果是在 324 道留出问题上的绝对成功率,按意图细分并在三种意图上取平均:
图 4:按购物意图划分,在留出的 ShoppingBench 测试集上的绝对成功率(ASR)。蒸馏后的 4B 学生模型(橙色)平均比基础模型提升 24 个百分点,达到生成其数据的 GLM-5.3 教师模型以及一个中等规模前沿模型的水平
| 模型 | 产品 | 店铺 | 优惠券 | 平均 |
| 基础 Qwen3.5-4B(未微调) | 0.456 | 0.314 | 0.121 | 0.297 |
| 蒸馏学生模型(SeqKD) | 0.563 | 0.638 | 0.422 | 0.541 |
| GLM-5.3 教师模型(采样,温度 1.0) | 0.553 | 0.619 | 0.362 | 0.511 |
| GLM-5.3 教师模型(贪心解码,温度 0) | 0.544 | 0.638 | 0.439 | 0.540 |
| 中等规模前沿模型 | 0.583 | 0.667 | 0.388 | 0.546 |
总之,一个 4B 模型基本恢复了其前沿教师模型的全部任务质量,与教师的差距在运行间波动范围内。我们以两种方式报告教师模型的结果,因为它们回答的是不同的问题。贪心解码返回单条最可能的轨迹,当每个模型对每道题只有一次尝试机会时(如本次评估),这是更强的选择,因此贪心教师模型得分为 0.540,而采样教师模型为 0.511。我们以另一种方式生成训练数据,即在温度 1.0 下采样,使重复尝试探索不同的路径,评分器则可以保留那些解出任务的轨迹。这种过滤也正是学生模型能够略胜其所学习的采样教师模型的原因:拒绝采样只让它在教师模型成功解出的运行上训练,而非其平均输出,因此它模仿的是最佳切片,而非全部尝试的分布。
目标中的推理很重要
学生学习的是教师的推理过程以及它的工具调用,这一选择在很大程度上决定了最终结果。用同样的轨迹训练同一个学生模型,但去掉推理过程、只保留裸露的工具调用,平均 ASR 会从 0.541 降到 0.415:
图 5:基础模型、仅在工具调用上训练的学生模型,以及在教师推理加工具调用上训练的学生模型的绝对成功率。保留推理过程使平均 ASR 提升约 13 个百分点,且这一增益在每个意图上都成立
两个学生模型都超过了基础模型,因此即便是仅训练工具调用的版本也是实打实的提升,但真正让学生模型一路追平其教师模型的,是保留推理过程。动作之间的推理并不是可以丢弃的叙述,而是解决问题行为的一部分。
保持通用能力
ASR 衡量的是我们蒸馏出的技能。与之相权衡的是,学生模型是否保留了其通用指令遵循能力,还是过拟合于购物任务。我们通过 IFEval,这是一个通用指令遵循基准,用于对可验证的格式指令(例如“答案必须恰好为三个要点”)进行评分。请将其视为回归检查,而非排行榜。学生模型的提示严格分数从基础模型的 0.819 降至 0.741,指令严格分数从 0.872 降至 0.813,下降了约 6 到 8 分(图 6)。这是专业化带来的预期代价。学生模型牺牲了一部分通用指令遵循能力,换来了在购物任务上的大幅提升。这是一个通用指令遵循基准,用于对可验证的格式指令(例如“答案必须恰好为三个要点”)进行评分。请将其视为回归检查,而非排行榜。学生模型的提示严格分数从基础模型的 0.819 降至 0.741,指令严格分数从 0.872 降至 0.813,下降了约 6 到 8 分(图 6)。这是专业化带来的预期代价。学生模型牺牲了一部分通用指令遵循能力,换来了在购物任务上的大幅提升。如果降幅大得多,那才是警示信号,说明蒸馏把模型收窄得过头了。
图 6:基础模型与蒸馏后学生模型在 IFEval 通用指令遵循上的表现,分别按提示严格和指令严格指标衡量。小幅下降是在购物任务上专业化所预期的代价
收益与权衡
单看成功率本身并不足以证明切换的合理性。进行蒸馏的理由在于,你能以较小模型的成本和速度获得较强模型的大部分质量。这里最清晰的证据是,4B 学生模型达到了其前沿教师模型和一个中等规模前沿模型的任务质量,同时其规模小到足以在单个 G 系列实例(本例中为 NVIDIA A10G/L40/L40S)上提供服务。
成本差异既与规模有关,也与形态有关。通过 API 访问的前沿模型按 token 计费,在每项任务的每个步骤上都要计费,因此成本随使用量增长而不会停止,引言中提到的智能体循环每次都要支付这笔费用。蒸馏出的学生模型运行在你自行托管的端点上,该端点按小时对实例计费,无论它服务一项任务还是数千项任务,因此在横向扩展之前,增加的流量不会带来额外成本。你是在用按请求计费的变动成本换取按小时计费的固定成本,而你能采用的最小模型越小,这项固定成本就越低。
我们在此不公布每项任务的成本或延迟数字,因为这取决于你的流量:G 系列实例所能维持的吞吐量,以及由此摊到每项任务上的成本,都是你的请求模式、序列长度和并发度的函数。请在你自己的工作负载上,对照当前的 SageMaker AI 定价 以及你的前沿提供商的 token 定价,而不是依赖单一的头条数字。
局限性与最佳实践
任务是否适合该学生模型(容量墙)
蒸馏传递的是学生模型已经能够表示的分布,而不是它原本缺乏的原始能力。当差距过大时,学生模型只会复制教师的格式,而学不会如何解题。用大型教师的思维链训练的小模型,会产生冗长、结构良好的推理,但答案始终是错的;用同样方式训练的智能体则会发出格式正确的工具调用,却毫无进展,用改写的查询反复搜索,一直跑到步数上限也满足不了约束。两者的症状相同:轨迹变长,而准确率保持不变。一旦看到这种情况,就应停止添加数据,要么扩大学生模型,要么降低任务难度。
决定这一差距的是任务与学生模型之间的差距,而不是学生模型与教师模型之间的差距。只要任务本身在学生模型的能力范围之内,一个小模型就能恢复出大得多的教师模型的质量。这就是这里的购物智能体能够奏效的原因。教师模型是一个混合专家模型,约有 400 亿活跃参数,大约是学生模型 40 亿参数的十倍,但购物任务完全在 4B 模型的能力范围之内,因此学生模型恢复的是教师模型的质量,而不只是把轨迹拉长。
SeqKD 在何处不再有帮助
SeqKD 能提升简单和中等难度的问题,但在最难的问题上会陷入停滞,因为在这些问题上学生必须从自己的错误中恢复。对智能体而言,这一点比单步任务更为突出。在多步循环中,早期的错误会不断累积,因为学生看不到自己的错误会导致哪些状态。弥合这一差距正是 on-policy 方法的目标,这也是为什么 SeqKD 通常只是第一阶段,而不是终点。SeqKD 对基础设施的要求也最低,因为它只使用教师的文本,而 logit 蒸馏需要教师的逐 token 分布,on-policy 方法则要对学生自己的 rollout 打分。先用纯文本的 SeqKD 做热启动、再接一个 on-policy 阶段,是一种常见的做法,因此本文介绍第一阶段,而突破这一停滞则是自然的下一步。
决定结果的两个选择
首先,让训练目标与期望的行为相匹配。在这里,这意味着在每次工具调用时保留教师的推理过程,而不是压缩成最终答案,因为推理正是解决问题的一部分(“结果”部分量化了这一差距),并且用智能体在服务时所用的聊天模板来渲染每一个目标。其次,诚实地评估。要在模型未见过的数据上衡量,因为泄露的基准会抬高分数;并且要与不止一个未训练的基座模型进行比较——学生模型注定会胜过它——而是要将它与教师模型和一个独立的前沿模型并列,这样数字才能显示出学生模型距离其蒸馏来源有多近。
清理
实时推理端点只要在运行就会计费,因此在你获得结果之后,请拆除你为提供服务而搭建的端点。由于学生模型是作为基础推理组件之上的适配器推理组件来提供服务的,所以各部分必须按依赖顺序拆除:先拆适配器组件,再拆基础组件,最后拆端点及其配置。随附的 kd_serve 辅助脚本会为你完成这项工作:它以结构化方式查找适配器,因此无论组件的确切名称是什么,它都能正确清理。
import kd_serve
# Deletes the adapter IC, then the base IC, then the endpoint and its configuration
kd_serve.teardown_endpoint(endpoint_name, region=region)
训练作业本身无需清理。无服务器训练会自动配置计算资源、运行并释放它,因此已完成的作业不会持续计费。注册的模型包及其模型组也不产生计算成本,所以你可以保留它们用于记录血缘关系,并在之后重新部署。唯一持续产生费用的是你在 S3 中的数据(编码后的轨迹和输出产物),只要保留就会产生存储费用,因此请删除不再需要的数据,或将其迁移到成本更低的存储类别。
结论
在这篇博文中,我们展示了如何使用序列级知识蒸馏(SeqKD)将一个智能体购物助手蒸馏为一个更小、成本更低的学生模型。我们在 Bedrock 上使用前沿教师模型生成了高质量的单轮、多步轨迹,并仅保留那些使用 ShoppingBench 自带评分器成功解决任务的轨迹,将每一步编码为聊天格式的消息,在保留教师推理过程的同时保留其工具调用,使用 SageMaker AI 模型定制在无服务器基础设施上对学生模型进行微调,使用 MLflow 跟踪训练过程,使用 SageMaker AI 模型评估运行质量门禁检查,并将结果部署在一个运行工具循环的推理端点之后。回报是一个 4B 的学生模型,它基本恢复了其前沿教师模型的全部任务质量。由于 SeqKD 只需要教师的文本输出,因此即使教师是通过 API 访问的封闭前沿模型,它也能发挥作用;并且在学生模型接管之后,你可以不再为每一步支付前沿模型的价格。
SeqKD 是坚实的第一步,但也有其实际局限。让学生模型的规模与任务相匹配,在留出测试集上对照强参考基准进行评估,而不只是与未经训练的基座模型比较;当最难的长尾问题难以通过模仿学习解决时,可在 SeqKD 之后采用本文未涉及的 on-policy 方法,让学生模型从自身的 rollout 中学习。
要自行尝试,请参阅 配套代码仓库 以及 SageMaker AI 文档.