Ottimizzazione distribuita di Qwen2-0.5B con LoRA

Questo notebook illustra come affinare in modo efficiente il grande modello linguistico Qwen2-0.5B utilizzando tecniche di ottimizzazione dei parametri efficienti nel Serverless GPU Compute. Si apprenderà come:

• Applicare LoRA (Adattamento a Basso Rango) per ridurre i parametri addestrabili di circa il 99% mantenendo la qualità del modello.
• Utilizzare i kernel Liger per un training efficiente in termini di memoria con kernel Triton ottimizzati.
• Sfruttare il TRL (Transformer Reinforcement Learning) per l'ottimizzazione con supervisione
• Registrare il modello ottimizzato in Unity Catalog per la governance e la distribuzione

Concetti chiave:

• LoRA: una tecnica che blocca il modello di base e esegue il training di livelli di adattatori di piccole dimensioni, riducendo notevolmente i requisiti di memoria e il tempo di training
• Kernel Liger: kernel ottimizzati per GPU che riducono l'utilizzo della memoria fino all'80% tramite operazioni di fusione
• TRL: una libreria per l'addestramento di modelli linguistici con apprendimento per rinforzo e affinamento supervisionato
• Calcolo GPU serverless: calcolo gestito di Databricks che ridimensiona automaticamente le risorse GPU

LoRA e matrice di decisioni di ottimizzazione completa

LoRA (Low-Rank Adaptation) blocca il modello di base ed esegue il training solo di piccoli livelli adattatori, riducendo i parametri sottoponibili a training di circa 99%. Questo rende l'addestramento più veloce e più efficiente dal punto di vista della memoria.

Vantaggi del kernel Liger

I kernel Liger sono operazioni ottimizzate per GPU che combinano più passaggi in singoli kernel, riducendo i trasferimenti di memoria e migliorando l'efficienza. Il documento tecnico fornisce benchmark dettagliati che mostrano miglioramenti significativi delle prestazioni.

• Operazioni fuse: combina le operazioni (ad esempio, lineare + perdita) per ridurre il sovraccarico di memoria fino a 80%
• Kernels Triton: Kernel GPU personalizzati di Triton ottimizzati per operazioni di trasformazione (RMSNorm, RoPE, SwiGLU, CrossEntropy)
• Efficienza della memoria: consente dimensioni di batch o modelli di dimensioni maggiori che altrimenti non rientrano nella memoria GPU
• Ottimizzazione gpu singola: particolarmente efficace per scenari di training A10/A100 a GPU singola

Questo notebook usa la libreria TRL per semplificare la configurazione di training e applicare automaticamente queste ottimizzazioni.

Connettersi al calcolo GPU serverless

Questo notebook richiede l'ambiente di calcolo GPU serverless. Per connettersi:

1. Fare clic sul selettore di calcolo del notebook in alto a destra e selezionare GPU serverless
2. Sul lato destro fare clic sul pulsante ambiente
3. Selezionare 8xH100 come acceleratore
4. Scegliere intelligenza artificiale v4 dall'ambiente di base
5. Fare clic su Applica

La funzione di training eseguirà automaticamente il provisioning di 8 GPU H100 per il training distribuito.

Installare le librerie necessarie

La cella successiva installa i pacchetti Python necessari per l'ottimizzazione distribuita:

Librerie di training di base:

• trl: libreria di apprendimento per rinforzo basato su transformer per ottimizzazione supervisionata e RLHF
• peft: libreria Parameter-Efficient Fine-Tuning che fornisce l'implementazione LoRA
• liger-kernel: kernel GPU ottimizzati per l'uso della memoria per un allenamento efficiente dei trasformatori

Librerie di supporto:

• hf_transfer: download accelerati da Hugging Face Hub tramite trasferimento basato su Rust
• mlflow>=3.0: Monitoraggio degli esperimenti e integrazione del registro dei modelli

Il comando %restart_python riavvia l'interprete Python per assicurarsi che i pacchetti appena installati vengano caricati correttamente.

%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

Impostazione della configurazione

Integrazione del catalogo Unity

La cella successiva configura la posizione in cui verrà archiviato e registrato il modello ottimizzato:

• Catalogo e schema: organizzare i modelli all'interno dello spazio dei nomi del catalogo Unity (impostazione predefinita: main.default)
• Nome modello: nome del modello registrato nel catalogo Unity per la governance e la distribuzione
• Volume: Volume di Unity Catalog per l'archiviazione dei checkpoint del modello durante l'addestramento

Questi widget consentono di personalizzare la posizione di archiviazione senza modificare il codice. Il modello verrà registrato come {catalog}.{schema}.{model_name} per semplificare l'accesso e il controllo della versione.

