Otimização completa do Qwen3-4B

Otimização completa do modelo de linguagem Qwen3-4B em uma única GPU H100. Este passo a passo mostra como:

• Execute o ajuste fino completo, que atualiza todos os parâmetros do modelo para máxima adaptação aos seus dados.
• Usar o ambiente de IA do Databricks v5 sem instalar nenhuma biblioteca adicional
• Aproveitar o TRL (Aprendizado de Reforço do Transformador) para ajuste refinado supervisionado
• Registrar o modelo ajustado no Catálogo do Unity para governança e implantação

Principais conceitos:

• Ajuste fino completo: atualiza todos os pesos do modelo, dando ao modelo a maior capacidade de aprender com seu conjunto de dados, ao custo de maior memória e computação do que os métodos com uso eficiente de parâmetros.
• TRL: uma biblioteca para o treinamento de modelos de linguagem com aprendizado por reforço e ajuste fino supervisionado
• Treinamento com uso eficiente de memória: usa precisão mista BF16 e checkpoint de gradiente para ajustar um modelo com 4 bilhões de parâmetros em uma única GPU H100.

Ajuste fino completo vs. matriz de decisão LoRA.

Este notebook usa ajuste fino completo, que atualiza todos os parâmetros de modelo. A alternativa, LoRA (Adaptação de Baixa Ordem), congela o modelo base e treina apenas pequenas camadas adaptadoras.

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

Conectar-se à computação de GPU sem servidor

Para se conectar à computação de GPU sem servidor:

1. Clique no menu suspenso Conectar no notebook e selecione GPU sem servidor.
2. Escolha uma GPU 1x H100 como o acelerador.
3. Abra o painel Ambiente e escolha IA v5 como o ambiente base.
4. Clique em Aplicar.

Para obter mais informações, consulte a documentação de computação da GPU.

Importar bibliotecas

O ambiente de IA do Databricks v5 já inclui todas as bibliotecas necessárias para este exemplo (comotrl, , transformersdatasetse 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 conjunto de dados e acompanhamento de MLflow.

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 do Catálogo do Unity

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

• Catálogo & Esquema: Organizar modelos no Namespace do Catálogo do Unity (padrão: main.default)
• Nome do modelo: o nome do modelo registrado no Catálogo do Unity para governança e implantação
• Volume: Volume do Catálogo do Unity para armazenar pontos de verificação de 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}, visando facilitar o acesso e o controle de versão.

Hiperparâmetros de treinamento

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

• Modelo & Conjunto de dados: Qwen3-4B com o conjunto de dados conversacional Capybara
• Tamanho do Lote (1): número de exemplos por GPU por etapa de treinamento, mantido pequeno para caber um ajuste fino completo na memória
• Acúmulo de gradiente (8): acumula gradientes em 8 lotes para um tamanho de lote efetivo de 8
• Taxa de Aprendizagem (2e-5): taxa conservadora apropriada para ajuste completo
• Máximo de etapas (50): limita o treinamento a 50 etapas para uma execução de demonstração rápida
• Registro em log & Ponto de verificação: salva o progresso a cada 25 etapas, registra métricas a cada 10 etapas

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 conjunto de dados

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

• Conjunto de dados: trl-lib/Capybara - dados de conversação de alta qualidade otimizados para seguimento de instruções
• Divisão de treinamento/validação: cria uma divisão de 90/10 se nenhum conjunto de testes existir
• Validação de dados: garante a formatação correta para ajuste fino de conversas

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 o modelo e o tokenizador

A próxima célula carrega o modelo base e o tokenizador e, em seguida, configura-os para ajuste fino de conversa:

• Carregamento do modelo: baixa o Qwen3-4B da Hugging Face com precisão BF16
• Configuração do tokenizador: configura o tokenizador rápido com preenchimento adequado
• Formatação de chat: aplica um modelo de chat para conversas estruturadas se o tokenizer ainda não definir um
• Configuração de token: define um token de preenchimento como o token EOS para lidar adequadamente com sequências

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")

Treinar o modelo

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

Configuração de treinamento

• Configuração do lote: 1 exemplo por dispositivo com 8 etapas de acúmulo de gradiente (tamanho efetivo do lote: 8)
• Otimização: etapas de aquecimento, decadência de peso e melhor seleção de modelo com base na perda de avaliação
• Registro em log: relata métricas ao MLflow para acompanhamento de experimentos

Otimizações principais habilitadas

• Precisão mista BF16: computação mais rápida com menor volume de memória, adequado para GPUs H100
• Checkpoint de gradiente: troca computação 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
• Acúmulo de gradiente: simula tamanhos de lote maiores para treinamento estável
• Checkpointing: salva o modelo a cada 25 etapas com um limite de 2 checkpoints

O loop de treinamento registra o progresso a cada 10 etapas e realiza uma avaliação a cada 25 etapas.

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 de modelo

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

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

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

Registrar modelo no Catálogo do Unity

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

Fluxo de trabalho de registro de modelo

1. Carregar modelo treinado: carrega o modelo de pesos completos salvo e o tokenizador
2. Preparar para registro: cria um dicionário de modelo de transformadores com o modelo e o tokenizador
3. Registrar no Catálogo do Unity: registra no MLflow e registra no Catálogo do Unity
4. Adicionar metadados: inclui o tipo de tarefa, a família de modelos e as informações de tamanho

Benefícios do registro do Catálogo do Unity

• Governança: Registro de modelo centralizado com controle de acesso e rastreamento de linhagem
• Versionamento: gerenciamento automático de versão para o ciclo de vida do modelo
• Implantação: implantação fácil para pontos de extremidade do serviço de modelo
• Capacidade de descoberta: os modelos são pesquisáveis e documentados no Catálogo do Unity

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={
                "task": task,
                "pretrained_model_name": MODEL_NAME,
                "databricks_model_family": "Qwen3ForCausalLM",
                "databricks_model_size_parameters": 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óximas Etapas 

Seu modelo Qwen3-4B foi ajustado com sucesso usando o ajuste fino de peso completo e registrado no Catálogo do Unity. Em seguida, você pode:

• Implantar o modelo: servir modelos com o Model Serving
• Saiba mais sobre o treinamento distribuído: treinamento distribuído com várias GPUs e vários nós
• Acompanhar experimentos e monitorar GPUs: Acompanhamento e observabilidade de experimentos
• Solucionar problemas: solucionar problemas na computação de GPU sem servidor

Notebook de exemplo

Otimização completa do Qwen3-4B

Obter laptop