Ajuste fino completo do Qwen3-4B

Open in Databricks

Realize o ajuste fino completo do modelo de linguagem grande Qwen3-4B em uma única GPU H100. Este passo a passo mostra como:

Realize a execução do ajuste fino completo , que atualiza cada parâmetro do modelo para adaptação máxima aos seus dados
Use o ambiente Databricks AI v5 sem instalar nenhuma biblioteca adicional
Aproveite TRL (Transformer Reinforcement Learning) para ajuste fino supervisionado
Registre o modelo ajustado no Unity Catalog para governança e implantação

key conceitos:

Ajuste fino completo : Atualiza todos os pesos do modelo, dando ao modelo a maior capacidade de aprender com seu dataset ao custo de maior memória e compute do que os métodos eficientes em parâmetros
TRL: Uma biblioteca para treinamento de modelos de linguagem com aprendizado por reforço e ajuste fino supervisionado
Treinamento com eficiência de memória : Usa precisão mista BF16 e checkpointing de gradiente para ajustar um ajuste fino completo de 4 bilhões de parâmetros em uma única GPU H100

Matriz de decisão de ajuste fino completo vs LoRA

Este Notebook usa ajuste fino completo , que atualiza todos os parâmetros do modelo. A alternativa, LoRA (Low-Rank Adaptation), congela o modelo base e ensina apenas pequenas camadas adaptadoras.

Cenário	Recomendação	Motivo
Grande mudança no comportamento do modelo	Ajuste fino completo	Atualiza todos os parâmetros para mudanças fundamentais no comportamento do modelo
A mais alta qualidade possível em uma única tarefa	Ajuste fino completo	Sem aproximação de baixa classificação, para que o modelo tenha capacidade total de adaptação
Memória da GPU limitada	LoRA	Acomoda modelos maiores na memória com o treinamento de apenas aproximadamente 1% dos parâmetros
Vários adaptadores específicos para tarefas	LoRA	Fazer swap de diferentes adaptadores no mesmo modelo base

O ajuste fino completo de um modelo de 4 bilhões de parâmetros requer significativamente mais memória de GPU do que LoRA porque o estado do otimizador e os gradientes são mantidos para cada parâmetro. Este Notebook usa precisão mista BF16 e checkpointing de gradiente para que o treinamento caiba em uma única GPU H100 (80 GB).

Conectar-se ao compute de GPU serverless

Para se conectar ao compute de GPU serverless:

Clique no menu suspenso Conectar no Notebook e selecione GPU Serverless .
Escolha uma GPU H100 de 1x como acelerador.
Abra o painel Ambiente e escolha AI v5 como o ambiente base.
Clique em Aplicar .

Para mais informações, consulte a documentação de compute de GPU.

Importar bibliotecas

O ambiente Databricks AI v5 já inclui todas as bibliotecas necessárias para este exemplo (como trl, transformers, datasets e mlflow), portanto, nenhuma instalação adicional é necessária.

A próxima célula importa as bibliotecas necessárias para treinamento de modelo, manipulação de dataset e acompanhamento do MLflow.

Python
from datasets import load_dataset
from transformers import AutoModelForCausalLM, AutoTokenizer
from trl import (
    SFTConfig,
    SFTTrainer,
    setup_chat_format
)
import torch
import mlflow

Configuração

Integração com o Unity Catalog

A próxima célula configura onde seu modelo ajustado será armazenado e registrado:

Catálogo e esquema : Organize modelos no seu espaço de nomes do Unity Catalog (default: main.default)
Nome do Modelo : O nome do modelo registrado no Unity Catalog para governança e implantação
Volume : Volume do Unity Catalog para armazenar checkpoints do modelo durante o treinamento

Esses widgets permitem personalizar o local de armazenamento sem editar o código. O modelo será registrado como {catalog}.{schema}.{model_name} para fácil acesso e controle de versão.

Hiperparâmetros de treinamento

A célula também define os key parâmetros de treinamento:

