Ajuste fino distribuído de Qwen2-0.5B com LoRA

Este notebook demonstra como ajustar de forma eficiente o modelo de linguagem grande Qwen2-0.5B utilizando técnicas eficientes em parâmetros numa computação de GPU sem servidor. Você aprenderá a:

• Aplicar LoRA (Low-Rank Adaptation) para reduzir os parâmetros treináveis em ~99%, mantendo a qualidade do modelo.
• Use Liger Kernels para treino com eficiência de memória com kernels Triton otimizados
• Aproveite o TRL (Aprendizagem por Reforço por Transformadores) para ajuste fino supervisionado
• Registe o modelo finamente ajustado no Catálogo Unity para governação e implementação

Conceitos-chave:

• LoRA: Uma técnica que congela o modelo base e treina pequenas camadas adaptadoras, reduzindo drasticamente os requisitos de memória e o tempo de treino
• Kernels Liger: Kernels otimizados para GPU que reduzem o uso de memória até 80% através de operações de fusão
• TRL: Uma biblioteca para treinar modelos de linguagem com aprendizagem por reforço e ajuste fino supervisionado
• AI Runtime: Computação gerida com Databricks que escala automaticamente os recursos da GPU

LoRA vs matriz de decisão de ajuste fino completo

LoRA (Adaptação de Baixo Rank) congela o modelo base e treina apenas pequenas camadas adaptadoras, reduzindo os parâmetros treináveis em cerca de 99%. Isto torna o treino mais rápido e eficiente em termos de memória.

Benefícios do Liger Kernel

Os kernels Liger são operações otimizadas para GPU que fundem múltiplos passos em kernels únicos, reduzindo as transferências de memória e melhorando a eficiência. O documento técnico fornece referências detalhadas que mostram melhorias significativas de desempenho.

• Operações fundidas: Combina operações (por exemplo, linear + perda) para reduzir a sobrecarga de memória até 80%
• Kernels Triton: Kernels personalizados de GPU otimizados para operações de transformer (RMSNorm, RoPE, SwiGLU, CrossEntropy)
• Eficiência de memória: Permite lotes maiores ou modelos que de outra forma não cabiam na memória da GPU
• Otimização de GPU única: Particularmente eficaz para cenários de treino A10/A100 com GPU única

Este notebook utiliza a biblioteca TRL para simplificar a configuração de treino e aplicar estas otimizações automaticamente.

Ligar-se à computação GPU sem servidor

Este portátil requer computação com GPU serverless para executar o treino distribuído. A computação sem servidor de GPU fornece e gere automaticamente os recursos da GPU para o seu trabalho.

Para se ligar, clique no menu suspenso Conectar no portátil e selecione GPU Serverless.

Para mais informações, consulte a documentação do AI Runtime.

Instalar as bibliotecas necessárias

A célula seguinte instala os pacotes Python necessários para o ajuste fino distribuído:

Bibliotecas de formação básica:

• trl==0.18.1: Biblioteca de Aprendizagem por Reforço de Transformadores para ajuste fino supervisionado e RLHF
• Peft: Parameter-Efficient Fine-Tuning biblioteca que fornece implementação LoRA
• liger-kernel: Kernels GPU otimizados em memória para treino eficiente de transformadores

Bibliotecas de apoio:

• hf_transfer: Downloads acelerados do Hugging Face Hub usando transferência baseada em Rust
• mlflow>=3.6.0: Rastreio de experiências e integração do registo de modelos

O %restart_python comando reinicia o interpretador Python para garantir que os pacotes recém-instalados são devidamente carregados.

%pip install trl==0.18.1 peft hf_transfer  liger-kernel
%pip install mlflow>=3.6.0
%restart_python

Configuração

Integração com o Unity Catalog

A célula seguinte configura onde o seu modelo finamente ajustado será armazenado e registado:

• Catálogo & Esquema: Organize os modelos dentro do seu namespace do Unity Catalog (por padrão: main.default)
• Nome do Modelo: O nome do modelo registado no Catálogo Unity para governação e implementação
• Volume: Volume do Unity Catalog para armazenar pontos de verificação do modelo durante o treino

