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

Este Notebook demonstra como ajustar de forma eficiente o modelo de linguagem Qwen2-0.5B de grande porte usando técnicas de otimização de parâmetros em compute GPU serverless . Você aprenderá como:

Aplique LoRA (Adaptação de Baixa Classificação) para reduzir os parâmetros treináveis em aproximadamente 99%, mantendo a qualidade do modelo.
Utilize kernels Liger para treinamento com uso eficiente de memória, utilizando kernels Triton otimizados.
Utilize TRL (Transformer Reinforcement Learning) para ajuste fino supervisionado.
Registre o modelo otimizado no Unity Catalog para governança e implantação.

conceitos-chave:

LoRa: Uma técnica que congela o modelo base e ensina pequenas camadas adaptadoras, reduzindo drasticamente os requisitos de memória e o tempo de treinamento.
Kernels Liger: Kernels otimizados para GPU que reduzem o uso de memória em até 80% por meio de operações combinadas.
TRL: Uma biblioteca para treinamento de modelos de linguagem com aprendizado por reforço e ajuste fino supervisionado.
AI Runtime: Databricks gerenciam compute que escala automaticamente recurso de GPU

Matriz de decisão LoRa versus ajuste fino completo

LoRA (Low-Rank Adaptation) congela o modelo base e ensina apenas pequenas camadas adaptadoras, reduzindo os parâmetros treináveis em cerca de 99%. Isso torna o treinamento mais rápido e com melhor aproveitamento da memória.

Cenário	Recomendação	Razão
Memória de GPU limitada	LoRA	Ajusta modelos maiores na memória treinando apenas 1% dos parâmetros.
Adaptação específica da tarefa	LoRA	Troque os adaptadores no mesmo modelo base para várias tarefas.
Mudança comportamental importante do modelo	Ajuste fino completo	Atualiza todos os parâmetros para alterações fundamentais no comportamento do modelo.
Implantação em produção	LoRA	Arquivos menores (MB vs GB), carregamento mais rápido, controle de versão mais fácil.

Benefícios do Liger Kernel

Os kernels Liger são operações otimizadas para GPU que fundem várias etapas do sistema operacional em um único kernel, reduzindo as transferências de memória e melhorando a eficiência. O documento técnico fornece parâmetros de comparação detalhados que demonstram melhorias significativas de desempenho.

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

Este notebook utiliza a biblioteca TRL para simplificar a configuração do treinamento e aplicar automaticamente essas otimizações.

Conecte-se à computeGPU serverless

Este notebook requer compute GPU serverless para executar o treinamento distribuído. compute de GPU sem servidor, provisionamento automático e gerenciamento de recursos de GPU para sua carga de trabalho.

Para conectar, clique no menu suspenso Conectar no Notebook e selecione GPU sem servidor .

Para obter mais informações, consulte a documentaçãoAI Runtime.

Instale a biblioteca necessária.

A próxima célula instala o pacote Python necessário para o ajuste fino distribuído:

Biblioteca principal do Windows:

trl==0.18.1 : Biblioteca de aprendizado por reforço Transformer para ajuste fino supervisionado e RLHF
peft : Biblioteca de ajuste fino com otimização de parâmetros que fornece implementação LoRA.
liger-kernel : Kernels de GPU otimizados para memória para treinamento eficiente de transformadores

Biblioteca de apoio:

hf_transfer : downloads acelerados do Hugging Face Hub usando transferência baseada em Rust
mlflow>=3.6.0 : Envio de experimento e registro de integração de modelo

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

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

Configuração

IntegraçãoUnity Catalog

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

Catálogo e Esquema : Organize os modelos dentro do seu namespace Unity Catalog (default: main.default)
Nome do modelo : O nome do modelo registrado no Unity Catalog para fins de governança e implantação.
Volume : Volume Unity Catalog para armazenar pontos de verificação 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 facilitar o acesso e o controle de versão.

Hiperparâmetros de treinamento

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

Modelo e conjunto de dados : Qwen2-0.5B com datasetconversacionais Capybara
Tamanho dos lotes (8) : Número de exemplos por GPU por treinamento o passo
Acumulação de gradiente (4) : Acumula gradientes em 4 lotes para um tamanho efetivo de lote de 32
Taxa de Aprendizagem (1e-4) : Taxa conservadora, automaticamente escalada 10 vezes maior para treinamento LoRa.
Épocas (1) : Passagem única pelo dataset para evitar sobreajuste
Logging & Checkpointing : salva o progresso a cada 100 passos, logs estatísticas a cada 25 passos