Modelo & Dataset : Qwen3-4B com o dataset conversacional Capybara
Tamanho do lote (1) : número de exemplos por GPU por o passo de treinamento, mantido pequeno para caber um ajuste fino completo na memória
Acúmulo de Gradientes (8) : Acumula gradientes em 8 lotes para um tamanho de lote efetivo de 8
Taxa de Aprendizado (2e-5) : Taxa conservadora apropriada para ajuste fino completo
Máximo de os passos (50) : Limita o treinamento a 50 os passos para uma execução de demonstração rápida
Registro & Checkpoint : Salva o progresso a cada 25 os passos, log métricas a cada 10 os passos

Python
dbutils.widgets.text("uc_catalog", "main")
dbutils.widgets.text("uc_schema", "default")
dbutils.widgets.text("uc_model_name", "qwen3_4b_assistant")
dbutils.widgets.text("uc_volume", "checkpoints")

UC_CATALOG = dbutils.widgets.get("uc_catalog")
UC_SCHEMA = dbutils.widgets.get("uc_schema")
UC_MODEL_NAME = dbutils.widgets.get("uc_model_name")
UC_VOLUME = dbutils.widgets.get("uc_volume")

print(f"UC_CATALOG: {UC_CATALOG}")
print(f"UC_SCHEMA: {UC_SCHEMA}")
print(f"UC_MODEL_NAME: {UC_MODEL_NAME}")
print(f"UC_VOLUME: {UC_VOLUME}")

# MLflow and Unity Catalog configuration

# Model selection
MODEL_NAME = "Qwen/Qwen3-4B"
DATASET_NAME = "trl-lib/Capybara"
OUTPUT_DIR = f"/Volumes/{UC_CATALOG}/{UC_SCHEMA}/{UC_VOLUME}/{UC_MODEL_NAME}"

# Training hyperparameters
BATCH_SIZE = 1
GRADIENT_ACCUMULATION_STEPS = 8
LEARNING_RATE = 2e-5
MAX_STEPS = 50
EVAL_STEPS = 25
LOGGING_STEPS = 10
SAVE_STEPS = 25

Carregar e preparar o dataset

A próxima célula carrega o dataset de treinamento e o prepara para o ajuste fino:

Dataset : trl-lib/Capybara - dados conversacionais de alta qualidade otimizados para seguir instruções
Divisão de ensinar/validação : cria uma divisão de 90/10 se nenhum conjunto de teste existir
Validação de dados : garante a formatação adequada para ajuste fino conversacional

Python
dataset = load_dataset(DATASET_NAME)
print(f"✓ Dataset loaded: {dataset}")

if "test" not in dataset:
    print("Creating validation split from training data...")
    dataset = dataset["train"].train_test_split(test_size=0.1, seed=42)
    print("✓ Data split: 90% train, 10% validation")

Inicializar modelo e tokenizador

A próxima célula carrega o modelo base e o tokenizador, e então os configura para o ajuste fino conversacional:

Carregamento do modelo : Faz o download de Qwen3-4B do Hugging Face em precisão BF16
Configuração do tokenizador : Configura o tokenizador rápido com preenchimento adequado
Formatação de chat : Aplica um padrão de chat para conversas estruturadas se o tokenizador ainda não definir um
Configuração de tokens : define o tokens de preenchimento como token EOS para o tratamento adequado da sequência

Python
model = AutoModelForCausalLM.from_pretrained(
    MODEL_NAME,
    torch_dtype=torch.bfloat16,
    trust_remote_code=True,
)

tokenizer = AutoTokenizer.from_pretrained(
    MODEL_NAME,
    trust_remote_code=True,
    use_fast=True
)

# Chat template formatting for conversational fine-tuning
if tokenizer.chat_template is None:
    print("Adding chat template for proper conversation formatting...")
    model, tokenizer = setup_chat_format(model, tokenizer, format="chatml")
    print("✓ ChatML format applied for structured conversations")

if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token
    print("✓ Padding token set to EOS token")

print("✓ Model and tokenizer loaded successfully")

Ensinar o modelo

A próxima célula configura e executa o processo de ajuste fino completo:

Configuração de treinamento

Configuração de lotes : 1 amostra por dispositivo com 8 os passos de acumulação de gradiente (tamanho de lote efetivo: 8)
Otimização : os passos de aquecimento, decaimento de peso e seleção do melhor modelo com base na perda de avaliação
Registro : Relata métricas ao MLflow para o acompanhamento de experimentos

