Pular para o conteúdo principal
Última atualização em

Ajuste fino do Llama 3.2 1B com LoRa usando o AI Runtime.

Este notebook demonstra como ajustar um modelo de linguagem de grande porte usando o ajuste fino supervisionado (SFT) com adaptação de baixa classificação (LoRA) no Databricks AI Runtime. O Notebook utiliza a biblioteca Transformers Reinforcement Learning (TRL) com a otimização DeepSpeed ZeRO Stage 3 para treinar eficientemente Llama 3.2 1B em um único nó com 8 GPUs H100.

conceitos-chave:

  • LoRA (Low-Rank Adaptation) : Uma técnica de ajuste fino com uso eficiente de parâmetros que reduz o número de parâmetros treináveis adicionando pequenas matrizes de decomposição de classificação treináveis às camadas do modelo.
  • TRL (Transformers Reinforcement Learning) : Uma biblioteca que fornece ferramentas para treinar modelos de linguagem com aprendizado por reforço e ajuste fino supervisionado.
  • DeepSpeed ZeRO Stage 3 : Uma técnica de otimização de memória que particiona parâmetros do modelo, gradientes e estados do otimizador entre as GPUs para permitir o treinamento de modelos grandes.
  • AI Runtime : Databricks-gerenciar compute de GPU que provisiona e escala automaticamente recursos de GPU para cargas de trabalho de treinamento.

Para obter mais informações, consulte AI Runtime.

Requisitos​

Este notebook requer o seguinte:

  • RuntimeAI : O Notebook utiliza Databricks AI Runtime com 8 GPUs H100 para treinamento distribuído. Não é necessária nenhuma configuração de cluster.
  • Unity Catalog : Um catálogo e esquema Unity Catalog para armazenar pontos de verificação do modelo e registrar o modelo treinado.
  • Tokens HuggingFace : Um access token HuggingFace armazenado nos segredos Databricks para download o modelo base e dataset.
  • PacotePython : Os pacotes necessários (peft, trl, deepspeed, mlflow, hf_transfer) são instalados na seção de configuração abaixo.

Instale o pacote necessário​

Instale o pacote Python necessário para o ajuste fino:

  • peft: Fornece implementação LoRA para ajuste fino de parâmetros com eficiência.
  • trlBiblioteca de aprendizado por reforço Transformers para ajuste fino supervisionado
  • deepspeedPermite treinamento distribuído com uso eficiente de memória através da otimização ZeRO.
  • mlflow: Monitora experimentos e logs modelos treinados
  • hf_transferAcelera downloads de modelos do HuggingFace Hub.

Após a instalação, reinicie o kernel Python para garantir que todos os pacotes sejam carregados corretamente.

Python
%pip install --upgrade transformers==4.56.1
%pip install peft==0.17.1
%pip install trl==0.18.1
%pip install deepspeed>=0.15.4
%pip install mlflow>=3.6.0
%pip install hf_transfer==0.1.9
Python
%restart_python

Configurar Unity Catalog e a variável de ambiente​

Configure os locais Unity Catalog para armazenar os pontos de verificação do modelo e registre o modelo treinado. O Notebook utiliza parâmetros de consulta para configurar:

  • Catálogo e esquema : Espaço de nomes Unity Catalog para organizar modelos e pontos de verificação.
  • Nome do modelo : Nome do modelo registrado no Unity Catalog
  • Volume : Volume Unity Catalog para armazenar pontos de verificação do modelo durante o treinamento.

A configuração também recupera os tokens do HuggingFace dos segredos Databricks e configura o experimento MLflow para acompanhamento de treinamentos.

Python
dbutils.widgets.text("uc_catalog", "main")
dbutils.widgets.text("uc_schema", "default")
dbutils.widgets.text("uc_model_name", "llama3_2-1b")
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")

# Get HuggingFace token and username
hf_token = dbutils.secrets.get(scope="sgc-nightly-notebook", key="hf_token")
username = spark.sql("SELECT session_user()").collect()[0][0]

