Vollständiges Fine-Tuning von Qwen3-4B

Optimieren Sie das große Qwen3-4B-Sprachmodell auf einer einzelnen H100-GPU vollständig. Diese Anleitung zeigt, wie Sie:

• Führen Sie die vollständige Feinabstimmung aus, die jeden Modellparameter aktualisiert, um maximale Anpassung an Ihre Daten zu erhalten.
• Verwenden Sie die Databricks AI v5-Umgebung , ohne zusätzliche Bibliotheken zu installieren
• 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:

• Vollständige Feinabstimmung: Aktualisiert alle Modellgewichte und verleiht dem Modell damit die größte Kapazität, aus Ihrem Datensatz zu lernen – auf Kosten eines höheren Speicher- und Rechenaufwands als bei parametereffizienten Methoden
• TRL: Eine Bibliothek für Schulungssprachenmodelle mit Verstärkungslernen und überwachter Feinabstimmung
• Speichereffizientes Training: Verwendet BF16 mit gemischter Präzision und Gradient Checkpointing, um eine vollständige Feinabstimmung eines Modells mit 4B Parametern auf einer einzelnen H100-GPU zu ermöglichen

Matrix für vollständige Feinabstimmung vs LoRA-Entscheidungsmatrix

Dieses Notebook verwendet vollständiges Fine-Tuning, bei dem alle Modellparameter aktualisiert werden. Die Alternative LoRA (Low-Rank Adaption) fixiert das Basismodell und trainiert nur kleine Adapterschichten.

Das vollständige Fine-Tuning eines Modells mit 4 Milliarden Parametern erfordert deutlich mehr GPU-Speicher als LoRA, da für jeden Parameter Optimiererzustände und Gradienten gespeichert werden. Dieses Notebook verwendet gemischte BF16-Präzision und Gradient Checkpointing, sodass das Training auf einer einzelnen H100-GPU (80 GB) ausgeführt werden kann.

Verbindung zu Serverless GPU-Compute herstellen

So stellen Sie eine Verbindung mit serverlosem GPU-Compute her:

1. Klicken Sie im Notizbuch auf das Dropdownmenü "Verbinden ", und wählen Sie "Serverlose GPU" aus.
2. Wählen Sie als Beschleuniger eine 1x H100-GPU aus.
3. Öffnen Sie den Bereich "Umgebung ", und wählen Sie "AI v5 " als Basisumgebung aus.
4. Klicken Sie auf Anwenden.

Weitere Informationen finden Sie in der GPU-Computedokumentation.

Importieren von Bibliotheken

Die Databricks AI v5-Umgebung enthält bereits alle Bibliotheken, die für dieses Beispiel erforderlich sind (z. B. trl, transformers, datasets und mlflow), sodass keine zusätzliche Installation erforderlich ist.

In der nächsten Zelle werden die erforderlichen Bibliotheken für Modellschulungen, Datasetverarbeitung und MLflow-Nachverfolgung importiert.

from datasets import load_dataset
from transformers import AutoModelForCausalLM, AutoTokenizer
from trl import (
    SFTConfig,
    SFTTrainer,
    setup_chat_format
)
import torch
import mlflow

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: Qwen3-4B mit dem Capybara-Konversationsdatensatz
• Batchgröße (1): Anzahl der Beispiele pro GPU pro Schulungsschritt, klein gehalten, um eine vollständige Feinabstimmung im Arbeitsspeicher zu erhalten
• Gradient Accumulation (8): Akkumuliert Gradienten über 8 Batches für eine effektive Batchgröße von 8
• Lernrate (2e-5): Konservative Rate, die für die vollständige Feinabstimmung geeignet ist
• Maximale Schritte (50): Begrenzt das Training auf 50 Schritte für einen kurzen Demonstrationslauf
• Logging & Checkpointing: Speichert den Fortschritt alle 25 Schritte, protokolliert Metriken alle 10 Schritte.

