在线AI教程 / AI上云
[Amazon Bedrock] 在 Amazon SageMaker AI 上用知识蒸馏构建高效 agent
人工智能 在 Amazon SageMaker AI 上用知识蒸馏构建高效 agent 你基于一个前沿模型(frontier model)搭了一个购物助手,它确实能跑通。有顾客提出
人工智能
在 Amazon SageMaker AI 上用知识蒸馏构建高效 agent
你基于一个前沿模型(frontier model)搭了一个购物助手,它确实能跑通。有顾客提出「找一个 20 美元以下、来自评分至少 4.5 星店铺的黑色腰包」,agent 就会搜索商品目录、逐个查看候选商品、核对约束条件,并最终给出一个明确的推荐。但这一次交互的实际成本是多少?一个 agentic 助手不是一次模型调用,而是一个循环:搜索、读取结果、细化、验证、决策。一个完成的任务可能是对那个前沿模型的多次调用,而每一次都要付费,每一位顾客都要付费。到了生产规模,这种按任务累积的成本会不断放大、侵蚀项目的 ROI,最终让项目无法上线。
让一个 agentic 应用更高效的一条路是使用更小的模型——它响应更快,托管起来也更划算。难点在质量:小模型通常在这类任务上达不到要求,而这正是当初选择更大模型的原因。知识蒸馏(knowledge distillation)补上了这个缺口:它把更大的「教师」模型的行为迁移进更小的「学生」模型,训练学生去复现教师的输出,从而在小模型的速度和成本下保住大部分质量。
在这篇文章里,你将了解如何使用知识蒸馏,特别是 sequence-level knowledge distillation(SeqKD,序列级知识蒸馏),在 Amazon SageMaker AI 上微调一个购物 agent。我们会讲到:如何用前沿的教师模型生成高质量的多步轨迹(trajectory)并只保留解开任务的那些,如何把它们编码以用于监督微调(SFT),如何在 serverless 基础设施上运行训练,以及如何在留出的 benchmark 上评测学生模型。

图 1:一个 agentic 任务是一个循环,不是单次调用——agent 搜索、读取结果、细化、验证、决策,而每一步都是对前沿模型的一次单独计费调用。一个完成的任务可能是多次调用,而且发生在每一位顾客身上
什么是 sequence-level knowledge distillation(SeqKD)?
最简单可用的蒸馏形式就是 sequence-level knowledge distillation(SeqKD)。它用教师模型生成高质量回答,然后用监督微调(SFT)在这些回答上训练学生模型。SeqKD 只需要教师的文本,既不需要访问教师的 logits,也不需要强化学习基础设施,所以即使教师是闭源的、或者只能通过 API 访问,它也照样可用。和 logit-based、on-policy 等其他方案相比,它处在蒸馏谱系里最具成本效益、也最容易移植的那一端。教师写下答案,学生在这些答案上做微调。因为这些答案会以数据集的形式留存下来,你可以反复用它们来训练或测试不同的学生,而不必重新生成数据。
原则上,你会希望学生去匹配教师在整个可能输出空间上的完整分布。这里的问题是:你做不到。可能的序列空间大得离谱,精确匹配并不可行。SeqKD 的绕法是用教师实际产出的样本来近似那个分布。
对一个 agent 来说,教师产出的「序列」不是一段成文的回答,而是一条轨迹:推理、工具调用(tool call)以及这些调用返回的 observation 按顺序连成的一整段运行,其中每一步都以它之前的全部内容作为条件。把学生放在这整段序列上训练,教给它的既有导向每个动作的推理,也有多步行动方案本身,而不只是一段答案文字。这让 agentic 任务成为序列级 KD 最干净的例证之一,因为被复制的那段序列本身就是行为。

