Llama 3.2 1B verfijnen met BEHULP van AI Runtime

Notitieblokversie van deze pagina openen

Een groot taalmodel volledig finetunen met supervised fine-tuning (SFT) op Databricks AI Runtime. Het notebook maakt gebruik van de TRL-bibliotheek (Transformers Reinforcement Learning) met DeepSpeed ZeRO Fase 3-optimalisatie om Llama 3.2 1B efficiënt te trainen op één knooppunt met 8 H100 GPU's.

Sleutelbegrippen:

  • TRL (Transformers Reinforcement Learning): een bibliotheek die hulpprogramma's biedt voor het trainen van taalmodellen met versterking van het leren en afstemmen onder supervisie.
  • DeepSpeed ZeRO Fase 3: Een techniek voor geheugenoptimalisatie waarmee modelparameters, kleurovergangen en optimalisatiestatussen tussen GPU's worden gepartitionaliseerd om het trainen van grote modellen mogelijk te maken.
  • AI Runtime: Door Databricks beheerde GPU-rekenkracht waarmee GPU-resources automatisch worden uitgevoerd en geschaald voor trainingsworkloads.

Zie AI Runtime voor meer informatie.

Note

Dit voorbeeld vereist de Databricks AI-omgeving versie 5 of hoger.

Requirements

Voor dit notebook is het volgende vereist:

  • AI Runtime: Het notebook maakt gebruik van Databricks AI Runtime met 8 H100 GPU's voor gedistribueerde training. Er is geen clusterconfiguratie nodig.
  • Unity Catalog: Een Unity Catalog-catalogus en -schema voor het opslaan van modelcontrolepunten en het registreren van het getrainde model.
  • HuggingFace-token: een HuggingFace-toegangstoken dat is opgeslagen in Databricks-geheimen om het basismodel en de basisgegevensset te downloaden.
  • Python pakketten: AI Runtime installeert de meeste vereiste bibliotheken vooraf. In het onderstaande installatiegedeelte wordt deepspeed geïnstalleerd.

Verbinding maken met serverloze GPU-rekenkracht

Voor dit notebook is serverloze GPU-rekenkracht vereist. Verbinding maken:

  1. Klik op de rekenkiezer van het notebook in de rechterbovenhoek en selecteer serverloze GPU.
  2. Klik aan de rechterkant op de knop Omgeving.
  3. Selecteer 8xH100 als de Accelerator.
  4. Kies een AI v5-omgeving in het rechterdeelvenster met alle vereiste bibliotheken om dit notebookvoorbeeld uit te voeren.
  5. Klik op Toepassen.

De trainingsfunctie richt automatisch 8 H100 GPU's in voor gedistribueerde training.

De vereiste pakketten installeren

AI Runtime bevat al de meeste vereiste bibliotheken die vooraf zijn geïnstalleerd. Voor dit voorbeeld moet u alleen deepspeedinstalleren.

%pip install deepspeed==0.19.1
%restart_python

Unity Catalog en omgevingsvariabelen configureren

Stel Unity Catalog-locaties in voor het opslaan van modelcontrolepunten en het registreren van het getrainde model. In het notebook worden queryparameters gebruikt om configuraties in te stellen.

  • Catalogus en schema: Unity Catalog-naamruimte voor het ordenen van modellen en controlepunten
  • Modelnaam: naam voor het geregistreerde model in Unity Catalog
  • Volume: Unity Catalog-volume voor het opslaan van modelcontrolepunten tijdens de training

De configuratie haalt ook het HuggingFace-token op uit Databricks-geheimen en stelt het MLflow-experiment in voor het bijhouden van metrische trainingsgegevens.

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

DeepSpeed ZeRO Fase 3-configuratie maken

DeepSpeed ZeRO (Zero Redundanty Optimizer) Fase 3 partitioneert modelparameters, kleurovergangen en optimizer statussen voor alle GPU's om het geheugenverbruik per GPU te verminderen. Dit maakt het mogelijk om grote modellen te trainen die niet in het geheugen van één GPU passen.

