Verteiltes Feintuning von Qwen2-0.5B mit LoRA

Dieses Notizbuch veranschaulicht, wie Sie das große Qwen2-0.5B-Sprachmodell mithilfe von parametereffizienten Techniken auf Serverless GPU Compute effizient optimieren. Hier erfahren Sie, wie Sie:

• Wenden Sie LoRA (Low-Rank Anpassung) an, um trainierbare Parameter um ~99% zu reduzieren und dabei die Modellqualität aufrechtzuerhalten
• Verwenden von Liger Kernels für eine speichereffiziente Schulung mit optimierten Triton-Kernels
• Nutzen von TRL (Transformer Reinforcement Learning) für überwachte Feinabstimmung
• Registrieren Sie das feingetunte Modell im Unity-Katalog für Governance und Bereitstellung.

Schlüsselkonzepte:

• LoRA: Eine Technik, die das Basismodell einfriert und kleine Adapterschichten trainiert, wodurch die Speicheranforderungen und die Trainingszeit erheblich reduziert werden.
• Liger Kernels: GPU-optimierte Kernel, die die Speicherauslastung durch fused-Vorgänge um bis zu 80% reduzieren
• TRL: Eine Bibliothek für Schulungssprachenmodelle mit Verstärkungslernen und überwachter Feinabstimmung
• Serverless GPU Compute: Vom Databricks verwaltetes Compute, das GPU-Ressourcen automatisch skaliert

LoRA vs vollständige Feinabstimmungsentscheidungsmatrix

LoRA (Low-Rank Adaption) fixiert das Basismodell und trainiert nur kleine Adapterschichten, wodurch trainierbare Parameter um ~99%reduziert werden. Dadurch wird das Training schneller und speichereffizienter.

Liger Kernel-Vorteile

Liger Kernels sind GPU-optimierte Vorgänge, die mehrere Schritte in einzelne Kernels zusammenführten, die Speicherübertragungen reduzieren und die Effizienz verbessern. Das technische Dokument enthält detaillierte Benchmarks, die erhebliche Leistungsverbesserungen zeigen.

• ** Fused-Vorgänge: Kombiniert Vorgänge (z. B. linear + Fehler), um den Speicherbedarf um bis zu 80% zu reduzieren
• Triton-Kernel: Benutzerdefinierte GPU-Kernel, die für Transformatorvorgänge optimiert sind (RMSNorm, RoPE, SwiGLU, CrossEntropy)
• Speichereffizienz: Ermöglicht größere Batchgrößen oder Modelle, die sonst nicht in den GPU-Speicher passen
• Optimierung einzelner GPU: Besonders effektiv für A10/A100 Einzel-GPU-Schulungsszenarien

Dieses Notizbuch verwendet die TRL-Bibliothek , um die Schulungskonfiguration zu vereinfachen und diese Optimierungen automatisch anzuwenden.

Herstellen einer Verbindung mit Serverless GPU Compute

Für dieses Notizbuch ist eine serverlose GPU-Berechnung erforderlich. So stellen Sie eine Verbindung her:

1. Klicken Sie oben rechts auf die Computeauswahl des Notizbuchs, und wählen Sie serverlose GPU aus.
2. Klicken Sie auf der rechten Seite auf die Schaltfläche "Umgebung".
3. Wählen Sie 8xH100 als Beschleuniger aus.
4. Wählen Sie AI v4 aus der Basisumgebung aus.
5. Klicken Sie auf Anwenden.

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

Installieren erforderlicher Bibliotheken

In der nächsten Zelle werden die Python Pakete installiert, die für die verteilte Feinabstimmung erforderlich sind:

Kernschulungsbibliotheken:

• trl: Transformer Reinforcement Learning-Bibliothek für überwachte Feinabstimmung und RLHF
• peft: Parameter-Efficient Fine-Tuning Bibliothek, die die LoRA-Implementierung bereitstellt
• liger-kernel: Speicheroptimierte GPU-Kernel für eine effiziente Transformatorschulung