key otimizações ativadas

Precisão mista BF16 : Computação mais rápida com menor consumo de memória, bem adequada para GPUs H100
Checkpointing de gradiente : troca compute extra por uma grande redução na memória de ativação, o que permite que um ajuste fino completo de 4B caiba em um único H100
Acumulação de gradiente : simula tamanhos de lote maiores para treinamento estável
Pontos de verificação : Salva o modelo a cada 25 os passos com um limite de 2 pontos de verificação

O loop de treinamento log o progresso a cada 10 os passos e avalia a cada 25 os passos.

Python
with mlflow.start_run(run_name=f"{MODEL_NAME}_full-fine-tuning", log_system_metrics=True):
    try:
        print(f"Learning rate: {LEARNING_RATE}")

        training_args_dict = {
            "output_dir": OUTPUT_DIR,
            "per_device_train_batch_size": BATCH_SIZE,
            "per_device_eval_batch_size": BATCH_SIZE,
            "gradient_accumulation_steps": GRADIENT_ACCUMULATION_STEPS,
            "learning_rate": LEARNING_RATE,
            "max_steps": MAX_STEPS,
            "eval_steps": EVAL_STEPS,
            "logging_steps": LOGGING_STEPS,
            "save_steps": SAVE_STEPS,
            "save_total_limit": 2,
            "report_to": "mlflow",  # Log to MLflow
            "warmup_steps": 10,
            "weight_decay": 0.01,
            "metric_for_best_model": "eval_loss",
            "greater_is_better": False,
            "eval_strategy": "steps",  # Run evaluation every eval_steps
            "save_strategy": "steps",  # Checkpoint on the same cadence as eval
            "load_best_model_at_end": True,  # Register the best-eval checkpoint, not the last
            "dataloader_pin_memory": False,
            "remove_unused_columns": False,
            "bf16": True,  # Mixed precision training
            "gradient_checkpointing": True,  # Reduce activation memory for full fine-tuning
            "gradient_checkpointing_kwargs": {"use_reentrant": False},
        }

        training_args = SFTConfig(**training_args_dict)

        trainer = SFTTrainer(
            model=model,
            args=training_args,
            train_dataset=dataset["train"],
            eval_dataset=dataset["test"],
            processing_class=tokenizer,
        )

        print("\n" + "="*50)
        print("STARTING TRAINING")
        print("="*50)

        print("🚀 Full fine-tuning Qwen3-4B on a single H100 GPU")

        trainer.train()
        print("\n✓ Training completed successfully!")

    except Exception as e:
        print(f"✗ Training failed: {e}")
        raise

Salvar artefatos do modelo

A próxima célula salva o modelo treinado e o tokenizador no volume do Unity Catalog:

Pesos completos do modelo : Salva o modelo ajustado completo, pronto para carregar diretamente para inferência
Tokenizer : Salva a configuração do tokenizer para inferência
Localização de armazenamento : Salva em /Volumes/{catalog}/{schema}/{volume}/{model_name}

Python
try:
    print("\nSaving trained model...")

    trainer.save_model(training_args.output_dir)
    print("✓ Full model weights saved")

    tokenizer.save_pretrained(training_args.output_dir)
    print("✓ Tokenizer saved with model")
    print(f"\n🎉 All artifacts saved to: {training_args.output_dir}")

except Exception as e:
    print(f"✗ Model saving failed: {e}")
    raise

Faça o registro do modelo no Unity Catalog

A próxima célula registra o modelo ajustado no Unity Catalog para governança e implantação:

Fluxo de trabalho de registro de modelo

Carregar modelo treinado : Carrega o modelo de peso completo salvo e o tokenizador
Preparar para registro : cria um dicionário de modelo transformers com o modelo e o tokenizador
Registro no Unity Catalog : faz log no MLflow e registro no Unity Catalog
Adicionar metadados : Inclui o tipo de tarefa, a família do modelo e as informações de tamanho

Benefícios do registro do Unity Catalog

