Ottimizzare Llama 3.2 1B con LoRA usando il runtime di intelligenza artificiale

Apri la versione notebook di questa pagina

Questo notebook illustra come affinare un modello linguistico di grandi dimensioni usando l'addestramento supervisionato (SFT) con Low-Rank Adaptation (LoRA) sulla piattaforma Databricks AI Runtime. Il notebook usa la libreria TRL (Transformers Reinforcement Learning) con l'ottimizzazione DeepSpeed ZeRO Stage 3 per eseguire in modo efficiente il training di Llama 3.2 1B in un singolo nodo con 8 GPU H100.

Concetti chiave:

  • LoRA (Low-Rank Adaptation): tecnica di ottimizzazione efficiente dei parametri che riduce il numero di parametri sottoponibili a training aggiungendo matrici di classificazione sottoponibili a training di piccole dimensioni ai livelli del modello.
  • TRL (Transformers Reinforcement Learning): libreria che fornisce strumenti per il training di modelli linguistici con apprendimento per rinforzo e ottimizzazione supervisionata.
  • DeepSpeed ZeRO Fase 3: tecnica di ottimizzazione della memoria che partiziona i parametri del modello, le sfumature e gli stati dell'ottimizzatore tra gpu per abilitare il training di modelli di grandi dimensioni.
  • Runtime AI: calcolo GPU gestito da Databricks che provisiona e ridimensiona automaticamente le risorse GPU per i carichi di lavoro di addestramento.

Per altre informazioni, vedere Runtime di intelligenza artificiale.

Requisiti

Questo notebook richiede quanto segue:

  • AI Runtime di Databricks: il notebook usa Databricks AI Runtime con 8 GPU H100 per il training distribuito. Non è necessaria alcuna configurazione del cluster.
  • Catalogo Unity: Un catalogo e uno schema del catalogo di Unity Catalog per archiviare i checkpoint del modello e registrare il modello addestrato.
  • Token HuggingFace: token di accesso HuggingFace archiviato nei segreti di Databricks per scaricare il modello di base e il set di dati.
  • Pacchetti Python: i pacchetti necessari (peft, trl, deepspeed, mlflow, hf_transfer) vengono installati nella sezione di installazione seguente.

Installare i pacchetti necessari

Installare i pacchetti Python necessari per l'ottimizzazione:

  • peft: fornisce l'implementazione LoRA per l'ottimizzazione dei parametri in modo efficiente
  • trl: libreria di apprendimento per rinforzo dei trasformatori per l'ottimizzazione con supervisione
  • deepspeed: permette il training distribuito a memoria efficiente con l'ottimizzazione ZeRO
  • mlflow: monitora gli esperimenti e registra i modelli addestrati
  • hf_transfer: accelera i download dei modelli dall'HuggingFace Hub

Dopo l'installazione, riavviare il kernel Python per assicurarsi che tutti i pacchetti siano caricati correttamente.

%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

%restart_python

Configurare il catalogo unity e le variabili di ambiente

Configurare i percorsi di Unity Catalog per archiviare i checkpoint del modello e registrare il modello addestrato. Il notebook usa i parametri di query per configurare:

  • Catalogo e schema: namespace del catalogo Unity per organizzare i modelli e i checkpoint
  • Nome modello: nome del modello registrato nel catalogo unity
  • Volume: Volume di Unity Catalog per l'archiviazione dei checkpoint del modello durante l'addestramento

La configurazione recupera anche il token HuggingFace dai segreti di Databricks e configura l'esperimento MLflow per tenere traccia delle metriche di training.

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


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)

Creare la configurazione di DeepSpeed ZeRO Fase 3

DeepSpeed ZeRO (Zero Redundancy Optimizer) Fase 3 partiziona i parametri del modello, i gradienti e gli stati dell'ottimizzatore su tutte le GPU per ridurre il consumo di memoria per ciascuna GPU. In questo modo è possibile eseguire il training di modelli di grandi dimensioni che non si adattano alla memoria di una singola GPU.

Impostazioni di configurazione chiave:

  • bf16 abilitato: utilizza la precisione bfloat16 per un addestramento più veloce e una riduzione dell'uso della memoria
  • Ottimizzazione della fase 3: partiziona tutti gli stati del modello tra GPU
  • Nessun offload della CPU: mantiene tutti i dati sulle GPU per ottenere prestazioni massime sull'hardware H100
  • Comunicazione sovrapposta: integra la comunicazione dei gradienti con il calcolo per migliorare l'efficienza
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")

