{"slug": "build-efficient-agents-with-knowledge-distillation-on-amazon-sagemaker-ai", "title": "Build efficient agents with knowledge distillation on Amazon SageMaker AI", "summary": "Amazon Web Services published a technical guide on using sequence-level knowledge distillation (SeqKD) to fine-tune a smaller shopping agent on Amazon SageMaker AI, training the student model on multi-step trajectories generated by a frontier teacher model. The approach requires only the teacher's text output, with no access to teacher logits and no reinforcement-learning infrastructure, so it works with closed or API-only teacher models. AWS positions SeqKD as the most cost-effective and portable distillation method for cutting the per-task cost of agentic loops that make multiple frontier-model calls per completed task.", "body_md": "## [Artificial Intelligence](https://aws.amazon.com/blogs/machine-learning/)\n\n# Build efficient agents with knowledge distillation on Amazon SageMaker AI\n\nYou built a shopping assistant on a frontier model, and it works. A shopper asks for “a black waist bag under twenty dollars from a shop with at least 4.5 stars,” and the agent searches the catalog, inspects candidates, checks the constraints, and commits to a recommendation. But what does that interaction actually cost? An agentic assistant is not one model call, it is a loop: search, read results, refine, verify, decide. A single completed task can be multiple calls to that frontier model, and you pay for every one, on every shopper. At production scale, that per-task cost compounds, erodes the project’s ROI, and stops the project from shipping.\n\nOne way to make an agentic application more efficient is to use a smaller model, which is faster to respond and more cost-effective to host. The catch is quality: a smaller model usually falls short on the task, which is why the larger one was chosen in the first place. Knowledge distillation closes that gap by transferring the behavior of the larger “teacher” model into a smaller “student”, training the student to reproduce the teacher’s outputs, so it keeps much of the quality at the smaller model’s speed and cost.\n\nIn this post you will learn how to use knowledge distillation, and specifically sequence-level knowledge distillation (SeqKD), to fine-tune a shopping agent on [Amazon SageMaker AI](https://aws.amazon.com/sagemaker/ai/). We cover how to generate high-quality multi-step trajectories with a frontier teacher and keep the ones that solve the task, how to encode them for supervised fine-tuning, how to run the training on serverless infrastructure, and how to evaluate the student against a held-out benchmark.\n\n## What is sequence-level knowledge distillation (SeqKD)?\n\nThe simplest applicable form of distillation is sequence-level knowledge distillation (SeqKD). It uses a teacher model to generate high-quality responses, then trains a student model on those responses with supervised fine-tuning (SFT). SeqKD needs only the teacher’s *text*, with no access to teacher logits and no reinforcement-learning infrastructure, so it works even with a closed or API-only teacher. Compared to other approaches such as logit-based or on-policy, it sits at the most cost-effective and portable end of the distillation spectrum. The teacher writes the answers, and the student is fine-tuned on them. Because those answers persist as a dataset, you can reuse them to train or test different students without regenerating the data.\n\nIn principle, you would want the student to match the teacher’s full distribution over possible outputs. The caveat here is that you can’t. The space of possible sequences is astronomically large, so matching it exactly isn’t feasible. SeqKD sidesteps that by approximating it with samples the teacher actually produces.\n\nFor an agent, the sequence the teacher produces is not a block of prose but a trajectory, an ordered run of reasoning, tool calls, and the observations they return, where each step is conditioned on everything before it. Training the student on that whole sequence teaches it the reasoning that leads to each action and the multi-step course of action itself, not a single paragraph of answer text. That makes an agentic task one of the cleanest illustrations of sequence-level KD, because the sequence being copied is the behavior itself.\n\n## When SeqKD is the right choice\n\nSeqKD fits a specific set of situations:\n\n**You have the inputs but no high-quality target outputs**\n\nThis is the core case for SeqKD, which is fundamentally data generation. You supply the inputs, and the teacher writes the targets, so you can fine-tune on a task or domain that has no existing high-quality labels. The targets can be as simple as single answers or as involved as multi-turn tool-use trajectories. This post takes the demanding end of that spectrum, agentic trajectories, which are especially hard to source. A single task branches into many partially correct runs, and telling a solved run from a failed one takes a reliable grader. A benchmark with a programmatic scorer supplies exactly that, so you can let the teacher generate trajectories and keep only the ones that solve the task.\n\n**Your teacher is a frontier model you can only reach through an API**\n\nThis is the normal case for frontier models. You can’t access the weights to host it yourself, you can’t read its logits, and you can’t fine-tune it. That rules out logit-based distillation and on-policy methods, which all need that internal access. SeqKD needs only the teacher’s text responses, so it is the one approach that still works, and it frees you from paying frontier prices on every request after the student takes over.\n\n**You want a strong, cost-effective first step**\n\nSome of the strongest small reasoning models were built in significant part with SeqKD (supervised fine-tuning on teacher reasoning traces), so it is an effective starting point for reasoning transfer when the setup is right. The [DeepSeek-R1 distilled models](https://arxiv.org/abs/2501.12948) are a well-known example: the reasoning of a large teacher was distilled into smaller Qwen and Llama students by fine-tuning them on the teacher’s generated reasoning traces.\n\n## Prerequisites\n\nTo run this example, you need:\n\n- An AWS account with an [AWS Identity and Access Management (IAM) role for SageMaker AI](https://docs.aws.amazon.com/sagemaker/latest/dg/security-iam.html) that has read and write access to an Amazon Simple Storage Service (Amazon S3) bucket and permission to invoke your teacher model.\n- A frontier teacher model to generate the trajectories (this walkthrough uses one on [Amazon Bedrock](https://aws.amazon.com/bedrock/) ).\n- Access to [SageMaker AI model customization](https://docs.aws.amazon.com/sagemaker/latest/dg/customize-model.html) in a supported AWS Region.\n- [Amazon SageMaker Studio](https://docs.aws.amazon.com/sagemaker/latest/dg/studio-updated.html) (or a local IDE with the[AWS CLI configured](https://docs.aws.amazon.com/cli/latest/userguide/cli-configure-files.html) ), running a Python 3.12 or higher kernel, with the[SageMaker Python SDK v3](https://sagemaker.readthedocs.io/) and`boto3` (pin your versions). Use an[ml.m5.4xlarge instance](https://docs.aws.amazon.com/sagemaker/latest/dg/notebooks-available-instance-types.html) at minimum for the notebook environment, since the data preparation and local scoring steps run on its CPU.\n- Working knowledge of Python, supervised fine-tuning, and tool-using agents.\n- The [accompanying code repository](https://github.com/aws/sagemaker-hyperpod-recipes/tree/main/examples/sequence-distillation) and a task with a programmatic scorer (this walkthrough uses the ShoppingBench test set. Substitute your own task and grader).\n\n## Solution overview\n\nThis section walks the pipeline from end to end. A frontier teacher, GLM-5.3 on Amazon Bedrock, acts as the ShoppingBench agent to generate trajectories. The benchmark’s own scorer keeps only the ones that solve the task, and those solved trajectories fine-tune a Qwen3.5-4B student with serverless supervised fine-tuning on SageMaker AI model customization. The distilled student is then evaluated on a held-out ShoppingBench split and served behind a SageMaker AI Inference real-time endpoint that runs the tool loop.\n\n### Preparing the dataset\n\nSourcing trajectories worth imitating is the hard part, and the ShoppingBench scorer is what separates the successful runs from the failed ones. ShoppingBench publishes 750 synthesized shopping problems, 250 each for the product, shop, and voucher intents, each with a hidden gold answer and the constraints that define success. We split them once into 426 training problems and 324 held-out evaluation problems (product 103, shop 105, voucher 116), as fixed indices so the two sets don’t overlap.\n\nFor each training problem, a frontier teacher, GLM-5.3 on Amazon Bedrock with reasoning on and temperature 1.0, acts as the ShoppingBench agent. It searches the catalog of about 2.7 million products, inspects candidates, and recommends the ones that meet the constraints. We sample the teacher twice per problem, which gives 850 trajectories, then let ShoppingBench’s own scorer keep only the ones that solve the task. Temperature 1.0 is deliberate. Sampling rather than greedy decoding makes the two attempts explore different queries and reasoning paths, which gives more problems at least one solved trajectory and a more diverse set to learn from. This rejection-sampling step turns a capable-but-imperfect teacher into a clean training set, since the teacher only has to be right often enough that the solved runs are worth imitating.\n\nEach solved trajectory is a complete episode: the shopper request, the teacher’s reasoning and tool calls at each step (search, inspect, recommend, terminate), and the observation the environment returned. We keep the reasoning alongside the actions, which the Results show is what carries the student most of the way to its teacher.\n\n### Encoding trajectories for training\n\nSequence-level distillation is supervised learning over trajectories that already exist, so nothing calls a tool during training, and the application only ever sees text. Because these trajectories are multi-step tool-use conversations, we encode each as a `messages` list of chat turns and let the student’s own chat template render it, so the tokens the trainer sees match the tokens the agent produces at inference.\n\nEach trajectory becomes one training row per step. The `user` turn carries the system prompt and the conversation so far. The `assistant` turn carries that step’s reasoning and the tool call it leads to, a … block followed by a [{…}] in the harness’s native form. Loss falls on the `assistant` turn only, so the model learns to produce the next reasoning-and-action, not to predict the observations the environment returned. A trajectory with several steps expands into several rows. Here the solved trajectories expand to about 2,600 training and 130 validation rows after filtering.\n\nA single row pairs the conversation so far (user) with the one next step to learn (assistant):\n\nEach later step becomes its own row, with the new observation appended to the user content and the next reasoning-and-tool-call as the `assistant` target, through to the final `recommend_product` and `terminate` calls. Render every row with your model’s exact chat template.\n\nSplit the data by problem, not by row. The teacher is sampled twice per problem, and each trajectory expands into several steps, so many rows share one problem, and a row-level split would leak near-identical context across training and validation. Holding out whole problems (about five percent) keeps the validation set honest.\n\nThe cost of per-step encoding is repeated context, because each row restates the growing conversation, and serverless fine-tuning truncates anything past `dataset_max_len`, so a row that crosses the limit loses part of what it needs. Rather than let the limit decide what gets cut, keep rows within budget. The most effective fix is to compress tool observations, which dominate the token count. For product search results, keep only the fields the next decision depends on (id, price, rating, and the constraints being checked), drop the raw listing text, and cap how many results you retain per search. If trajectories are still over budget, cap the steps per trajectory. Here `dataset_max_len` is 16,384 tokens, which fits most trajectories after observations are trimmed.\n\nAfter the JSONL for each split is in S3, register each as a customization [dataset asset](https://docs.aws.amazon.com/sagemaker/latest/dg/model-customize-open-weight-create-assets-ui.html) with `DataSet.create` from the SageMaker Python SDK v3, which returns the `train_dataset` and `val_dataset` the training job consumes.\n\n### Fine-tuning the student\n\nTraining uses [SageMaker AI model customization](https://aws.amazon.com/sagemaker/ai/model-customization/) with serverless SFT. You select a base model from the [model hub](https://docs.aws.amazon.com/sagemaker/latest/dg/studio-jumpstart.html), point the job at the encoded datasets, and the managed service provisions the compute, runs the job, and returns the customized model, with metrics tracked in [Amazon SageMaker with MLflow](https://docs.aws.amazon.com/sagemaker/latest/dg/mlflow.html). Because the data is already in a supported format, there is no training script, container image, or instance to manage.\n\nThe job writes its output into a [model package group](https://docs.aws.amazon.com/sagemaker/latest/dg/model-registry-model-group.html) that versions every model it produces, so create the group first with `ModelPackageGroup.create` (reusing it if the name already exists) and pass it to the trainer. Then configure the [customization job](https://docs.aws.amazon.com/sagemaker/latest/dg/model-customize-open-weight-job.html). Choose the base model and training type, point it at the train and validation datasets plus the model package group, then name an MLflow experiment so metrics stream to MLflow while the job runs.\n\nThe trainer starts from a documented set of default hyperparameters that you can inspect with `sft_trainer.hyperparameters.to_dict()` and override before launching. For this run we raise the sequence length to fit the multi-step trajectories, use a Low-Rank Adaptation (LoRA) rank of 64, and train for three epochs over the small dataset:\n\nValidate the configuration with a dry run, then launch the managed job and read the registered Model Package ARN when it finishes. That ARN is the deploy-ready artifact the evaluation and serving steps consume:\n\n### Evaluating the student\n\nEvaluate in two stages, lowest cost first. The full end-to-end evaluation is both time and resource intensive. It stands up a search server over millions of products, deploys the endpoint, and runs a multi-step agent loop on every held-out problem. Before running the full evaluation, run a low-cost managed pre-check using SageMaker AI model evaluation, which scores one generated step at a time against the teacher’s ground-truth action.\n\nThe pre-check needs a scoring function. Model evaluation supports a", "url": "https://wpnews.pro/news/build-efficient-agents-with-knowledge-distillation-on-amazon-sagemaker-ai", "canonical_source": "https://aws.amazon.com/blogs/machine-learning/build-efficient-agents-with-knowledge-distillation-on-amazon-sagemaker-ai/", "published_at": "2026-11-02 17:01:31+00:00", "updated_at": "2026-11-02 17:20:33.893994+00:00", "lang": "en", "topics": ["artificial-intelligence", "machine-learning", "large-language-models", "ai-agents", "ai-tools"], "entities": ["Amazon SageMaker AI", "Amazon Web Services", "SeqKD"], "also_reported_by": [], "alternates": {"html": "https://wpnews.pro/news/build-efficient-agents-with-knowledge-distillation-on-amazon-sagemaker-ai", "markdown": "https://wpnews.pro/news/build-efficient-agents-with-knowledge-distillation-on-amazon-sagemaker-ai.md", "text": "https://wpnews.pro/news/build-efficient-agents-with-knowledge-distillation-on-amazon-sagemaker-ai.txt", "jsonld": "https://wpnews.pro/news/build-efficient-agents-with-knowledge-distillation-on-amazon-sagemaker-ai.jsonld"}}