LoRA による Qwen2-0.5B の分散ファインチューニング

このノートブックでは、サーバーレス GPU コンピュートでの最も効率的な手法を使用して、Qwen2-0.5B 大規模言語モデルを効率的に微調整する方法を示します。あなたは以下のことを学びます：

LoRA (低ランク適応) を適用して、モデルの品質を維持しながらトレーニング可能な論点を最大 99% 削減します
最適化されたTritonカーネルを使用して、メモリ効率の良いトレーニングを行うには、 Ligerカーネルを 使用してください。
TRL（Transformer Reinforcement Learning） を活用して、教師ありファインチューニングを行う
ガバナンスとデプロイメントのための微調整されたモデルをUnity Catalogに登録する

主要概念：

LoRA : 基本モデルとトレーニングする小さなアダプター層をフリーズし、メモリ要件とトレーニング時間を大幅に削減する技術
Liger Kernels ：GPUに最適化されたカーネルで、融合演算によりメモリ使用量を最大80%削減します。
TRL ：強化学習と教師ありファインチューニングによる言語モデルのトレーニング用ライブラリ
AI Runtime : GPU リソースを自動的にスケールするDatabricksマネージドコンピュート

LoRAとフルファインチューニングの決定マトリックス

LoRA (Low-Rank Adaptation)は、ベースモデルをフリーズし、小さなアダプター層のみをトレーニングすることで、トレーニング可能な懸念を最大 99% 削減します。これにより、トレーニングが高速化され、メモリ効率も向上します。

シナリオ	推奨事項	理由:
GPUメモリが制限されている	ロラ	わずか 1% のトレーニングにより、より大きなモデルをメモリに適合させる
タスク固有の適応	ロラ	複数のタスクに対して同じ基本モデル上の異なるアダプタをスワップする
主要なモデル動作の変更	フルファインチューニング	モデルの動作に対する根本的な変更をすべて更新します
本番運用展開	ロラ	ファイルサイズが小さくなり（MB単位対GB単位）、読み込み速度が速くなり、バージョン管理が容易になる。

ライガーカーネルの効能

Ligerカーネルは、複数の処理ステップを単一のカーネルに統合することでメモリ転送を削減し、効率を向上させる、GPUに最適化された演算です。この技術論文では、大幅な性能向上を示す詳細なベンチマークデータが掲載されている。

融合演算 ：演算（例：線形演算＋損失演算）を組み合わせることで、メモリオーバーヘッドを最大80%削減します。
Tritonカーネル ：トランスフォーマー演算（RMSNorm、RoPE、SwiGLU、CrossEntropy）に最適化されたカスタムGPUカーネル
メモリ効率 ：GPUメモリに収まらないような大きなバッチサイズやモデルの処理を可能にする
シングル GPU の最適化 : A10/A100 シングル GPU トレーニングシナリオに特に効果的

このノートブックはTRL ライブラリを使用してトレーニング構成を簡素化し、これらの最適化を自動的に適用します。

サーバレスGPUコンピュートに接続する

このノートブックは分散トレーニングを実行するためにサーバレス GPU コンピュートを必要とします。サーバレス GPU コンピュートは自動的にプロビジョニングを行い、ワークロードの GPU リソースを管理します。

接続するには、ノートブックの [接続] ドロップダウンメニューをクリックし、 [サーバレス GPU] を選択します。

詳細については、 AI Runtimeドキュメントを参照してください。

必要なライブラリをインストールします

次のセルでは、分散ファインチューニングに必要なPythonパッケージをインストールします。

コアトレーニングライブラリ：

trl==0.18.1 :教師ありファインチューニングとRLHFのためのTransformer強化学習ライブラリ
peft : LoRA 実装を提供する、欠点-Efficient ファインチューニングライブラリ
liger-kernel : トランスフォーマーの効率的なトレーニングのためのメモリ最適化された GPU カーネル

サポートライブラリ:

hf_transfer : Rustベースの転送を使用してHugging Face Hubからのダウンロードを高速化
mlflow>=3.6.0 :エクスペリメント追跡とモデルレジストリの統合

%restart_pythonコマンドは、新しくインストールされたパッケージが正しくロードされるように、Python インタープリタを再起動します。

Python
%pip install trl==0.18.1 peft hf_transfer  liger-kernel
%pip install mlflow>=3.6.0
%restart_python

設定構成

Unity Catalogの連携

次のセルでは、微調整したモデルの保存場所と登録場所を設定します。

カタログとスキーマ : Unity Catalog名前空間内でモデルを整理します (確実: main.default )
モデル名 ：ガバナンスとデプロイメントのためにUnity Catalogに登録されたモデル名
Volume : トレーニング中にモデルのチェックポイントを保存するためのUnity Catalogボリューム

これらのウィジェットを使用すると、コードを編集することなく保存場所をカスタマイズできます。このモデルは、アクセスとバージョン管理を容易にするため、 {catalog}.{schema}.{model_name}として登録されます。

ハイパーパラメータのトレーニング

このセルでは、主要なトレーニングの問題も定義します。

