AWS 演示用 SeqKD 在 SageMaker AI 上将购物 Agent 蒸馏到 Qwen3.5-4B
Build efficient agents with knowledge distillation on Amazon SageMaker AI
AWS Machine Learning Blog 发布教程,讲解如何用序列级知识蒸馏(SeqKD)在 Amazon SageMaker AI 上把购物 Agent 从 frontier 教师模型 GLM-5.3 蒸馏到 Qwen3.5-4B 学生模型。
AWS 官方给出了完整的 SeqKD 智能体蒸馏流程与实测数字,包括保留推理带来的约13个点提升,方法可直接复用。
你在前沿模型上构建了一个购物助手,而且它能正常工作。购物者询问“一家至少 4.5 星的店铺里 20 美元以下的黑色腰包”,智能体会搜索商品目录、检查候选商品、核对约束条件,并最终给出推荐。但这样一次交互的实际成本是多少?一个智能体化助手并不是一次模型调用,而是一个循环:搜索、读取结果、细化、验证、决策。一个完成的任务可能需要多次调用那个前沿模型,而每一次调用你都要付费,且面向每一位购物者。在生产规模下,这种单任务成本会不断累积,侵蚀项目的投资回报率(ROI),甚至导致项目无法上线。
让智能体化应用更高效的一种方法是使用更小的模型,它响应更快、托管成本更低。问题在于质量:更小的模型通常无法胜任该任务,而这也正是当初选择大模型的原因。知识蒸馏通过把更大的“教师”模型的行为迁移到更小的“学生”模型中来弥合这一差距——训练学生模型复现教师模型的输出,从而以更小模型的速度和成本保留大部分质量。
在这篇文章中,你将学习如何使用知识蒸馏,特别是序列级知识蒸馏(SeqKD),在 Amazon SageMaker AI 上微调一个购物智能体。我们会介绍如何用前沿教师模型生成高质量的多步轨迹并保留那些能解决任务的轨迹,如何为监督微调对它们进行编码,如何在无服务器基础设施上运行训练,以及如何针对保留的基准测试集评估学生模型。
图 1:一个智能体化任务是一个循环,而不是单次调用:智能体搜索、读取结果、细化、验证并决策,而每一步都是对前沿模型的一次独立计费调用。一个完成的任务可能是多次调用,且面向每一位购物者
什么是序列级知识蒸馏(SeqKD)?
最简单且可用的蒸馏形式是序列级知识蒸馏(SeqKD)。它使用教师模型生成高质量响应,然后通过监督微调(SFT)在这些响应上训练学生模型。SeqKD 只需要教师模型的文本输出,无需访问教师模型的 logits,也无需强化学习基础设施,因此即使教师模型是闭源的或只能通过 API 访问也能使用。与基于 logits 或策略内(on-policy)等其他方法相比,它处于蒸馏光谱中成本效益最高、可移植性最强的一端。教师模型撰写答案,学生模型则在答案上进行微调。由于这些答案会以数据集的形式保存下来,你可以复用它们来训练或测试不同的学生模型,而无需重新生成数据。
原则上,你会希望学生模型匹配教师模型在所有可能输出上的完整分布。但这里的问题是:你做不到。可能的序列空间极其庞大,精确匹配并不现实。SeqKD 通过用教师模型实际生成的样本来近似这一分布,从而绕过了这个难题。
对智能体而言,教师生成的不是一段散文,而是一条轨迹:一次有序的运行,包含推理、工具调用以及它们返回的观察结果,其中每一步都以前面的所有内容为条件。让学生在整个序列上训练,教会它引出每个动作的推理过程以及多步骤的行动路径本身,而不是单段答案文本。这使得智能体任务成为序列级知识蒸馏最清晰的例证之一,因为被复制的序列正是行为本身。
图 2:SeqKD 仅通过教师的文本输出来传递行为:教师生成轨迹,学生在其上训练,结果在小模型的速度和成本下保留了大量教师的质量。训练期间不需要教师权重、logits 或梯度,因此它就像普通的监督微调一样运行。
何时选择 SeqKD
SeqKD 适用于一组特定场景:
你有输入数据但没有高质量的目标输出
这是 SeqKD 的核心场景,其本质是数据生成。你提供输入,教师来编写目标输出,因此你可以在没有现成高质量标签的任务或领域上进行微调。目标可以简单到单个答案,也可以复杂到多轮工具使用轨迹。本文选取了这一谱系中要求最高的端——智能体轨迹,这类数据尤其难以获取。单个任务会分叉出许多部分正确的运行,而区分一次成功的运行和一次失败的运行需要一个可靠的评分器。带有程序化评分器的基准测试恰好提供了这一点,因此你可以让教师生成轨迹,并只保留那些成功解决任务的轨迹。
你的教师是只能通过 API 访问的前沿模型
这是前沿模型的常见情况。你无法访问权重来自行托管,无法读取其 logits,也无法对其进行微调。这就排除了基于 logits 的蒸馏和 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 上的一个模型)。
- 在受支持的 AWS 区域中访问 SageMaker AI 模型自定义。
- Amazon SageMaker Studio(或已配置好 AWS CLI 的本地 IDE),运行 Python 3.12 或更高版本的内核,并安装 SageMaker Python SDK v3 和
boto3(请固定你的版本号)。笔记本环境至少使用 ml.m5.4xlarge 实例,因为数据准备和本地评分步骤在其 CPU 上运行。 - 具备 Python、监督微调以及使用工具的智能体的基本知识。
- 随附的代码仓库,以及一个带有程序化评分器的任务(本演练使用 ShoppingBench 测试集。请替换为你自己的任务和评分器)。
解决方案概览
本节从头到尾介绍整个流水线。前沿教师模型——Amazon Bedrock 上的 GLM-5.3——充当 ShoppingBench 智能体来生成轨迹。基准测试自带的评分器只保留成功解决任务的轨迹,这些成功轨迹再通过 SageMaker AI 模型自定义的无服务器监督微调来微调 Qwen3.5-4B 学生模型。蒸馏出的学生模型随后在留出的 ShoppingBench 划分数据集上评估,并通过运行工具循环的 SageMaker AI Inference 实时端点提供服务。
图 3:在 Amazon SageMaker AI 上蒸馏购物智能体:Amazon Bedrock 上的前沿教师模型在训练问题上充当 ShoppingBench 智能体,ShoppingBench 自带的评分器只保留解决任务的轨迹(拒绝采样),每条成功轨迹按步骤编码,学生模型通过 Amazon SageMaker AI 模型自定义进行无服务器微调,结果在留出的 ShoppingBench 测试划分上评估
准备数据集
获取值得模仿的轨迹是难点,而 ShoppingBench 评分器正是区分成功与失败运行的关键。ShoppingBench 发布了 750 个合成的购物问题,商品、店铺和优惠券三类意图各 250 个,每个问题都有隐藏的标准答案和定义成功的约束条件。我们将它们一次性划分为 426 个训练问题和 324 个留出的评估问题(商品 103、店铺 105、优惠券 116),并固定索引以使两个集合不重叠。
对于每个训练问题,一个前沿教师模型——开启推理、temperature 为 1.0 的 Amazon Bedrock 上的 GLM-5.3——充当 ShoppingBench 智能体。它在约 270 万件商品的目录中搜索、检查候选商品,并推荐满足约束条件的商品。我们对每个问题采样教师两次,得到 850 条轨迹,然后让 ShoppingBench 自带的评分器只保留成功解决任务的轨迹。temperature 1.0 是有意为之。相比贪婪解码,采样使两次尝试探索不同的查询和推理路径,这让更多问题至少拥有一条成功轨迹,并提供更多样的学习素材。这个拒绝采样步骤将一个能力强但并不完美的教师转化为干净的训练集,因为教师只需足够频繁地做对,其成功运行就值得模仿。
每条成功轨迹都是一个完整的回合:购物请求、教师每一步的推理和工具调用(搜索、检查、推荐、终止),以及环境返回的观察结果。我们保留了推理和动作,结果显示正是这些内容让学生模型大部分程度地达到了其教师的水准。
为训练编码轨迹
序列级蒸馏是对已经存在的轨迹进行监督学习,因此训练期间不会有任何工具调用,应用程序只会看到文本。由于这些轨迹是多步骤的工具使用对话,我们把每条轨迹编码为聊天轮次的messages列表,并让学生模型自己的聊天模板来渲染它,使训练器看到的 token 与智能体在推理时产生的 token 一致。
每条轨迹在每个步骤都变成一行训练数据。user轮携带系统提示词和到目前为止的对话。assistant轮携带该步骤的推理及其引发的工具调用,即一个 … 块,后跟工具调用框架原生格式的 [{…}]。损失只落在assistant轮上,因此模型学习产生下一步的推理和行动,而不是预测环境返回的观察结果。包含多个步骤的轨迹会扩展为多行。这里的已解决轨迹在过滤后扩展为大约 2,600 行训练数据和 130 行验证数据。
每行数据将到目前为止的对话(user)与要学习的一个下一步(assistant)配对:
{
"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>"}
]
}后续的每个步骤都成为单独的一行,把新的观察结果追加到 user 内容中,下一步的推理和工具调用作为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中。由于数据已经是受支持的格式,因此无需管理训练脚本、容器镜像或实例。
该任务会将其输出写入一个 模型包组,该组会对其产出的每个模型进行版本管理,因此请先用 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 的低秩适配(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 模型评估运行一个低成本的托管预检,它会将生成的每一步与教师模型的真实动作进行逐一评分。
预检需要一个评分函数。模型评估支持你提供的 Custom Scorer 奖励函数,它对每个样本运行一次,将学生模型生成的动作与教师模型的动作进行比较。对于工具调用蒸馏,应评判工具选择(模型是否按正确顺序调用了正确的工具,忽略参数格式),因为这是蒸馏最需要迁移的行为。为了让重点放在选择而非语法上,该奖励函数从学生模型的动作和教师模型的参考中读取工具名称,而不要求格式完全匹配。
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.create() 将该函数注册为带版本的 Evaluator,然后调用 evaluate() 为每个任务运行一个 Custom Scorer 作业,同时对学生模型和基础模型进行评分,使每个结果都有对比。
决策规则很简单。学生模型必须在这个低成本预检中胜过基础模型。如果做到了,说明蒸馏已生效,值得运行更全面的端到端评估。如果没有做到,则应先修复训练。
服务化蒸馏后的智能体
蒸馏后的模型不会自己运行工具。微调教会它在正确的时刻以文本形式发出正确的工具调用,但仍然需要某种东西来执行该调用并反馈结果。在推理时,你需要将模型包裹在一个小型的确定性框架中:提示模型、解析它发出的工具调用、运行该工具、追加观察结果,然后再次提示模型,循环往复,直到模型产出最终答案。这只是普通的应用程序代码,或现成的服务框架,而不是强化学习,也不是在策略训练循环。
学生模型是开源权重基座模型上的低秩适配(LoRA)适配器,这决定了它的部署方式。与其把适配器合并后另起一份完整的权重副本,不如只部署一次基座模型,再把学生模型作为适配器挂载其上,作为同一实时端点上的两个推理组件。调用时指定推理组件名称即可选择由谁响应,因此同一个端点可以同时服务基座模型和蒸馏出的学生模型,方便进行并排对比。
我们使用 Deep Java Library Large Model Inference(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 上。我们同时报告其采样模式(temperature 1.0,即生成数据时的设置)和贪婪模式(temperature 0)的结果。
- 中等规模前沿模型:一个作为智能体运行的强大专有模型,作为实际的上限参考。
每个模型都在三种购物者意图(product、shop 和 voucher)上运行相同的、未经修改的 ShoppingBench 智能体循环。我们报告绝对成功率(ASR),即智能体解决的留出问题所占的比例,取三种意图的平均值,缺失的轨迹计为失败。当推荐的一个或多个产品满足所有给定约束时,该问题才算解决。以下数字来自 LoRA 学生模型的单次运行,因此预期会有几个百分点的运行间波动。
任务成功率
核心结果是在 324 个留出问题上的绝对成功率(ASR),按意图细分并取三者平均:
图 4:按购物者意图在留出的 ShoppingBench 测试集上的绝对成功率(ASR)。蒸馏后的 4B 学生模型(橙色)平均比基座模型提高 24 个百分点,达到生成其数据的 GLM-5.3 教师模型和中等规模前沿模型的水平
| 模型 | product | shop | voucher | 平均 |
| 基座 Qwen3.5-4B(无微调) | 0.456 | 0.314 | 0.121 | 0.297 |
| 蒸馏学生模型(SeqKD) | 0.563 | 0.638 | 0.422 | 0.541 |
| GLM-5.3 教师模型(采样,temperature 1.0) | 0.553 | 0.619 | 0.362 | 0.511 |
| GLM-5.3 教师模型(贪婪,temperature 0) | 0.544 | 0.638 | 0.439 | 0.540 |
| 中等规模前沿模型 | 0.583 | 0.667 | 0.388 | 0.546 |
总而言之,这个 4B 模型几乎完全恢复了其前沿教师模型的任务质量,在运行间噪声范围内与之持平。我们以两种方式报告教师模型的结果,因为它们回答的是不同的问题。贪心解码返回单条最可能的轨迹,当每个模型对每个问题只有一次尝试机会时(如本次评估),这是更强的选择,这也是为什么贪心教师得分为 0.540,而采样教师为 0.511。我们以另一种方式生成训练数据:在 temperature 1.0 下采样,让重复尝试探索不同的路径,评分器可以保留成功解决任务的那些。这种过滤也是学生模型能略胜其学习对象——采样教师——的原因,因为拒绝采样只在教师成功解决任务的运行上训练它,而不是其平均输出,所以它模仿的是最佳尝试切片,而非完整的尝试分布。
目标中的推理很重要
学生模型不仅从教师模型的工具调用中学习,还从其推理中学习,这一选择是结果的重要组成部分。让同一个学生在相同的轨迹上训练但剥离推理内容,只留下裸的工具调用,平均 ASR 从 0.541 降至 0.415:
图 5:基础模型、仅用工具调用训练的学生模型、以及用教师推理加工具调用训练的学生模型的绝对成功率。保留推理带来约 13 个点的平均 ASR 提升,且该提升在每种意图上都成立
两个学生模型都超过了基础模型,因此即使是仅用工具调用的版本也是真正的改进,但保留推理才是让学生模型走完剩余路程、追上教师模型的关键。动作之间的推理并不是可以丢弃的叙述,它是解决任务的行为的一部分。
保持通用能力
ASR 衡量的是我们为之蒸馏的技能。对应的权衡是学生模型是否保留了其通用的指令遵循能力,还是过拟合到了购物任务上。我们用 IFEval 来检查,这是一个通用指令遵循基准,评估可验证的格式指令,例如“恰好用三个要点回答”。把它视为回归检查,而不是排行榜。学生模型的 prompt-strict 得分从基础模型的 0.819 降至 0.741,instruction-strict 得分从 0.872 降至 0.813,下降约 6 到 8 个点(图 6)。这是专业化的预期代价。学生模型用一些通用指令遵循能力换取了在购物任务上的大幅提升。如果降幅大得多,那才是蒸馏把模型收窄过度的警告信号。
图 6:基础模型与蒸馏学生模型在 prompt-strict 和 instruction-strict 指标上的 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)方法,让学生模型从自身的轨迹中学习。
要亲自尝试,请参阅随附的代码仓库和SageMaker AI 文档。
关于作者
来源: AWS Machine Learning Blog · aws.amazon.com