图 2:SeqKD 只通过教师的文本输出来迁移行为:教师生成轨迹,学生在这些轨迹上训练,结果是在小模型的速度和成本下保住教师的大部分质量。训练期间不需要教师的权重、logits 或梯度,所以它就以普通的监督微调方式运行
什么时候 SeqKD 是正确选择
SeqKD 适配下面这一组特定情形:
你有输入,但没有高质量的目标输出
这是 SeqKD 的核心场景,它本质上就是数据生成:你来提供输入,教师来写 targets,于是你可以在一个没有现成高质量标签的任务或领域上做微调。targets 可以简单到单个回答,也可以复杂到多轮工具使用轨迹。本文取的是这个谱系里要求最高的一端——agentic 轨迹——它尤其难以获得。同一个任务会派生出许多部分正确的运行,而要把解开的运行和失败的运行区分开,需要一个可靠的 grader。一个带程序化打分器(scorer)的 benchmark 正好提供这一点,于是你可以让教师去生成轨迹,只保留解开任务的那些。
你的教师是只能通过 API 访问的前沿模型
这是前沿模型的常态。你访问不到它的权重、无法自己托管它,读不到它的 logits,也无法对它做微调。这就排除了 logit-based 蒸馏和 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 Region 中拥有 SageMaker AI model customization 的访问权限。 Amazon SageMaker Studio(或一个本地 IDE,并配置好 AWS CLI),运行 Python 3.12 或更高版本的 kernel,并装有 SageMaker Python SDK v3 和 boto3(请锁定你的版本)。notebook 环境至少要用一个 ml.m5.4xlarge 实例,因为数据准备和本地打分步骤都跑在它的 CPU 上。 具备 Python、监督微调,以及使用工具的 agent 方面的工作知识。 配套的代码仓库,以及一个带程序化打分器的任务(本文的走查使用 ShoppingBench 测试集。请换成你自己的任务和 grader)。
方案总览
本节把整条 pipeline 从头走到尾。一个前沿教师——Amazon Bedrock 上的 GLM-5.3——充当 ShoppingBench 的 agent 来生成轨迹。这个 benchmark 自带的打分器只保留解开任务的那些,而这些解开的轨迹用来微调一个 Qwen3.5-4B 学生,采用 SageMaker AI model customization 上的 serverless 监督微调。随后在留出的 ShoppingBench 划分(split)上评测蒸馏出的学生,并通过一个运行工具循环的 SageMaker AI Inference 实时端点(endpoint)把它对外服务。

