Build efficient agents with knowledge distillation on Amazon SageMaker AI

You 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.

One 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.

In 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. 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.

Diagram of an agentic task as a loop: search, read results, refine, verify, and decide, with each step a separate billed call to the frontier model

Figure 1: One agentic task is a loop, not a single call: the agent searches, reads results, refines, verifies, and decides, and each step is a separate billed call to the frontier model. A single completed task can be multiple calls, on every shopper

What is sequence-level knowledge distillation (SeqKD)?

The 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.

In 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.

For 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.

Diagram of SeqKD: the teacher generates trajectories and the student trains on them as ordinary supervised fine-tuning

Figure 2: SeqKD transfers behavior through the teacher’s text outputs alone: the teacher generates trajectories, the student trains on them, and the result keeps much of the teacher’s quality at the small model’s speed and cost. No teacher weights, logits, or gradients are required during training, so it runs as ordinary supervised fine-tuning

When SeqKD is the right choice

SeqKD fits a specific set of situations:

You have the inputs but no high-quality target outputs

This 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.

Your teacher is a frontier model you can only reach through an API

This 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.

You want a strong, cost-effective first step

Some 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 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.

Prerequisites

To run this example, you need:

Solution overview

This 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.

Architecture diagram of distilling a shopping agent on SageMaker AI: the teacher generates trajectories, the scorer keeps the solved ones, the student is fine-tuned, and results are evaluated on the held-out split

Figure 3: Distilling a shopping agent on Amazon SageMaker AI: a frontier teacher on Amazon Bedrock acts as the ShoppingBench agent on the training problems, ShoppingBench’s own scorer keeps only the trajectories that solve the task (rejection sampling), each solved trajectory is encoded per step, the student is fine-tuned serverlessly with Amazon SageMaker AI model customization, and the result is evaluated on the held-out ShoppingBench test split

Preparing the dataset

Sourcing 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.

For 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.

Each 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.

Encoding trajectories for training

Sequence-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.

Each 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.

A single row pairs the conversation so far (user) with the one next step to learn (assistant):

{
  "messages": [
    {"role": "user", "content": "You are a shopping assistant. ...n# TasknFind a black waist bag under $20 from a shop rated 4.5 or higher.n# Observations so farn(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>"}
  ]
}

Each 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.

Split 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.

The 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.

After the JSONL for each split is in S3, register each as a customization dataset asset with DataSet.create from the SageMaker Python SDK v3, which returns the train_dataset and val_dataset the training job consumes.

Fine-tuning the student

Training uses SageMaker AI model customization with serverless SFT. You select a base model from the model hub, 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. Because the data is already in a supported format, there is no training script, container image, or instance to manage.

The job writes its output into a model package group 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. 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.

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,
)

The 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:

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

Validate 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:

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

Evaluating the student

Evaluate 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.

The pre-check needs a scoring function. Model evaluation supports a Custom Scorer reward function you supply that runs once per sample, comparing the student’s generated action against the teacher’s. For tool-call distillation, grade tool selection (did the model call the right tools, in order, ignoring argument formatting), because that is the behavior distillation most needs to transfer. To keep the focus on selection rather than syntax, the reward function reads the tool names out of both the student’s action and the teacher’s reference without requiring an exact format match.

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"},
        ],
    }

Register that function as a versioned Evaluator with Evaluator.create(), then call evaluate() to run one Custom Scorer job per task, scoring the base model alongside the student so every result has a comparison.

The decision rule is straightforward. The student must beat the base model on this low-cost pre-check. If it does, the distillation took and it is worth running the more comprehensive end-to-end evaluation. If it does not, fix training first.

Serving the distilled agent

The distilled model doesn’t run tools on its own. Fine-tuning teaches it to emit the right tool call at the right moment as text, but something still has to execute that call and feed back the result. At inference you wrap the model in a small deterministic harness that prompts the model, parses the tool call it emits, runs the tool, appends the observation, and prompts again, looping until the model produces a final answer. This is ordinary application code, or an off-the-shelf serving framework, not reinforcement learning and not an on-policy training loop.