Governança : Registro de modelo centralizado com controle de acesso e acompanhamento da linhagem
Versionamento : Gerenciamento automático de versões para o ciclo de vida do modelo
Implantação : Implantação fácil para endpoint de servindo modelo
Capacidade de descoberta : Os modelos são pesquisáveis e documentados no Unity Catalog

Python
mlflow_run_id = mlflow.last_active_run().info.run_id
print("\nRegistering model with MLflow and Unity Catalog...")

with mlflow.start_run(run_id=mlflow_run_id):
    try:
        # Load the trained full-weight model for registration
        print("Loading fine-tuned model for registration...")
        trained_model = AutoModelForCausalLM.from_pretrained(
            training_args.output_dir,
            torch_dtype=torch.bfloat16,
            trust_remote_code=True
        )
        tokenizer = AutoTokenizer.from_pretrained(training_args.output_dir)
        model_type = "Full fine-tuning"
        size_params = "4b"

        # Prepare transformers model dictionary
        transformers_model = {
            "model": trained_model,
            "tokenizer": tokenizer
        }

        # Create Unity Catalog model name
        full_model_name = f"{UC_CATALOG}.{UC_SCHEMA}.{UC_MODEL_NAME}"

        print(f"Registering model as: {full_model_name}")

        # Start MLflow run and log model
        task = "llm/v1/chat"
        model_info = mlflow.transformers.log_model(
            transformers_model=transformers_model,
            task=task,
            registered_model_name=full_model_name,
            metadata={
                &quot;task&quot;: task,
                &quot;pretrained_model_name&quot;: MODEL_NAME,
                &quot;databricks_model_family&quot;: &quot;Qwen3ForCausalLM&quot;,
                &quot;databricks_model_size_parameters&quot;: size_params,
            },
            repo_type="local",  # Fix: specify repo_type for local path
        )

        print(f"✓ Model successfully registered in Unity Catalog: {full_model_name}")
        print(f"✓ MLflow model URI: {model_info.model_uri}")

        # Print deployment information
        print(f"\n📦 Model Registration Complete!")
        print(f"Unity Catalog Path: {full_model_name}")
        print(f"Model Type: {model_type}")

    except Exception as e:
        print(f"✗ Model registration failed: {e}")
        print("Model is still saved locally and can be registered manually")
        print(f"Local model path: {training_args.output_dir}")
        raise

Próximos os passos

Seu modelo Qwen3-4B foi ajustado com sucesso com ajuste fino de peso completo e registrado no Unity Catalog. Em seguida, você pode:

Implante o modelo : Implante modelos com o Model Serving
Saiba mais sobre treinamento distribuído : Treinamento distribuído multi-GPU e multinós
Acompanhe experimentos e monitore GPUs : Acompanhamento de experimentos e observabilidade
Solucionar problemas : Solucionar problemas no compute GPU serverless

Notebook de exemplo

Ajuste fino completo do Qwen3-4B

Abrir notebook em uma nova aba Open in Databricks

Matriz de decisão de ajuste fino completo vs LoRA​

Conectar-se ao compute de GPU serverless​

Importar bibliotecas​

Configuração​

Integração com o Unity Catalog​

Hiperparâmetros de treinamento​

Carregar e preparar o dataset​

Inicializar modelo e tokenizador​

Ensinar o modelo​

Configuração de treinamento​

key otimizações ativadas​

Salvar artefatos do modelo​

Faça o registro do modelo no Unity Catalog​

Fluxo de trabalho de registro de modelo​

Benefícios do registro do Unity Catalog​

Próximos os passos​

Notebook de exemplo​

Ajuste fino completo do Qwen3-4B

Matriz de decisão de ajuste fino completo vs LoRA

Conectar-se ao compute de GPU serverless

Importar bibliotecas

Configuração

Integração com o Unity Catalog

Hiperparâmetros de treinamento

Carregar e preparar o dataset

Inicializar modelo e tokenizador

Ensinar o modelo

Configuração de treinamento

key otimizações ativadas

Salvar artefatos do modelo

Faça o registro do modelo no Unity Catalog

Fluxo de trabalho de registro de modelo

Benefícios do registro do Unity Catalog

Próximos os passos

Notebook de exemplo