Fine-tuning de LLM com Databricks Unity Catalog e Amazon SageMaker AI

O desafio de governança no fine-tuning de LLMs

Realizar fine-tuning de Modelos de Linguagem de Grande Escala (LLMs) em ambientes que combinam serviços de diferentes provedores traz um desafio específico: como manter governança rigorosa sobre os dados sem abrir mão das melhores ferramentas de Aprendizado de Máquina (ML) disponíveis?

Quando se usa o Amazon SageMaker AI em conjunto com o Databricks Unity Catalog, esse problema fica ainda mais evidente. O Unity Catalog gerencia metadados e permissões, enquanto os dados em si ficam armazenados no Amazon Simple Storage Service (Amazon S3) — quando o ambiente de nuvem escolhido é a AWS. O risco é que os jobs de treinamento do SageMaker AI acessem esses dados no S3 contornando o modelo de autorização granular do Unity Catalog, criando lacunas de auditoria e exposição a riscos de conformidade.

Em indústrias reguladas, saber exatamente quais dados treinaram quais modelos não é opcional — é um requisito crítico. Sem um padrão de integração estruturado, perde-se visibilidade sobre essa linhagem, e a conformidade fica comprometida.

Visão geral da solução

A solução demonstrada pela AWS aborda esse problema integrando três serviços principais: o Unity Catalog para governança centralizada, o Amazon EMR Serverless para pré-processamento com Apache Spark, e o SageMaker AI para o treinamento do modelo. O fluxo completo realiza as seguintes etapas:

  • Leitura dos dados de treinamento a partir de uma tabela gerenciada pelo Unity Catalog, com controles de governança aplicados
  • Pré-processamento dos dados usando EMR Serverless com Apache Spark
  • Fine-tuning do modelo Ministral-3-3B-Instruct via jobs de treinamento do SageMaker AI
  • Rastreamento da linhagem de dados no Unity Catalog, desde os dados brutos até o modelo treinado

Os componentes e suas funções na arquitetura são os seguintes:

  • Amazon SageMaker AI Studio – JupyterLab Space: orquestração do fluxo e treinamento do modelo
  • Amazon EMR Serverless: pré-processamento de dados com Spark, sem gerenciamento de clusters
  • Databricks Unity Catalog: catálogo de metadados, governança e rastreamento de linhagem
  • Hugging Face: acesso ao modelo pré-treinado
  • Amazon S3: armazenamento de dados e artefatos do modelo
  • AWS Secrets Manager: gerenciamento seguro de credenciais

No fluxo operacional, o usuário acessa o SageMaker AI Studio, inicia o pré-processamento via job do EMR Serverless — que acessa os dados no S3 gerenciado pelo Unity Catalog usando OAuth armazenado no Secrets Manager — e, após o processamento, o job de treinamento do SageMaker AI busca o modelo no Hugging Face, realiza o fine-tuning e devolve os artefatos ao bucket S3. Por fim, o modelo é registrado no Unity Catalog com linhagem de dados completa.

Pré-requisitos

Antes de iniciar, é necessário ter acesso e permissões configuradas para os seguintes serviços AWS: SageMaker AI, EMR Serverless, Amazon S3, AWS Secrets Manager, AWS Identity and Access Management (IAM), Amazon Virtual Private Cloud (Amazon VPC), Amazon CloudWatch Logs e Amazon Elastic Container Registry (Amazon ECR). Também é necessário ter uma VPC com acesso à internet configurado.

No lado do Databricks, os requisitos são: configurar o Unity Catalog para o Workspace, habilitar acesso externo a dados no Metastore (desativado por padrão) e criar credenciais OAuth (Client ID e secret) para acesso programático ao Databricks.

Passo a passo da implementação

O tutorial completo está disponível no notebook LLM_Finetunig_SageMaker_AI_Unity_Catalog.ipynb, que deve ser executado no SageMaker AI Studio. Para isso, acesse o Console do Amazon SageMaker AI, crie um domínio Studio se necessário (usando a configuração rápida), e crie um JupyterLab Space com a instância ml.m5.2xlarge, imagem Sagemaker Distribution 3.8.0 e 5 GB de armazenamento.

Etapa 1: Configuração AWS

Nesta etapa, são criados os recursos base: bucket S3 gerenciado pelo Unity Catalog com upload do dataset, secret no AWS Secrets Manager para as credenciais OAuth do Databricks, e as roles de Gerenciamento de Identidade e Acesso (IAM) necessárias para o EMR Serverless e o SageMaker AI.

O dataset utilizado são os arquivos SEC EDGAR (Sistema Eletrônico de Coleta, Análise e Recuperação de Dados da Comissão de Valores Mobiliários dos EUA) — especificamente formulários 10-K e 10-Q de empresas do S&P 500 de 2023 a 2024, com foco na seção de Fatores de Risco. Os dados são armazenados em formato JSON no S3, com a seguinte estrutura:

