Distribuerad finjustering av Qwen2-0.5B med LoRA

Den här notebook-filen visar hur du effektivt finjusterar den stora språkmodellen Qwen2-0.5B med hjälp av parametereffektiva tekniker på serverlös GPU-beräkning. Du får lära dig att:

• Använd LoRA (Low-Rank Adaptation) för att minska träningsbara parametrar med ~99% samtidigt som modellkvaliteten bibehålls
• Använda Liger-kärnor för minneseffektiv träning med optimerade Triton-kernels
• Utnyttja TRL (Transformer Reinforcement Learning) för övervakad finjustering
• Registrera den finjusterade modellen i Unity Catalog för styrning och distribution

Viktiga begrepp:

• LoRA: En teknik som fryser basmodellen och tränar små adapterlager, vilket dramatiskt minskar minneskraven och träningstiden
• Liger-kärnor: GPU-optimerade kärnor som minskar minnesanvändningen med upp till 80% genom sammansvetsade åtgärder
• TRL: Ett bibliotek för att träna språkmodeller med förstärkningsinlärning och övervakad finjustering
• Serverlös GPU-beräkning: Databricks-hanterad beräkning som automatiskt skalar GPU-resurser

LoRA jämfört med fullständig beslutsmatris för finjustering

LoRA (Low-Rank Adaptation) fryser basmodellen och tränar endast små adapterskikt, vilket minskar träningsbara parametrar med ~99%. Detta gör träningen snabbare och mer minneseffektiv.

Fördelar med Liger Kernel

Liger Kernels är GPU-optimerade åtgärder som kombinerar flera steg i enskilda kärnor, vilket minskar minnesöverföringar och förbättrar effektiviteten. Det tekniska dokumentet innehåller detaljerade riktmärken som visar betydande prestandaförbättringar.

• Sammansvetsade åtgärder: Kombinerar åtgärder (t.ex. linjär + förlust) för att minska minneskostnaderna med upp till 80%
• Tritonkärnor: Anpassade GPU-kernels optimerade för transformeringsåtgärder (RMSNorm, RoPE, SwiGLU, CrossEntropy)
• Minneseffektivitet: Tillåter större batchstorlekar eller modeller som annars inte skulle få plats i GPU-minnet
• Enkel GPU-optimering: Särskilt effektivt för träningsscenarier med en enda GPU för A10/A100

Den här notebook-filen använder TRL-biblioteket för att förenkla träningskonfigurationen och automatiskt tillämpa dessa optimeringar.

Ansluta till serverlös GPU-beräkning

Den här anteckningsboken kräver beräkningsresurser med serverlös GPU. Så här ansluter du:

1. Klicka på notebook-filens beräkningsväljare längst upp till höger och välj Serverlös GPU
2. Klicka på miljöknappen till höger
3. Välj 8xH100 som accelerator
4. Välj AI v4 från basmiljön
5. Klicka på Använd

Träningsfunktionen etablerar automatiskt 8 H100 GPU:er för distribuerad träning.

Installera nödvändiga bibliotek

Nästa cell installerar de Python paket som behövs för distribuerad finjustering:

Grundläggande utbildningsbibliotek:

• trl: Transformer Reinforcement Learning-bibliotek för övervakad finjustering och RLHF
• peft: Parameter-Efficient Fine-Tuning bibliotek som tillhandahåller LoRA-implementering
• liger-kernel: Minnesoptimerade GPU-kärnor för effektiv transformeringsträning

Bibliotek för stöd:

• hf_transfer: Accelererade nedladdningar från Hugging Face Hub med hjälp av Rust-baserad överföring
• mlflow>=3.0: Experimentspårning och modellregisterintegrering

Kommandot %restart_python startar om Python tolken för att säkerställa att de nyligen installerade paketen läses in korrekt.

%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

Konfigurationsinställning

Integration av Unity-katalog

Nästa cell konfigurerar var din finjusterade modell ska lagras och registreras:

• Katalog och schema: Ordna modeller i unity-katalogens namnområde (standard: main.default)
• Modellnamn: Det registrerade modellnamnet i Unity Catalog för styrning och distribution
• Volym: Unity Catalog-volym för lagring av modellkontrollpunkter under träning

Med de här widgetarna kan du anpassa lagringsplatsen utan att redigera kod. Modellen registreras som {catalog}.{schema}.{model_name} för enkel åtkomst och versionskontroll.

