サーバーレス GPU を使用して LoRA でLlama 3.2 1B を微調整する
このノートブックでは、 Databricksサーバレス GPU で低ランク適応 (LoRA) を備えた教師付きファインチューニング (SFT) を使用して大規模言語モデルを微調整する方法を示します。 ノートブックは、DeepSpeed ZeRO Stage 3 最適化を備えた Transformers Reinforcement Learning (TRL) ライブラリを使用して、4 つの A10 GPU を備えた単一ノードでLlama 3.2 1B を効率的にトレーニングします。
重要な概念:
- LoRA (Low-Rank Adaptation) : 小規模でトレーニング可能なランク分解行列をモデル レイヤーに追加することで、トレーニング可能なタンクの数を減らす、効率的なファインチューニング手法。
- TRL (Transformers Reinforcement Learning) : 強化学習と教師付きファインチューニングを使用して言語モデルをトレーニングするためのツールを提供するライブラリ。
- DeepSpeed ZeRO Stage 3 : モデルの欠点、勾配、およびオプティマイザーの状態を GPU 全体に分割して、大規模なモデルのトレーニングを可能にするメモリ最適化手法。
- サーバーレス GPU : Databricksが管理する GPU コンピュート。トレーニング ワークロードのプロビジョニングと GPU リソースのスケーリングを自動的に行います。
詳細については、 「サーバレス GPU コンピュート」を参照してください。
要件
このノートブックには次のものが必要です。
- サーバレス GPU コンピュート : ノートブックは、分散トレーニング用に 4 つの A10 GPU を備えたDatabricksサーバレス GPU を使用します。 クラスター構成は必要ありません。
- Unity Catalog : モデルのチェックポイントを保存し、トレーニング済みモデルを登録するためのUnity Catalogカタログとスキーマ。
- HuggingFace ウイルス : ベース モデルとデータセットをダウンロードするための HuggingFace アクセス マラソンはDatabricksシークレットに保存されます。
- Python パッケージ : 必要なパッケージ (peft、trl、deepspeed、mlflow、hf_transfer) は、以下のセットアップ セクションでインストールされます。
必要なパッケージをインストールする
ファインチューニングに必要なPythonパッケージをインストールします。
peft: 効率的なファインチューニングのための LoRA 実装を提供しますtrl: 教師ありファインチューニング用トランスフォーマー強化学習ライブラリdeepspeed: ZeRO最適化によるメモリ効率の高い分散トレーニングを実現mlflow: 経験を追跡し、トレーニングされたモデルを記録しますhf_transfer: HuggingFace Hubからのモデルのダウンロードを高速化
インストール後、すべてのパッケージが適切にロードされていることを確認するために Python カーネルを再起動します。
Python
%pip install --upgrade transformers==4.56.1
%pip install peft==0.17.1
%pip install trl==0.18.1
%pip install deepspeed>=0.15.4
%pip install mlflow>=3.6.0
%pip install hf_transfer==0.1.9
Python
%restart_python
Unity Catalogと環境変数を構成する
モデルのチェックポイントを保存し、トレーニング済みモデルを登録するためのUnity Catalog場所を設定します。 ノートブックはクエリを使用して次の設定を行います。
- カタログとスキーマ : モデルとチェックポイントを整理するための Unity Catalog 名前空間
- モデル名 : Unity Catalogに登録されたモデルの名前
- Volume : トレーニング中にモデルのチェックポイントを保存するためのUnity Catalogボリューム
この構成では、 Databricksシークレットから HuggingFace を取得し、トレーニング メトリクスを追跡するためのMLflowエクスペリメントも設定します。
Python
dbutils.widgets.text("uc_catalog", "main")
dbutils.widgets.text("uc_schema", "default")
dbutils.widgets.text("uc_model_name", "llama3_2-1b")
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")
# Get HuggingFace token and username
hf_token = dbutils.secrets.get(scope="sgc-nightly-notebook", key="hf_token")
username = spark.sql("SELECT session_user()").collect()[0][0]
REGISTERED_MODEL_NAME = f"{UC_CATALOG}.{UC_SCHEMA}.{UC_MODEL_NAME}"
CHECKPOINT_DIR = f"/Volumes/{UC_CATALOG}/{UC_SCHEMA}/{UC_VOLUME}/{UC_MODEL_NAME}"
FINE_TUNED_MODEL_PATH = f"{CHECKPOINT_DIR}/fine-tuned-peft-model"
MLFLOW_EXPERIMENT_NAME = f"/Users/{username}/{UC_MODEL_NAME}"
# Create the Unity Catalog volume if it doesn't exist
spark.sql(f"CREATE VOLUME IF NOT EXISTS {UC_CATALOG}.{UC_SCHEMA}.{UC_VOLUME}")
print(f"👤 Username: {username}")
print("🔑 HuggingFace token configured")
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}")
print(f"CHECKPOINT_DIR: {CHECKPOINT_DIR}")
print(f"MLFLOW_EXPERIMENT_NAME: {MLFLOW_EXPERIMENT_NAME}")
Python
import os
import json
import tempfile
import torch
import mlflow
from huggingface_hub import constants
from datasets import load_dataset
from transformers import AutoTokenizer, AutoModelForCausalLM
from trl import SFTTrainer, SFTConfig, ModelConfig, ScriptArguments, setup_chat_format
from peft import LoraConfig, get_peft_model, PeftModel
if mlflow.get_experiment_by_name(MLFLOW_EXPERIMENT_NAME) is None:
mlflow.create_experiment(name=MLFLOW_EXPERIMENT_NAME)
mlflow.set_experiment(MLFLOW_EXPERIMENT_NAME)
DeepSpeed ZeRO Stage 3構成を作成する
DeepSpeed ZeRO (Zero Redundancy Optimizer) ステージ 3 は、モデル カーネル、勾配、およびオプティマイザーの状態をすべての GPU に分割し、GPU あたりのメモリ消費量を削減します。 これにより、単一の GPU のメモリに収まらない大規模なモデルのトレーニングが可能になります。
主要な構成設定:
- bf16 有効 : bfloat16 精度を使用してトレーニングを高速化し、メモリ使用量を削減します
- ステージ3の最適化 : すべてのモデル状態をGPUに分割する
- CPUオフロードなし ：すべてのデータをGPU上に保持し、A10ハードウェアで最大のパフォーマンスを実現します。
- オーバーラップ通信 : 効率化のため勾配通信と計算をオーバーラップする
Python
def create_deepspeed_config():
"""Create DeepSpeed ZeRO Stage 3 configuration for single node A10 training."""
deepspeed_config = {
"fp16": {
"enabled": False
},
"bf16": {
"enabled": True
},
"zero_optimization": {
"stage": 3,
"offload_optimizer": {
"device": "none"
},
"offload_param": {
"device": "none"
},
"overlap_comm": True,
"contiguous_gradients": True,
"sub_group_size": 1e9,
"reduce_bucket_size": "auto",
"stage3_prefetch_bucket_size": "auto",
"stage3_param_persistence_threshold": "auto",
"stage3_max_live_parameters": 1e9,
"stage3_max_reuse_distance": 1e9,
"stage3_gather_16bit_weights_on_model_save": True
},
"gradient_accumulation_steps": "auto",
"gradient_clipping": "auto",
"steps_per_print": 2000,
"train_batch_size": "auto",
"train_micro_batch_size_per_gpu": "auto",
"wall_clock_breakdown": False
}
return deepspeed_config
# Create DeepSpeed configuration
deepspeed_config = create_deepspeed_config()
print("⚙️ DeepSpeed ZeRO Stage 3 configuration created")
トレーニング問題と LoRA 構成を定義する
監視付きファインチューニングを設定します。
- モデル : Llama 3.2 1B Instruct、A10 GPUに適したコンパクトモデル
- データセット : 会話型AIトレーニング用の TRL ライブラリからの Capybara データセット
- バッチ サイズ : デバイスあたり 2、有効バッチ サイズ 64 の 4 勾配累積ステップ
- 学習率 : コサインスケジューラとウォームアップ付き 2e-4
- トレーニング ステップ : デモンストレーション用 60 ステップ (完全なトレーニング用に増加)
- LoRA 害 : ランク 16、アルファ 32、注目層と MLP プロジェクション層をターゲット
この構成では、bfloat16 精度と勾配チェックポイントを使用してメモリ使用量を最適化します。
Python
def create_training_config():
"""Create training configuration for TRL SFT with LoRA."""
# Model and dataset configuration (not part of TrainingArguments)
model_config = {
"model_name": "meta-llama/Llama-3.2-1B-Instruct", # Small Llama model for A10
"dataset_name": "trl-lib/Capybara"
}
# Training arguments that will be passed directly to TrainingArguments
training_args_config = {
"output_dir": CHECKPOINT_DIR,
"per_device_train_batch_size": 2,
"per_device_eval_batch_size": 2,
"gradient_accumulation_steps": 4,
"learning_rate": 2e-4,
"max_steps": 60, # TO DO remove when fine-tuning on full dataset. Demo purposes only.
# "num_train_epochs": 1, # TO DO update to >= 1 when fine-tuning on full dataset
"logging_steps": 10,
"save_steps": 30,
"eval_steps": 30,
"eval_strategy": "steps",
"warmup_steps": 10,
"lr_scheduler_type": "cosine",
"gradient_checkpointing": True,
"fp16": False,
"bf16": True,
"optim": "adamw_torch",
"remove_unused_columns": False,
"run_name": f"llama3.2-1b-lora",
"report_to": "mlflow",
"save_total_limit": 2,
"load_best_model_at_end": True,
"metric_for_best_model": "eval_loss",
"greater_is_better": False,
}
# LoRA configuration
lora_config = {
"r": 16,
"lora_alpha": 32,
"target_modules": ["q_proj", "v_proj", "k_proj", "o_proj", "gate_proj", "up_proj", "down_proj"],
"lora_dropout": 0.1,
"bias": "none",
"task_type": "CAUSAL_LM"
}
return model_config, training_args_config, lora_config
# Create training configuration
model_config, training_args_config, lora_config = create_training_config()
print("📊 Training Configuration:")
print(f" 🤖 Model: {model_config['model_name']}")
print(f" 📚 Dataset: {model_config['dataset_name']}")
print(f" 🎯 Batch size: {training_args_config['per_device_train_batch_size']}")
print(f" 📈 Learning rate: {training_args_config['learning_rate']}")
print(f" 🧠 LoRA rank: {lora_config['r']}")
分散トレーニング関数を定義する
serverless_gpuライブラリの@distributedデコレータにより、 Databricksサーバレス GPU での GPU ワークロードのシームレスな実行が可能になります。 デコレータは 4 つの A10 GPU をプロビジョニングし、分散トレーニングのセットアップを自動的に処理します。
鍵となる問題:
- gpus=4 : 分散トレーニング用に4つのGPUを要求する
- gpu_type='A10' : A10 GPUハードウェアを指定します
- Remote=True : リモートサーバレスGPUコンピュートで実行
トレーニング機能:
- HuggingFaceからベースモデルとトークナイザーをロードします
- 会話型AIのチャットフォーマットを設定する
- 効率的なファインチューニングのために LoRA を構成します
- トレーニングデータセットをロードする
- DeepSpeed最適化でTRL SFTTrainerを初期化します
- モデルをトレーニングしてチェックポイントを保存します
- トレーニング結果とMLflow実行IDを返します。
詳細については、サーバレス GPU APIドキュメントを参照してください。
Python
from serverless_gpu import distributed
os.environ['MLFLOW_EXPERIMENT_NAME'] = MLFLOW_EXPERIMENT_NAME
@distributed(
gpus=4,
gpu_type='A10',
remote=True, # Set to False to run locally, True for remote GPUs
)
def run_distributed_trl_sft():
"""
Distributed TRL SFT training function using serverless GPU.
This function will be executed on the A10 GPU with DeepSpeed optimization.
"""
# Set up environment variables for remote jobs
import os
import tempfile
import json
from huggingface_hub import constants
from datasets import load_dataset
from transformers import AutoTokenizer, AutoModelForCausalLM, TrainingArguments
from trl import SFTTrainer, setup_chat_format
from peft import LoraConfig, get_peft_model
# HuggingFace configuration
os.environ["HUGGING_FACE_HUB_TOKEN"] = hf_token
os.environ['HF_TOKEN'] = hf_token
constants.HF_HUB_ENABLE_HF_TRANSFER = True
# Set up temporary directories
temp_dir = tempfile.mkdtemp()
print("🚀 Starting TRL SFT training on A10 GPU...")
try:
# Load tokenizer and model
print(f"📥 Loading model: {model_config['model_name']}")
tokenizer = AutoTokenizer.from_pretrained(model_config['model_name'])
# Add pad token if not present
if tokenizer.pad_token is None:
tokenizer.pad_token = tokenizer.eos_token
model = AutoModelForCausalLM.from_pretrained(
model_config['model_name'],
dtype="auto"
)
# Setup chat template if the model doesn't have one
# This is crucial for conversational AI models and TRL SFTTrainer
if tokenizer.chat_template is None:
print("🗨️ Setting up chat template...")
model, tokenizer = setup_chat_format(model, tokenizer, format="chatml")
# Configure LoRA
print("🔧 Setting up LoRA configuration...")
peft_config = LoraConfig(**lora_config)
# Load dataset
print(f"📚 Loading dataset: {model_config['dataset_name']}")
dataset = load_dataset(model_config['dataset_name'])
# Create temporary DeepSpeed config file
deepspeed_config_path = os.path.join(temp_dir, "deepspeed_config.json")
with open(deepspeed_config_path, "w") as f:
json.dump(deepspeed_config, f, indent=2)
# Training arguments - dynamically pass all config parameters
training_args = TrainingArguments(
**training_args_config,
deepspeed=deepspeed_config_path, # Override deepspeed with the config file path
)
# Initialize SFT Trainer
print("🏋️ Initializing SFT Trainer with DeepSpeed...")
trainer = SFTTrainer(
model=model,
args=training_args,
train_dataset=dataset["train"],
eval_dataset=dataset["test"] if "test" in dataset else None,
processing_class=tokenizer,
peft_config=peft_config
)
# Start training
print("🎯 Starting training...")
trainer.train()
# Save the model
print("💾 Saving trained model...")
trainer.save_model()
# Get training results
train_results = trainer.state.log_history
final_loss = train_results[-1].get('train_loss', 'N/A') if train_results else 'N/A'
print("✅ Training completed successfully!")
print(f"📊 Final training loss: {final_loss}")
mlflow_run_id = None
if mlflow.last_active_run() is not None:
mlflow_run_id = mlflow.last_active_run().info.run_id
return {
"status": "success",
"final_loss": final_loss,
"output_dir": training_args_config['output_dir'],
"model_name": model_config['model_name'],
"mlflow_run_id": mlflow_run_id,
}
except Exception as e:
print(f"❌ Training failed: {e}")
import traceback
traceback.print_exc()
return {
"status": "failed",
"error": str(e)
}
分散トレーニングジョブを実行する
装飾された関数で.distributed()を呼び出してトレーニング関数を実行します。これは、サーバーレス GPU リソースをプロビジョニングし、DeepSpeed 最適化を使用して 4 つの A10 GPU にわたってトレーニングを実行し、結果を返します。
トレーニングプロセス:
- 4つのA10 GPUを自動的にプロビジョニング
- HuggingFaceからモデルとデータセットをダウンロードします
- LoRAファインチューニングでモデルをトレーニングする
- チェックポイントをUnity Catalogボリュームに保存します
- メトリクスのログをMLflowに記録します
- トレーニング ステータス、最終損失、 MLflow実行 ID を返します。
Python
# Execute the distributed training
results = run_distributed_trl_sft.distributed()
print("🏁 Training execution completed!")
print(f"📊 Results: {results}")
if results and results[0].get('status') == 'success':
print("✅ Training completed successfully!")
print(f"💾 Model saved to: {results[0].get('output_dir', 'N/A')}")
print(f"📈 Final loss: {results[0].get('final_loss', 'N/A')}")
print(f"🎉 MLflow run ID: {results[0].get('mlflow_run_id', 'N/A')}")
else:
print("❌ Training failed!")
if results and 'error' in results:
print(f"🔍 Error: {results['error']}")
微調整したモデルを保存し、推論をテストする
このオプションのステップでは、トレーニング済みの LoRA アダプターを読み込み、それをベース モデルとマージして、微調整された完全なモデルを保存します。 マージされたモデルは、サンプルプロンプトを使用してテストされ、ファインチューニングの結果を検証できます。
プロセス:
- ベースとなるLlama 3.2 1Bモデルをロードします
- トレーニング済みのLoRAアダプタの重みを適用する
- アダプタをベースモデルにマージします
- マージされたモデルをUnity Catalogボリュームに保存します
- サンプル会話プロンプトでモデルをテストする
Python
%pip install hf_transfer
Python
def save_and_load_trained_model():
"""Load the trained model from the Unity Catalog volume."""
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import PeftModel
# Load base model and tokenizer
base_model = AutoModelForCausalLM.from_pretrained(
model_config['model_name'],
dtype=torch.bfloat16,
token=hf_token,
trust_remote_code=True,
device_map="auto"
)
tokenizer = AutoTokenizer.from_pretrained(model_config['model_name'], token=hf_token, trust_remote_code=True)
# Load LoRA weights
model = PeftModel.from_pretrained(base_model, training_args_config['output_dir'])
# Merge LoRA weights into base model
model = model.merge_and_unload()
# Save the merged model
model.save_pretrained(FINE_TUNED_MODEL_PATH)
tokenizer.save_pretrained(FINE_TUNED_MODEL_PATH)
# Return the merged model and tokenizer
return model, tokenizer
def test_trained_model(model, tokenizer):
"""Test the trained model with simple inference."""
try:
import torch
# Test prompt
# Create a conversation following the schema
conversation = [
{
"content": "What is machine learning?",
"role": "user"
}
]
# Convert conversation to chat format
prompt = ""
for message in conversation:
if message["role"] == "user":
prompt += f"### User: {message['content']}\n### Response:"
else:
prompt += f" {message['content']}\n\n"
# Tokenize
inputs = tokenizer(prompt, return_tensors="pt").to("cuda")
# Generate
with torch.no_grad():
outputs = model.generate(
**inputs,
max_new_tokens=500,
temperature=0.7,
do_sample=True,
pad_token_id=tokenizer.eos_token_id
)
# Decode
response = tokenizer.decode(outputs[0], skip_special_tokens=True)
print("🤖 Model Response:")
print(response)
return response
except Exception as e:
print(f"❌ Model testing failed: {e}")
# Save and load the trained model
model, tokenizer = save_and_load_trained_model()
# Test the trained model
response = test_trained_model(model, tokenizer)
Unity Catalogにモデルを登録する
微調整されたモデルをMLflowに記録し、デプロイと提供のためにUnity Catalogに登録します。 モデルは次のようにログに記録されます。
- モデルとトークナイザー : 推論に必要な両方のコンポーネント
- タスクタイプ : 会話型 AI 用に
llm/v1/chatとして構成されています - 入力例 : テスト用のサンプルチャットメッセージ形式
- Unity Catalog登録 : 設定したカタログとスキーマにモデルを自動的に登録します
登録が完了すると、モデルをモデルサービングエンドポイントにデプロイしたり、バッチ推論に使用したりできます。
Python
run_id = results[0].get('mlflow_run_id')
mlflow.set_registry_uri("databricks-uc")
# log the model to mlflow using the latest run id and register to Unity Catalog
with mlflow.start_run(run_id=run_id) as run:
components = {
"model": model,
"tokenizer": tokenizer
}
logged_model = mlflow.transformers.log_model(
transformers_model=components,
name="model",
task="llm/v1/chat",
input_example={
"messages": [
{"role": "user", "content": "What is machine learning?"}
]
},
registered_model_name=REGISTERED_MODEL_NAME
)
print(f"🔍 Model logged to: {logged_model}")
次のステップ
- サーバレスGPUコンピュート
- サーバレス GPU コンピュートのベスト プラクティス
- サーバーレス GPU コンピュートの問題のトラブルシューティング
- マルチGPUおよびマルチノード分散トレーニング
- MLflowでモデルをトレーニングする
- モデルサービングでモデルをデプロイする