User's blog

Como construir agentes eficientes com destilação de conhecimento no Amazon SageMaker AI

Written by

Make.com Service User

in

O problema de custo em aplicações agênticas

Imagine um assistente de compras construído sobre um modelo frontier. O usuário pede “uma pochete preta por menos de vinte dólares de uma loja com pelo menos 4,5 estrelas” e o agente busca no catálogo, examina candidatos, verifica restrições e faz uma recomendação. Parece simples, mas o que acontece por baixo dos panos é um loop de múltiplas chamadas: buscar, ler resultados, refinar, verificar, decidir. Cada etapa é uma chamada separada ao modelo frontier — e você paga por todas elas, para cada usuário.

Em escala de produção, esse custo por tarefa se acumula rapidamente, corrói o retorno sobre o investimento e pode inviabilizar o projeto. Uma alternativa natural é usar um modelo menor, que responde mais rápido e custa menos para hospedar. O problema é que modelos menores geralmente entregam qualidade inferior — é exatamente por isso que o modelo maior foi escolhido.

É aqui que entra a destilação de conhecimento: a técnica transfere o comportamento do modelo maior (“professor”) para um modelo menor (“aluno”), treinando o aluno para reproduzir as saídas do professor. O resultado é um modelo que mantém boa parte da qualidade do professor, mas com a velocidade e o custo do modelo menor.

A AWS publicou um guia completo mostrando como aplicar essa abordagem — especificamente a Destilação de Conhecimento em Nível de Sequência (SeqKD, do inglês Sequence-level Knowledge Distillation) — para treinar um agente de compras no Amazon SageMaker AI.

O que é SeqKD e quando usá-la

A SeqKD é a forma mais simples e portátil de destilação. O funcionamento é direto: o modelo professor gera respostas de alta qualidade, e o modelo aluno é treinado nessas respostas via Ajuste Fino Supervisionado (SFT, do inglês Supervised Fine-Tuning). Não é necessário acessar os pesos internos do professor, seus logits ou qualquer infraestrutura de aprendizado por reforço — apenas as saídas de texto já são suficientes.

Isso significa que a SeqKD funciona mesmo quando o professor é um modelo frontier fechado, acessível apenas por API. Como as respostas geradas persistem como um conjunto de dados, elas podem ser reutilizadas para treinar ou testar diferentes alunos sem precisar regenerar os dados.

Para agentes, a sequência que o professor produz não é um parágrafo de texto, mas uma trajetória: uma sequência ordenada de raciocínios, chamadas de ferramentas e observações retornadas, onde cada etapa depende de tudo que veio antes. Treinar o aluno nessa sequência completa ensina tanto o raciocínio que leva a cada ação quanto o curso de ação em múltiplas etapas.

A SeqKD é a escolha certa em três situações principais:

  • Você tem as entradas, mas não tem saídas de alta qualidade: o professor gera os alvos de treinamento para você, mesmo para trajetórias agênticas complexas.
  • Seu professor é um modelo frontier acessível apenas por API: como a SeqKD precisa apenas do texto gerado, é a única abordagem que funciona sem acesso aos pesos internos.
  • Você quer um primeiro passo sólido e econômico: os modelos destilados DeepSeek-R1 são um exemplo conhecido dessa abordagem, onde o raciocínio de um modelo grande foi destilado em modelos Qwen e Llama menores.

Visão geral da solução

O pipeline descrito pela AWS funciona assim: um modelo frontier (GLM-5.3 no Amazon Bedrock) atua como agente no benchmark ShoppingBench para gerar trajetórias. O próprio scorer do benchmark filtra apenas as trajetórias que resolvem a tarefa, e essas trajetórias aprovadas são usadas para fazer o ajuste fino de um modelo Qwen3.5-4B (aluno) com SFT serverless via customização de modelos do SageMaker AI. O aluno destilado é então avaliado em um conjunto de teste separado e servido por um endpoint de inferência em tempo real.