REGISTERED_MODEL_NAME = f"{UC_CATALOG}.{UC_SCHEMA}.{UC_MODEL_NAME}"
CHECKPOINT_DIR = f"/Volumes/{UC_CATALOG}/{UC_SCHEMA}/{UC_VOLUME}/{UC_MODEL_NAME}"
FINE_TUNED_MODEL_PATH = f"{CHECKPOINT_DIR}/fine-tuned-peft-model"
MLFLOW_EXPERIMENT_NAME = f"/Users/{username}/{UC_MODEL_NAME}"

# Create the Unity Catalog volume if it doesn't exist
spark.sql(f"CREATE VOLUME IF NOT EXISTS {UC_CATALOG}.{UC_SCHEMA}.{UC_VOLUME}")

print(f"👤 Username: {username}")
print("🔑 HuggingFace token configured")
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}")
print(f"CHECKPOINT_DIR: {CHECKPOINT_DIR}")
print(f"MLFLOW_EXPERIMENT_NAME: {MLFLOW_EXPERIMENT_NAME}")

Python
import os
import json
import tempfile
import torch
import mlflow
from huggingface_hub import constants
from datasets import load_dataset
from transformers import AutoTokenizer, AutoModelForCausalLM
from trl import SFTTrainer, SFTConfig, ModelConfig, ScriptArguments, setup_chat_format
from peft import LoraConfig, get_peft_model, PeftModel

if mlflow.get_experiment_by_name(MLFLOW_EXPERIMENT_NAME) is None:
    mlflow.create_experiment(name=MLFLOW_EXPERIMENT_NAME)
mlflow.set_experiment(MLFLOW_EXPERIMENT_NAME)

Criar configuração DeepSpeed ZeRO Estágio 3​

O DeepSpeed ZeRO (Zero Redundancy Optimizer) Stage 3 divide os parâmetros do modelo, os gradientes e os estados do otimizador entre todas as GPUs para reduzir o consumo de memória por GPU. Isso possibilita o treinamento de modelos grandes que não caberiam na memória de uma única GPU.

Configurações principais:

  • bf16 ativado : Utiliza precisão bfloat16 para treinamento mais rápido e menor uso de memória.
  • Otimização de estágio 3 : Particiona todos os estados do modelo entre as GPUs.
  • Sem descarregamento de CPU : mantém todos os dados nas GPUs para obter o máximo desempenho no hardware H100.
  • Comunicação sobreposta : sobrepõe a comunicação do gradiente com a computação para maior eficiência.
Python
def create_deepspeed_config():
    """Create DeepSpeed ZeRO Stage 3 configuration for single node A10 training."""

    deepspeed_config = {
        "fp16": {
            "enabled": False
        },
        "bf16": {
            "enabled": True
        },
        "zero_optimization": {
            "stage": 3,
            "offload_optimizer": {
                "device": "none"
            },
            "offload_param": {
                "device": "none"
            },
            "overlap_comm": True,
            "contiguous_gradients": True,
            "sub_group_size": 1e9,
            "reduce_bucket_size": "auto",
            "stage3_prefetch_bucket_size": "auto",
            "stage3_param_persistence_threshold": "auto",
            "stage3_max_live_parameters": 1e9,
            "stage3_max_reuse_distance": 1e9,
            "stage3_gather_16bit_weights_on_model_save": True
        },
        "gradient_accumulation_steps": 1,
        "gradient_clipping": "auto",
        "steps_per_print": 2000,
        "train_batch_size": "auto",
        "train_micro_batch_size_per_gpu": "auto",
        "wall_clock_breakdown": False
    }

    return deepspeed_config


# Create DeepSpeed configuration
deepspeed_config = create_deepspeed_config()
print("⚙️  DeepSpeed ZeRO Stage 3 configuration created")

Defina os parâmetros de treinamento e a configuração do LoRa.​

Configure os parâmetros de ajuste fino supervisionado:

  • Modelo : Llama 3.2 1B Instruct, um modelo compacto adequado para GPUs H100.
  • conjunto de dados : dataset Capivara da biblioteca TRL para treinamento AI conversacional
  • Tamanho dos lotes : 2 por dispositivo com 4 etapas de acumulação de gradiente para um tamanho efetivo de lote de 64
  • Taxa de aprendizado : 2e-4 com programador de cosseno e aquecimento.
  • treinamento os passos : 60 os passos para demonstração (aumento para treinamento completo)
  • Parâmetros LoRa : Rank 16 com alfa 32, direcionado para camadas de atenção e projeção MLP.

