Feinabstimmung von Llama 3.2 1B mit AI Runtime

Notizbuchversion dieser Seite öffnen

Optimieren Sie ein großes Sprachmodell vollumfänglich mit überwachter Feinabstimmung (SFT) auf Databricks AI Runtime. Das Notizbuch verwendet die TRL-Bibliothek (Transformers Reinforcement Learning) mit DeepSpeed ZeRO Stage 3 Optimierung, um Llama 3.2 1B auf einem einzigen Knoten mit 8 H100 GPUs effizient zu trainieren.

Schlüsselkonzepte:

TRL (Transformers Reinforcement Learning): Eine Bibliothek, die Tools zum Training von Sprachmodellen mit Verstärkungslernen und überwachter Feinabstimmung bereitstellt.
DeepSpeed ZeRO Phase 3: Eine Speicheroptimierungstechnik, die Modellparameter, Farbverläufe und Optimiererzustände über GPUs partitioniert, um die Schulung großer Modelle zu ermöglichen.
AI-Runtime: Databricks-verwaltete GPU-Compute, die GPU-Ressourcen automatisch für Schulungsworkloads bereitstellt und skaliert.

Weitere Informationen finden Sie unter AI Runtime.

Note

Dieses Beispiel benötigt die Databricks KI-Umgebung Version 5 oder höher.

Anforderungen

Dieses Notizbuch erfordert Folgendes:

AI Runtime: Das Notizbuch verwendet Databricks AI Runtime mit 8 H100 GPUs für verteilte Schulungen. Es ist keine Clusterkonfiguration erforderlich.
Unity-Katalog: Ein Unity-Katalog und -Schema zum Speichern von Modellprüfpunkten und Registrieren des trainierten Modells.
HuggingFace-Token: Ein HuggingFace-Zugriffstoken, das in Databricks geheimen Schlüsseln gespeichert ist, um das Basismodell und das Dataset herunterzuladen.
Python Pakete: KI-Runtime vorinstalliert die meisten erforderlichen Bibliotheken. Im folgenden Abschnitt „Einrichtung“ wird deepspeed installiert.

Verbindung zu Serverless GPU-Compute herstellen

Für dieses Notebook ist ein serverloses GPU-Compute erforderlich. So stellen Sie eine Verbindung her:

Klicken Sie oben rechts auf die Compute-Auswahl des Notebooks und wählen Sie Serverlose GPU aus.
Klicken Sie auf der rechten Seite auf die Umgebungsschaltfläche.
Wählen Sie 8xH100 als Beschleuniger aus.
Wählen Sie die AI v5-Umgebung aus dem rechten Bereich aus, der alle erforderlichen Bibliotheken enthält, um dieses Notizbuchbeispiel auszuführen.
Klicken Sie auf Anwenden.

Die Schulungsfunktion stellt automatisch 8 H100 GPUs für verteilte Schulungen bereit.

Installieren Sie die erforderlichen Pakete

DIE KI-Runtime verfügt bereits über die meisten erforderlichen Bibliotheken vorinstalliert. In diesem Beispiel müssen Sie nur deepspeedinstallieren.

%pip install deepspeed==0.19.1
%restart_python

Konfigurieren von Unity-Katalog- und Umgebungsvariablen

Richten Sie Die Speicherorte des Unity-Katalogs zum Speichern von Modellprüfpunkten und zum Registrieren des trainierten Modells ein. Das Notizbuch verwendet Abfrageparameter, um Folgendes zu konfigurieren:

Katalog und Schema: Unity Catalog-Namespace zum Organisieren von Modellen und Prüfpunkten
Modellname: Name für das registrierte Modell im Unity-Katalog
Volume: Unity Catalog Volume zum Speichern von Modellprüfpunkten während der Schulung

Die Konfiguration ruft auch das HuggingFace-Token aus Databricks-Geheimnissen ab und richtet das MLflow-Experiment für die Nachverfolgung von Schulungsmetriken ein.

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

Erstellen der DeepSpeed ZeRO Phase 3-Konfiguration

DeepSpeed ZeRO (Zero Redundanzoptimierer) Stufe 3 partitioniert Modellparameter, Gradienten und Optimiererzustände über alle GPUs, um den Speicherverbrauch pro GPU zu reduzieren. Dies ermöglicht die Schulung großer Modelle, die nicht in den Speicher einer einzelnen GPU passen.

Wichtige Konfigurationseinstellungen:

bf16 aktiviert: Verwendet bfloat16-Genauigkeit für schnellere Schulungen und reduzierte Arbeitsspeicherauslastung
Phase 3-Optimierung: Partitioniert alle Modellzustände über GPUs hinweg
Kein CPU-Offloading: Hält alle Daten auf GPUs für maximale Leistung auf H100-Hardware
Überlappende Kommunikation: Überlappt Gradientenkommunikation mit Berechnung für effizientere Abläufe

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

