Effiziente Agenten durch Knowledge Distillation auf Amazon SageMaker AI
Der Beitrag zeigt, wie sich mit sequence-level knowledge distillation das Agentenverhalten eines teuren Frontier-Modells auf ein kleineres, guenstigeres Modell uebertragen laesst. Das senkt die Kosten mehrstufiger Tool-Loops in der Produktion.
Du hast einen Shopping-Assistenten auf einem Frontier-Modell aufgebaut, und er funktioniert. Eine Käuferin fragt nach „einer schwarzen Bauchtasche unter zwanzig Dollar von einem Shop mit mindestens 4,5 Sternen“, und der Agent durchsucht den Katalog, prüft Kandidaten, kontrolliert die Bedingungen und schlägt eine Empfehlung vor. Aber was kostet diese Interaktion tatsächlich?
Ein agentischer Assistent ist kein einzelner Modellaufruf, sondern eine Schleife: suchen, Ergebnisse lesen, verfeinern, verifizieren, entscheiden. Eine einzelne abgeschlossene Aufgabe kann mehrere Aufrufe des Frontier-Modells bedeuten, und du zahlst für jeden einzelnen davon, bei jedem Käufer. Im Produktionsmaßstab summieren sich diese Kosten pro Aufgabe, sie zehren die ROI des Projekts auf und verhindern, dass das Projekt ausgeliefert wird.
Ein Weg, eine agentische Anwendung effizienter zu machen, ist der Einsatz eines kleineren Modells, das schneller antwortet und kostengünstiger bereitzustellen ist. Der Haken ist die Qualität: Ein kleineres Modell schneidet bei der Aufgabe meist schlechter ab, weshalb ursprünglich das größere gewählt wurde. Knowledge Distillation schließt diese Lücke, indem das Verhalten des größeren „Teacher“-Modells in ein kleineres „Student“-Modell übertragen wird. Der Student wird darauf trainiert, die Ausgaben des Teachers zu reproduzieren, und behält so einen Großteil der Qualität bei Geschwindigkeit und Kosten des kleineren Modells.
In diesem Beitrag lernst du, wie du Knowledge Distillation, genauer gesagt Sequence-Level Knowledge Distillation (SeqKD), nutzt, um einen Shopping-Agenten auf Amazon SageMaker AI zu fine-tunen. Wir behandeln, wie man hochwertige mehrstufige Trajektorien mit einem Frontier-Teacher generiert und diejenigen behält, die die Aufgabe lösen, wie man sie für überwachtes Fine-Tuning kodiert, wie man das Training auf serverelloser Infrastruktur ausführt und wie man den Studenten gegen einen zurückgehaltenen Benchmark evaluiert.
Abbildung 1: Eine agentische Aufgabe ist eine Schleife, kein einzelner Aufruf: Der Agent sucht, liest Ergebnisse, verfeinert, verifiziert und entscheidet, und jeder Schritt ist ein separater abgerechneter Aufruf des Frontier-Modells. Eine einzelne abgeschlossene Aufgabe kann mehrere Aufrufe bedeuten, bei jedem Käufer
Die einfachste praktikable Form der Destillation ist Sequence-Level Knowledge Distillation (SeqKD). Dabei nutzt man ein Teacher-Modell, um hochwertige Antworten zu generieren, und trainiert dann ein Student-Modell auf diesen Antworten mittels überwachtem Fine-Tuning (SFT). SeqKD benötigt nur den Text des Teachers, ohne Zugriff auf Teacher-Logits und ohne Reinforcement-Learning-Infrastruktur, funktioniert also auch mit einem geschlossenen oder nur über eine API erreichbaren Teacher.
Verglichen mit anderen Ansätzen wie logit-basierten oder On-Policy-Methoden liegt es am kosteneffizientesten und portabelsten Ende des Destillationsspektrums. Der Teacher schreibt die Antworten, und der Student wird darauf fine-getuned. Da diese Antworten als Datensatz erhalten bleiben, kannst du sie wiederverwenden, um verschiedene Studenten zu trainieren oder zu testen, ohne die Daten neu zu generieren.
Im Grunde würde man wollen, dass der Student die vollständige Verteilung des Teachers über mögliche Ausgaben abbildet. Das Problem dabei ist: Das geht nicht. Der Raum möglicher Sequenzen ist astronomisch groß, eine exakte Abbildung ist also nicht machbar. SeqKD umgeht das, indem es diese Verteilung durch Stichproben approximiert, die der Teacher tatsächlich erzeugt.
Bei einem Agenten ist die Sequenz, die der Teacher erzeugt, kein Prosa-Block, sondern eine Trajektorie: eine geordnete Abfolge von Reasoning, Tool-Aufrufen und den dazugehörigen Beobachtungen, wobei jeder Schritt von allem vorher Bedingung ist. Den Studenten auf dieser Gesamtsequenz zu trainieren, vermittelt ihm das Reasoning, das zu jeder Aktion führt, und zugleich den mehrstufigen Handlungsverlauf selbst, nicht nur einen einzelnen Absatz Antworttext. Das macht eine agentische Aufgabe zu einer der klarsten Illustrationen von Sequence-Level KD, denn die kopierte Sequenz ist das Verhalten selbst.
Abbildung 2: SeqKD überträgt Verhalten allein über die Textausgaben des Teachers: Der Teacher generiert Trajektorien, der Student trainiert darauf, und das Ergebnis behält einen Großteil der Qualität des Teachers bei Geschwindigkeit und Kosten des kleinen Modells. Während des Trainings werden weder Teacher-Gewichte noch Logits oder Gradienten benötigt, es läuft also als ganz normales überwachtes Fine-Tuning
SeqKD passt auf eine bestimmte Reihe von Situationen:
Du hast die Eingaben, aber keine hochwertigen Ziel-Ausgaben
Das ist der Kernfall für SeqKD, das im Wesentlichen Datengenerierung ist. Du lieferst die Eingaben, der Teacher schreibt die Ziele, du kannst also auf einer Aufgabe oder Domäne fine-tunen, für die es keine hochwertigen Labels gibt. Die Ziele können so einfach wie einzelne Antworten oder so aufwendig wie mehrstufige Tool-Nutzungs-Trajektorien sein.
Dieser Beitrag nimmt das anspruchsvolle Ende dieses Spektrums, agentische Trajektorien, die besonders schwer zu beschaffen sind. Eine einzelne Aufgabe verzweigt sich in viele teilweise korrekte Läufe, und einen gelösten Lauf von einem gescheiterten zu unterscheiden, erfordert einen verlässlichen Bewerter. Ein Benchmark mit programmatischem Scorer liefert genau das, sodass du den Teacher Trajektorien generieren lassen und nur diejenigen behalten kannst, die die Aufgabe lösen.
Dein Teacher ist ein Frontier-Modell, das du nur über eine API erreichst
Das ist der Normalfall bei Frontier-Modellen. Du kannst nicht auf die Gewichte zugreifen, um es selbst zu hosten, du kannst seine Logits nicht lesen, und du kannst es nicht fine-tunen. Das schließt logit-basierte Destillation und On-Policy-Methoden aus, die alle diesen internen Zugriff benötigen. SeqKD benötigt nur die Textantworten des Teachers, ist also der eine Ansatz, der weiterhin funktioniert, und es befreit dich davon, nach der Übernahme durch den Studenten bei jedem Request Frontier-Preise zu zahlen.
Du willst einen starken, kosteneffizienten ersten Schritt
Einige der stärksten kleinen Reasoning-Modelle wurden zu einem erheblichen Teil mit SeqKD gebaut (überwachtes Fine-Tuning auf Reasoning-Spuren des Teachers), es ist also ein effektiver Ausgangspunkt für Reasoning-Transfer, wenn das Setup passt. Die destillierten DeepSeek-R1-Modelle sind ein bekanntes Beispiel: Das Reasoning eines großen Teachers wurde in kleinere Qwen- und Llama-Studenten destilliert, indem man sie auf den generierten Reasoning-Spuren des Teachers fine-tunte.
Dieser Abschnitt durchläuft die Pipeline von Ende zu Ende. Ein Frontier-Teacher, GLM-5.3 auf Amazon Bedrock, agiert als ShoppingBench-Agent und generiert Trajektorien. Der Scorer des Benchmarks behält nur diejenigen, die die Aufgabe lösen, und diese gelösten Trajektorien fine-tunen einen Qwen3.5-4B-Studenten mit serverlosem überwachtem Fine-Tuning über SageMaker AI Model Customization. Der destillierte Student wird anschließend auf einem zurückgehaltenen ShoppingBench-Split evaluiert und hinter einem SageMaker AI Inference-Real-Time-Endpoint bereitgestellt, der die Tool-Schleife ausführt.
Abbildung 3: Destillation eines Shopping-Agenten auf Amazon SageMaker AI: Ein Frontier-Teacher auf Amazon Bedrock agiert als ShoppingBench-Agent auf den Trainingsproblemen, der Scorer von ShoppingBench selbst behält nur die Trajektorien, die die Aufgabe lösen (Rejection Sampling), jede gelöste Trajektorie wird pro Schritt kodiert, der Student wird serverlos mit Amazon SageMaker AI Model Customization fine-getuned, und das Ergebnis wird auf dem zurückgehaltenen ShoppingBench-Test-Split evaluiert
Trajektorien zu beschaffen, die nachahmenswert sind, ist der schwierige Teil, und der ShoppingBench-Scorer trennt die erfolgreichen von den gescheiterten Läufen. ShoppingBench veröffentlicht 750 synthetisierte Shopping-Probleme, je 250 für die Intents Produkt, Shop und Voucher, jedes mit einer verborgenen Gold-Antwort und den Bedingungen, die Erfolg definieren. Wir teilen sie einmal in 426 Trainingsprobleme und 324 zurückgehaltene Evaluationsprobleme auf (Produkt 103, Shop 105, Voucher 116), mit festen Indizes, damit sich die beiden Mengen nicht überschneiden.
Für jedes Trainingsproblem agiert ein Frontier-Teacher, GLM-5.3 auf Amazon Bedrock mit aktiviertem Reasoning und Temperatur 1.0, als ShoppingBench-Agent. Er durchsucht den Katalog mit rund 2,7 Millionen Produkten, prüft Kandidaten und empfiehlt diejenigen, die die Bedingungen erfüllen. Wir samplen den Teacher zweimal pro Problem, was 850 Trajektorien ergibt, und lassen dann den Scorer von ShoppingBench selbst nur diejenigen behalten, die die Aufgabe lösen.
Temperatur 1.0 ist Absicht. Sampling statt Greedy-Decoding lässt die beiden Versuche unterschiedliche Queries und Reasoning-Pfade erkunden, was mehr Problemen mindestens eine gelöste Trajektorie verschafft und eine diversere Lernmenge ergibt. Dieser Rejection-Sampling-Schritt verwandelt einen fähigen, aber unvollkommenen Teacher in einen sauberen Trainingsdatensatz, da der Teacher nur oft genug richtig liegen muss, damit sich die gelösten Läufe nachahmen lassen.
Jede gelöste Trajektorie ist eine vollständige Episode: die Anfrage des Käufers, das Reasoning und die Tool-Aufrufe des Teachers bei jedem Schritt (suchen, prüfen, empfehlen, abbrechen) sowie die Beobachtung, die die Umgebung zurückgibt. Wir behalten das Reasoning zusammen mit den Aktionen, was laut den Ergebnissen den Studenten den Großteil des Weges zu seinem Teacher trägt.
Destillation auf Sequenzebene ist überwachtes Lernen auf bereits vorhandenen Trajektorien, während des Trainings wird also kein Tool aufgerufen, und die Anwendung sieht ausschließlich Text. Da diese Trajektorien mehrstufige Tool-Nutzung-Gespräche sind, kodieren wir jede als Messages-Liste aus Chat-Turns und lassen das eigene Chat-Template des Studenten sie rendern, sodass die Tokens, die der Trainer sieht, den Tokens entsprechen, die der Agent bei der Inferenz erzeugt.
Jede Trajektorie wird zu einer Trainingszeile pro Schritt. Der User-Turn enthält den System-Prompt und die bisherige Konversation. Der Assistant-Turn enthält das Reasoning des jeweiligen Schritts und den daraus resultierenden Tool-Aufruf, einen …-Block gefolgt von einem [{…}] in der nativen Form des Harness. Der Loss fällt nur auf den Assistant-Turn, sodass das Modell lernt, das nächste Reasoning und die nächste Aktion zu erzeugen, nicht aber die Beobachtungen vorherzusagen, die die Umgebung zurückgegeben hat. Eine Trajektorie mit mehreren Schritten wird zu mehreren Zeilen expandiert. Hier werden die gelösten Trajektorien nach der Filterung zu etwa 2.600 Trainings- und 130 Validierungszeilen expandiert.
Eine einzelne Zeile paart die bisherige Konversation (User) mit dem einen zu lernenden nächsten Schritt (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\n[{\"name\": \"find_product\", \"parameters\": {\"q\": \"black waist bag\", \"price\": \"0-20\"}}]"} ] } Jeder spätere Schritt wird zu einer eigenen Zeile, wobei die neue Beobachtung an den User-Content angehängt wird und das nächste Reasoning samt Tool-Aufruf als Assistant-Ziel dient, bis hin zu den abschließenden recommend_product- und terminate-Aufrufen. Rendern Sie jede Zeile mit dem exakten Chat-Template Ihres Modells.
Teilen Sie die Daten nach Problem, nicht nach Zeile auf. Der Teacher wird pro Problem zweimal gesampelt, und jede Trajektorie wird zu mehreren Schritten expandiert, sodass viele Zeilen sich ein Problem teilen und ein zeilenbasierter Split nahezu identische Kontexte über Training und Validierung hinweg lecken würde. Ganze Probleme zurückzuhalten (etwa fünf Prozent) hält das Validierungsset ehrlich.
Der Preis der Kodierung pro Schritt ist der wiederholte Kontext, da jede Zeile die wachsende Konversation wiederholt, und serverloses Fine-tuning schneidet alles ab, was über dataset_max_len hinausgeht, sodass eine Zeile, die das Limit überschreitet, einen Teil dessen verliert, was sie braucht. Anstatt das Limit entscheiden zu lassen, was abgeschnitten wird, halten Sie die Zeilen im Budget. Die wirksamste Korrektur besteht darin, die Tool-Beobachtungen zu komprimieren, die den Löwenanteil der Token ausmachen.
Bei Produktsuchergebnissen behalten Sie nur die Felder, von denen die nächste Entscheidung abhängt (id, price, rating und die geprüften Bedingungen), lassen den rohen Listing-Text weg und begrenzen Sie, wie viele Ergebnisse Sie pro Suche behalten. Falls Trajektorien weiterhin über dem Budget liegen, begrenzen Sie die Schritte pro Trajektorie. Hier beträgt dataset_max_len 16.384 Token, was nach dem Kürzen der Beobachtungen für die meisten Trajektorien ausreicht.
Nachdem sich die JSONL-Dateien für jeden Split in S3 befinden, registrieren Sie jeden als Customization-Dataset-Asset mit DataSet.create aus dem SageMaker Python SDK v3, das die train_dataset und val_dataset zurückgibt, die der Trainingsjob verarbeitet.
Das Training nutzt SageMaker AI Model Customization mit serverlosem SFT. Sie wählen ein Basismodell aus dem Model Hub, lassen den Job auf die kodierten Datensätze zeigen, und der verwaltete Service stellt die Compute-Ressourcen bereit, führt den Job aus und liefert das angepasste Modell zurück, wobei die Metriken in Amazon SageMaker mit MLflow nachverfolgt werden. Da die Daten bereits in einem unterstützten Format vorliegen, gibt es kein Trainingsskript, kein Container-Image und keine Instanz zu verwalten.
Der Job schreibt seine Ausgabe in eine Model Package Group, die jedes erzeugte Modell versioniert. Erstellen Sie die Gruppe daher zuerst mit ModelPackageGroup.create (unter Wiederverwendung, falls der Name bereits existiert) und übergeben Sie sie an den Trainer. Konfigurieren Sie dann den Customization-Job. Wählen Sie das Basismodell und den Trainingstyp, verweisen Sie auf die Trainings- und Validierungsdatensätze sowie die Model Package Group, und benennen Sie dann ein MLflow-Experiment, damit die Metriken während des Job-Laufs an MLflow gestreamt werden.
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, ) Der Trainer startet mit einem dokumentierten Satz von Standard-Hyperparametern, den Sie mit sft_trainer.hyperparameters.to_dict() einsehen und vor dem Start überschreiben können. Für diesen Lauf erhöhen wir die Sequenzlänge, um die mehrstufigen Trajektorien unterzubringen, verwenden einen LoRA-Rang (Low-Rank Adaptation) von 64 und trainieren drei Epochen über den kleinen Datensatz:
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 Validieren Sie die Konfiguration mit einem Dry Run, starten Sie dann den verwalteten Job und lesen Sie die ARN des registrierten Model Packages aus, sobald er abgeschlossen ist. Diese ARN ist das deploymentfertige Artefakt, das die Evaluierungs- und Serving-Schritte verarbeiten:
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 Evaluieren Sie in zwei Stufen, die günstigste zuerst. Die vollständige End-to-End-Evaluierung ist zeit- und ressourcenintensiv. Sie bringt einen Suchserver über Millionen von Produkten hoch, deployt den Endpoint und führt für jedes zurückgehaltene Problem eine mehrstufige Agent-Schleife aus.
Bevor Sie die vollständige Evaluierung durchführen, führen Sie einen kostengünstigen verwalteten Pre-Check mit SageMaker AI Model Evaluation aus, der jeweils einen generierten Schritt gegen die Ground-Truth-Aktion des Teachers bewertet.
Der Pre-Check benötigt eine Scoring-Funktion. Model Evaluation unterstützt eine Custom Scorer Reward Function, die Sie bereitstellen und die einmal pro Sample läuft und die generierte Aktion des Studenten mit der des Teachers vergleicht. Für die Tool-Call-Destillation bewerten Sie die Tool-Auswahl (hat das Modell die richtigen Tools in der richtigen Reihenfolge aufgerufen, unabhängig von der Argumentformatierung), denn genau dieses Verhalten muss die Destillation am stärksten übertragen. Um den Fokus auf die Auswahl statt auf die Syntax zu legen, liest die Reward Function die Tool-Namen sowohl aus der Aktion des Studenten als auch aus der Referenz des Teachers, ohne ein exaktes Format-Match zu verlangen.
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 <think 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"}, ], } Registrieren Sie diese Funktion als versionierten Evaluator mit Evaluator.create() und rufen Sie dann evaluate() auf, um pro Task einen Custom Scorer Job auszuführen, der das Basismodell neben dem Studenten bewertet, sodass jedes Ergebnis einen Vergleich liefert.
Die Entscheidungsregel ist einfach. Der Student muss das Basismodell bei diesem kostengünstigen Pre-Check schlagen. Wenn ja, hat die Destillation gegriffen und es lohnt sich, die umfassendere End-to-End-Evaluierung durchzuführen. Wenn nicht, beheben Sie zuerst das Training.
Das destillierte Modell führt nicht von selbst Tools aus. Das Fine-tuning bringt ihm bei, im richtigen Moment den richtigen Tool-Aufruf als Text zu erzeugen, aber etwas muss diesen Aufruf weiterhin ausführen und das Ergebnis zurückspielen. Beim Inferenz wickeln Sie das Modell in einen kleinen deterministischen Harness, der das Modell promptet, den ausgegebenen Tool-Aufruf parst, das Tool ausführt, die Beobachtung anhängt und erneut promptet, in einer Schleife, bis das Modell eine finale Antwort liefert. Das ist gewöhnlicher Anwendungscode oder ein Serving-Framework von der Stange, kein Reinforcement Learning und keine On-Policy-Trainingsschleife.
Der Student ist ein LoRA-Adapter (Low-Rank Adaptation) über einem Basismodell mit offenen Gewichten, was die Art des Deployments prägt. Anstatt den Adapter zu mergen und eine zweite vollständige Kopie der Gewichte zu serven, serven Sie das Basismodell einmal und hängen den Studenten als Adapter darüber an, als zwei Inference Components auf einem Real-Time-Endpoint. Die Benennung der Inference Component zum Aufrufzeitpunkt bestimmt, welche antwortet, sodass derselbe Endpoint sowohl das Basismodell als auch den destillierten Studenten für einen direkten Vergleich bedient.
Wir serven das Modell mit offenen Gewichten mit dem Deep Java Library Large Model Inference Container (DJL LMI), der die vLLM-Engine für hochdurchsatzfähige Generierung bündelt, auf eine bestimmte Version statt auf latest gepinnt, damit Redeploys reproduzierbar sind. Das begleitende Notebook kapselt die Inference-Component-Aufrufe in einem kleinen kd_serve-Helfer, sodass das Deployen der Basiskomponente mit aktiviertem LoRA-Serving und das anschließende Anhängen des Studenten-Adapters reduziert wird auf:
from sagemaker.core import image_uris import kd_serve # DJL LMI container (vLLM engine enthalten); Version fixieren für reproduzierbare erneute Deployments image_uri = image_uris.retrieve("djl-lmi", region=region, version="0.36.0") # 1. Das Basismodell als Inference Component deployieren, mit aktiviertem LoRA-Serving. 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. Die Adapter-Artefakte des Studenten aus seinem Model Package auflösen und als zweite Inference Component an derselben Basis anhängen. 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, ) Die Liste instance_types ist eine geordnete Präferenz.
SageMaker AI Inference deployt auf dem ersten Typ mit verfügbarer Kapazität und fällt die Liste nach unten ab, sodass der Student je nach Verfügbarkeit auf einer ml.g5.2xlarge, ml.g6.2xlarge oder ml.g6e.2xlarge läuft. Die Basalgewichte werden einmal geladen und von beiden Inference Components gemeinsam genutzt, sodass das Hinzufügen des Adapters adapter-großen Speicher kostet statt einer zweiten Kopie des Modells.
Wir evaluieren den destillierten Studenten auf dem zurückgehaltenen ShoppingBench-Test-Split, gemessen an einer Reihe von Referenzpunkten, die einordnen, wie ein gutes Ergebnis aussieht:
Jedes Modell durchläuft dieselbe unveränderte ShoppingBench-Agent-Schleife über die drei Shopper-Intents (product, shop und voucher). Wir berichten die Absolute Success Rate (ASR), den Anteil der gelösten zurückgehaltenen Probleme, gemittelt über die drei Intents, wobei fehlende Trajektorien als Fehler zählen. Ein Problem gilt als gelöst, wenn das oder die empfohlenen Produkte jede genannte Einschränkung erfüllen. Die folgenden Zahlen stammen aus einem einzelnen Durchlauf mit einem LoRA-Studenten, rechne also mit einigen Punkten Schwankung von Durchlauf zu Durchlauf.
Das Kernergebnis ist die Absolute Success Rate auf den 324 zurückgehaltenen Problemen, aufgeschlüsselt nach Intent und über die drei gemittelt:
Abbildung 4: Absolute Success Rate (ASR) auf dem zurückgehaltenen ShoppingBench-Test-Split, nach Shopper-Intent. Der destillierte 4B-Student (orange) verbessert das Basismodell um durchschnittlich 24 Punkte, erreicht das GLM-5.3-Lehrermodell, das seine Daten erzeugt hat, und ein mittelgroßes Frontier-Modell
Zusammenfassend hat ein 4B-Modell im Wesentlichen die gesamte Aufgabenqualität seines Frontier-Lehrers zurückgewonnen und es innerhalb des Rauschens zwischen Durchläufen erreicht. Wir berichten das Lehrermodell auf zwei Arten, weil sie unterschiedliche Fragen beantworten. Greedy Decoding liefert die wahrscheinlichste einzelne Trajektorie, die stärkere Wahl, wenn jedes Modell einen Versuch pro Problem erhält (wie in dieser Evaluation), weshalb der Greedy-Lehrer 0,540 gegenüber 0,511 des gesampelten Lehrers erzielt.
Die Trainingsdaten haben wir auf die andere Art erzeugt, mit Sampling bei Temperatur 1,0, sodass wiederholte Versuche verschiedene Pfade erkundeten und der Scorer diejenigen behalten konnte, die die Aufgabe lösten. Diese Filterung ist auch der Grund, warum der Student den gesampelten Lehrer, von dem er lernte, übertreffen kann, denn Rejection Sampling trainiert ihn nur auf den gelösten Durchläufen des Lehrers, nicht auf dessen durchschnittlicher Ausgabe, sodass er ein Best-of-Anteil nachahmt statt der gesamten Streuung der Versuche.
Der Student lernt aus dem Reasoning des Lehrers ebenso wie aus dessen Tool-Aufrufen, und diese Entscheidung ist für einen großen Teil des Ergebnisses verantwortlich. Trainiert man denselben Studenten auf denselben Trajektorien mit entferntem Reasoning, sodass nur die nackten Tool-Aufrufe bleiben, fällt die durchschnittliche ASR von 0,541 auf 0,415:
Abbildung 5: Absolute Success Rate für das Basismodell, den Studenten, der nur auf Tool-Aufrufen trainiert wurde, und den Studenten, der auf Reasoning plus Tool-Aufrufen des Lehrers trainiert wurde. Das Beibehalten des Reasonings bringt etwa 13 Punkte durchschnittliche ASR zusätzlich, und der Zugewinn hält auf jedem Intent an
Beide Studenten übertreffen das Basismodell, sodass selbst die reine Tool-Aufrufe-Version eine echte Verbesserung ist, aber das Beibehalten des Reasonings ist es, was den Studenten den Rest des Weges zu seinem Lehrer trägt. Das Reasoning zwischen den Aktionen ist keine Erzählung zum Wegwerfen, es ist Teil des Verhaltens, das die Aufgabe löst.
ASR misst die Fähigkeit, die wir destilliert haben. Die Gegenfrage ist, ob der Student seine allgemeine Fähigkeit zum Befolgen von Anweisungen behalten hat oder auf die Shopping-Aufgabe überpasst hat. Wir prüfen das mit IFEval, einem allgemeinen Instruction-Following-Benchmark, der überprüfbare Format-Anweisungen bewertet, etwa „antworte in genau drei Stichpunkten“.
Betrachte es als Regressionsprüfung, nicht als Leaderboard. Der Prompt-strict-Score des Studenten fällt von 0,819 für das Basismodell auf 0,741, und sein Instruction-strict-Score von 0,872 auf 0,813, ein Rückgang von rund 6 bis 8 Punkten (Abbildung 6). Das ist die erwartete Kosten der Spezialisierung.
Der Student hat etwas allgemeines Instruction-Following gegen einen großen Zugewinn bei der Shopping-Aufufe eingetauscht. Ein deutlich größerer Rückgang wäre das Warnsignal, dass die Destillation das Modell zu stark eingeengt hat.
Abbildung 6: IFEval allgemeines Instruction-Following für das Basismodell und den destillierten Studenten, für die Metriken prompt-strict und instruction-strict. Der kleine Rückgang sind die erwarteten Kosten der Spezialisierung auf die Shopping-Aufgabe
Eine Erfolgsrate allein rechtfertigt den Wechsel nicht. Der Grund zu destillieren ist, dass man den größten Teil der Qualität des stärkeren Modells zu den Kosten und der Geschwindigkeit des kleinen Modells erhält. Der klarste Beleg hier ist, dass der 4B-Student die Aufgabenqualität seines Frontier-Lehrers und eines mittelgroßen Frontier-Modells erreicht, dabei aber klein genug ist, um auf einer einzigen G-Family-Instanz zu laufen (NVIDIA A10G/L40/L40S in diesem Beispiel).
Der Kostenunterschied ist dabei ebenso eine Frage der Form wie der Größe. Ein Frontier-Modell über eine API wird pro Token abgerechnet, bei jedem Schritt jeder Aufgabe, sodass die Kosten mit der Nutzung wachsen und nicht aufhören, und die Agent-Schleife aus der Einleitung zahlt diesen Tribut jedes Mal. Der destillierte Student läuft auf einem von dir gehosteten Endpoint, der die Instanz stündlich abrechnet, ob er eine oder Tausende Aufgaben bedient, sodass zusätzlicher Verkehr nichts weiter kostet, bis du horizontal skalierst. Man tauscht eine variable Abrechnung pro Anfrage gegen eine feste stündliche, und je kleiner das Modell ist, mit dem man auskommt, desto kleiner sind diese Fixkosten.
Wir veröffentlichen hier keine Kosten- oder Latenzzahl pro Aufgabe, weil sie von deinem Verkehr abhängt: der Durchsatz, den eine G-Family-Instanz hält, und damit die pro Aufgabe amortisierten Kosten sind eine Funktion deines Anfragemusters, der Sequenzlängen und der Parallelität. Miss es an deiner eigenen Arbeitslast gegen die aktuellen SageMaker-AI-Preise und die Token-Preise deines Frontier-Anbieters, statt dich auf eine einzelne Schlagzahl zu verlassen.
Destillation überträgt eine Verteilung, die der Student bereits darstellen kann, nicht eine rohe Fähigkeit, die ihm fehlt. Über eine zu große Lücke hinweg kopiert der Student das Format des Lehrers, ohne lösen zu lernen. Ein kleines Modell, trainiert auf die Gedankenketten eines großen Lehrers, erzeugt langes, wohlstrukturiertes Reasoning, dessen Antworten trotzdem falsch bleiben, und ein auf dieselbe Weise trainierter Agent gibt wohlgeformte Tool-Aufrufe aus, die keinen Fortschritt machen, mit umformulierten Anfragen erneut suchen und bis zu seinem Schrittlimit laufen, ohne die Einschränkungen zu erfüllen.
Das Symptom ist in beiden Fällen dasselbe, die Traces werden länger, während die Genauigkeit flach bleibt. Wenn du das siehst, hör auf, Daten hinzuzufügen, und vergrößere entweder den Studenten oder senke die Aufgabenschwierigkeit.
Die Lücke, über die das entscheidet, liegt zwischen der Aufgabe und dem Studenten, nicht zwischen dem Studenten und dem Lehrer. Ein kleiner Student kann die Qualität eines viel größeren Lehrers zurückgewinnen, solange die Aufgabe selbst in seiner Reichweite liegt. Deshalb funktioniert der Shopping-Agent hier. Der Lehrer ist ein Mixture-of-Experts-Modell mit rund 40 Milliarden aktiven Parametern, etwa zehnmal so viel wie die 4 Milliarden des Studenten, aber die Shopping-Aufgabe liegt klar innerhalb der Kapazität eines 4B-Modells, sodass der Student die Qualität des Lehrers zurückgewinnt, statt nur seine Trajektorien zu verlängern.
SeqKD hebt leichte und mittelschwere Probleme an, erreicht aber bei den schwersten ein Plateau, wo der Student aus seinen eigenen Fehlern herausfinden muss. Das ist für einen Agenten schärfer als für eine Einzelschritt-Aufgabe. In einer mehrstufigen Schleife vervielfacht sich ein früher Fehler, weil der Student die Zustände nicht sieht, in die seine eigenen Fehler führen.
Diese Lücke zu schließen ist das Ziel von On-Policy-Methoden, weshalb SeqKD üblicherweise die erste Stufe ist und nicht die Ziellinie. SeqKD stellt außerdem die geringsten Anforderungen an deine Infrastruktur, da es nur den Text des Lehrers nutzt, während Logit-Destillation die Per-Token-Verteilung des Lehrers braucht und On-Policy-Methoden die eigenen Rollouts des Studenten bewerten. Ein rein textueller SeqKD-Warmstart gefolgt von einer On-Policy-Stufe ist ein gängiges Rezept, dieser Beitrag deckt die erste Stufe ab, und über das Plateau hinauszustoßen ist der natürliche nächste Schritt.
Erstens sollten die Trainingsziele an das gewünschte Verhalten angepasst werden. In diesem Fall bedeutete das, das Reasoning des Lehrermodells zusammen mit jedem Tool-Aufruf beizubehalten, statt auf die finalen Antworten zu komprimieren, da das Reasoning ein Teil dessen ist, was die Aufgabe löst (die Ergebnisse quantifizieren diese Lücke), und jedes Ziel mit dem Chat-Template zu rendern, das der Agent zur Serving-Zeit verwendet. Zweitens sollte man ehrlich evaluieren.
Man sollte auf Daten messen, die die Modelle noch nicht gesehen haben, da geleakte Benchmarks die Scores aufblähen, und man sollte nicht nur mit dem untrainierten Basismodell vergleichen, das der Schüler zwangsläufig schlägt, sondern ihn neben das Lehrermodell und ein unabhängiges Frontier-Modell stellen, damit die Zahl zeigt, wie nah der Schüler an das herankam, wovon er destilliert wurde.
Ein Echtzeit-Inference-Endpoint wird abgerechnet, solange er läuft. Man sollte daher den Endpoint, den man für das Serving eingerichtet hat, wieder abbauen, sobald man die Ergebnisse hat. Da der Schüler als Adapter-Inference-Komponente auf einer Basis-Inference-Komponente betrieben wird, müssen die Bestandteile in Abhängigkeitsreihenfolge abgebaut werden: zuerst die Adapter-Komponente, dann die Basiskomponente, dann der Endpoint samt seiner Konfiguration. Das beiliegende kd_serve-Hilfsprogramm übernimmt dies und findet den Adapter strukturell, sodass die Bereinigung unabhängig von den genauen Komponentennamen funktioniert:
import kd_serve # Löscht die Adapter-IC, dann die Basis-IC, dann den Endpoint samt Konfiguration kd_serve.teardown_endpoint(endpoint_name, region=region) Der Trainingsjob selbst erfordert keine Bereinigung. Serverless Training stellt Compute bereit, führt den Job aus und gibt die Ressourcen automatisch wieder frei, sodass ein abgeschlossener Job keine laufenden Kosten verursacht. Auch das registrierte Modelpaket und seine Gruppe verursachen keine Compute-Kosten, sodass man sie für die Herkunftsnachverfolgung behalten und später erneut bereitstellen kann.
Das Einzige, was weiter kostet, sind die Daten in S3 (die kodierten Trajektorien und Ausgabe-Artefakte), die Speicherkosten verursachen, solange man sie behält. Man sollte daher löschen, was man nicht mehr braucht, oder in eine günstigere Speicherklasse verschieben.
In diesem Beitrag haben wir gezeigt, wie man einen agentischen Shopping-Assistenten per sequence-level knowledge distillation (SeqKD) in einen kleineren, günstigeren Schüler destilliert. Wir haben hochwertige Single-Turn-Trajektorien mit mehreren Schritten mit einem Frontier-Lehrermodell auf Bedrock erzeugt und nur diejenigen behalten, die die Aufgabe mit dem eigenen Scorer von ShoppingBench lösten, jeden Schritt als chatformatierte Nachrichten kodiert, die das Reasoning des Lehrermodells zusammen mit seinen Tool-Aufrufen bewahren, den Schüler auf serverloser Infrastruktur mit SageMaker AI model customization fine-getunt, das Training mit MLflow verfolgt, eine Quality-Gate-Prüfung mit SageMaker AI model evaluation durchgeführt und das Ergebnis hinter einem Inference-Endpoint betrieben, der die Tool-Schleife ausführt. Das Ergebnis war ein 4B-Schüler, der im Wesentlichen die gesamte Aufgabenqualität seines Frontier-Lehrermodells wiederherstellt.
Da SeqKD nur die Textausgaben des Lehrermodells benötigt, funktioniert es auch dann, wenn das Lehrermodell ein geschlossenes Frontier-Modell ist, auf das man über eine API zugreift, und man kann aufhören, Frontier-Preise für jeden Schritt zu zahlen, sobald der Schüler übernimmt.
SeqKD ist ein starker erster Schritt mit realen Grenzen. Man sollte die Schülergröße an die Aufgabe anpassen, auf einem zurückgehaltenen Testset gegen starke Referenzpunkte evaluieren statt nur gegen das untrainierte Basismodell, und wenn der schwierigste Rest der Probleme der Imitation widersteht, sollte man SeqKD mit On-Policy-Methoden fortsetzen, die hier nicht behandelt wurden und es dem Schüler erlauben, aus seinen eigenen Rollouts zu lernen.
Um das selbst auszuprobieren, siehe das begleitende Code-Repository und die SageMaker-AI-Dokumentation.
Giuseppe ist Principal GenAI Specialist Solutions Architect bei AWS und hilft großen Unternehmen dabei, generative KI-Lösungen auf AWS zu entwerfen und bereitzustellen. Er verfügt über mehr als 20 Jahre Erfahrung als Full-Stack-Softwareentwickler und hat die letzten 7 Jahre bei AWS im Bereich KI verbracht.
Ankita ist Software Development Engineer bei AWS und arbeitet dort an verteiltem Training und Reinforcement Learning für Amazon SageMaker AI. Sie hat einen Master-Abschluss in Computer Engineering von der University of Southern California. In ihrer Freizeit reist sie gern.
Amin ist Senior Data Scientist und Forscher bei AWS und arbeitet dort an Post-Training und Modellanpassung: Reinforcement Learning, Fine-Tuning und Knowledge Distillation, um große Modelle anzupassen und ihre Fähigkeiten in kleinere, einsetzbare Modelle zu komprimieren. Er verfügt über mehr als acht Jahre Erfahrung im Aufbau und Einsatz von Machine-Learning-Systemen in den Bereichen Computer Vision, Natural Language Processing und statistischer Inferenz für Finanzsysteme. Er hat einen PhD in theoretischer Physik, was seine Herangehensweise an Modellverhalten und Evaluation prägt.
Michael ist Product Manager für Amazon SageMaker AI. Es treibt ihn an, Technologie und KI zur Lösung realer Probleme einzusetzen. Bei AWS hilft er Kunden aus verschiedenen Branchen, KI-/ML-Modelle im großen Maßstab zu bauen, zu trainieren und bereitzustellen. Vor AWS entwarf und betrieb Michael Offshore-Öl- und Gasbohrungen sowie Anlagen. In seiner Freizeit beschäftigt er sich gern mit Verhaltenswissenschaft und Psychologie durch Bücher und Podcasts.
Automatisch übersetzt. Das Original auf Englisch ist unten aufklappbar.
Original (Englisch) anzeigen
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.
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
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.
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
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.
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.
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
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.
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# 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>"} ] } 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.
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.
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.
We evaluate the distilled student on the held-out ShoppingBench test split, against a set of reference points that bracket what good looks like:
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.
The headline result is the Absolute Success Rate on the 324 held-out problems, broken down by intent and averaged across the three:
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
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.
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:
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.
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.
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
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.
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.
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.
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.
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.
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.
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 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 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 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.