Definire i parametri di addestramento e la configurazione LoRA

Configurare i parametri di ottimizzazione con supervisione:

  • Modello: Llama 3.2 1B Instruct, un modello compatto adatto per GPU H100
  • Set di dati: set di dati Capybara dalla libreria TRL per l'addestramento dell'intelligenza artificiale conversazionale
  • Dimensione batch: 2 per dispositivo con 4 fasi di accumulo gradiente per dimensione batch effettiva di 64
  • Frequenza di apprendimento: 2e-4 con programmatore coseno e riscaldamento
  • Passaggi di training: 60 passaggi per la dimostrazione (aumentare per la formazione completa)
  • Parametri LoRA: rank 16 con alfa 32, mirati agli strati di attenzione e proiezione MLP

La configurazione utilizza la precisione bfloat16 e il checkpointing del gradiente per ottimizzare l'utilizzo della memoria.

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']}")

Definire la funzione di training distribuita

Il @distributed decorator della libreria serverless_gpu consente l'esecuzione senza problemi dei carichi di lavoro GPU su Databricks AI Runtime. Il decoratore esegue il provisioning di 8 GPU H100 e gestisce automaticamente la configurazione dell'addestramento distribuito.

Parametri chiave:

  • gpus=8: richieste di 8 GPU per il training distribuito
  • gpu_type='H100': indica l'hardware GPU H100

La funzione di training:

  1. Carica il modello di base e il tokenizer da HuggingFace
  2. Configura la formattazione della chat per l'intelligenza artificiale conversazionale
  3. Configura LoRA per l'ottimizzazione efficiente dei parametri
  4. Carica il set di dati di training
  5. Inizializza TRL SFTTrainer con l'ottimizzazione DeepSpeed
  6. Esegue il training del modello e salva i checkpoint
  7. Restituisce i risultati dell'addestramento e l'ID dell'esecuzione MLflow

Per altre informazioni, vedere la documentazione dell'API runtime di intelligenza artificiale.

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

Eseguire il processo di training distribuito

Eseguire la funzione di addestramento chiamando .distributed() sulla funzione decorata. Questo esegue il provisioning delle risorse del runtime di intelligenza artificiale, esegue il training su 8 GPU H100 con l'ottimizzazione DeepSpeed e restituisce i risultati.

Processo di training:

  • Effettua automaticamente il provisioning di 8 GPU H100
  • Scarica il modello e il set di dati da HuggingFace
  • Esegue il training del modello con il fine-tuning di LoRA
  • Salva i checkpoint nel volume di Unity Catalog
  • Registra le metriche in MLflow
  • Restituisce lo stato del training, la perdita finale e l'ID esecuzione MLflow
# 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']}")

Salvare il modello ottimizzato e inferenziare i test

Questo passaggio facoltativo carica l'adattatore LoRA sottoposto a training, lo unisce al modello di base e salva il modello completo ottimizzato. Il modello unito può quindi essere testato con richieste di esempio per verificare i risultati dell'ottimizzazione.

Processo:

  1. Carica il modello Llama di base 3.2 1B
  2. Applica i pesi dell'adattatore LoRA addestrato
  3. Unisce l'adattatore nel modello di base
  4. Salva il modello unito nel volume del catalogo Unity
  5. Verifica il modello con un prompt di conversazione di esempio
%pip install hf_transfer

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)

Registrare il modello nel catalogo unity

Registrare il modello ottimizzato in MLflow e registrarlo nel catalogo Unity per la distribuzione e la gestione. Il modello viene registrato con:

  • Modello e tokenizer: entrambi i componenti necessari per l'inferenza
  • Tipo di attività: configurato come llm/v1/chat per l'intelligenza artificiale conversazionale
  • Esempio di input: formato di messaggio di chat di esempio per il test
  • Registrazione del catalogo Unity: registra automaticamente il modello nel catalogo e nello schema configurati

Dopo la registrazione, il modello può essere distribuito agli endpoint di gestione del modello o usato per l'inferenza batch.

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

Passaggi successivi

Notebook di esempio

Ottimizzare Llama 3.2 1B con LoRA usando il runtime di intelligenza artificiale

Ottieni il notebook

  • Last updated on