Python
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 LoRA (Adaptação de Baixa Classificação) que controlam como o modelo é ajustado. O LoRa congela os pesos do modelo base e ensina apenas pequenas matrizes adaptadoras, reduzindo drasticamente os requisitos de memória.

Seleção de parâmetros

Classificação (r=8) : Oferece um bom equilíbrio entre desempenho e parâmetros.
Alfa (32) : Fator de escala, tipicamente 2 a 4 vezes a classificação
Dropout (0.1) : Regularização para evitar sobreajuste

Módulos alvo para Qwen2

Este exemplo abrange todas as camadas de transformações key :

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

Python
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 treinamento

A próxima célula cria a função de treinamento distribuído que será executada em várias GPUs. Eis o que ele faz:

Configuração de treinamento distribuído

O decorador @distributed configura compute GPU serverless :

8 GPUs : Distribui o treinamento entre 8 GPUs H100 para um treinamento mais rápido.
Orquestração automática : Lida com provisionamento de GPU, distribuição de dados e sincronização

8 fluxo de trabalho

A função executa estes passos:

Carregar dataset : baixa e prepara o datasetconversacionais do Capybara.
Inicializar modelo : Carrega Qwen2-0.5B e tokenizador com formatação de chat.
Aplicar LoRa : Anexa camadas adaptadoras para reduzir os parâmetros treináveis em cerca de 99%.
Configurar treinamento : Define o tamanho dos lotes, a taxa de aprendizado e as otimizações do kernel Liger.
modelo de ensino : execução do loop de treinamento com checkpoint e registro automáticos
Salvar artefatos : Armazena adaptadores LoRa e tokenizador no volume Unity Catalog
Retorna o ID da execução do MLflow : Fornece o ID da execução para o registro do modelo.

principais otimizações ativadas

Kernels Liger : Operações de GPU integradas reduzem o uso de memória em até 80%.
Precisão mista (FP16) : Computação mais rápida com menor consumo de memória.
Checkpointing de gradiente : troca computação por memória para acomodar lotes maiores.
Acumulação de gradiente : Simula tamanhos de lotes maiores para treinamento estável.

Python
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

execução do treinamento distribuído

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

provisionamento recurso compute GPU serverless
Distribuir a carga de trabalho de treinamento entre várias GPUs.
Coletando o ID de execução do MLflow para registro do modelo.

O treinamento normalmente leva de 15 a 30 minutos, dependendo do tamanho dataset e da disponibilidade de compute .

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

Registro no catálogoMLflow e Unity Catalog

Estratégia de registro de modelo

AcompanhamentoMLflow : modelos de artefatos registrados e metadados
Unity Catalog : modelo de registro para governança e implantação
Versionamento de modelos : Versionamento automático para gerenciamento do ciclo de vida do modelo
Metadados : Informações completas do modelo para reprodutibilidade

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

Próximos passos

Agora que você ajustou e registrou seu modelo, você pode:

Implantei o modelo : Modelos de serviço com modelo de instalação
Saiba mais sobre treinamento distribuído : Treinamento distribuído com múltiplas GPUs
Otimize o uso de GPUs serverless : Melhores práticas para computeem GPU sem servidor.
Solução de problemas : Solucione problemas em computeGPU serverless

Exemplo de caderno

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

Abrir notebook em uma nova aba

Matriz de decisão LoRa versus ajuste fino completo​

Benefícios do Liger Kernel​

Conecte-se à computeGPU serverless​

Instale a biblioteca necessária.​

Configuração​

IntegraçãoUnity Catalog​

Hiperparâmetros de treinamento​

Configuração LoRa​

Seleção de parâmetros​

Módulos alvo para Qwen2​

Defina a função de treinamento​

Configuração de treinamento distribuído​

8 fluxo de trabalho​

principais otimizações ativadas​

execução do treinamento distribuído​

Registro no catálogoMLflow e Unity Catalog​

Estratégia de registro de modelo​

Próximos passos​

Exemplo de caderno​

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

Matriz de decisão LoRa versus ajuste fino completo

Benefícios do Liger Kernel

Conecte-se à computeGPU serverless

Instale a biblioteca necessária.

Configuração

IntegraçãoUnity Catalog

Hiperparâmetros de treinamento

Configuração LoRa

Seleção de parâmetros

Módulos alvo para Qwen2

Defina a função de treinamento

Configuração de treinamento distribuído

8 fluxo de trabalho

principais otimizações ativadas

execução do treinamento distribuído

Registro no catálogoMLflow e Unity Catalog

Estratégia de registro de modelo

Próximos passos

Exemplo de caderno