Como melhorar a precisão de chamada de ferramentas do seu agente com SFT e DPO no Amazon SageMaker AI

O problema com agentes que escolhem a ferramenta errada

Agentes de IA têm a capacidade de executar tarefas complexas de forma autônoma, mas tudo depende de uma condição básica: chamar a ferramenta certa no momento certo. Quando um agente seleciona a ferramenta errada, formata parâmetros incorretamente ou quebra uma cadeia de fluxo de trabalho, o impacto é imediato — tempos de execução aumentam, taxas de erro sobem e a experiência do usuário se deteriora.

À medida que mais organizações levam aplicações agênticas do piloto para a produção, garantir que o agente selecione a ferramenta adequada para cada requisição se torna essencial para uma automação confiável. A AWS publicou um guia técnico mostrando como usar duas técnicas de ajuste fino em conjunto — Ajuste Fino Supervisionado (SFT) e Otimização Direta de Preferências (DPO) — para melhorar significativamente essa precisão em um modelo de linguagem pequeno (SLM).

As duas metodologias de ajuste fino

Ajuste Fino Supervisionado (SFT)

O SFT envolve a curadoria de um conjunto de dados de alta qualidade que se alinha estreitamente com a função pretendida do modelo. Ele fornece exemplos explícitos de como o modelo deve executar determinadas tarefas ou interagir com ferramentas específicas. Esse método é particularmente eficaz para ensinar o modelo a reconhecer as nuances da linguagem, comandos e restrições específicos de cada ferramenta.

Otimização Direta de Preferências (DPO)

O DPO refina as interações incorporando feedback humano ou objetivos predefinidos diretamente no ciclo de treinamento. Ele alinha a saída do modelo com os resultados desejados ao enfatizar preferências por certos tipos de resposta ou comportamento. Os dados de treinamento no DPO seguem um formato de “prefira isso, não aquilo”, o que otimiza os mesmos objetivos do aprendizado por reforço, mas sem funções de recompensa ou modelos de recompensa. Isso reduz os requisitos de recursos e o tempo de treinamento, mantendo a qualidade.

A biblioteca TRL da HuggingFace para DPO, por exemplo, aceita amostras de treinamento no seguinte formato:

{
  "prompt": ["<array of input samples>"],
  "chosen": "<complete preferred response (j)>",  # rated better than k
  "rejected": "<complete non-preferred response (k)>"  # rated worse than j
}

Esse formato é baseado na pesquisa publicada em arXiv:2305.18290 [cs.LG]. Juntos, SFT e DPO formam um framework robusto para ajustar modelos de linguagem para interagir com uma ampla variedade de ferramentas digitais.

Para entender os custos associados ao Amazon SageMaker Studio e aos jobs de treinamento do SageMaker AI, consulte a página de preços do SageMaker AI.

Visão geral da solução

O exemplo demonstra como ajustar o modelo Qwen3 1.7B usando jobs de treinamento do Amazon SageMaker AI, um serviço totalmente gerenciado que suporta configurações distribuídas com múltiplas GPUs e múltiplos nós. Com esses jobs, é possível provisionar clusters de alto desempenho sob demanda, treinar modelos com bilhões de parâmetros mais rapidamente e desligar os recursos automaticamente ao final do job. As métricas de infraestrutura e do loop de treinamento são enviadas para o MLflow no SageMaker AI para análise posterior.

Pré-requisitos

Para executar o ajuste fino de modelos de chamada de função no SageMaker AI, são necessários:

Preparação do dataset

A escolha e criação do dataset correto é um passo fundamental no ajuste fino de modelos de fundação (FMs). O exemplo usa o dataset When2Call, publicado pela NVIDIA — um benchmark projetado para avaliar a tomada de decisão de chamada de ferramentas em FMs. Ele cobre cenários como: quando gerar uma chamada de ferramenta, quando fazer perguntas de acompanhamento, quando indicar que a pergunta não pode ser respondida com as ferramentas disponíveis, e o que fazer quando uma chamada de ferramenta parece necessária mas não pode ser feita. Os scripts de geração de dados sintéticos estão no repositório GitHub da NVIDIA.

O dataset é composto por três partes:

  • Dataset para SFT: contém 15.000 amostras.
  • Dataset para alinhamento de preferências (DPO): contém 9.000 amostras com colunas chosen_response e rejected_response.
  • Dataset de teste: inclui avaliação por Questões de Múltipla Escolha (MCQ) com 3.652 amostras e avaliação por LLM-como-juiz com 300 amostras.