The student is a Low-Rank Adaptation (LoRA) adapter over an open-weight base, which shapes how you deploy it. Rather than merge the adapter and serve a second full copy of the weights, serve the base model once and attach the student as an adapter on top, as two inference components on one real-time endpoint. Naming the inference component at invoke time selects which one answers, so the same endpoint serves both the base and the distilled student for a side-by-side comparison.

We serve the open-weight model with the Deep Java Library Large Model Inference (DJL LMI) container, which bundles the vLLM engine for high-throughput generation, pinned to a specific version rather than latest, so redeploys are reproducible. The accompanying notebook wraps the inference-component calls in a small kd_serve helper, so deploying the base component with LoRA serving enabled and then attaching the student adapter reduces to:

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,
)

The instance_types list is an ordered preference. SageMaker AI Inference deploys on the first type with available capacity and falls back down the list, so the student serves on an ml.g5.2xlarge, ml.g6.2xlarge, or ml.g6e.2xlarge depending on availability. The base weights load once, and both inference components share them, so adding the adapter costs adapter-sized memory instead of a second copy of the model.

Results

We evaluate the distilled student on the held-out ShoppingBench test split, against a set of reference points that bracket what good looks like:

  • Base Qwen3.5-4B: the off-the-shelf student with no fine-tuning. This is the floor.
  • Distilled student (SeqKD): the same 4B model fine-tuned on the teacher trajectories, the approach this post describes. This is the result that matters.
  • GLM-5.3 teacher: the frontier model that generated the training data, on Bedrock. We report it both sampled (temperature 1.0, the setting that produced the data) and greedy (temperature 0).
  • Medium-sized frontier model: a strong proprietary model run as an agent, included as a practical ceiling.

Every model runs the same unmodified ShoppingBench agent loop over the three shopper intents (product, shop, and voucher). We report Absolute Success Rate (ASR), the share of held-out problems the agent solves, averaged over the three intents, with missing trajectories counted as failures. A problem counts as solved when the recommended product or products meet every stated constraint. The following numbers come from a single run with a LoRA student, so expect a few points of run-to-run variation.

Task success

The headline result is the Absolute Success Rate on the 324 held-out problems, broken down by intent and averaged across the three:

Bar chart of Absolute Success Rate by shopper intent, showing the distilled 4B student matching the GLM-5.3 teacher and a medium-sized frontier model

Figure 4: Absolute Success Rate (ASR) on the held-out ShoppingBench test split, by shopper intent. The distilled 4B student (orange) improves on the base model by 24 points on average, reaches the GLM-5.3 teacher that generated its data and a medium-sized frontier model

Model product shop voucher Average
Base Qwen3.5-4B (no fine-tuning) 0.456 0.314 0.121 0.297
Distilled student (SeqKD) 0.563 0.638 0.422 0.541
GLM-5.3 teacher (sampled, temperature 1.0) 0.553 0.619 0.362 0.511
GLM-5.3 teacher (greedy, temperature 0) 0.544 0.638 0.439 0.540
Medium-sized frontier model 0.583 0.667 0.388 0.546

In summary, a 4B model recovered essentially all of its frontier teacher’s task quality, matching it within run-to-run noise. We report the teacher two ways because they answer different questions. Greedy decoding returns the single most likely trajectory, the stronger choice when each model gets one attempt per problem (as in this evaluation) which is why the greedy teacher scores 0.540 against the sampled teacher’s 0.511. We generated the training data the other way, sampling at temperature 1.0 so repeated attempts explored different paths, and the scorer could keep the ones that solved the task. That filtering is also why the student can edge out the sampled teacher it learned from, since rejection sampling trains it only on the teacher’s solved runs, not its average output, so it imitates a best-of slice rather than the full spread of attempts.

Reasoning in the targets matters

The student learns from the teacher’s reasoning as well as its tool calls, and that choice is responsible for a large part of the result. Training the same student on the same trajectories with the reasoning stripped, leaving only the bare tool calls, drops average ASR from 0.541 to 0.415:

Bar chart of Absolute Success Rate for the base model, the student trained on tool calls only, and the student trained on reasoning plus tool calls