Definieren von Schulungsparametern

Konfigurieren Sie die überwachten Feinabstimmungseinstellungen.

Modell: Llama 3.2 1B Instruct, ein kompaktes Modell geeignet für H100 GPUs
Dataset: Capybara-Dataset aus der TRL-Bibliothek für die Ausbildung konversationaler KI
Batchgröße: 2 pro Gerät mit 4 Gradientenakkumulationsschritten für eine effektive Batchgröße von 64
Lernrate: 2e-4 mit Kosinus-Scheduler und Aufwärmung
Trainingsschritte: 60 Schritte zur Präsentation (Förderung für die vollständige Schulung)

Die Konfiguration verwendet bfloat16-Präzision und Gradienten-Checkpointing, um die Speicherauslastung zu optimieren.

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

Definieren der verteilten Schulungsfunktion

Der @distributed-Dekorator aus der Bibliothek serverless_gpu ermöglicht das Ausführen von GPU-Workloads auf Databricks AI Runtime. Der Decorator stellt 8 H100 GPUs bereit und verwaltet das verteilte Trainings-Setup automatisch.

Schlüsselparameter:

gpus=8: Fordert 8 GPUs für verteilte Schulungen an
gpu_type='H100': Gibt H100 GPU-Hardware an

Die Schulungsfunktion:

Lädt das Basismodell und den Tokenizer von HuggingFace
Einrichten der Chatformatierung für Konversations-KI
Lädt das Schulungsdatenset
Initialisiert den TRL SFTTrainer mit DeepSpeed-Optimierung
Trainiert das Modell und speichert Prüfpunkte
Gibt Schulungsergebnisse und MLflow-Ausführungs-ID zurück.

Weitere Informationen finden Sie in der AI-Runtime-API-Dokumentation.

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

Den verteilten Trainingsjob ausführen

Führen Sie die Trainingsfunktion aus, indem Sie .distributed() für die dekorierte Funktion aufrufen. Dies stellt die AI-Runtime-Ressourcen bereit, führt die Schulung über 8 H100 GPUs mit DeepSpeed-Optimierung aus und gibt die Ergebnisse zurück.

Der Schulungsvorgang:

Automatische Bereitstellung von 8 H100 GPUs
Lädt das Modell und das Dataset von HuggingFace herunter.
Trainiert das Modell mit voller Feinabstimmung
Speichert Prüfpunkte im Unity-Katalogvolume
Protokolliert Metriken bei MLflow
Gibt den Schulungsstatus, den endgültigen Verlust und die MLflow-Ausführungs-ID zurück.

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

Speichern des fein abgestimmten Modells und Testen der Ableitung

Dieser optionale Schritt lädt das fein abgestimmte Modell und testet es mit einer Beispielaufforderung, um die Ergebnisse zu überprüfen.

Der Prozess:

Lädt das gespeicherte Modell
Das Modell wird mit einer Beispielkonversationsanfrage getestet.

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)

Registrieren des Modells im Unity-Katalog

Protokollieren Sie das fein abgestimmte Modell bei MLflow und registrieren Sie es im Unity-Katalog für die Bereitstellung und Servierung. Das Modell wird mit protokolliert:

Modell- und Tokenizer: Beide Komponenten, die für die Ableitung erforderlich sind
Aufgabentyp: Konfiguriert für llm/v1/chat konversationelle KI
Eingabebeispiel: Beispiel für chatnachrichtenformat zum Testen
Unity-Katalogregistrierung: Registriert das Modell automatisch im konfigurierten Katalog und Schema.

Nach der Registrierung kann das Modell für Modellbereitstellungsendpunkte bereitgestellt oder für die Batch-Ableitung verwendet werden.

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

Nächste Schritte

Beispiel-Notebook

Feinabstimmung von Llama 3.2 1B mit AI Runtime

Notebook abrufen

Feedback

War diese Seite hilfreich?

Last updated on 2026-08-18

Feinabstimmung von Llama 3.2 1B mit AI Runtime

Anforderungen

Verbindung zu Serverless GPU-Compute herstellen

Installieren Sie die erforderlichen Pakete

Konfigurieren von Unity-Katalog- und Umgebungsvariablen

Erstellen der DeepSpeed ZeRO Phase 3-Konfiguration

Definieren von Schulungsparametern

Definieren der verteilten Schulungsfunktion

Den verteilten Trainingsjob ausführen

Speichern des fein abgestimmten Modells und Testen der Ableitung

Registrieren des Modells im Unity-Katalog

Nächste Schritte

Beispiel-Notebook

Feinabstimmung von Llama 3.2 1B mit AI Runtime

Feedback

Zusätzliche Ressourcen