Affinement distribué de Qwen2-0.5B avec LoRA

Ouvrir dans Databricks

Affinez le grand modèle linguistique Qwen2-0.5B efficacement sur AI Runtime en utilisant des techniques à parameter efficace. Vous apprendrez comment :

Appliquer LoRA (Low-Rank Adaptation) pour réduire les paramètres entraînables d'environ 99 % tout en maintenant la qualité du modèle
Utilisez Liger Kernels pour un entraînement à mémoire efficace avec des noyaux Triton optimisés
Tirer parti de **TRL (Transformer Reinforcement Learning)** pour l'affinement supervisé
Enregistrez le modèle affiné dans Unity Catalog pour la gouvernance et le déploiement

Concepts clés :

LoRA: Une technique qui fige le modèle de base et entraîne de petites couches d’adaptateur, réduisant considérablement les exigences de mémoire et le temps d’entraînement.
Liger Kernels: noyaux optimisés pour GPU qui réduisent l'utilisation de la mémoire jusqu'à 80 % grâce à des opérations fusionnées.
TRL: Une bibliothèque pour l'entraînement de modèles de langage avec l'apprentissage par renforcement et l'affinement supervisé
Compute GPU Serverless: compute géré par Databricks qui gère la montée en charge automatique des ressources GPU

Matrice de décision LoRA ou affinement complet

LoRA (Low-Rank Adaptation) gèle le modèle de base et entraîne uniquement de petites couches d'adaptateur, réduisant les parameters entraînables d'environ 99 %. Cela rend l'entraînement plus rapide et plus économe en mémoire.

Scénario	Recommandation	Raison
Mémoire GPU limitée	LoRA	Adapte des modèles plus grands en mémoire en entraînant seulement 1 % des paramètres
Adaptation spécifique aux tâches	LoRA	Échanger différents adaptateurs sur le même modèle de base pour plusieurs tâches
Changement majeur de comportement du modèle	Affinement complet	Met à jour tous les paramètres pour les modifications fondamentales du comportement du modèle
Déploiement en production	LoRA	Fichiers plus petits (Mo vs Go), chargement plus rapide, contrôle de version plus facile

Scénario	Recommandation	Raison
Mémoire GPU limitée	LoRA	Adapte des modèles plus grands en mémoire en entraînant seulement 1 % des paramètres
Adaptation spécifique aux tâches	LoRA	Échanger différents adaptateurs sur le même modèle de base pour plusieurs tâches
Changement majeur de comportement du modèle	Affinement complet	Met à jour tous les paramètres pour les modifications fondamentales du comportement du modèle
Déploiement en production	LoRA	Fichiers plus petits (Mo vs Go), chargement plus rapide, contrôle de version plus facile

Avantages du noyau Liger

Liger Kernels sont des operations optimisées pour le GPU qui fusionnent plusieurs étapes en des noyaux uniques, réduisant les transferts de mémoire et améliorant l'efficacité. Le document technique fournit des benchmarks détaillés montrant des améliorations significatives des performances.

Opérations fusionnées : combine les opérations (par exemple, linéaire + perte) pour réduire la surcharge de mémoire jusqu'à 80 %
Noyaux Triton : Noyaux GPU personnalisés optimisés pour les opérations de transformateur (RMSNorm, RoPE, SwiGLU, CrossEntropy)
Efficacité de la mémoire : Permet des tailles de batch plus importantes ou des modèles qui ne tiendraient pas autrement dans la mémoire GPU
Optimisation d'un seul GPU : particulièrement efficace pour les scénarios d'entraînement sur GPU unique A10/A100

Ce Notebook utilise la bibliothèque TRL pour simplifier la configuration d'entraînement et appliquer automatiquement ces optimisations.

Se connecter au compute GPU serverless

Ce notebook nécessite un compute GPU Serverless. Connecter :