Estes widgets permitem-lhe personalizar a localização de armazenamento sem precisar de editar código. O modelo será registado como {catalog}.{schema}.{model_name} para fácil acesso e controlo de versões.

Treino de hiperparâmetros

A célula também define parâmetros chave de treino:

• Modelo & Conjunto de Dados: Qwen2-0.5B com conjunto de dados conversacionais Capybara
• Tamanho do lote (8): Número de exemplos por GPU por passo de treino
• Acumulação de gradiente (4): Acumula gradientes ao longo de 4 lotes para um tamanho efetivo de lote de 32
• Taxa de Aprendizagem (1e-4): Taxa conservadora, aumentada automaticamente 10 vezes para formação LoRA
• Épocas (1): Passagem única pelo conjunto de dados para prevenir ajuste excessivo
• Registo e Checkpointing: Guarda o progresso a cada 100 passos, regista métricas a cada 25 passos

dbutils.widgets.text("uc_catalog", "main")
dbutils.widgets.text("uc_schema", "default")
dbutils.widgets.text("uc_model_name", "qwen2_liger_lora_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 - Choose based on your compute constraints
MODEL_NAME = "Qwen/Qwen2-0.5B"
DATASET_NAME = "trl-lib/Capybara"
OUTPUT_DIR = f"/Volumes/{UC_CATALOG}/{UC_SCHEMA}/{UC_VOLUME}/qwen2-0.5b-lora"

# Training hyperparameters
BATCH_SIZE = 8
GRADIENT_ACCUMULATION_STEPS = 4
LEARNING_RATE = 1e-4
NUM_EPOCHS = 1
EVAL_STEPS = 100
LOGGING_STEPS = 25
SAVE_STEPS = 100

Configuração LoRA

A próxima célula configura os parâmetros de Adaptação de Baixa Ordem (LoRA) que controlam como o modelo é ajustado finamente. O LoRA congela os pesos do modelo base e treina apenas pequenas matrizes adaptadoras, reduzindo drasticamente os requisitos de memória.

Seleção de parâmetros

• Classificação (r=8): Proporciona um bom equilíbrio entre desempenho e parâmetros
• Alfa (32): Fator de escala, tipicamente 2-4 vezes o ranking
• Dropout (0.1): Regularização para evitar sobreajustes

Módulos-alvo para Qwen2

Este exemplo dirige-se a todas as camadas chave de transformação:

• Atenção: q_proj, k_proj, v_proj, o_proj
• MLP: gate_proj, up_proj, down_proj

LORA_R = 8
LORA_ALPHA = 32
LORA_DROPOUT = 0.1
LORA_TARGET_MODULES = [
    "q_proj", "k_proj", "v_proj", "o_proj",
    "gate_proj", "up_proj", "down_proj"
]

Defina a função de treino

A célula seguinte cria a função de treino distribuído que irá correr em várias GPUs. Veja o que ele faz:

Configuração de treino distribuído

O decorador @distributed configura o processamento GPU sem servidor:

• 8 GPUs: Distribui o treino entre 8 GPUs H100 para um treino mais rápido
• Orquestração automática: Gere o provisionamento da GPU, distribuição de dados e sincronização

Fluxo de trabalho de treinamento

A função executa estes passos:

1. Carregar o conjunto de dados: Descarrega e prepara o conjunto de dados conversacional Capybara
2. Inicializar modelo: Carrega Qwen2-0.5B e o tokenizador com a formatação de chat
3. Aplicar LoRA: Anexa camadas adaptadoras para reduzir parâmetros treináveis em ~99%
4. Configurar treino: Define o tamanho do lote, a taxa de aprendizagem e as otimizações do kernel Liger
5. Treinar modelo: Executa o ciclo de treino com armazenamento de pontos de verificação e registo automáticos
6. Guardar artefactos: Armazena adaptadores LoRA e tokenizador no volume do Unity Catalog
7. Return MLflow run ID: Fornece o ID de execução para registo de modelos

Principais otimizações ativadas

• Kernels Liger: Operações de GPU fundidas reduzem o uso de memória até 80%
• Precisão mista (FP16): Computação mais rápida com menor área de memória
• Checkpoint de gradiente: Troca computação por memória para se ajustar a lotes maiores
• Acumulação de gradiente: Simula lotes maiores para treino estável

from serverless_gpu import distributed
from serverless_gpu import runtime as rt

@distributed(gpus=8, gpu_type="H100")
def run_train(use_lora=True):
    import logging
    from datasets import load_dataset
    from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer
    from peft import LoraConfig, TaskType, get_peft_model
    from trl import (
        SFTConfig,
        SFTTrainer,
        setup_chat_format
    )
    import json
    import os
    import mlflow

    dataset = load_dataset(DATASET_NAME)
    logging.info(f"✓ Dataset loaded: {dataset}")

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

    # model and tokenizer initialization
    model = AutoModelForCausalLM.from_pretrained(
        MODEL_NAME,
        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:
        logging.info("Adding chat template for proper conversation formatting...")
        model, tokenizer = setup_chat_format(model, tokenizer, format="chatml")
        logging.info("✓ ChatML format applied for structured conversations")

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

    logging.info("✓ Model and tokenizer loaded successfully")

    # PEFT
    peft_config = None
    if use_lora:
        try:
            logging.info("Configuring LoRA for parameter-efficient fine-tuning...")

            peft_config = LoraConfig(
                task_type=TaskType.CAUSAL_LM,
                inference_mode=False,
                r=LORA_R,
                lora_alpha=LORA_ALPHA,
                lora_dropout=LORA_DROPOUT,
                target_modules=LORA_TARGET_MODULES,
                bias="none",
                use_rslora=False,
                modules_to_save=None,
            )

            logging.info(f"LoRA configuration: rank={LORA_R}, alpha={LORA_ALPHA}, dropout={LORA_DROPOUT}")
            logging.info(f"Target modules: {', '.join(LORA_TARGET_MODULES)}")

            original_params = model.num_parameters()
            model = get_peft_model(model, peft_config)

            trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
            total_params = sum(p.numel() for p in model.parameters())
            efficiency_ratio = 100 * trainable_params / total_params

            logging.info(f"✓ LoRA applied successfully:")
            logging.info(f"  • Original parameters: {original_params:,}")
            logging.info(f"  • Trainable parameters: {trainable_params:,}")
            logging.info(f"  • Training efficiency: {efficiency_ratio:.2f}% of parameters")
            logging.info(f"  • Memory savings: ~{100-efficiency_ratio:.1f}% reduction in gradient memory")

        except Exception as e:
            logging.info(f"✗ LoRA configuration failed: {e}")
            logging.info("Falling back to full fine-tuning...")
            peft_config = None
    else:
        logging.info("Full fine-tuning mode selected")
        trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
        logging.info(f"Trainable parameters: {trainable_params:,} (100% of model)")

    # Learning rate adjustment for LoRA
    adjusted_lr = LEARNING_RATE * 10 if use_lora else LEARNING_RATE
    logging.info(f"Learning rate: {adjusted_lr} ({'LoRA-adjusted' if use_lora else 'standard'})")

    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": adjusted_lr,
        "num_train_epochs": NUM_EPOCHS,
        "eval_steps": EVAL_STEPS,
        "logging_steps": LOGGING_STEPS,
        "save_steps": SAVE_STEPS,
        "save_total_limit": 2,
        "report_to": "mlflow",
        "run_name": f"{MODEL_NAME}_fine-tuning",
        "warmup_steps": 50,
        "weight_decay": 0.01,
        "metric_for_best_model": "eval_loss",
        "greater_is_better": False,
        "dataloader_pin_memory": False,
        "remove_unused_columns": False,
        "use_liger_kernel": True,  # Enable Liger kernel optimizations
        "fp16": True,  # Mixed precision training
        "gradient_checkpointing": True,
        "gradient_checkpointing_kwargs": {"use_reentrant": False}, # Required for LORA with DDP
    }

    logging.info("✓ Liger kernel optimizations enabled")

    training_args = SFTConfig(**training_args_dict)

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

    logging.info("\n" + "="*50)
    logging.info("STARTING TRAINING")
    logging.info("="*50)

    logging.info("🚀 Training with Liger kernels for memory-efficient single GPU training")
    if use_lora:
        logging.info("🎯 Using LoRA for parameter-efficient fine-tuning")

    trainer.train()
    logging.info("\n✓ Training completed successfully!")
    if rt.get_global_rank() == 0:
        logging.info("\nSaving trained model...")

        logging.info("Saving LoRA adapter weights...")
        trainer.save_model(training_args.output_dir)
        logging.info("✓ LoRA adapters saved - use with base model for inference")
        tokenizer.save_pretrained(training_args.output_dir)
        logging.info("✓ Tokenizer saved with model")
        logging.info(f"\n🎉 All artifacts saved to: {training_args.output_dir}")

    mlflow_run_id = None
    if mlflow.last_active_run() is not None:
        mlflow_run_id = mlflow.last_active_run().info.run_id

    return mlflow_run_id

Executar o treinamento distribuído

Esta célula executa a função de treino em 8 GPUs H100. O distributed() método trata:

• Provisão de recursos de computação de GPU serverless
• Distribuição da carga de trabalho de treino entre várias GPUs
• Recolha do ID da execução MLflow para registo de modelos

O treino normalmente demora entre 15 a 30 minutos, dependendo do tamanho do conjunto de dados e da disponibilidade de computação.

mlflow_run_id = run_train.distributed(use_lora=True)[0]
print(mlflow_run_id)

Registo do MLflow e do Catálogo Unity

Estratégia de registo de modelos

• MLflow Tracking: Registo de artefactos e metadados do modelo
• Unity Catalog: Modelo de registo para governação e implementação
• Versionamento de Modelos: Versionamento automático para a gestão do ciclo de vida do modelo
• Metadados: Informação completa do modelo para reprodutibilidade

print("\nRegistering model with MLflow and Unity Catalog...")

from transformers import AutoTokenizer, AutoModelForCausalLM
from peft import PeftModel
import mlflow
from mlflow import transformers as mlflow_transformers

try:
    # Load the trained model for registration
    print("Loading LoRA model for registration...")
    # For LoRA models, we need both base model and adapter
    base_model = AutoModelForCausalLM.from_pretrained(
        MODEL_NAME,
        trust_remote_code=True
    )
    # Load tokenizer
    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
    adapter_dir = OUTPUT_DIR
    peft_model = PeftModel.from_pretrained(base_model, adapter_dir)
    # Merge LoRA into base and drop PEFT wrappers
    merged_model = peft_model.merge_and_unload()

    components = {
        "model": merged_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"
    with mlflow.start_run(run_id=mlflow_run_id):
        model_info = mlflow.transformers.log_model(
            transformers_model=components,
            artifact_path="model",
            task=task,
            registered_model_name=full_model_name,
            metadata={
                "task": task,
                "pretrained_model_name": MODEL_NAME,
                "databricks_model_family": "QwenForCausalLM",
            },
        )

    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}")
    print(f"Optimization: Liger Kernels + LoRA")

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: {OUTPUT_DIR}")

Passos seguintes

Agora que afinou e registou o seu modelo, pode:

• Implemente o modelo: Sirva modelos com o Model Serving
• Saiba mais sobre treino distribuído: Treino distribuído multi-GPU
• Otimizar o uso de GPUs serverless: Boas práticas para computação de GPUs serverless
• Solucionar problemas: Solucionar problemas na computação com GPU sem servidor

Exemplo de bloco de notas

Ajuste fino distribuído de Qwen2-0.5B com LoRA

Obter bloco de notas