Ottimizzare Llama 3.2 1B usando il runtime di intelligenza artificiale

Apri la versione notebook di questa pagina

Ottimizzare completamente un modello linguistico di grandi dimensioni con ottimizzazione con supervisione (SFT) in 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:

  • 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: runtime di intelligenza artificiale preinstalla le librerie più richieste. La sezione relativa all'installazione seguente installa deepspeed.

Connessione alla computazione GPU senza server

Questo notebook richiede l'elaborazione GPU in modalità serverless. Per connettersi:

  1. Fare clic sul selettore di calcolo del notebook in alto a destra e selezionare GPU serverless.
  2. Sul lato destro fare clic sul pulsante ambiente.
  3. Selezionare 8xH100 come acceleratore.
  4. Scegliere Ambiente di intelligenza artificiale v5 nel pannello destro che contiene tutte le librerie necessarie per eseguire questo esempio di notebook.
  5. Fare clic su Applica.

La funzione di training effettua automaticamente il provisioning di 8 GPU H100 per il training distribuito.

Installare i pacchetti necessari

Il runtime di intelligenza artificiale include già la maggior parte delle librerie necessarie preinstallate. Per questo esempio, è necessario installare solo deepspeed.

%pip install deepspeed==0.19.1
%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}"
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

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(stage: int):
    """Create a DeepSpeed ZeRO configuration for single-node 8xH100 training."""

    deepspeed_config = {
        "fp16": {
            "enabled": False
        },
        "bf16": {
            "enabled": True
        },
        "zero_optimization": {
            "stage": stage,
            "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": 0,
            "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
zero_stage = 3
deepspeed_config = create_deepspeed_config(zero_stage)
print(f"⚙️  DeepSpeed ZeRO Stage {zero_stage} configuration created")

Definire i parametri di training

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)

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."""

    # Model and dataset configuration (not part of TrainingArguments)
    model_config = {
        "model_name": "meta-llama/Llama-3.2-1B-Instruct",  # Small Llama model suitable for 8xH100
        "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.
        "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-fft-zero3",
        "report_to": "mlflow",
        "save_total_limit": 2,
        "load_best_model_at_end": True,
        "metric_for_best_model": "eval_loss",
        "greater_is_better": False,
    }

    return model_config, training_args_config

# Create training configuration
model_config, training_args_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']}")

Definire la funzione di training distribuita

Il decoratore @distributed della libreria serverless_gpu consente di eseguire 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. Carica il set di dati di training
  4. Inizializza TRL SFTTrainer con l'ottimizzazione DeepSpeed
  5. Esegue il training del modello e salva i checkpoint
  6. 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

mlflow.set_experiment(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
    from trl import SFTTrainer, SFTConfig

    # 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
        print(f"📥 Loading tokenizer from 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

        # 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 = SFTConfig(
            **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_config["model_name"],
            args=training_args,
            train_dataset=dataset["train"],
            eval_dataset=dataset["test"] if "test" in dataset else None,
            processing_class=tokenizer,
        )

        # 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

Esegui 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 ottimizzazione completa
  • 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 il modello ottimizzato e lo testa con una richiesta di esempio per verificare i risultati.

Processo:

  1. Carica il modello salvato
  2. Verifica il modello con un prompt di conversazione di esempio
def save_and_load_trained_model():
    """Load the fully fine-tuned model from the Unity Catalog volume."""

    import torch
    from transformers import AutoModelForCausalLM, AutoTokenizer

    print(f"📥 Loading fine-tuned model from: {training_args_config['output_dir']}")

    # Load the fully fine-tuned model directly from the checkpoint directory
    model = AutoModelForCausalLM.from_pretrained(
        training_args_config['output_dir'],
        torch_dtype=torch.bfloat16,
        trust_remote_code=True,
        device_map={"":0}
    )
    tokenizer = AutoTokenizer.from_pretrained(training_args_config['output_dir'], trust_remote_code=True)

    print("✅ Model loaded successfully!")
    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
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 usando il runtime di intelligenza artificiale

Ottieni il notebook

  • Last updated on