Cliquez sur le sélecteur de compute du notebook en haut à droite et sélectionnez Serverless GPU .
Sur le côté droit, cliquez sur le bouton d'environnement.
Sélectionnez 8xH100 comme Accélérateur
Choisissez AI v5 dans l'environnement de base
Cliquez sur Appliquer

La fonction d'entraînement provisionnera automatiquement 8 GPU H100 pour l'entraînement distribué.

Installez les bibliothèques requises

L'environnement Databricks AI v5 comprend déjà la plupart des bibliothèques requises pour cet exemple (telles que trl, peft, transformers, hf_transfer et mlflow). La cellule suivante installe liger-kernel, qui ne fait pas encore partie de l'environnement.

La commande %restart_python redémarre l'interpréteur Python pour s'assurer que le nouveau package installé est correctement chargé.

Python
%pip install liger-kernel==0.8.0
%restart_python

Configuration de l'installation

Intégration de Unity Catalog

La cellule suivante configure l'emplacement où votre modèle ajusté sera stocké et enregistré :

**Catalogue et Schéma** : Organisez les modèles dans votre espace de noms Unity Catalog (par défautmain.default:)
Nom du modèle : nom du modèle enregistré dans Unity Catalog pour la gouvernance et le déploiement.
**Volume** : volume Unity Catalog pour le stockage des points de contrôle du modèle pendant l'entraînement

Ces widgets vous permettent de personnaliser l'emplacement de stockage sans modifier le code. Le modèle sera enregistré en tant que {catalog}.{schema}.{model_name} pour un accès facile et le contrôle de version.

Hyperparamètres d'entraînement

La cellule définit également les principaux paramètres de formation :

Modèle et Dataset : Qwen2-0.5B avec dataset conversationnel Capybara
**Taille de batch (8)** : Nombre d'exemples par GPU et par étape d'entraînement.
Accumulation de gradient (4) : Accumule les gradients sur 4 batches pour une taille de batch effective de 32
Taux d'apprentissage (1e-4) : Taux conservateur, mis à l'échelle automatiquement 10 fois plus élevé pour l'entraînement LoRA
**Époques (1)** : passage unique à travers le dataset pour éviter le surapprentissage
Logging & Checkpointing : enregistre la progression toutes les 100 étapes, logs les métriques toutes les 25 étapes

Python
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

Configuration LoRA

La cellule suivante configure les paramètres LoRA (adaptation de bas rang) qui contrôlent la façon dont le modèle est affiné. LoRA gèle les poids du modèle de base et entraîne uniquement de petites matrices d’adaptateurs, réduisant considérablement les besoins en mémoire.

Sélection de paramètres

Classement (r=8) : Offre un bon équilibre entre les performances et les parameters
**Alpha (32)** : facteur d'échelle, généralement 2 à 4 fois le rang
Dropout (0,1) : Régularisation pour prévenir le surajustement

Modules cibles pour Qwen2

Cet exemple cible toutes les couches de transformations clés :

Attention : q_proj, k_proj, v_proj, o_proj
**MLP** gate_proj``up_proj:,, down_proj

Python
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"
]

Définissez la fonction d’entraînement

La cellule suivante crée la fonction d'entraînement distribué qui s'exécutera sur plusieurs GPU. Voici ce qu’il fait :

Configuration de la formation distribuée

Le décorateur @distributed configure le compute GPU Serverless :

8 GPU : Distribue l'entraînement sur 8 GPU H100 pour un entraînement plus rapide.
**Orchestration automatique** : Gère le provisionnement des GPU, la distribution des données et la synchronisation.

Workflow d'entraînement

La fonction exécute les étapes suivantes :

Charger le jeu de données : download et prépare le jeu de données conversationnel Capybara
Initialiser le modèle : charge Qwen2-0.5B et le tokenizer avec le formatage du chat
**Appliquer LoRA **: Attache des couches d'adaptateurs pour réduire les paramètres entraînables d'environ 99 %
Configurer l'entraînement : Définit la taille du batch, le taux d'apprentissage et les optimisations du noyau Liger.
Entraîner le modèle : Exécute la boucle d’entraînement avec des points de contrôle et une journalisation automatiques
Enregistrer les artefacts : Stocke les adaptateurs LoRA et le tokenizer dans le volume Unity Catalog
ID d'exécution MLflow renvoyé : Fournit l'ID d'exécution pour l'enregistrement du modèle.