s3://aws-blog-smai-uc-bucket-ACCOUNTID/
├── raw/
│   └── risk_factors/
│       ├── form_type=10-K/
│       │   └── fiscal_year=2024/
│       │       └── cik=0000320193/
│       │           └── risk_factors.json
│       └── form_type=10-Q/
│           └── fiscal_year=2024/
│               └── quarter=1/
│                   └── cik=0000320193/
│                       └── risk_factors.json
├── curated/
│   └── ml/

Para autenticação, a solução usa OAuth para service principals (OAuth M2M), que gera tokens de curta duração. As credenciais são criadas seguindo a documentação do Databricks e armazenadas no AWS Secrets Manager.

As políticas de IAM para o EMR Serverless e para o SageMaker AI concedem acesso ao bucket S3 gerenciado pelo Unity Catalog, ao Secrets Manager e ao CloudWatch Logs. A role do SageMaker AI inclui também permissões para o ECR. Veja os exemplos de política abaixo:

Role do EMR Serverless:

emr_policy = {
    "Version": "2012-10-17",
    "Statement": [
        {
            "Effect": "Allow",
            "Action": [
                "s3:GetObject",
                "s3:PutObject",
                "s3:DeleteObject",
                "s3:ListBucket"
            ],
            "Resource": [
                f"arn:aws:s3:::{UC_MANAGED_BUCKET}/*",
                f"arn:aws:s3:::{UC_MANAGED_BUCKET}"
            ]
        },
        {
            "Effect": "Allow",
            "Action": [
                "secretsmanager:GetSecretValue"
            ],
            "Resource": [
                f"arn:aws:secretsmanager:{AWS_REGION}:{AWS_ACCOUNT_ID}:secret:databricks/*"
            ]
        },
        {
            "Effect": "Allow",
            "Action": [
                "logs:CreateLogGroup",
                "logs:CreateLogStream",
                "logs:PutLogEvents"
            ],
            "Resource": f"arn:aws:logs:{AWS_REGION}:{AWS_ACCOUNT_ID}:*"
        }
    ]
}

Role de execução do SageMaker AI:

sagemaker_policy = {
    "Version": "2012-10-17",
    "Statement": [
        {
            "Effect": "Allow",
            "Action": [
                "s3:GetObject",
                "s3:PutObject",
                "s3:DeleteObject",
                "s3:ListBucket"
            ],
            "Resource": [
                f"arn:aws:s3:::{UC_MANAGED_BUCKET}/*",
                f"arn:aws:s3:::{UC_MANAGED_BUCKET}"
            ]
        },
        {
            "Effect": "Allow",
            "Action": [
                "secretsmanager:GetSecretValue"
            ],
            "Resource": [
                f"arn:aws:secretsmanager:{AWS_REGION}:{AWS_ACCOUNT_ID}:secret:databricks/*"
            ]
        },
        {
            "Effect": "Allow",
            "Action": [
                "logs:CreateLogGroup",
                "logs:CreateLogStream",
                "logs:PutLogEvents",
                "logs:DescribeLogStreams"
            ],
            "Resource": f"arn:aws:logs:{AWS_REGION}:{AWS_ACCOUNT_ID}:*"
        },
        {
            "Effect": "Allow",
            "Action": [
                "ecr:GetAuthorizationToken",
                "ecr:BatchCheckLayerAvailability",
                "ecr:GetDownloadUrlForLayer",
                "ecr:BatchGetImage"
            ],
            "Resource": "*"
        }
    ]
}

Etapa 2: Configuração do Databricks Unity Catalog

Nesta etapa, o bucket S3 é registrado como External Location no Unity Catalog. Em seguida, são criados os objetos de banco de dados: catálogo, schemas e tabela externa. As permissões necessárias são concedidas ao Service Principal — incluindo USE CATALOG, USE SCHEMA, SELECT na tabela de dados brutos e CREATE TABLE nos schemas de curadoria e ML.

Após a configuração, o cliente Databricks é inicializado e a conexão com o Unity Catalog é testada:

# Initialize Databricks client
from databricks.sdk import WorkspaceClient

w = WorkspaceClient(
    host=UNITY_WORKSPACE_URL,
    client_id=DATABRICKS_CLIENT_ID,
    client_secret=DATABRICKS_CLIENT_SECRET
)

table_info = w.tables.get(f"{UNITY_CATALOG_NAME}.{UNITY_SCHEMA_DATA}.{UNITY_TABLE_NAME}")
print(f"Table: {table_info.name}")
print(f"Storage Location: {table_info.storage_location}")
print(f"Table Format: {table_info.data_source_format}")

A saída esperada confirma a tabela, sua localização no S3 e o formato Delta:

Table: risk_factors
Storage Location: s3:// aws-blog-smai-uc-bucket-ACCOUNTID/raw/risk_factors
Table Format: DataSourceFormat.DELTA

Etapa 3: Configuração da aplicação EMR Serverless