É necessário pré-processar o dataset para corresponder aos formatos esperados pelo SFTTrainer e pelo DPOTrainer da biblioteca TRL. Isso inclui construir um prompt de sistema com a lista de ferramentas disponíveis e adicioná-lo às listas de mensagens do dataset original.

def generate_and_tokenize_prompt(data_point):
    """
    Generates a tool using prompt based on patient information.
    Args:
        data_point (dict): Dictionary containing target and meaning_representation keys
    Returns:
        dict: Dictionary containing the formatted prompt
    """
    full_prompt = f"""
You are a helpful assistant with access to the following tools or function calls.
Your task is to produce a sequence of tools or function calls necessary to generate response to the user utterance.
Use the following tools or function calls as required:
{data_point["tools"]}
"""
    return {"system_prompt": full_prompt.strip()}

Para o DPO, o DPOTrainer do TRL exige colunas chamadas chosen e rejected, então é necessário renomear as colunas chosen_response e rejected_response do dataset original. Após o pré-processamento, os datasets de SFT e DPO são salvos no Amazon Simple Storage Service (Amazon S3) para ficarem disponíveis durante o treinamento.

Ajuste Fino Supervisionado (SFT) no modelo base

O repositório contém a receita de treinamento no diretório scripts, onde é possível modificar o modelo base e os parâmetros de treinamento para SFT. O exemplo usa uma receita de ajuste fino baseada em Spectrum, mas outras técnicas de Ajuste Fino Eficiente de Parâmetros (PEFT), como LoRA ou QLoRA, também podem ser utilizadas. A receita define, entre outros parâmetros:

# Model arguments
model_name_or_path: Qwen/Qwen3-1.7B
tokenizer_name_or_path: Qwen/Qwen3-1.7B
model_revision: main
torch_dtype: bfloat16
attn_implementation: flash_attention_2
bf16: true
tf32: true
output_dir: /opt/ml/model/Qwen3-1.7B-function-calling

# Dataset arguments
dataset_id_or_path: /opt/ml/input/data/dataset/dataset.json
max_seq_length: 2048
packing: true

# Training arguments
num_train_epochs: 10
per_device_train_batch_size: 4
gradient_accumulation_steps: 2
learning_rate: 5.0e-5
lr_scheduler_type: cosine
warmup_ratio: 0.1

Criando o job de treinamento com o ModelTrainer do SageMaker AI

As APIs ModelTrainer do SDK Python do SageMaker AI executam jobs de treinamento em infraestrutura totalmente gerenciada, cuidando da configuração do ambiente, do escalonamento e do gerenciamento de artefatos. É possível especificar scripts de treinamento, dados de entrada e recursos de computação sem provisionar servidores manualmente.

A seção Compute define os requisitos de infraestrutura, e a seção SourceCode define os caminhos locais para o código que será importado no job:

compute = Compute(
    instance_count=1,
    instance_type= "ml.p4d.24xlarge",
    volume_size_in_gb=96,
    keep_alive_period_in_seconds=3600,
)
source_code = SourceCode(
    source_dir="./scripts",
    requirements="requirements.txt",
    entry_script="run_training_sft.sh",
)

Para o treinamento distribuído em múltiplas GPUs, o exemplo usa o Hugging Face Accelerate com DeepSpeed ZeRO-3. O Accelerate simplifica o lançamento do treinamento distribuído, enquanto o ZeRO-3 reduz o uso de memória particionando estados do otimizador, gradientes e parâmetros entre as GPUs. O comando de lançamento tem a seguinte forma:

NUM_GPUS=$(nvidia-smi --list-gpus | wc -l)
echo "Detected ${NUM_GPUS} GPUs on the machine"
accelerate launch \
  --config_file accelerate_configs/deepspeed_zero3.yaml \
  --num_processes ${NUM_GPUS} run_sft.py \
  --config receipes/Qwen3-0.6B-spectrum.yaml

Otimização Direta de Preferências (DPO) no modelo ajustado com SFT

Com o artefato do modelo SFT pronto, ele é usado como modelo base para o treinamento DPO. A receita DPO é similar à do SFT, com algumas diferenças importantes nos hiperparâmetros:

  • beta: hiperparâmetro específico do DPO, geralmente entre 0 e 2, que controla quão agressivamente o modelo adota novas preferências. Um valor mais próximo de 0 é mais agressivo; mais próximo de 2 é mais conservador. Um ponto de partida típico é entre 0,1 e 0,5.
  • learning_rate: o DPO se beneficia de taxas de aprendizado mais baixas (por exemplo, 5e-7) com warmup_ratio para evitar sobreajuste. Esse valor contrasta com o 5e-5 usado no SFT.
  • batch_size: tamanhos de lote maiores tendem a ter melhor desempenho. O tamanho neste exemplo é intencionalmente pequeno para reduzir os requisitos de recursos.