Träna hyperparametrar

Cellen definierar även nyckelträningsparametrar:

• Modell och datauppsättning: Qwen2-0.5B med Capybara-konversationsdatauppsättning
• Batchstorlek (8): Antal exempel per GPU vid varje träningssteg
• Gradientansamling (4): Ackumulerar gradienter över 4 omgångar för en effektiv omgångsstorlek på 32
• Inlärningsfrekvens (1e-4): Konservativ hastighet, automatiskt skalad 10x högre för LoRA-utbildning
• Epoker (1): Enkel genomgång av datauppsättningen för att undvika överanpassning
• Loggning och kontrollpunkter: Sparar förlopp var 100:e steg, loggar mått var 25:e steg

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

Nästa cell konfigurerar LoRA-parametrar (Low-Rank Anpassning) som styr hur modellen finjusteras. LoRA fryser basmodellens vikter och tränar endast små adaptermatriser, vilket dramatiskt minskar minneskraven.

Parameterval

• Rangordning (r=8): Ger bra balans mellan prestanda och parametrar
• Alfa (32): Skalningsfaktor, vanligtvis 2–4 gånger rangordningen
• Dropout (0.1): Regularisering för att förhindra överträning

Målmoduler för Qwen2

Det här exemplet riktar sig till alla viktiga transformeringslager:

• Uppmärksamhet: 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"
]

Definiera träningsfunktionen

Nästa cell skapar den distribuerade träningsfunktionen som körs över flera GPU:er. Så här gör den:

Konfiguration av distribuerad utbildning

Dekoratören @distributed konfigurerar serverlös GPU-beräkning:

• 8 GPU:er: Distribuerar träning över 8 H100 GPU:er för snabbare träning
• Automatisk orkestrering: Hanterar GPU-etablering, datadistribution och synkronisering

Arbetsflöde för träning

Funktionen kör följande steg:

1. Ladda datamängden: Laddar ned och förbereder Capybara-konversationsdatasetet
2. Initiera modellen: Läser in Qwen2-0.5B och tokenizer med chattformatering
3. Använd LoRA: Fäster adapterskikt för att minska träningsbara parametrar med ~99%
4. Konfigurera träning: Konfigurerar batchstorlek, inlärningshastighet och optimering av Liger-kernel
5. Träningsmodell: Kör träningsloopen med automatisk kontrollpunkt och loggning
6. Spara artefakter: Lagrar LoRA-adaptorer och tokenizer på Unity Catalog-volymen
7. Returnera MLflow-kör-ID: Tillhandahåller kör-ID för modellregistrering

Nyckeloptimeringar aktiverade

• Liger Kernels: Sammansvetsade GPU-åtgärder minskar minnesanvändningen med upp till 80%
• Blandad precision (FP16): Snabbare beräkning med lägre minnesfotavtryck
• Gradient checkpointing: Byter beräkning mot minne för att få plats med större satser
• Gradientaccumulering: Simulerar större batchstorlekar för stabil träning

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

Kör den distribuerade träningen

Den här cellen kör träningsfunktionen på 8 H100 GPU:er. Metoden distributed() hanterar:

• Etablera serverlösa GPU-beräkningsresurser
• Distribuera träningsarbetsbelastningen över flera GPU:er
• Samla in MLflow-körnings-ID för modellregistrering

Träningen tar vanligtvis 15–30 minuter beroende på datamängdens storlek och beräkningstillgänglighet.

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

Registrering av MLflow- och Unity-katalog

Strategi för modellregistrering

• MLflow-spårning: Loggmodellartefakter och metadata
• Unity Catalog: Registrera modell för styrning och distribution
• Modellversionshantering: Automatisk versionshantering för modelllivscykelhantering
• Metadata: Fullständig modellinformation för reproducerbarhet

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ästa steg

Nu när du har finjusterat och registrerat din modell kan du:

• Distribuera modellen: Hantera modeller med modellservering
• Läs mer om distribuerad utbildning: Distribuerad träning med flera GPU:er
• Optimera serverlös GPU-beräkningsanvändning: Metodtips för serverlös GPU-beräkning
• Felsöka problem: Felsöka problem med serverlös GPU-beräkning

Exempelanteckningsbok

Distribuerad finjustering av Qwen2-0.5B med LoRA

Hämta anteckningsbok