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

Ajuste fino Llama 3.2 1B com LoRa usando GPU sem servidor

Este notebook demonstra como ajustar um modelo de linguagem complexo usando ajuste fino supervisionado (SFT) com adaptação de baixa classificação (LoRA) em uma GPU sem servidor Databricks . 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 4 GPUs A10.

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.
  • GPU sem servidor : Databricks gerenciam compute de GPU que provisiona e escala automaticamente recurso de GPU para cargas de trabalho de treinamento.

Para obter mais informações, consulte computeGPU sem servidor.

Requisitos​

Este notebook requer o seguinte:

  • computeGPU sem servidor : O Notebook utiliza a GPU sem servidor Databricks com 4 GPUs A10 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 A10.
  • 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": "auto",
        "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 A10.
  • 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": 4,
        "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": True,
        "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 servidor Databricks GPU sem servidor. O decorador provisiona 4 GPUs A10 e lida com a configuração de treinamento distribuído automaticamente.

parâmetros-chave:

  • gpus=4 : Solicita 4 GPUs para treinamento distribuído
  • gpu_type='A10' : Especifica o hardware da GPU A10
  • remote=True : Executa em computeGPU remota serverless .

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 GPU sem servidor.

Python
from serverless_gpu import distributed

os.environ['MLFLOW_EXPERIMENT_NAME'] = MLFLOW_EXPERIMENT_NAME
@distributed(
    gpus=4,
    gpu_type='A10',
    remote=True,  # Set to False to run locally, True for remote GPUs
)
def run_distributed_trl_sft():
    """
    Distributed TRL SFT training function using serverless GPU.

    This function will be executed on the A10 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 A10 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 disponibiliza o recurso de GPU serverless , executa o treinamento em 4 GPUs A10 com otimização DeepSpeed e retorna os resultados.

O processo de treinamento:

  • provisionamento 4 GPUs A10 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 Llama 3.2 1B com LoRa usando GPU sem servidor

Open notebook in new tab