Pré-requisitos

Para executar esse exemplo, são necessários:

Preparando o conjunto de dados

O ShoppingBench publica 750 problemas de compras sintéticos (250 para cada intenção: produto, loja e voucher), cada um com uma resposta correta oculta e as restrições que definem o sucesso. Os dados são divididos em 426 problemas de treinamento e 324 para avaliação.

Para cada problema de treinamento, o modelo professor (GLM-5.3 no Bedrock, com raciocínio ativado e temperatura 1.0) age como agente do ShoppingBench: busca no catálogo de cerca de 2,7 milhões de produtos, examina candidatos e faz recomendações. O professor é amostrado duas vezes por problema, gerando 850 trajetórias no total. O scorer do ShoppingBench então filtra apenas as que resolvem a tarefa.

A temperatura 1.0 é intencional: a amostragem em vez da decodificação greedy faz com que as duas tentativas explorem consultas e caminhos de raciocínio diferentes, aumentando a chance de que pelo menos uma trajetória por problema seja bem-sucedida. Essa etapa de rejection sampling transforma um professor capaz-mas-imperfeito em um conjunto de treinamento limpo.

Codificando as trajetórias para treinamento

Cada trajetória aprovada é codificada como uma lista de mensagens de chat e renderizada com o template de chat do próprio modelo aluno — assim os tokens que o treinador vê são os mesmos que o agente produz na inferência.

Cada trajetória se torna uma linha de treinamento por etapa. O turno do usuário carrega o prompt de sistema e a conversa até aquele ponto. O turno do assistente carrega o raciocínio daquela etapa e a chamada de ferramenta correspondente. A perda é calculada apenas sobre o turno do assistente, para que o modelo aprenda a produzir o próximo raciocínio-e-ação, não a prever as observações retornadas pelo ambiente.

Um exemplo de linha de treinamento:

{
  "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": "Start with a broad search and let the price filter work before I narrow by shop rating.\n[{\"name\": \"find_product\", \"parameters\": {\"q\": \"black waist bag\", \"price\": \"0-20\"}}]"}
  ]
}

Cada etapa posterior vira sua própria linha, com a nova observação adicionada ao conteúdo do usuário e o próximo raciocínio-e-chamada-de-ferramenta como alvo do assistente.

Um ponto importante: divida os dados por problema, não por linha. Como o professor é amostrado duas vezes por problema e cada trajetória se expande em várias etapas, muitas linhas compartilham o mesmo problema — uma divisão por linha vazaria contextos quase idênticos entre treinamento e validação.

Para controlar o tamanho dos tokens, comprima as observações das ferramentas, que dominam a contagem. Para resultados de busca de produtos, mantenha apenas os campos que a próxima decisão depende (id, preço, avaliação e as restrições sendo verificadas). Após o JSONL de cada divisão estar no S3, registre cada um como um ativo de conjunto de dados de customização com DataSet.create do SageMaker Python SDK v3.

Ajuste fino do modelo aluno

O treinamento usa a customização de modelos do SageMaker AI com SFT serverless. Basta selecionar um modelo base do hub de modelos, apontar o job para os conjuntos de dados codificados, e o serviço gerenciado provisiona o compute, executa o job e retorna o modelo customizado — com métricas rastreadas no Amazon SageMaker com MLflow. Não há script de treinamento, imagem de container ou instância para gerenciar.

O job grava sua saída em um grupo de pacotes de modelo que versiona cada modelo produzido. A configuração do job de customização é feita assim:

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:///seqkd/output/",
    accept_eula=True,
)

Os hiperparâmetros para este experimento:

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

Para validar a configuração sem lançar o job, use dry_run=True. Para lançar e aguardar:

from sagemaker.core.resources import TrainingJob

sft_trainer.train(dry_run=True)           # verifica a config sem lançar
training_job = sft_trainer.train(wait=True) # lança e aguarda
model_package_arn = TrainingJob.get(
    training_job_name=training_job.training_job_name
).output_model_package_arn