O EMR Serverless é criado em uma VPC com acesso à internet. Isso é necessário porque, como o EMR Serverless não inclui suporte nativo ao Delta Lake (pelo menos até fevereiro de 2026), a aplicação Spark precisa baixar os arquivos JAR do Delta Lake do repositório Maven Central durante a inicialização. Sem acesso à internet, esse download falha e o job não consegue resolver o formato de tabela Delta.

Etapa 4: Pré-processamento com EMR Serverless

Um script de pré-processamento lê os fatores de risco do SEC EDGAR da tabela Delta, filtra e limpa os dados, formata o conteúdo como prompt no estilo instrução para fine-tuning e grava os resultados como tabela Delta no S3. Os parâmetros Spark necessários para habilitar o Delta Lake são configurados assim:

jobDriver={
    'sparkSubmit': {
        'sparkSubmitParameters': ' '.join([
            '--packages io.delta:delta-spark_2.12:3.2.0',
            '--conf spark.sql.extensions=io.delta.sql.DeltaSparkSessionExtension',
            '--conf spark.sql.catalog.spark_catalog=org.apache.spark.sql.delta.catalog.DeltaCatalog'
        ])
    }

Após a conclusão do job, os dados pré-processados são registrados como tabela Delta no schema de curadoria do Unity Catalog via SDK do Databricks, usando um SQL warehouse serverless.

Etapa 5: Fine-tuning com SageMaker AI

Com o dataset preparado, o job de treinamento do SageMaker AI realiza o fine-tuning do modelo Ministral-3-3B-Instruct-2512, disponível no Hugging Face. Trata-se de um modelo compacto de 3 bilhões de parâmetros da Mistral AI, otimizado para seguir instruções com bom desempenho em tarefas de raciocínio.

O script de treinamento aplica duas técnicas de eficiência de memória: quantização em ponto flutuante de 8 bits (FP8) para reduzir o uso de memória, e Adaptação de Baixa Classificação (LoRA) para treinar apenas 1–2% dos parâmetros em vez dos 3 bilhões completos. Os dados são tokenizados com janela de contexto de 1.024 tokens. O treinamento usa uma instância ml.g5.16xlarge, e os artefatos do modelo são enviados ao bucket S3 ao final do job.

Etapa 6: Registro do modelo no Unity Catalog

Após o upload dos artefatos para o S3, o modelo fine-tunado é registrado no schema ML do Unity Catalog usando o MLflow gerenciado pelo Databricks via SDK. Um experimento MLflow é criado para registrar metadados de treinamento — hiperparâmetros, tabela de origem, detalhes do job — e o modelo é versionado com documentação que inclui linhagem dos dados de origem, serviço de treinamento e especificações do modelo.

Etapa 7: Criação de linhagem de dados no Unity Catalog

O Unity Catalog captura automaticamente a linhagem de dados para workloads executados dentro do Databricks. Para workloads externos — como jobs do EMR Serverless e jobs de treinamento do SageMaker AI — é possível adicionar metadados e linhagem através da API de Metadados Externos e da API de Linhagem Externa. Essa capacidade de “traga sua própria linhagem de dados” (atualmente em prévia pública desde abril de 2026) oferece uma visão completa da linhagem no Unity Catalog.

Para o fluxo completo, são criados os seguintes objetos de linhagem:

  • Metadados externos do job EMR Serverless
  • Linhagem upstream: tabela bruta → job EMR Serverless
  • Linhagem downstream: job EMR Serverless → tabela pré-processada
  • Metadados externos do job de treinamento do SageMaker AI
  • Linhagem upstream: tabela pré-processada → job de treinamento
  • Linhagem downstream: job de treinamento → versão do modelo

Com isso, é possível visualizar o grafo completo de linhagem — dos dados brutos ao modelo fine-tunado — diretamente no console do Unity Catalog.

Limpeza dos recursos

Após validar a solução, a AWS recomenda remover os seguintes recursos para evitar cobranças contínuas: aplicação EMR Serverless, buckets e objetos S3, roles e políticas IAM, secrets no AWS Secrets Manager, recursos de VPC e objetos do Unity Catalog (catálogo, schema, tabela e service principal).

Conclusão

A integração entre Databricks Unity Catalog e Amazon SageMaker AI apresentada pela AWS demonstra como é possível realizar fine-tuning de LLMs com governança contínua, sem comprometer segurança ou conformidade. Combinando o Unity Catalog para controle de acesso e linhagem, o EMR Serverless para pré-processamento escalável e o SageMaker AI para treinamento, o resultado é um pipeline de ML governado de ponta a ponta — adequado para ambientes de produção e indústrias reguladas.

O padrão pode ser reutilizado com outros LLMs e datasets maiores, servindo como base para workloads de IA generativa em ambientes multicloud com requisitos rigorosos de auditoria.

Fonte

Fine-tune LLM with Databricks Unity Catalog and Amazon SageMaker AI (https://aws.amazon.com/blogs/machine-learning/fine-tune-llm-with-databricks-unity-catalog-and-amazon-sagemaker-ai/)

Comments

Leave a Reply

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