A configuração utiliza precisão bfloat16 e checkpointing de gradiente para otimizar o uso da memória.

Python
def create_training_config():
    """Create training configuration for TRL SFT with LoRA."""

    # Model and dataset configuration (not part of TrainingArguments)
    model_config = {
        "model_name": "meta-llama/Llama-3.2-1B-Instruct",  # Small Llama model for A10
        "dataset_name": "trl-lib/Capybara"
    }

    # Training arguments that will be passed directly to TrainingArguments
    training_args_config = {
        "output_dir": CHECKPOINT_DIR,
        "per_device_train_batch_size": 2,
        "per_device_eval_batch_size": 2,
        "gradient_accumulation_steps": 1,
        "learning_rate": 2e-4,
        "max_steps": 60,   # TO DO remove when fine-tuning on full dataset. Demo purposes only.
        # "num_train_epochs": 1, # TO DO update to >= 1 when fine-tuning on full dataset
        "logging_steps": 10,
        "save_steps": 30,
        "eval_steps": 30,
        "eval_strategy": "steps",
        "warmup_steps": 10,
        "lr_scheduler_type": "cosine",
        "gradient_checkpointing": False,
        "fp16": False,
        "bf16": True,
        "optim": "adamw_torch",
        "remove_unused_columns": False,
        "run_name": f"llama3.2-1b-lora",
        "report_to": "mlflow",
        "save_total_limit": 2,
        "load_best_model_at_end": True,
        "metric_for_best_model": "eval_loss",
        "greater_is_better": False,
    }

    # LoRA configuration
    lora_config = {
        "r": 16,
        "lora_alpha": 32,
        "target_modules": ["q_proj", "v_proj", "k_proj", "o_proj", "gate_proj", "up_proj", "down_proj"],
        "lora_dropout": 0.1,
        "bias": "none",
        "task_type": "CAUSAL_LM"
    }

    return model_config, training_args_config, lora_config

# Create training configuration
model_config, training_args_config, lora_config = create_training_config()

print("📊 Training Configuration:")
print(f"  🤖 Model: {model_config['model_name']}")
print(f"  📚 Dataset: {model_config['dataset_name']}")
print(f"  🎯 Batch size: {training_args_config['per_device_train_batch_size']}")
print(f"  📈 Learning rate: {training_args_config['learning_rate']}")
print(f"  🧠 LoRA rank: {lora_config['r']}")

Defina a função de treinamento distribuído​

O decorador @distributed da biblioteca serverless_gpu permite a execução perfeita de cargas de trabalho de GPU no Databricks AI Runtime. O decorador provisiona 8 GPUs H100 e lida com a configuração de treinamento distribuído automaticamente.

parâmetros-chave:

  • gpus=8 : Solicita 8 GPUs para treinamento distribuído
  • gpu_type='H100' : Especifica o hardware da GPU H100

A função de treinamento:

  1. Carrega o modelo base e o tokenizador do HuggingFace.
  2. Configura a formatação de bate-papo para AIconversacional.
  3. Configura o LoRA para ajuste fino com otimização de parâmetros.
  4. Carrega o datasetde treinamento
  5. Inicializa o TRL SFTTrainer com otimização DeepSpeed.
  6. Ensina o modelo e salva os pontos de verificação.
  7. Retorna os resultados do treinamento e o ID da execução do MLflow.

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

Python
from serverless_gpu import distributed