Figure 5: Absolute Success Rate for the base model, the student trained on tool calls only, and the student trained on the teacher’s reasoning plus tool calls. Keeping the reasoning adds about 13 points of average ASR, and the gain holds on every intent

Both students clear the base model, so even the tool-calls-only version is a real improvement, but keeping the reasoning is what carries the student the rest of the way to its teacher. The reasoning between the actions isn’t narration to discard, it is part of the behavior that solves the task.

Keeping general ability

ASR measures the skill we distilled for. The counterweight is whether the student kept its general instruction-following ability or overfit to the shopping task. We check with IFEval, a general instruction-following benchmark that scores verifiable format instructions such as “answer in exactly three bullet points.” Treat it as a regression check, not a leaderboard. The student’s prompt-strict score falls from 0.819 for the base model to 0.741, and its instruction-strict score from 0.872 to 0.813, a drop of roughly 6 to 8 points (Figure 6). That is the expected cost of specialization. The student traded some general instruction-following for a large gain on the shopping task. A much larger drop would be the warning sign that distillation narrowed the model too far.

Bar chart of IFEval prompt-strict and instruction-strict scores for the base model and the distilled student

Figure 6: IFEval general instruction-following for the base model and the distilled student, on the prompt-strict and instruction-strict metrics. The small drop is the expected cost of specializing on the shopping task

Benefits and tradeoffs

A success rate on its own doesn’t justify the switch. The reason to distill is that you get most of the stronger model’s quality at the small model’s cost and speed. The clearest evidence here is that the 4B student reaches the task quality of its frontier teacher and a medium-sized frontier model, while being small enough to serve on a single G family instance (NVIDIA A10G/L40/L40S in this example).

The cost difference is as much about shape as size. A frontier model reached through an API is billed per token, on every step of every task, so the cost grows with usage and doesn’t stop, and the agentic loop from the introduction pays that toll every time. The distilled student runs on an endpoint you host, which bills for the instance by the hour whether it serves one task or thousands, so added traffic costs nothing more until you scale out. You’re trading a variable per-request charge for a fixed hourly one, and the smaller the model you can get away with, the smaller that fixed cost is.

We don’t publish a per-task cost or latency number here because it depends on your traffic: the throughput a G family instance sustains, and therefore the cost amortized per task, is a function of your request pattern, sequence lengths, and concurrency. Measure it on your own workload against the current SageMaker AI pricing and your frontier provider’s token pricing, rather than relying on a single headline figure.

Limitations and best practices

Will the task fit the student (the capacity wall)

Distillation transfers a distribution the student can already represent, not raw capability it lacks. Across too large a gap, the student copies the teacher’s format without learning to solve. A small model trained on a large teacher’s chains of thought produces long, well-structured reasoning whose answers stay wrong, and an agent trained the same way emits well-formed tool calls that make no progress, re-searches with reworded queries, and runs to its step limit without meeting the constraints. The symptom is the same in both, traces lengthen while accuracy stays flat. When you see it, stop adding data and either grow the student or lower task difficulty.

The gap that decides this is between the task and the student, not between the student and the teacher. A small student can recover a much larger teacher’s quality as long as the task itself is within its reach. That is why the shopping agent here works. The teacher is a mixture-of-experts model with roughly 40 billion active parameters, about ten times the student’s 4 billion, but the shopping task is well within a 4B model’s capacity, so the student recovers the teacher’s quality instead of only lengthening its trajectories.

Where SeqKD stops helping

SeqKD lifts easy and mid-difficulty problems but plateaus on the hardest ones, where the student must recover from its own mistakes. This is sharper for an agent than a single-step task. In a multi-step loop an early error compounds, because the student doesn’t see the states its own mistakes lead to. Closing that gap is what on-policy methods target, which is why SeqKD is usually the first stage rather than the finish line. SeqKD also asks the least of your infrastructure, since it uses only the teacher’s text, whereas logit distillation needs the teacher’s per-token distribution and on-policy methods score the student’s own rollouts. A text-only SeqKD warm start followed by an on-policy stage is a common recipe, so this post covers the first stage and pushing past the plateau is the natural next step.