Addestramento degli iperparametri

La cella definisce anche i parametri di training chiave:

• Modello e set di dati: Qwen2-0.5B con set di dati conversazionale Capybara
• Dimensioni batch (8): Numero di esempi per GPU per passaggio di training
• Accumulo dei gradienti (4): accumula gradienti su 4 batch per una dimensione effettiva del batch di 32
• Tasso di apprendimento (1e-4): tasso conservativo, automaticamente aumentato di 10 volte per la formazione LoRA
• Periodi (1): passaggio singolo del set di dati per impedire l'overfitting
• Registrazione e checkpoint: salva lo stato di avanzamento ogni 100 passaggi, registra le metriche ogni 25 passaggi

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

Configurazione di LoRA

La cella successiva configura i parametri LoRA (Low-Rank Adaptation) che controllano la modalità di ottimizzazione del modello. LoRA congela i pesi del modello di base ed esegue l'addestramento solo su piccole matrici adattatrici, riducendo notevolmente i requisiti di memoria.

Selezione dei parametri

• Rank (r=8): fornisce un buon equilibrio tra prestazioni e parametri
• Alfa (32): fattore di ridimensionamento, in genere 2-4 volte il rango
• Dropout (0.1): regolarizzazione per impedire l'overfitting

Moduli di destinazione per Qwen2

Questo esempio è destinato a tutti i livelli di trasformazione chiave:

• Attenzione: q_proj, k_proj, v_proj, o_proj
• MLP: gate_proj, up_proj, down_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"
]

Definire la funzione di training

La cella successiva crea la funzione di training distribuita che verrà eseguita tra più GPU. Ecco cosa fa:

Configurazione del training distribuito

Il decorator @distributed configura il calcolo GPU serverless.

• 8 GPU: distribuisce il training tra 8 GPU H100 per un training più rapido
• Orchestrazione automatica: Gestisce il provisioning delle GPU, la distribuzione dei dati e la sincronizzazione.

Flusso di lavoro di training

La funzione esegue questi passaggi:

1. Caricare il set di dati: scarica e prepara il set di dati conversazionale Capybara
2. Inizializza il modello: Carica Qwen2-0.5B e tokenizer con formattazione della chat
3. Applica LoRA: collega i livelli dell'adattatore per ridurre i parametri sottoponibili a training di circa 99%
4. Configurare il training: configura le dimensioni del batch, il tasso di apprendimento e le ottimizzazioni del kernel Liger
5. Modello di training: esegue il ciclo di training con checkpoint e registrazione automatici
6. Salva artefatti: archivia gli adattatori LoRA e il tokenizer nel volume del catalogo Unity
7. Restituisce l'ID di esecuzione MLflow: fornisce l'ID di esecuzione per la registrazione del modello

Ottimizzazioni principali abilitate

• Liger Kernels: operazioni GPU fuse riducono l'uso della memoria fino al 80%
• Precisione mista (FP16): calcolo più veloce con impronta di memoria inferiore
• Checkpointing di gradiente: Scambia calcolo per memoria per adattarsi a batch più grandi
• Accumulo dei gradienti: simula dimensioni batch maggiori per il training stabile

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

Eseguire la formazione distribuita

Questa cella esegue la funzione di training su 8 GPU H100. Il distributed() metodo gestisce:

• Calcolo Serverless con GPU
• Distribuzione del carico di lavoro di training tra più GPU
• Raccolta dell'ID di esecuzione MLflow per la registrazione del modello

Il training richiede in genere 15-30 minuti a seconda delle dimensioni del set di dati e della disponibilità di calcolo.

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

Registrazione del catalogo MLflow e Unity

Strategia per la registrazione dei modelli

• MLflow Tracking: registra artefatti e metadati del modello
• Unity Catalog: Registra il modello per la governance e la distribuzione
• Controllo delle versioni del modello: controllo delle versioni automatico per la gestione del ciclo di vita del modello
• Metadati: informazioni complete sul modello per la riproducibilità

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

Passaggi successivi

Dopo aver ottimizzato e registrato il modello, è possibile:

• Distribuire il modello: Gestire i modelli con Model Serving
• Altre informazioni sul training distribuito: Training distribuito multi-GPU
• Ottimizzare l'utilizzo del calcolo GPU serverless: Migliori pratiche per il calcolo GPU serverless
• Risolvere i problemi: Risolvere i problemi relativi all'ambiente di calcolo GPU serverless

Notebook di esempio

Ottimizzazione distribuita di Qwen2-0.5B con LoRA

Ottieni il notebook