モデルとデータセット : Qwen2-0.5B (Capybara 会話データセット付き)
バッチサイズ（8） ：トレーニングステップごとにGPUあたりで処理するサンプル数
グラジエント累積 (4) : 有効バッチサイズ 32 の場合、4 バッチにわたるグラジエントを累積します。
学習率 (1e-4) : 控えめな率、LoRA トレーニング用に自動的に 10 倍に調整されます。
エポック（1） ：データセットを1回通過させて過学習を防ぐ
ロギングとチェックポイント : 100 ステップごとに進行状況を保存し、25 ステップごとにメトリクスを記録します

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

LoRA設定

次のセルは、モデルの微調整方法を制御する LoRA (Low-Rank Adaptation) を構成します。 LoRA は、ベースモデルの重みを固定し、小さなアダプターマトリックスのみをトレーニングすることで、メモリ要件を大幅に削減します。

争点選択

ランク (r=8) : パフォーマンスとデメリットのバランスが取れています。
アルファ（32） ：スケーリング係数。通常はランクの2～4倍。
ドロップアウト（0.1） ：過学習を防ぐための正則化

Qwen2のターゲットモジュール

この例では、すべての主要な変換レイヤーを対象としています。

注意 : 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"
]

トレーニング関数を定義する

次のセルでは、複数のGPUにまたがって実行される分散トレーニング関数を作成します。その機能は以下のとおりです。

分散型トレーニングの設定

@distributedデコレータは、サーバレス GPU コンピュートを構成します。

8つのGPU ：トレーニングを8つのH100 GPUに分散することで、トレーニングを高速化します。
自動オーケストレーション : GPU プロビジョニング、データ配布、同期を処理します。

トレーニングワークフロー

この関数は次のステップを実行します。

データセットの読み込み ：カピバラの会話データセットをダウンロードして準備します。
モデルの初期化 ：Qwen2-0.5Bとチャットフォーマット付きトークナイザーをロードします
LoRA を適用 : アダプター層を接続して、トレーニング可能量を ~99% 削減します
トレーニングの設定 ：バッチサイズ、学習率、Ligerカーネルの最適化を設定します。
トレーニングするモデル : 自動チェックポイントとロギングを使用してトレーニングループを実行します。
アーティファクトの保存 : LoRA アダプターとトークナイザーをUnity Catalogボリュームに保存します
Return MLflow実行 ID : モデル登録の実行 ID を提供します。

主要な最適化が有効になりました

Liger Kernels ：GPU演算の融合によりメモリ使用量を最大80%削減
混合精度（FP16） ：メモリ使用量を抑えつつ、より高速な計算が可能
勾配チェックポイント ：計算処理とメモリ使用量をトレードオフして、より大きなバッチサイズに対応します。
勾配蓄積 ：安定したトレーニングのために、より大きなバッチサイズをシミュレートします。

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

分散トレーニングを実行する

このセルは、8基のH100 GPU上でトレーニング機能を実行します。distributed()メソッドは以下を処理します。

プロビジョニングサーバレス GPU コンピュートリソース
トレーニングワークロードを複数のGPUに分散させる
モデル登録用のMLflow実行IDの収集

データセットのサイズとコンピュートの可用性に応じて、トレーニングには通常 15 ～ 30 分かかります。

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

MLflowとUnity Catalog登録

モデル登録戦略

MLflow Tracking : 記録済みモデルのアーティファクトとメタデータ
Unity Catalog : ガバナンスとデプロイメントのための登録するモデル
モデルのバージョン管理 ：モデルライフサイクル管理のための自動バージョン管理
メタデータ ：再現性を確保するための完全なモデル情報

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

次のステップ

モデルの微調整と登録が完了したら、次の操作が可能になります。

モデルのデプロイ ：モデルサービングを使用してモデルを提供する
分散トレーニングの詳細については 、マルチGPU分散トレーニングをご覧ください。
サーバレス GPU の使用を最適化する :サーバレス GPU コンピュートのベストプラクティス
問題のトラブルシューティング :サーバレス GPU コンピュートの問題のトラブルシューティング

サンプルノートブック

LoRA による Qwen2-0.5B の分散ファインチューニング

ノートブックを新しいタブで開く

LoRAとフルファインチューニングの決定マトリックス​

ライガーカーネルの効能​

サーバレスGPUコンピュートに接続する​

必要なライブラリをインストールします​

設定構成​

Unity Catalogの連携​

ハイパーパラメータのトレーニング​

LoRA設定​

争点選択​

Qwen2のターゲットモジュール​

トレーニング関数を定義する​

分散型トレーニングの設定​

トレーニングワークフロー​

主要な最適化が有効になりました​

分散トレーニングを実行する​

MLflowとUnity Catalog登録​

モデル登録戦略​

次のステップ​

サンプルノートブック​

LoRA による Qwen2-0.5B の分散ファインチューニング

LoRAとフルファインチューニングの決定マトリックス

ライガーカーネルの効能

サーバレスGPUコンピュートに接続する

必要なライブラリをインストールします

設定構成

Unity Catalogの連携

ハイパーパラメータのトレーニング

LoRA設定

争点選択

Qwen2のターゲットモジュール

トレーニング関数を定義する

分散型トレーニングの設定

トレーニングワークフロー

主要な最適化が有効になりました

分散トレーニングを実行する

MLflowとUnity Catalog登録

モデル登録戦略

次のステップ

サンプルノートブック