Avaliando o modelo aluno

A avaliação acontece em dois estágios, do menor custo para o maior.

Primeiro, um pré-check gerenciado de baixo custo usando a avaliação de modelos do SageMaker AI, que pontua uma etapa gerada por vez contra a ação de referência do professor. O critério é a seleção de ferramentas: o modelo chamou as ferramentas certas, na ordem certa? A função de recompensa para o Custom Scorer é:

def reward_function(sample, index):
    pred = sample.get("model_response", "")
    ref  = sample.get("reference_answer", "") or sample.get("response", "")

    ref_names  = tool_names(ref)   # tool names from the teacher's [...]
    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,
        "metrics_list": [
            {"name": "toolset_match", "value": toolset, "type": "Reward"},
        ],
    }

A regra de decisão é simples: o aluno deve superar o modelo base nesse pré-check. Se superar, a destilação funcionou e vale a pena rodar a avaliação completa ponta a ponta. Se não superar, o treinamento precisa ser corrigido primeiro.

Servindo o agente destilado

O modelo destilado não executa ferramentas por conta própria. O ajuste fino ensina-o a emitir a chamada de ferramenta certa no momento certo como texto, mas ainda é necessário um harness que prompta o modelo, analisa a chamada emitida, executa a ferramenta, adiciona a observação e prompta novamente — em loop até o modelo produzir uma resposta final.

O aluno é um adaptador LoRA (Adaptação de Baixo Rank, do inglês Low-Rank Adaptation) sobre uma base de pesos abertos. Em vez de mesclar o adaptador e servir uma segunda cópia completa dos pesos, a abordagem recomendada é servir o modelo base uma vez e anexar o aluno como adaptador, usando dois componentes de inferência em um único endpoint em tempo real. Nomear o componente de inferência na invocação seleciona qual responde, permitindo comparação lado a lado entre o modelo base e o aluno destilado.

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

A lista instance_types é uma preferência ordenada: o SageMaker AI Inference implanta no primeiro tipo com capacidade disponível e recorre à lista conforme necessário. Os pesos base são carregados uma vez, e ambos os componentes de inferência os compartilham — adicionar o adaptador custa memória do tamanho do adaptador, não uma segunda cópia completa do modelo.

Resultados

A avaliação usa a Taxa de Sucesso Absoluta (ASR, do inglês Absolute Success Rate) — a proporção de problemas que o agente resolve — no conjunto de teste separado do ShoppingBench, com 324 problemas distribuídos entre três intenções (produto, loja e voucher).

Os pontos de referência comparados são: o Qwen3.5-4B base sem ajuste fino (piso), o aluno destilado via SeqKD (o resultado que importa), o professor GLM-5.3 (amostrado a temperatura 1.0 e greedy a temperatura 0) e um modelo frontier de médio porte como teto prático.

Os resultados de ASR médio por modelo:

  • Base Qwen3.5-4B (sem ajuste fino): 0,297
  • Aluno destilado (SeqKD): 0,541
  • Professor GLM-5.3 (amostrado, temperatura 1.0): 0,511
  • Professor GLM-5.3 (greedy, temperatura 0): 0,540
  • Modelo frontier de médio porte: 0,546

O resultado principal: um modelo de 4 bilhões de parâmetros recuperou essencialmente toda a qualidade de tarefa do seu professor frontier, igualando-o dentro da variação entre execuções.

O raciocínio nas saídas importa muito

Treinar o aluno apenas nas chamadas de ferramentas (sem o raciocínio do professor) derruba a ASR média de 0,541 para 0,415. Ambos os alunos superam o modelo base, mas manter o raciocínio é o que leva o aluno até o nível do seu professor. O raciocínio entre as ações não é narração descartável — é parte do comportamento que resolve a tarefa.

