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

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