os.environ['MLFLOW_EXPERIMENT_NAME'] = MLFLOW_EXPERIMENT_NAME
@distributed(
    gpus=8,
    gpu_type='H100',
)
def run_distributed_trl_sft():
    """
    Distributed TRL SFT training function using AI Runtime.

    This function will be executed on the H100 GPU with DeepSpeed optimization.
    """

    # Set up environment variables for remote jobs
    import os
    import tempfile
    import json
    from huggingface_hub import constants
    from datasets import load_dataset
    from transformers import AutoTokenizer, AutoModelForCausalLM, TrainingArguments
    from trl import SFTTrainer, setup_chat_format
    from peft import LoraConfig, get_peft_model

    # HuggingFace configuration
    os.environ["HUGGING_FACE_HUB_TOKEN"] = hf_token
    os.environ['HF_TOKEN'] = hf_token
    constants.HF_HUB_ENABLE_HF_TRANSFER = True

    # Set up temporary directories
    temp_dir = tempfile.mkdtemp()

    print("🚀 Starting TRL SFT training on H100 GPU...")

    try:
        # Load tokenizer and model
        print(f"📥 Loading model: {model_config['model_name']}")
        tokenizer = AutoTokenizer.from_pretrained(model_config['model_name'])

        # Add pad token if not present
        if tokenizer.pad_token is None:
            tokenizer.pad_token = tokenizer.eos_token

        model = AutoModelForCausalLM.from_pretrained(
            model_config['model_name'],
            dtype="auto"
        )

        # Setup chat template if the model doesn't have one
        # This is crucial for conversational AI models and TRL SFTTrainer
        if tokenizer.chat_template is None:
            print("🗨️  Setting up chat template...")
            model, tokenizer = setup_chat_format(model, tokenizer, format="chatml")

        # Configure LoRA
        print("🔧 Setting up LoRA configuration...")
        peft_config = LoraConfig(**lora_config)

        # Load dataset
        print(f"📚 Loading dataset: {model_config['dataset_name']}")
        dataset = load_dataset(model_config['dataset_name'])

        # Create temporary DeepSpeed config file
        deepspeed_config_path = os.path.join(temp_dir, "deepspeed_config.json")
        with open(deepspeed_config_path, "w") as f:
            json.dump(deepspeed_config, f, indent=2)

        # Training arguments - dynamically pass all config parameters
        training_args = TrainingArguments(
            **training_args_config,
            deepspeed=deepspeed_config_path,  # Override deepspeed with the config file path
        )

        # Initialize SFT Trainer
        print("🏋️ Initializing SFT Trainer with DeepSpeed...")
        trainer = SFTTrainer(
            model=model,
            args=training_args,
            train_dataset=dataset["train"],
            eval_dataset=dataset["test"] if "test" in dataset else None,
            processing_class=tokenizer,
            peft_config=peft_config
        )

        # Start training
        print("🎯 Starting training...")
        trainer.train()

        # Save the model
        print("💾 Saving trained model...")
        trainer.save_model()

        # Get training results
        train_results = trainer.state.log_history
        final_loss = train_results[-1].get('train_loss', 'N/A') if train_results else 'N/A'

        print("✅ Training completed successfully!")
        print(f"📊 Final training loss: {final_loss}")

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

        return {
            "status": "success",
            "final_loss": final_loss,
            "output_dir": training_args_config['output_dir'],
            "model_name": model_config['model_name'],
            "mlflow_run_id": mlflow_run_id,
        }

    except Exception as e:
        print(f"❌ Training failed: {e}")
        import traceback
        traceback.print_exc()
        return {
            "status": "failed",
            "error": str(e)
        }

execução do treinamento distribuído Job​

Execute a função de treinamento chamando .distributed() na função decorada. Este provisionamento do recurso AI Runtime , execução do treinamento em 8 GPUs H100 com otimização DeepSpeed, e retorna os resultados.

O processo de treinamento:

  • provisionamento 8 GPUs H100 automaticamente
  • baixa o modelo e dataset do HuggingFace
  • ensinar o modelo com ajuste fino LoRa
  • Salva pontos de verificação no volume Unity Catalog
  • logs números para MLflow
  • Retorna o status do treinamento, a perda final e o ID da execução do MLflow.
Python
# Execute the distributed training
results = run_distributed_trl_sft.distributed()

print("🏁 Training execution completed!")
print(f"📊 Results: {results}")

if results and results[0].get('status') == 'success':
    print("✅ Training completed successfully!")
    print(f"💾 Model saved to: {results[0].get('output_dir', 'N/A')}")
    print(f"📈 Final loss: {results[0].get('final_loss', 'N/A')}")
    print(f"🎉 MLflow run ID: {results[0].get('mlflow_run_id', 'N/A')}")
else:
    print("❌ Training failed!")
    if results and 'error' in results:
        print(f"🔍 Error: {results['error']}")

Salve o modelo ajustado e teste a inferência.​