Optimisations clés activées

Liger Kernels : les opérations GPU fusionnées réduisent l'utilisation de la mémoire jusqu'à 80 %.
**Précision mixte (FP16)** : Calcul plus rapide avec une empreinte mémoire plus faible.
Point de contrôle de gradient : Échange le calcul contre la mémoire pour s’adapter à des batchs plus grands
Accumulation de gradient : simule des tailles de batch plus grandes pour un entraînement stable

Python
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

Exécuter la formation distribuée

Cette cellule exécute la fonction d'entraînement sur 8 GPU H100. La méthode distributed() gère :

Provisionnement de ressources Serverless GPU Compute
Distribution de la charge de travail d'entraînement sur plusieurs GPU
Collecte de l'ID d'exécution MLflow pour l'enregistrement du modèle.

L'entraînement prend généralement de 15 à 30 minutes selon la taille du dataset et la disponibilité du compute.

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

Enregistrement de MLflow et Unity Catalog

Stratégie d’enregistrement du modèle

MLflow Tracking : Enregistrer les artefacts et les métadonnées du modèle
Unity Catalog : Enregistrer le modèle pour la gouvernance et le déploiement
Gestion des versions de modèle : Versioning automatique pour la gestion du cycle de vie du modèle
Métadonnées : Informations complètes sur le modèle pour la reproductibilité

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

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

try:
    # Load the trained model for registration
    print("Loading LoRA model for registration...")
    # For LoRA models, we need both base model and adapter
    model_type = "LoRA"
    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={
                &quot;task&quot;: task,
                &quot;pretrained_model_name&quot;: MODEL_NAME,
                &quot;databricks_model_family&quot;: &quot;QwenForCausalLM&quot;,
            },
        )

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

Étapes suivantes

Maintenant que vous avez affiné et enregistré votre modèle, vous pouvez :

Déployer le modèle : Diffuser les modèles avec Model Serving
En savoir plus sur la formation distribuée : Formation distribuée multi-GPU
Optimiser l'utilisation du compute GPU Serverless : Bonnes pratiques pour le compute GPU Serverless
Résolution des problèmes : Résoudre les problèmes sur Serverless GPU Compute

Exemple de Notebook

Affinement distribué de Qwen2-0.5B avec LoRA

Ouvrir le notebook dans un nouvel onglet Ouvrir dans Databricks

Matrice de décision LoRA ou affinement complet​

Avantages du noyau Liger​

Se connecter au compute GPU serverless​

Installez les bibliothèques requises​

Configuration de l'installation​

Intégration de Unity Catalog​

Hyperparamètres d'entraînement​

Configuration LoRA​

Sélection de paramètres​

Modules cibles pour Qwen2​

Définissez la fonction d’entraînement​

Configuration de la formation distribuée​

Workflow d'entraînement​

Optimisations clés activées​

Exécuter la formation distribuée​

Enregistrement de MLflow et Unity Catalog​

Stratégie d’enregistrement du modèle​

Étapes suivantes​

Exemple de Notebook​

Affinement distribué de Qwen2-0.5B avec LoRA

Matrice de décision LoRA ou affinement complet

Avantages du noyau Liger

Se connecter au compute GPU serverless

Installez les bibliothèques requises

Configuration de l'installation

Intégration de Unity Catalog

Hyperparamètres d'entraînement

Configuration LoRA

Sélection de paramètres

Modules cibles pour Qwen2

Définissez la fonction d’entraînement

Configuration de la formation distribuée

Workflow d'entraînement

Optimisations clés activées

Exécuter la formation distribuée

Enregistrement de MLflow et Unity Catalog

Stratégie d’enregistrement du modèle

Étapes suivantes

Exemple de Notebook