# Training arguments
beta: 0.1
max_length: 1536
max_prompt_length: 768
loss_type: sigmoid
num_train_epochs: 10
per_device_train_batch_size: 2
gradient_accumulation_steps: 8
learning_rate: 5.0e-7
lr_scheduler_type: constant
warmup_ratio: 0.03

Opcionalmente, é possível combinar diferentes valores de loss para realizar a Otimização de Preferência Mista (Mixed Preference Optimization), que permite combinar e ponderar múltiplos tipos de loss. Se apenas dados DPO estiverem disponíveis, é possível usar o tipo de loss sft para aproveitar a coluna de respostas aceitas do DPO para SFT também:

# MPO (Mixed Preference Optimization): Combines DPO (sigmoid) for preference and BCO (bco_pair) for quality
loss_type : ["sigmoid", "bco_pair", "sft"],  # Loss types to combine
loss_weights : [0.8, 0.2, 1.0]  # Corresponding weights, as used in the MPO paper

Resultados

O experimento foi executado para três modelos diferentes, usando o script de avaliação fornecido pela NVIDIA. Os resultados mostram ganhos expressivos:

  • Entre os modelos base, o Qwen3-0.6B foi o mais forte apesar de ser o menor, superando o Qwen3-1.7B em aproximadamente 6% e o Llama-3.2-3B-Instruct em aproximadamente 1%.
  • Após o SFT, o Qwen3-1.7B ganhou aproximadamente 19% de acurácia e passou a superar os outros em 4–7%.
  • Após o DPO, o Qwen3-1.7B adicionou mais aproximadamente 10,5% de acurácia, encerrando o experimento na liderança com 8–9% de vantagem sobre os demais modelos.
  • No total, o Qwen3-1.7B ganhou 30% de acurácia geral e performou 9% melhor que o modelo Llama-3.2-3B, que tem quase o dobro de parâmetros.

A tabela abaixo resume os resultados:

Modelo Técnica de Ajuste Acc-Norm
Llama 3.2 3B Instruct Base 46,50%
Llama 3.2 3B Instruct Spectrum SFT 53,41%
Llama 3.2 3B Instruct Spectrum SFT + DPO 62,67%
Qwen3-0.6B Base 47,64%
Qwen3-0.6B Spectrum SFT 56,10%
Qwen3-0.6B Spectrum SFT + DPO 62,02%
Qwen3-1.7B Base 41,57%
Qwen3-1.7B Spectrum SFT 60,43%
Qwen3-1.7B Spectrum SFT + DPO 71,06%

Alcançar desempenho similar ou superior com um modelo menor pode reduzir custos e melhorar o throughput na hora de hospedar o modelo.

Limpeza dos recursos

Para evitar cobranças por recursos que não são mais necessários, a AWS recomenda:

  • Excluir os jobs de treinamento do SageMaker AI lançados (jobs concluídos com sucesso não continuam gerando cobranças).
  • Remover os datasets enviados ao Amazon S3: aws s3 rm s3://<your-bucket>/datasets/nvidia_function_calling/ --recursive
  • Parar ou excluir a instância de notebook JupyterLab do SageMaker Studio para evitar cobranças por ociosidade.
  • Excluir os checkpoints do modelo armazenados no Amazon S3 que não são mais necessários.

Conclusão

A combinação de SFT e DPO no Amazon SageMaker AI representa uma abordagem robusta para melhorar a precisão de chamada de ferramentas em agentes de IA. O SFT usa datasets rotulados para refinar os parâmetros do modelo, desenvolvendo uma compreensão fundamental a partir de exemplos anotados por especialistas. O DPO então alinha as saídas do modelo com preferências humanas ou critérios de desempenho específicos por meio de feedback direto, sem a necessidade de definir funções de recompensa.

O resultado é um modelo mais preciso, mais relevante e mais alinhado com o comportamento desejado — e, como os experimentos mostram, é possível superar modelos maiores com modelos menores bem ajustados, o que representa uma vantagem real em termos de custo e latência em produção.

Para mais exemplos de ajuste fino de modelos de fundação, acesse o repositório GitHub de exemplos de IA generativa do SageMaker AI. Para mais informações sobre treinamento de modelos no SageMaker AI, consulte a documentação do SageMaker AI.

Fonte

Improve your agent’s tool-calling accuracy with SFT and DPO on Amazon SageMaker AI (https://aws.amazon.com/blogs/machine-learning/improve-your-agents-tool-calling-accuracy-with-sft-and-dpo-on-amazon-sagemaker-ai/)

Comments

Leave a Reply

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