Como acelerar o fine-tuning de LLMs com dados não estruturados usando SageMaker Unified Studio e Amazon S3

Em 2026, a AWS anunciou uma integração entre o Amazon SageMaker Unified Studio e os buckets de uso geral do Amazon S3, facilitando o uso de dados não estruturados armazenados no S3 para casos de uso em machine learning (ML) e análise de dados. Este artigo apresenta um guia prático para integrar buckets S3 com o Amazon SageMaker Catalog e realizar o fine-tuning do modelo Llama 3.2 11B Vision Instruct para a tarefa de Visual Question Answering (VQA) usando o SageMaker Unified Studio.

Contexto e objetivo do fine-tuning

O modelo base Llama 3.2 11B Vision Instruct, acessado via SageMaker JumpStart, já apresenta desempenho robusto com uma pontuação de Average Normalized Levenshtein Similarity (ANLS) de 85,3% no dataset DocVQA. O ANLS é uma métrica que avalia a similaridade entre a resposta prevista pelo modelo e a resposta correta em tarefas de VQA. No entanto, para aplicações que exigem maior precisão, o fine-tuning com dados específicos é fundamental.

Para isso, utilizamos o dataset DocVQA disponível no Hugging Face, que contém 39.500 exemplos compostos por imagens, perguntas e respostas. O objetivo é criar três versões do modelo fine-tuned utilizando conjuntos de dados com 1.000, 5.000 e 10.000 imagens para comparar o impacto do volume de dados na performance.

Pré-requisitos essenciais

Conta AWS ativa.
Domínio SageMaker Unified Studio configurado via quick setup, com Identity Center habilitado.
Dois projetos criados no SageMaker Unified Studio: um para o perfil produtor de dados e outro para o consumidor que realizará o fine-tuning. Veja mais em Create a project.
Aplicação serverless do MLflow gerenciada pelo SageMaker para rastreamento das experiências e métricas, configurada conforme criando aplicação serverless MLflow.
Bucket S3 previamente populado com o dataset DocVQA, obtido via API do Hugging Face.
Solicitação de aumento de cota para instâncias p4de.24xlarge para treinamento, conforme Requesting a quota increase.

Arquitetura da solução e fluxo de trabalho

A arquitetura contempla seis etapas principais:

Criação de uma role IAM com permissão de leitura para o bucket S3 contendo o dataset DocVQA.
Projeto produtor de dados usa essa role para descobrir e catalogar o dataset no SageMaker Catalog.
Enriquecimento opcional do dataset com metadados empresariais, como README e glossário.
Publicação do dataset no SageMaker Catalog para consumo.
Projeto consumidor assina o dataset publicado, tornando-o disponível para o time de desenvolvimento.
Pré-processamento do dataset e fine-tuning do modelo em três versões, com acompanhamento via MLflow e avaliação pelo ANLS.

Preparação e ingestão dos dados no Amazon S3

Para iniciar, utilizamos a API Datasets do Hugging Face para baixar e salvar os splits de treino e validação do DocVQA localmente, conforme o código abaixo:

import os
from datasets import load_dataset

os.makedirs("data", exist_ok=True)
train_data = load_dataset("HuggingFaceM4/DocumentVQA", split="train[:10000]", cache_dir="./data")
train_data.save_to_disk("data/train")
val_data = load_dataset("HuggingFaceM4/DocumentVQA", split="validation[:100]", cache_dir="./data")
val_data.save_to_disk("data/validation")

Em seguida, sincronizamos esses dados para o bucket S3 que será usado durante o treinamento.

Catalogação e publicação do dataset no SageMaker Unified Studio

No projeto produtor, adicionamos o bucket S3 via Data → Add data → Add S3 location, selecionando a role IAM criada para garantir acesso seguro. O dataset aparece então no catálogo do projeto.

Opcionalmente, adicionamos metadados empresariais para facilitar a governança e documentação, como um arquivo README. Após isso, publicamos o dataset no SageMaker Catalog para que outros projetos possam consumi-lo.

Consumo do dataset e ambiente de desenvolvimento

No projeto consumidor, assinamos o dataset publicado para que ele fique disponível para o time de ML. Criamos um espaço com JupyterLab dentro do SageMaker Unified Studio para desenvolver, configurar e executar o fine-tuning.

É recomendado configurar o espaço com 100 GB de armazenamento e definir o tempo de inatividade para 6 horas, pois o treinamento pode levar até 4 horas usando a instância ml.p4de.24xlarge.

Fine-tuning do modelo Llama 3.2 11B Vision Instruct

O processo de fine-tuning é realizado em três conjuntos de dados com tamanhos variados (1k, 5k e 10k imagens) para avaliar o impacto do volume de dados na performance do modelo.

O dataset é transformado em estrutura própria para o fine-tuning, contendo subdiretórios images e arquivos metadata.jsonl com exemplos no formato:

{"file_name": "images/img_0.jpg", "prompt": "qual é a data mencionada nesta carta?", "completion": "1/8/93"}
{"file_name": "images/img_1.jpg", "prompt": "qual o nome da pessoa de contato mencionada na carta?", "completion": "P. Carter"}

O treinamento é orquestrado por uma função parametrizada que utiliza o SageMaker JumpStart para acessar o modelo base, ajustando hiperparâmetros e configurando a instância de treinamento. O MLflow é integrado para registrar automaticamente hiperparâmetros, metadados do job e resultados.

Avaliação dos modelos fine-tuned com a métrica ANLS

A métrica ANLS é usada para avaliar a qualidade das respostas geradas pelo modelo nas tarefas de VQA, considerando variações e pequenos erros textuais. O MLflow facilita a comparação entre as versões treinadas.

O pipeline de avaliação inclui funções para codificação das imagens, formatação do payload para inferência, cálculo da métrica ANLS e registro dos resultados.

Dicas práticas e limitações

O uso da instância ml.p4de.24xlarge é recomendado devido ao tamanho do modelo e necessidade de memória.
Solicite aumento de cota antecipadamente para evitar bloqueios no treinamento.
Configure o tempo de inatividade e armazenamento do espaço JupyterLab para evitar interrupções e falta de espaço.
O processo de sincronização de dados via AWS CLI com access grants pode ser adaptado conforme a linguagem e SDK preferidos.
O fine-tuning pode ser aplicado a outros tipos de dados não estruturados além de imagens, como documentos, logs e textos.