Unterstützende Bibliotheken:

• hf_transfer: Beschleunigte Downloads von Hugging Face Hub mit Rust-basierter Übertragung
• mlflow>=3.0: Experiment-Tracking und Modellregistrierungsintegration

Der Befehl %restart_python startet den Python-Interpreter neu, um sicherzustellen, dass die neu installierten Pakete ordnungsgemäß geladen werden.

%pip install --upgrade peft==0.17.1
%pip install --upgrade hf_transfer==0.1.9
%pip install --upgrade transformers==4.56.1
%pip install trl==0.18.1
%pip install liger-kernel
%pip install mlflow==3.7.0
%restart_python

Konfiguration

Integration des Unity Catalog

Die nächste Zelle konfiguriert, wo Ihr fein abgestimmtes Modell gespeichert und registriert wird:

• Katalog & Schema: Organisieren von Modellen innerhalb Ihres Unity-Katalognamespaces (Standard: main.default)
• Modellname: Der name des registrierten Modells im Unity-Katalog für Governance und Bereitstellung
• Volume: Unity Catalog Volume zum Speichern von Modellprüfpunkten während der Schulung

Mit diesen Widgets können Sie den Speicherort anpassen, ohne Code zu bearbeiten. Das Modell wird als {catalog}.{schema}.{model_name} für einfachen Zugriff und Versionsverwaltung registriert.

Trainieren von Hyperparametern

Die Zelle definiert außerdem wichtige Schulungsparameter:

• Model & Dataset: Qwen2-0.5B mit dem Capybara-Gesprächsdatensatz
• Batchgröße (8): Anzahl der Beispiele pro GPU pro Schulungsschritt
• Gradientenakkumulation (4): Akkumuliert Gradienten über 4 Batches für eine effektive Batchgröße von 32
• Lernrate (1e-4): Konservative Rate, automatisch skaliert 10x höher für LoRA-Schulungen
• Epochen (1): Einzelner Durchlauf durch den Datensatz, um Überanpassung zu vermeiden
• Logging & Checkpointing: Speichert den Fortschritt alle 100 Schritte, protokolliert Metriken alle 25 Schritte.

dbutils.widgets.text("uc_catalog", "main")
dbutils.widgets.text("uc_schema", "default")
dbutils.widgets.text("uc_model_name", "qwen2_liger_lora_assistant")
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")

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

# MLflow and Unity Catalog configuration

# Model selection - Choose based on your compute constraints
MODEL_NAME = "Qwen/Qwen2-0.5B"
DATASET_NAME = "trl-lib/Capybara"
OUTPUT_DIR = f"/Volumes/{UC_CATALOG}/{UC_SCHEMA}/{UC_VOLUME}/qwen2-0.5b-lora"

# Training hyperparameters
BATCH_SIZE = 8
GRADIENT_ACCUMULATION_STEPS = 4
LEARNING_RATE = 1e-4
NUM_EPOCHS = 1
EVAL_STEPS = 100
LOGGING_STEPS = 25
SAVE_STEPS = 100

LoRA-Konfiguration

In der nächsten Zelle werden LoRA (Low-Rank Adaptionsparameter) konfiguriert, die steuern, wie das Modell fein abgestimmt wird. LoRA friert die Basismodellgewichte ein und trainiert nur kleine Adaptermatrizen, wodurch die Speicheranforderungen erheblich reduziert werden.

Parameterauswahl

• Rang (r=8): Bietet eine gute Balance zwischen Leistung und Parametern.
• Alpha (32): Skalierungsfaktor, in der Regel 2-4x der Rang
• Dropout (0.1): Regularisierung zur Vermeidung von Überanpassung

Zielmodule für Qwen2

In diesem Beispiel werden alle Schlüsseltransformationsebenen anvisiert.

• Achtung: q_proj, , k_projv_projo_proj
• MLP: gate_proj, up_projdown_proj