Two choices that carry the result

First, match the training targets to the behavior you want. Here that meant keeping the teacher’s reasoning with each tool call rather than compressing to final answers, since the reasoning is part of what solves the task (the Results quantify the gap), and rendering every target with the chat template the agent uses at serving time. Second, evaluate honestly. Measure on data the models have not seen, since leaked benchmarks inflate scores, and compare against more than the untrained base model the student is bound to beat, setting it beside the teacher and an independent frontier model so the number shows how close the student came to what it was distilled from.

Cleaning up

A real-time inference endpoint bills for as long as it runs, so tear down the endpoint that you stood up for serving after you have your results. Because the student is served as an adapter inference component on top of a base inference component, the pieces must come down in dependency order: the adapter component first, then the base component, then the endpoint and its configuration. The accompanying kd_serve helper does this for you, finding the adapter structurally so it cleans up regardless of the exact component names:

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)

The training job itself needs no cleanup. Serverless training provisions compute, runs, and releases it automatically, so a completed job isn’t an ongoing charge. The registered model package and its group carry no compute cost either, so you can keep them for lineage and redeploy later. The one thing that does keep costing you is your data in S3 (the encoded trajectories and output artifacts), which incurs storage cost for as long as you keep it, so delete what you no longer need or move it to a lower-cost storage class.

Conclusion

In this post, we showed how to distill an agentic shopping assistant into a smaller, lower-cost student with sequence-level knowledge distillation (SeqKD). We generated high-quality single-turn, multi-step trajectories with a frontier teacher on Bedrock and kept only the ones that solved the task using ShoppingBench’s own scorer, encoded each step as chat-formatted messages that preserve the teacher’s reasoning alongside its tool calls, fine-tuned the student on serverless infrastructure with SageMaker AI model customization, tracked training with MLflow, ran a quality gate check with SageMaker AI model evaluation, and served the result behind an inference endpoint that runs the tool loop. The payoff was a 4B student that recovers essentially all of its frontier teacher’s task quality. Because SeqKD needs only the teacher’s text outputs, it works even when the teacher is a closed frontier model you reach through an API, and you can stop paying frontier prices on every step after the student takes over.

SeqKD is a strong first step with real limits. Match the student size to the task, evaluate on a held-out suite against strong reference points rather than only the untrained base model, and when the hardest tail of problems resists imitation, follow SeqKD with on-policy methods not covered here that let the student learn from its own rollouts.

To try this yourself, see the accompanying code repository and the SageMaker AI documentation.


About the authors

Giuseppe Zappia

Giuseppe Zappia

Giuseppe is a Principal GenAI Specialist Solutions Architect at AWS, focused on helping large enterprises design and deploy generative AI solutions on AWS. He has over 20 years of experience as a full stack software engineer and has spent the past 7 years at AWS focused on the field of AI.

Ankita Agarwal

Ankita Agarwal

Ankita is a Software Development Engineer at AWS, where she works on distributed training and reinforcement learning for Amazon SageMaker AI. She holds a master’s degree in computer engineering from the University of Southern California. Outside of work, she likes to travel.

Amin Dashti

Amin Dashti

Amin is a Senior Data Scientist and researcher at AWS, where he works on post-training and model customization: reinforcement learning, fine-tuning, and knowledge distillation to adapt large models and compress their capability into smaller, deployable ones. He has over eight years of experience building and deploying machine learning systems across computer vision, natural language processing, and statistical inference for financial systems. He holds a PhD in theoretical physics, which shapes how he approaches model behavior and evaluation.

Michael Oguike

Michael Oguike

Michael is a Product Manager for Amazon SageMaker AI. He is driven by the application of technology and AI to solve real-world problems. At AWS, he helps customers across industries build, train, and deploy AI/ML models at scale. Before AWS, Michael designed and operated offshore oil and gas wells and facilities. Outside of work, he enjoys exploring behavioral science and psychology through books and podcasts.

Leave a Comment

This site uses Akismet to reduce spam. Learn how your comment data is processed.