dbutils.widgets.text("uc_catalog", "main")
dbutils.widgets.text("uc_schema", "default")
dbutils.widgets.text("uc_model_name", "qwen3_4b_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
MODEL_NAME = "Qwen/Qwen3-4B"
DATASET_NAME = "trl-lib/Capybara"
OUTPUT_DIR = f"/Volumes/{UC_CATALOG}/{UC_SCHEMA}/{UC_VOLUME}/{UC_MODEL_NAME}"

# Training hyperparameters
BATCH_SIZE = 1
GRADIENT_ACCUMULATION_STEPS = 8
LEARNING_RATE = 2e-5
MAX_STEPS = 50
EVAL_STEPS = 25
LOGGING_STEPS = 10
SAVE_STEPS = 25

Laden und Vorbereiten des Datasets

Die nächste Zelle lädt den Trainingsdatensatz und bereitet ihn auf das Feintuning vor:

• Dataset: trl-lib/Capybara - qualitativ hochwertige Gesprächsdaten, die für das Befolgen von Anweisungen optimiert sind
• Aufteilung Training/Validierung: Erstellt eine 90/10-Aufteilung, wenn kein Testdatensatz vorhanden ist
• Datenüberprüfung: Stellt eine ordnungsgemäße Formatierung für die Optimierung von Unterhaltungen sicher.

dataset = load_dataset(DATASET_NAME)
print(f"✓ Dataset loaded: {dataset}")

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

Initialisieren von Modell und Tokenizer

Die nächste Zelle lädt das Basismodell und den Tokenizer und konfiguriert sie dann für die Feinabstimmung für Konversationen.

• Laden von Modellen: Lädt Qwen3-4B von Hugging Face mit BF16-Präzision herunter
• Tokenizer-Setup: Konfiguriert einen schnellen Tokenizer mit ordnungsgemäßen Abständen.
• Chatformatierung: Wendet eine Chatvorlage für strukturierte Unterhaltungen an, wenn der Tokenizer noch keines definiert
• Tokenkonfiguration: Legt das Padding-Token auf das EOS-Token fest, um eine ordnungsgemäße Verarbeitung der Sequenz zu gewährleisten.

model = AutoModelForCausalLM.from_pretrained(
    MODEL_NAME,
    torch_dtype=torch.bfloat16,
    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:
    print("Adding chat template for proper conversation formatting...")
    model, tokenizer = setup_chat_format(model, tokenizer, format="chatml")
    print("✓ ChatML format applied for structured conversations")

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

print("✓ Model and tokenizer loaded successfully")

Trainieren des Modells

Die nächste Zelle konfiguriert und führt den vollständigen Feinabstimmungsprozess aus:

Schulungskonfiguration

• Batchkonfiguration: 1 Sample pro Gerät mit 8 Schritten der Gradientenakkumulation (effektive Batchgröße: 8)
• Optimierung: Aufwärmschritte, Gewichtsabfall und beste Modellauswahl basierend auf Auswertungsverlust
• Protokollierung: Meldet Metriken an MLflow zum Nachverfolgen von Experimenten

Schlüsseloptimierungen aktiviert

• BF16 gemischte Genauigkeit: Schnellere Berechnung mit geringerem Speicherbedarf, gut geeignet für H100 GPUs
• Gradient Checkpointing: Tauscht zusätzliches Compute gegen eine deutliche Reduzierung des Aktivierungsspeichers, wodurch eine vollständige Feinabstimmung eines 4B-Modells auf eine einzelne H100 passt
• Gradientakkumulation: Simuliert größere Batchgrößen für stabiles Training
• Checkpointing: Speichert das Modell alle 25 Schritte, wobei maximal 2 Checkpoints gespeichert werden.

Die Schulungsschleife protokolliert den Fortschritt alle 10 Schritte und wertet alle 25 Schritte aus.

with mlflow.start_run(run_name=f"{MODEL_NAME}_full-fine-tuning", log_system_metrics=True):
    try:
        print(f"Learning rate: {LEARNING_RATE}")

        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": LEARNING_RATE,
            "max_steps": MAX_STEPS,
            "eval_steps": EVAL_STEPS,
            "logging_steps": LOGGING_STEPS,
            "save_steps": SAVE_STEPS,
            "save_total_limit": 2,
            "report_to": "mlflow",  # Log to MLflow
            "warmup_steps": 10,
            "weight_decay": 0.01,
            "metric_for_best_model": "eval_loss",
            "greater_is_better": False,
            "eval_strategy": "steps",  # Run evaluation every eval_steps
            "save_strategy": "steps",  # Checkpoint on the same cadence as eval
            "load_best_model_at_end": True,  # Register the best-eval checkpoint, not the last
            "dataloader_pin_memory": False,
            "remove_unused_columns": False,
            "bf16": True,  # Mixed precision training
            "gradient_checkpointing": True,  # Reduce activation memory for full fine-tuning
            "gradient_checkpointing_kwargs": {"use_reentrant": False},
        }

        training_args = SFTConfig(**training_args_dict)

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

        print("\n" + "="*50)
        print("STARTING TRAINING")
        print("="*50)

        print("🚀 Full fine-tuning Qwen3-4B on a single H100 GPU")

        trainer.train()
        print("\n✓ Training completed successfully!")

    except Exception as e:
        print(f"✗ Training failed: {e}")
        raise

Speichern von Modellartefakten

Die nächste Zelle speichert das trainierte Modell und den Tokenizer im Unity-Katalog-Volume.

• Vollständige Modellgewichte: Speichert das komplette Feinabstimmungsmodell, das direkt für die Inferenz geladen werden kann
• Tokenizer: Speichert die Tokenizerkonfiguration für die Ableitung
• Speicherort: Speichert in /Volumes/{catalog}/{schema}/{volume}/{model_name}

try:
    print("\nSaving trained model...")

    trainer.save_model(training_args.output_dir)
    print("✓ Full model weights saved")

    tokenizer.save_pretrained(training_args.output_dir)
    print("✓ Tokenizer saved with model")
    print(f"\n🎉 All artifacts saved to: {training_args.output_dir}")

except Exception as e:
    print(f"✗ Model saving failed: {e}")
    raise

Registrieren des Modells im Unity-Katalog

In der nächsten Zelle wird das fein abgestimmte Modell im Unity Catalog für Governance und Bereitstellung registriert:

Modellregistrierungsworkflow

1. Trainiertes Modell laden: Lädt das gespeicherte Modell mit vollständigen Gewichtungen und den Tokenisierer
2. Vorbereiten der Protokollierung: Erstellt ein Transformatorenmodellverzeichnis mit dem Modell und Tokenizer
3. Registrieren im Unity-Katalog: Protokolliert in MLflow und registriert im Unity-Katalog
4. Hinzufügen von Metadaten: Enthält Aufgabentyp, Modellfamilie und Größeninformationen

Vorteile der Unity-Katalogregistrierung

• Governance: Zentrale Modellregistrierung mit Zugriffssteuerung und Lineage-Tracking
• Versionsverwaltung: Automatische Versionsverwaltung für den Modelllebenszyklus
• Bereitstellung: Einfache Bereitstellung zu Endpunkten für die Modellbereitstellung
• Auffindbarkeit: Modelle sind durchsuchbar und werden im Unity-Katalog dokumentiert.

mlflow_run_id = mlflow.last_active_run().info.run_id
print("\nRegistering model with MLflow and Unity Catalog...")

with mlflow.start_run(run_id=mlflow_run_id):
    try:
        # Load the trained full-weight model for registration
        print("Loading fine-tuned model for registration...")
        trained_model = AutoModelForCausalLM.from_pretrained(
            training_args.output_dir,
            torch_dtype=torch.bfloat16,
            trust_remote_code=True
        )
        tokenizer = AutoTokenizer.from_pretrained(training_args.output_dir)
        model_type = "Full fine-tuning"
        size_params = "4b"

        # Prepare transformers model dictionary
        transformers_model = {
            "model": trained_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"
        model_info = mlflow.transformers.log_model(
            transformers_model=transformers_model,
            task=task,
            registered_model_name=full_model_name,
            metadata={
                "task": task,
                "pretrained_model_name": MODEL_NAME,
                "databricks_model_family": "Qwen3ForCausalLM",
                "databricks_model_size_parameters": size_params,
            },
            repo_type="local",  # Fix: specify repo_type for local path
        )

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

    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: {training_args.output_dir}")
        raise

Nächste Schritte

Ihr Qwen3-4B-Modell wurde erfolgreich mit vollständiger Feinabstimmung optimiert und im Unity Catalog registriert. Als Nächstes haben Sie folgende Möglichkeiten:

• Modell bereitstellen:Dienste für Modelle mit Model Serving bereitstellen
• Erfahren Sie mehr über verteilte Schulungen: Verteilte Mehr-GPU- und Mehrknotenschulung
• Verfolgen von Experimenten und Überwachen von GPUs: Nachverfolgen und Beobachten von Experimenten
• Behandeln von Problemen: Behandeln von Problemen beim Serverlosen GPU-Compute

Beispiel-Notebook

Vollständiges Fine-Tuning von Qwen3-4B

Notebook abrufen