LORA_R = 8
LORA_ALPHA = 32
LORA_DROPOUT = 0.1
LORA_TARGET_MODULES = [
    "q_proj", "k_proj", "v_proj", "o_proj",
    "gate_proj", "up_proj", "down_proj"
]

Definieren der Schulungsfunktion

Die nächste Zelle erstellt die verteilte Schulungsfunktion, die über mehrere GPUs hinweg ausgeführt wird. Hier erfahren Sie, was dies tut:

Verteiltes Schulungssetup

Der @distributed Dekorateur konfiguriert Serverless GPU Compute:

• 8 GPUs: Verteilt Schulungen über 8 H100 GPUs für schnellere Schulungen
• Automatische Orchestrierung: Behandelt gpu-Bereitstellung, Datenverteilung und Synchronisierung

Trainingsworkflow

Die Funktion führt die folgenden Schritte aus:

1. Laden des Datasets: Lädt das Capybara-Unterhaltungsdatenset herunter und bereitet es vor.
2. Initialisieren des Modells: Lädt Qwen2-0.5B und Tokenizer mit Chatformatierung
3. Anwenden von LoRA: Fügt Adapterschichten an, um trainierbare Parameter um ~99% zu reduzieren.
4. Schulung konfigurieren: Richtet Batchgröße, Lernrate und Liger-Kerneloptimierungen ein
5. Trainingsmodell: Führt die Trainingsschleife mit automatischer Prüfpunkterstellung und Protokollierung aus.
6. Speichern von Artefakten: Speichert LoRA-Adapter und Tokenizer im Unity-Katalogvolume
7. Zurückgeben der MLflow-Ausführungs-ID: Stellt die Ausführungs-ID für die Modellregistrierung bereit.

Schlüsseloptimierungen aktiviert

• Liger Kernels: Fused GPU-Vorgänge reduzieren die Speichernutzung um bis zu 80%
• Gemischte Genauigkeit (FP16): Schnellere Berechnung mit geringerem Speicherbedarf
• Gradienten-Checkpoint: Tauscht Rechenleistung gegen Speicherplatz, um größere Batches zu verarbeiten.
• Gradientakkumulation: Simuliert größere Batchgrößen für stabiles Training

from serverless_gpu import distributed
from serverless_gpu import runtime as rt