图 3:在 Amazon SageMaker AI 上蒸馏一个购物 agent:Bedrock 上的一个前沿教师在训练题目上充当 ShoppingBench 的 agent,ShoppingBench 自带的打分器只保留解开任务的轨迹(rejection sampling),每条解开的轨迹按步骤编码,学生在 Amazon SageMaker AI model customization 上以 serverless 方式微调,最终在留出的 ShoppingBench 测试划分上评测结果
准备数据集
找到值得模仿的轨迹是最难的一环,而 ShoppingBench 的打分器正是把成功的运行和失败的运行分开的那把尺子。ShoppingBench 发布了 750 道合成的购物题目,product、shop、voucher 三种意图各 250 道,每一题都带一个隐藏的 gold answer 和定义「成功」的约束条件。我们把它们一次性划分为 426 道训练题和 324 道留出的评测题(product 103、shop 105、voucher 116),用的是固定索引,所以两个集合不会重叠。
对每一道训练题,一个前沿教师——Amazon Bedrock 上开启推理、temperature 设为 1.0 的 GLM-5.3——充当 ShoppingBench 的 agent。它在约 270 万商品的目录里搜索、查看候选商品,并推荐那些满足约束的。每一题我们把教师采样两次,得到 850 条轨迹,然后让 ShoppingBench 自带的打分器只保留解开任务的那些。temperature 取 1.0 是有意的。用采样而不是 greedy 解码,会让两次尝试去探索不同的查询和推理路径,从而让更多题目至少有一条解开的轨迹,可学习的样本集也更丰富。这个 rejection sampling 步骤把一个「能干但不完美」的教师变成一份干净的训练集,因为教师只要对得足够频繁,解开的运行就值得模仿。
每条解开的轨迹都是一个完整的 episode:顾客请求、教师在每一步的推理和工具调用(search、inspect、recommend、terminate),以及环境返回的 observation。我们把推理和行动一起保留下来,后文的 Results 表明,正是这一点把学生送完了通往教师的大半程。
为训练编码轨迹
序列级蒸馏就是在已经存在的轨迹上做监督学习,所以训练期间不会有任何工具被真正调用,应用看到的只有文本。因为这些轨迹是多步的工具使用对话,我们把每条轨迹编码成由若干 chat 轮次组成的 messages 列表,并让学生自己的 chat template 去渲染它,这样训练器看到的 token 就和 agent 在推理时产出的 token 一致。
每条轨迹按步骤展开,每一步生成一行训练数据。user 轮带着 system prompt 和此前的对话。assistant 轮带着这一步的推理以及它所导向的那个工具调用,即一个 … 块后面跟一个 [{…}],按 harness 的原生形式来。loss 只落在 assistant 轮上,所以模型学的是产出下一个「推理加动作」,而不是去预测环境返回的 observation。一条轨迹有若干步就展开成若干行。这里,解开的轨迹经过滤后展开成约 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": "
之后的每一步都各自成一行:把新的 observation 追加到 user 内容里,把下一个「推理加工具调用」作为 assistant 的 target,一直到最后的 recommend_product 和 terminate 调用。每一行都要用你模型的确切 chat template 来渲染。
划分数据要按题目划分,而不是按行划分。教师每题采样两次,每条轨迹又展开成多步,所以许多行共用同一道题目;按行划分会让几乎相同的上下文泄漏到训练集和验证集之间。整题留出(约百分之五)才能让验证集保持诚实。
按步编码的代价是重复的上下文,因为每一行都把不断变长的对话重述一遍,而 serverless 微调会把超过 dataset_max_len 的内容截断,所以越过上限的那一行会丢掉它需要的一部分内容。与其让上限决定切掉什么,不如把每一行控制在预算之内。最有效的办法是压缩工具的 observation,它们在 token 总数里占大头。对商品搜索结果,只保留下一步决策要用的那些字段(id、price、rating,以及正在核对的约束),丢掉原始的商品文案,并限制每次搜索保留多少条结果。如果轨迹仍然超出预算,就限制每条轨迹的步数。这里 dataset_max_len 是 16,384 个 token,在把 observation 修剪之后能装下大多数轨迹。
每个划分的 JSONL 进入 S3 之后,用 SageMaker Python SDK v3 的 DataSet.create 把每个文件注册为一个 customization dataset asset,它返回训练任务要消费的 train_dataset 和 val_dataset。
微调学生模型
训练使用 SageMaker AI model customization 搭配 serverless SFT。你从 model hub 里选一个基础模型,把任务指向编码好的数据集,托管服务就会准备好算力、运行任务并返回定制后的模型,指标记录在 Amazon SageMaker with MLflow 里。因为数据已经是受支持的格式,你不用管理训练脚本、容器镜像或实例。
这个任务会把输出写进一个 model package group,它为它产出的每一个模型做版本管理,所以先用 ModelPackageGroup.create 创建这个组(如果同名已经存在就直接复用它),再把它传给 trainer。然后配置 customization job:选定基础模型和训练类型,指向训练集与验证集以及 model package group,再指定一个 MLflow experiment 的名字,这样任务运行期间指标就会流向 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://
trainer 从一组有文档记录的默认超参数出发,你可以用 sft_trainer.hyperparameters.to_dict() 查看它们,并在启动任务之前覆盖它们。本次运行我们把序列长度调高以容纳多步轨迹,Low-Rank Adaptation(LoRA)的 rank 用 64,并在这个不大的数据集上训练三个 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
先用一次 dry run 校验配置,然后启动托管任务,结束后读出注册好的 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
评测学生模型
评测分两个阶段,成本低的先做。完整的端到端评测既耗时又耗资源:它要在数百万商品之上架起一个搜索服务、部署端点,并在每一道留出题目上跑一个多步的 agent 循环。在跑完整评测之前,先用 SageMaker AI model evaluation 跑一次低成本的托管预检,它每次只针对一个生成的步骤打分,和教师的 ground-truth 动作做比对。
预检需要一个打分函数。Model evaluation 支持你自己提供的一个 Custom Scorer reward function,它对每个样本运行一次,把学生生成的动作和教师的动作做比较。对工具调用蒸馏,要给工具选择打分(模型是否按顺序调用了正确的工具,忽略参数的格式),因为那正是蒸馏最需要迁移的行为。为了让关注点留在「选择」而不是「语法」上,这个 reward function 分别从学生的动作和教师的参考动作里读出工具名,并不要求格式完全匹配。
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
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 任务,并且把基础模型和学生模型一起打分,这样每个结果都有一份对照。
决策规则很直接。学生必须在这次低成本的预检上打败基础模型。做到了,说明蒸馏生效了,值得再跑更全面的端到端评测;做不到,就先回去修训练。
服务蒸馏后的 agent
蒸馏出来的模型自己不会执行工具。微调教会它在正确的时机以文本形式发出正确的工具调用,但仍然要有东西去执行这个调用并把结果喂回去。在推理时,你把模型包进一个小的、确定性的 harness 里:给模型提示、解析它发出的工具调用、运行工具、追加 observation、再次提示,如此循环直到模型给出最终答案。这就是普通的应用代码,或者一个现成的 serving 框架,不是强化学习,也不是 on-policy 训练循环。
学生是一个建在开放权重基础模型之上的 Low-Rank Adaptation(LoRA)adapter,这会影响你部署它的方式。与其把 adapter 合并进去、再托管第二份完整的权重,不如只托管一次基础模型,把学生作为 adapter 挂在它上面——也就是同一个实时端点上的两个 inference component。调用时指定 inference component 的名字来选择由哪一个作答,所以同一个端点既能为基础模型服务,也能为蒸馏出的学生服务,方便你并排比较。
我们用 Deep Java Library Large Model Inference(DJL LMI)容器来服务这个开放权重模型,它内置打包了 vLLM 引擎以实现高吞吐生成,并锁定到具体版本而不是 latest,这样重新部署就是可复现的。配套的 notebook 把 inference component 的调用包进了一个小的 kd_serve helper,于是「先启用 LoRA serving 部署基础 component,再挂上学生的 adapter」这一步就简化为:
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 上。基础权重只加载一次,两个 inference component 共享它们,因此多挂一个 adapter 花掉的是 adapter 量级的内存,而不是模型的第二份副本。
结果
我们在留出的 ShoppingBench 测试划分上评测蒸馏出的学生,对照的是一组能框定「什么才算好」的参照点:
基础 Qwen3.5-4B: 开箱即用、没有做过微调的同一个学生模型。这是下限。 蒸馏后的学生(SeqKD): 同一个 4B 模型,在教师的轨迹上做了微调,也就是本文描述的方法。这才是重点结果。 GLM-5.3 教师: 生成了训练数据的那个前沿模型,跑在 Bedrock 上。我们既报告它采样时的结果(temperature 1.0,也就是产出那批数据的设置),也报告 greedy 时的结果(temperature 0)。 中等尺寸的前沿模型: 一个强的专有模型作为 agent 运行,作为一个实际可达的上限纳入对照。
每个模型都跑同一个未做修改的 ShoppingBench agent 循环,覆盖三种顾客意图(product、shop 和 voucher)。我们报告 Absolute Success Rate(ASR,绝对成功率),即 agent 解开的留出题目占比,对三种意图取平均,缺失的轨迹按失败计入。当推荐的一个或多个商品满足所有给出的约束时,这道题算解开。下面的数字来自 LoRA 学生的单次运行,所以运行之间的波动会有几个点。
任务成功率
核心结果是 324 道留出题目上的 Absolute Success Rate,按意图拆分,并对三者取平均:

图 4:留出的 ShoppingBench 测试划分上的 Absolute Success Rate(ASR),按顾客意图拆分。蒸馏后的 4B 学生(橙色)平均比基础模型提高 24 个点,达到了为它生成数据的 GLM-5.3 教师以及一个中等尺寸前沿模型的水平
模型productshopvoucher平均值 基础 Qwen3.5-4B(未微调)0.4560.3140.1210.297 蒸馏后的学生(SeqKD)0.5630.6380.4220.541 GLM-5.3 教师(采样,temperature 1.0)0.5530.6190.3620.511 GLM-5.3 教师(greedy,temperature 0)0.5440.6380.4390.540 中等尺寸的前沿模型0.5830.6670.3880.546
概括来说,一个 4B 模型几乎把它前沿教师的全部任务质量都恢复了出来,在运行之间的噪声范围内与教师打平。我们用两种方式报告教师,因为它们回答的是不同的问题。Greedy 解码返回单条最可能的轨迹,在每个模型每道题只有一次尝试的场合(本次评测就是这样)它是更强的选择,这也是 greedy 教师拿到 0.540、而采样教师只有 0.511 的原因。我们生成训练数据时用的是另一种方式:在 temperature 1.0 下采样,让重复的尝试去探索不同的路径,再由打分器把解开任务的那些保留下来。这层过滤也正是学生能略微超过它所学习的采样教师的原因——rejection sampling 让学生只在教师的解开运行上训练,而不是它的平均输出,所以它模仿的是一个 best-of 的切片,而不是全部尝试的完整分布。
保留 target 里的推理很重要
学生既从教师的推理里学习,也从它的工具调用里学习,而这个选择承担了结果里很大的一部分。把同一个学生放在同样的轨迹上训练、但把推理剥掉,只留下赤裸裸的工具调用,平均 ASR 就从 0.541 掉到 0.415:

图 5:基础模型、只用工具调用训练的学生,以及用教师的推理加工具调用训练的学生三者的 Absolute Success Rate。保留推理能带来约 13 个点的平均 ASR 提升,而且这个增益在每一种意图上都成立
两个学生都超过了基础模型,所以即使只用工具调用那一版也是实实在在的改进;但把推理保留下来,才是把学生送到教师那个水平的最后一段路。动作之间的推理不是可以丢掉的旁白,它就是解开任务的那个行为的一部分。
保住通用能力
ASR 衡量的是我们为之做蒸馏的那项技能。而与之制衡的问题是:学生到底保住了它通用的 instruction-following 能力,还是对购物任务过拟合了。我们用 IFEval 来检查——这是一个通用指令遵循 benchmark,对「用恰好三个 bullet 来回答」这类可验证的格式指令打分。把它当成一次回归检查,而不是排行榜。学生的 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 计费,在每个任务的每一步上都计费,所以成本随用量增长、永不停止,而开头提到的那个 agentic 循环每一次都要付这笔过路费。蒸馏出的学生跑在你自己托管的端点上,按小时为实例计费,无论它服务一个任务还是一千个任务都一样,所以在扩容之前额外的流量不需要再多付一分钱。你是在把一笔可变的按请求费用,换成一笔固定的按小时费用,而你能够将就用的模型越小,这笔固定成本就越小。
这里我们不公布单任务的成本或延迟数字,因为它取决于你的流量:一个 G 系列实例能撑起的吞吐、以及由此摊到每个任务上的成本,都是你的请求模式、序列长度和并发度的函数。请在你自己真实的工作负载上,对照当前的 SageMaker AI 定价 和你那家前沿模型供应商的 token 定价来测量,而不是依赖一个笼统的宣传数字。
局限与最佳实践
任务是否装得下学生(那堵容量墙)
蒸馏迁移的是学生本来就能表示的那个分布,而不是它本来就缺的能力。当差距过大时,学生只会照抄教师的格式,却学不会解题。一个小模型在大教师的思维链上训练,会产出很长、结构很规整的推理,但答案仍然是错的;一个用同样方式训练的 agent 会发出格式正确的工具调用,却毫无进展——换个措辞反复搜索,一路跑到步数上限也没能满足约束。两种情况的症状是一样的:轨迹越来越长,准确率原地不动。看到这种现象就停止继续加数据,要么把学生做大,要么降低任务难度。
决定这件事的那个差距,是在任务和模型之间,而不是在学生和教师之间。只要任务本身够得着,一个小学生完全可以恢复一个大得多的教师的质量。这也正是这里的购物 agent 能成的原因:教师是一个 mixture-of-experts 模型,激活参数约 400 亿,大约是学生 40 亿的十倍,但购物这个任务远在 4B 模型的能力范围之内,所以学生恢复的是教师的质量,而不是只把自己的轨迹拉长。
SeqKD 在哪里就不再有帮助了
SeqKD 能抬高简单题目和中等难度题目的表现,但在最难的题目上会进入平台期——那些需要学生从自己的错误里恢复过来的题。对 agent 来说,这一点比单步任务更明显:在多步循环里,一个早期的错误会不断复利放大,因为学生看不到「自己的错误会通向的那些状态」。而补上这个缺口正是 on-policy 方法所针对的目标,所以 SeqKD 通常是第一阶段,而不是终点线。SeqKD 对你基础设施的要求也是最低的,因为它只用教师的文本;而 logit 蒸馏需要教师的逐 token 分布,on-policy 方法则要给学生的自我 rollout 打分。先用纯文本的 SeqKD 做一个 warm start,再接一个 on-policy 阶段,是一种常见的配方,所以本文讲的是第一阶段;而要越过平台期,下一步自然就是它。
两个决定成败的选择
第一,把训练 target 和你想要的行为对齐。在这里,这意味着把教师的推理和每个工具调用一起保留,而不是压缩成最终答案——因为推理本身就是解开任务的一部分(Results 量化了这个差距),并且每一个 target 都要用 agent 在服务时使用的那个 chat template 来渲染。第二,诚实地评测。在模型没见过的数据上测量,因为泄漏的 benchmark 会把分数抬高;而且不要只跟没训练过的基础模型比——那是学生必然超过的对象——要把教师和一个独立的前沿模型摆在旁边一起比,这样这个数字才能说明学生离它的蒸馏来源有多近。
清理资源
实时推理端点只要还在运行就持续计费,所以在拿到你的结果之后,把为服务而架起来的那个端点拆掉。因为学生是作为 adapter inference component 挂在基础 inference component 之上的,这些部件必须按依赖顺序往下拆:先 adapter component,再基础 component,然后才是端点和它的配置。配套的 kd_serve helper 会替你做这件事,它按结构找出那个 adapter,所以无论 component 的具体名字是什么都能清理干净:
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)
训练任务本身不需要清理。Serverless 训练会准备好算力、运行、然后自动释放,所以一个已完成的任务不会留下持续的费用。注册好的 model package 和它的组也不产生计算成本,你可以为了 lineage 把它们留着,以后再重新部署。真正一直在花钱的只有你在 S3 里的数据(编码后的轨迹和输出产物),只要你保留着,就会有存储成本,所以删掉不再需要的部分,或者把它们移到成本更低的存储类别。
结论
在这篇文章里,我们展示了如何用 sequence-level knowledge distillation(SeqKD)把一个 agentic 购物助手蒸馏进一个更小、成本更低的学生模型。我们用 Bedrock 上的前沿教师生成了高质量的单轮、多步轨迹,并用 ShoppingBench 自带的打分器只保留解开任务的那些;把每一步编码成 chat 格式的消息,让教师的推理和它的工具调用一起保留;在 SageMaker AI model customization 的 serverless 基础设施上对学生做微调,用 MLflow 追踪训练,用 SageMaker AI model evaluation 跑一次质量门禁检查,最后把结果服务在一个运行工具循环的推理端点之后。回报是一个 4B 学生,它恢复了自己前沿教师几乎全部的任务质量。因为 SeqKD 只需要教师的文本输出,即使教师是你只能通过 API 访问的闭源前沿模型,它也照样可用;而一旦学生接管,你就不必再在每一步上支付前沿模型的价格。
SeqKD 是一个很强的第一步,也有真实的边界。让学生的尺寸和任务相匹配;在留出的测试套件上、对照强的参照点来评测,而不是只跟没训练过的基础模型比;当最难的那批题目难以靠模仿拿下时,就在 SeqKD 之后接上本文没有覆盖的 on-policy 方法,让学生从自己的 rollout 里学习。
想自己动手试试,参见配套的代码仓库和 SageMaker AI 文档。
关于作者