Belangrijkste configuratie-instellingen:

  • bf16 ingeschakeld: maakt gebruik van bfloat16-precisie voor snellere training en verminderd geheugengebruik
  • Fase 3-optimalisatie: Partitioneert alle modelstatussen tussen GPU's
  • Geen CPU-offloading: bewaart alle gegevens op GPU's voor maximale prestaties op H100-hardware
  • Overlappende communicatie: overlappen gradiëntcommunicatie met berekening voor efficiëntie
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")

Trainingsparameters definiëren

Configureer de parameters voor het afstemmen onder supervisie:

  • Model: Llama 3.2 1B Instruct, een compact model geschikt voor H100 GPU's
  • Gegevensset: Capybara-gegevensset uit de TRL-bibliotheek voor training voor conversationele AI
  • Batchgrootte: 2 per apparaat met 4 stappen voor gradiëntaccumulatie met een effectieve batchgrootte van 64
  • Learnsnelheid: 2e-4 met cosinusscheduler en warmup
  • Trainingsstappen: 60 stappen voor demonstratie (toename voor volledige training)

De configuratie maakt gebruik van bfloat16-nauwkeurigheid en controlepunten voor gradiënten om het geheugengebruik te optimaliseren.

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

De gedistribueerde trainingsfunctie definiëren

De @distributed decorator van de serverless_gpu-bibliotheek maakt het mogelijk om GPU-workloads uit te voeren op Databricks AI Runtime. De decorator voorziet in 8 H100 GPU's en verwerkt automatisch de gedistribueerde trainingsopstelling.

Belangrijkste parameters:

  • gpus=8: aanvragen 8 GPU's voor gedistribueerde training
  • gpu_type='H100': Hiermee geeft u H100 GPU-hardware op

De trainingsfunctie:

  1. Laadt het basismodel en de tokenizer van HuggingFace
  2. Hiermee stelt u de chatopmaak in voor conversationele AI
  3. De trainingsgegevensset laden
  4. Initialiseert de TRL SFTTrainer met DeepSpeed-optimalisatie
  5. Traint het model en slaat controlepunten op
  6. Retourneert trainingsresultaten en MLflow-uitvoerings-id

Zie de documentatie voor de AI Runtime-API voor meer informatie.

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

De gedistribueerde trainingstaak uitvoeren

Voer de trainingsfunctie uit door .distributed() aan te roepen op de gedecoreerde functie. Hiermee worden de AI Runtime-resources geïmplementeerd, wordt de training uitgevoerd op 8 H100 GPU's met DeepSpeed-optimalisatie en worden de resultaten geretourneerd.

Het trainingsproces:

  • Richt automatisch 8 H100 GPU's in
  • Downloadt het model en de gegevensset van HuggingFace
  • Traint het model met volledige afstemming
  • Hiermee worden controlepunten opgeslagen op het Unity Catalog-volume
  • Registreert metrische gegevens naar MLflow
  • Retourneert de trainingsstatus, het uiteindelijke verlies en de uitvoerings-id van 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']}")

Het nauwkeurig afgestemde model en testdeductie opslaan

Met deze optionele stap wordt het nauwkeurig afgestemde model geladen en getest met een voorbeeldprompt om de resultaten te controleren.

Het proces:

  1. Het opgeslagen model laden
  2. Test het model met een voorbeeld van een gespreksprompt
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)

Het model registreren in Unity Catalog

Registreer het nauwkeurig afgestemde model bij MLflow en registreer het in Unity Catalog voor implementatie en bediening. Het model wordt geregistreerd met:

  • Model en tokenizer: Beide onderdelen die nodig zijn voor deductie
  • Taaktype: geconfigureerd als llm/v1/chat voor conversationele AI
  • Voorbeeld van invoer: Voorbeeld van chatberichtindeling voor testen
  • Unity Catalog-registratie: het model automatisch registreren in de geconfigureerde catalogus en het geconfigureerde schema

Zodra het model is geregistreerd, kan het worden ingezet voor modelbedieningseindpunten of worden gebruikt voor batch-inferentie.

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

Volgende stappen 

Voorbeeld van notebook

Llama 3.2 1B verfijnen met BEHULP van AI Runtime

Notebook krijgen

  • Last updated on