@distributed(gpus=8, gpu_type="H100")
def run_train(use_lora=True):
    import logging
    from datasets import load_dataset
    from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer
    from peft import LoraConfig, TaskType, get_peft_model
    from trl import (
        SFTConfig,
        SFTTrainer,
        setup_chat_format
    )
    import json
    import os
    import mlflow

    dataset = load_dataset(DATASET_NAME)
    logging.info(f"✓ Dataset loaded: {dataset}")

    if "test" not in dataset:
        logging.info("Creating validation split from training data...")
        dataset = dataset["train"].train_test_split(test_size=0.1, seed=42)
        logging.info("✓ Data split: 90% train, 10% validation")

    # model and tokenizer initialization
    model = AutoModelForCausalLM.from_pretrained(
        MODEL_NAME,
        trust_remote_code=True,
    )

    tokenizer = AutoTokenizer.from_pretrained(
        MODEL_NAME,
        trust_remote_code=True,
        use_fast=True
    )

    # Chat template formatting for conversational fine-tuning
    if tokenizer.chat_template is None:
        logging.info("Adding chat template for proper conversation formatting...")
        model, tokenizer = setup_chat_format(model, tokenizer, format="chatml")
        logging.info("✓ ChatML format applied for structured conversations")

    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token
        logging.info("✓ Padding token set to EOS token")

    logging.info("✓ Model and tokenizer loaded successfully")

    # PEFT
    peft_config = None
    if use_lora:
        try:
            logging.info("Configuring LoRA for parameter-efficient fine-tuning...")

            peft_config = LoraConfig(
                task_type=TaskType.CAUSAL_LM,
                inference_mode=False,
                r=LORA_R,
                lora_alpha=LORA_ALPHA,
                lora_dropout=LORA_DROPOUT,
                target_modules=LORA_TARGET_MODULES,
                bias="none",
                use_rslora=False,
                modules_to_save=None,
            )

            logging.info(f"LoRA configuration: rank={LORA_R}, alpha={LORA_ALPHA}, dropout={LORA_DROPOUT}")
            logging.info(f"Target modules: {', '.join(LORA_TARGET_MODULES)}")

            original_params = model.num_parameters()
            model = get_peft_model(model, peft_config)

            trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
            total_params = sum(p.numel() for p in model.parameters())
            efficiency_ratio = 100 * trainable_params / total_params

            logging.info(f"✓ LoRA applied successfully:")
            logging.info(f"  • Original parameters: {original_params:,}")
            logging.info(f"  • Trainable parameters: {trainable_params:,}")
            logging.info(f"  • Training efficiency: {efficiency_ratio:.2f}% of parameters")
            logging.info(f"  • Memory savings: ~{100-efficiency_ratio:.1f}% reduction in gradient memory")

        except Exception as e:
            logging.info(f"✗ LoRA configuration failed: {e}")
            logging.info("Falling back to full fine-tuning...")
            peft_config = None
    else:
        logging.info("Full fine-tuning mode selected")
        trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
        logging.info(f"Trainable parameters: {trainable_params:,} (100% of model)")

    # Learning rate adjustment for LoRA
    adjusted_lr = LEARNING_RATE * 10 if use_lora else LEARNING_RATE
    logging.info(f"Learning rate: {adjusted_lr} ({'LoRA-adjusted' if use_lora else 'standard'})")

    training_args_dict = {
        "output_dir": OUTPUT_DIR,
        "per_device_train_batch_size": BATCH_SIZE,
        "per_device_eval_batch_size": BATCH_SIZE,
        "gradient_accumulation_steps": GRADIENT_ACCUMULATION_STEPS,
        "learning_rate": adjusted_lr,
        "num_train_epochs": NUM_EPOCHS,
        "eval_steps": EVAL_STEPS,
        "logging_steps": LOGGING_STEPS,
        "save_steps": SAVE_STEPS,
        "save_total_limit": 2,
        "report_to": "mlflow",
        "run_name": f"{MODEL_NAME}_fine-tuning",
        "warmup_steps": 50,
        "weight_decay": 0.01,
        "metric_for_best_model": "eval_loss",
        "greater_is_better": False,
        "dataloader_pin_memory": False,
        "remove_unused_columns": False,
        "use_liger_kernel": True,  # Enable Liger kernel optimizations
        "fp16": True,  # Mixed precision training
        "gradient_checkpointing": True,
        "gradient_checkpointing_kwargs": {"use_reentrant": False}, # Required for LORA with DDP
    }

    logging.info("✓ Liger kernel optimizations enabled")

    training_args = SFTConfig(**training_args_dict)

    trainer = SFTTrainer(
        model=model,
        args=training_args,
        train_dataset=dataset["train"],
        eval_dataset=dataset["test"],
        processing_class=tokenizer,
        peft_config=peft_config,
    )

    logging.info("\n" + "="*50)
    logging.info("STARTING TRAINING")
    logging.info("="*50)

    logging.info("🚀 Training with Liger kernels for memory-efficient single GPU training")
    if use_lora:
        logging.info("🎯 Using LoRA for parameter-efficient fine-tuning")

    trainer.train()
    logging.info("\n✓ Training completed successfully!")
    if rt.get_global_rank() == 0:
        logging.info("\nSaving trained model...")

        logging.info("Saving LoRA adapter weights...")
        trainer.save_model(training_args.output_dir)
        logging.info("✓ LoRA adapters saved - use with base model for inference")
        tokenizer.save_pretrained(training_args.output_dir)
        logging.info("✓ Tokenizer saved with model")
        logging.info(f"\n🎉 All artifacts saved to: {training_args.output_dir}")

    mlflow_run_id = None
    if mlflow.last_active_run() is not None:
        mlflow_run_id = mlflow.last_active_run().info.run_id

    return mlflow_run_id