Esta etapa opcional carrega o adaptador LoRa treinado, mescla-o com o modelo base e salva o modelo completo ajustado. O modelo de fusão pode então ser testado com exemplos de comandos para verificar os resultados do ajuste fino.

O processo:

  1. Carrega o modelo básico Llama 3.2 1B
  2. Aplica os pesos do adaptador LoRA treinados
  3. Incorpore o adaptador ao modelo base.
  4. Salva o modelo de mesclagem no volume do Unity Catalog .
  5. Testa o modelo com um exemplo de diálogo conversacional.
Python
%pip install hf_transfer
Python
def save_and_load_trained_model():
    """Load the trained model from the Unity Catalog volume."""

    import torch
    from transformers import AutoModelForCausalLM, AutoTokenizer
    from peft import PeftModel

    # Load base model and tokenizer
    base_model = AutoModelForCausalLM.from_pretrained(
        model_config['model_name'],
        dtype=torch.bfloat16,
        token=hf_token,
        trust_remote_code=True,
        device_map="auto"
    )
    tokenizer = AutoTokenizer.from_pretrained(model_config['model_name'], token=hf_token, trust_remote_code=True)

    # Load LoRA weights
    model = PeftModel.from_pretrained(base_model, training_args_config['output_dir'])

    # Merge LoRA weights into base model
    model = model.merge_and_unload()

    # Save the merged model
    model.save_pretrained(FINE_TUNED_MODEL_PATH)
    tokenizer.save_pretrained(FINE_TUNED_MODEL_PATH)

    # Return the merged model and tokenizer
    return model, tokenizer

def test_trained_model(model, tokenizer):
    """Test the trained model with simple inference."""

    try:
        import torch
        # Test prompt
        # Create a conversation following the schema
        conversation = [
            {
                "content": "What is machine learning?",
                "role": "user"
            }
        ]

        # Convert conversation to chat format
        prompt = ""
        for message in conversation:
            if message["role"] == "user":
                prompt += f"### User: {message['content']}\n### Response:"
            else:
                prompt += f" {message['content']}\n\n"

        # Tokenize
        inputs = tokenizer(prompt, return_tensors="pt").to("cuda")

        # Generate
        with torch.no_grad():
            outputs = model.generate(
                **inputs,
                max_new_tokens=500,
                temperature=0.7,
                do_sample=True,
                pad_token_id=tokenizer.eos_token_id
            )

        # Decode
        response = tokenizer.decode(outputs[0], skip_special_tokens=True)
        print("🤖 Model Response:")
        print(response)
        return response

    except Exception as e:
        print(f"❌ Model testing failed: {e}")

# Save and load the trained model
model, tokenizer = save_and_load_trained_model()

# Test the trained model
response = test_trained_model(model, tokenizer)

Registre o modelo no Unity Catalog​

Registre o modelo ajustado no MLflow e, em seguida, no Unity Catalog para implantação e disponibilização. O modelo é composto por registros com:

  • Modelo e tokenizador : ambos os componentes são necessários para a inferência.
  • Tipo de tarefa : Configurado como llm/v1/chat para AIconversacional
  • Exemplo de entrada : Formato de mensagem de chat de exemplo para teste
  • RegistroUnity Catalog : Registra automaticamente o modelo no catálogo e esquema configurados.

Uma vez registrado, o modelo pode ser implantado no endpoint do modelo de serviço ou usado para inferência em lotes.

Python
run_id = results[0].get('mlflow_run_id')
mlflow.set_registry_uri("databricks-uc")

# log the model to mlflow using the latest run id and register to Unity Catalog
with mlflow.start_run(run_id=run_id) as run:
    components = {
        "model": model,
        "tokenizer": tokenizer
    }
    logged_model = mlflow.transformers.log_model(
        transformers_model=components,
        name="model",
        task="llm/v1/chat",
        input_example={
            "messages": [
                {"role": "user", "content": "What is machine learning?"}
            ]
        },
        registered_model_name=REGISTERED_MODEL_NAME
        )
    print(f"🔍 Model logged to: {logged_model}")

Próximos passos​

Exemplo de caderno​

Ajuste fino do Llama 3.2 1B com LoRa usando o AI Runtime.

Abrir notebook em uma nova aba
Nesta página