Capacidade geral preservada

Para verificar se o aluno manteve sua capacidade geral de seguir instruções, o guia usa o IFEval, um benchmark de seguimento de instruções. A pontuação prompt-strict cai de 0,819 (modelo base) para 0,741, e a instruction-strict de 0,872 para 0,813 — uma queda de cerca de 6 a 8 pontos. Isso é o custo esperado da especialização: o aluno trocou alguma capacidade geral por um grande ganho na tarefa de compras.

Benefícios, limitações e boas práticas

A justificativa econômica é clara: um modelo frontier acessado via API é cobrado por token, em cada etapa de cada tarefa, e o custo cresce com o uso sem parar. O aluno destilado roda em um endpoint que você hospeda, cobrado por hora independente do volume — você troca uma cobrança variável por requisição por uma fixa por hora. Quanto menor o modelo que você consegue usar, menor esse custo fixo. A AWS recomenda medir o custo no seu próprio workload comparando a precificação do SageMaker AI com os preços de tokens do seu provedor frontier.

Há limitações importantes a considerar:

  • O muro de capacidade: a destilação transfere uma distribuição que o aluno já consegue representar. Se a lacuna entre a tarefa e o aluno for grande demais, o aluno copia o formato do professor sem aprender a resolver. O sintoma: as trajetórias ficam mais longas enquanto a acurácia permanece estagnada. Quando isso acontece, aumente o aluno ou reduza a dificuldade da tarefa.
  • Onde a SeqKD para de ajudar: a SeqKD melhora problemas fáceis e de dificuldade média, mas atinge um platô nos mais difíceis, onde o aluno precisa se recuperar dos próprios erros. Métodos on-policy são o próximo passo natural para fechar essa lacuna — a SeqKD é um forte ponto de partida, não o fim da linha.

Duas escolhas que fazem a diferença: primeiro, alinhe os alvos de treinamento ao comportamento desejado (mantenha o raciocínio do professor junto a cada chamada de ferramenta). Segundo, avalie honestamente — use dados que os modelos não viram, e compare contra mais do que apenas o modelo base não treinado.

Limpeza de recursos

Um endpoint de inferência em tempo real é cobrado enquanto estiver ativo. Para encerrar, os componentes devem ser removidos na ordem de dependência: o adaptador primeiro, depois o componente base, depois o endpoint e sua configuração. O helper kd_serve faz isso automaticamente:

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)

O job de treinamento não precisa de limpeza — o treinamento serverless provisiona compute, executa e libera automaticamente. O grupo de pacotes de modelo também não gera custo de compute. O que continua gerando custo é o armazenamento no S3 (trajetórias codificadas e artefatos de saída), então delete o que não for mais necessário ou mova para uma classe de armazenamento mais econômica.

Conclusão

O guia publicado pela AWS demonstra como a SeqKD permite destilar um assistente de compras agêntico em um modelo menor e mais barato que mantém essencialmente toda a qualidade de tarefa do professor frontier. O pipeline completo cobre: geração de trajetórias de alta qualidade com rejeição das que falham, codificação por etapa preservando o raciocínio do professor, ajuste fino serverless com rastreamento via MLflow, verificação de qualidade com avaliação de modelos do SageMaker AI, e serviço do resultado em um endpoint de inferência.

Como a SeqKD precisa apenas das saídas de texto do professor, ela funciona mesmo quando o professor é um modelo frontier fechado acessível apenas por API — e após o aluno assumir, você para de pagar preços frontier em cada etapa. Para explorar o pipeline completo, consulte o repositório de código do exemplo e a documentação do SageMaker AI.

Fonte

Build efficient agents with knowledge distillation on Amazon SageMaker AI (https://aws.amazon.com/blogs/machine-learning/build-efficient-agents-with-knowledge-distillation-on-amazon-sagemaker-ai/)

Comments

Leave a Reply

Your email address will not be published. Required fields are marked *

More posts