Führen Sie das verteilte Training aus

Diese Zelle führt die Trainingsfunktion auf 8 H100 GPUs aus. Die distributed() Methode behandelt:

• Bereitstellen serverloser GPU-Computeressourcen
• Verteilen der Schulungsworkloads über mehrere GPUs
• Sammeln der MLflow-Ausführungs-ID für die Modellregistrierung

Die Schulung dauert in der Regel 15-30 Minuten, je nach Datasetgröße und Berechnungsverfügbarkeit.

mlflow_run_id = run_train.distributed(use_lora=True)[0]
print(mlflow_run_id)

MLflow- und Unity-Katalogregistrierung

Modellregistrierungsstrategie

• MLflow-Nachverfolgung: Protokollmodellartefakte und Metadaten
• Unity-Katalog: Registrieren des Modells für Governance und Bereitstellung
• Modellversionsverwaltung: Automatische Versionsverwaltung für die Modelllebenszyklusverwaltung
• Metadaten: Vollständige Modellinformationen zur Reproduzierbarkeit

print("\nRegistering model with MLflow and Unity Catalog...")

from transformers import AutoTokenizer, AutoModelForCausalLM
from peft import PeftModel
import mlflow
from mlflow import transformers as mlflow_transformers

try:
    # Load the trained model for registration
    print("Loading LoRA model for registration...")
    # For LoRA models, we need both base model and adapter
    base_model = AutoModelForCausalLM.from_pretrained(
        MODEL_NAME,
        trust_remote_code=True
    )
    # Load tokenizer
    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
    adapter_dir = OUTPUT_DIR
    peft_model = PeftModel.from_pretrained(base_model, adapter_dir)
    # Merge LoRA into base and drop PEFT wrappers
    merged_model = peft_model.merge_and_unload()

    components = {
        "model": merged_model,
        "tokenizer": tokenizer,
    }

    # Create Unity Catalog model name
    full_model_name = f"{UC_CATALOG}.{UC_SCHEMA}.{UC_MODEL_NAME}"

    print(f"Registering model as: {full_model_name}")

    # Start MLflow run and log model
    task = "llm/v1/chat"
    with mlflow.start_run(run_id=mlflow_run_id):
        model_info = mlflow.transformers.log_model(
            transformers_model=components,
            artifact_path="model",
            task=task,
            registered_model_name=full_model_name,
            metadata={
                "task": task,
                "pretrained_model_name": MODEL_NAME,
                "databricks_model_family": "QwenForCausalLM",
            },
        )

    print(f"✓ Model successfully registered in Unity Catalog: {full_model_name}")
    print(f"✓ MLflow model URI: {model_info.model_uri}")

    # Print deployment information
    print(f"\n📦 Model Registration Complete!")
    print(f"Unity Catalog Path: {full_model_name}")
    print(f"Model Type: {model_type}")
    print(f"Optimization: Liger Kernels + LoRA")

except Exception as e:
    print(f"✗ Model registration failed: {e}")
    print("Model is still saved locally and can be registered manually")
    print(f"Local model path: {OUTPUT_DIR}")

Nächste Schritte

Nachdem Sie Ihr Modell optimiert und registriert haben, haben Sie folgende Möglichkeiten:

• Modell bereitstellen:Dienste für Modelle mit Model Serving bereitstellen
• Erfahren Sie mehr über verteilte Schulungen: Verteilte Multi-GPU-Schulung
• Optimieren der Serverless GPU Compute-Nutzung: Bewährte Methoden für serverlose GPU-Compute
• Behandeln von Problemen: Behandeln von Problemen beim Serverless GPU Compute

Beispiel-Notebook

Verteiltes Feintuning von Qwen2-0.5B mit LoRA

Notebook abrufen