AI ランタイムを使用して Llama 3.2 1B を微調整する

このページのノートブックバージョンを開く

Databricks AI Runtime で監視対象の微調整 (SFT) を使用して、大規模な言語モデルを完全に微調整します。このノートブックでは、DeepSpeed ZeRO Stage 3 最適化を備えたトランスフォーマー強化学習 (TRL) ライブラリを使用して、8 個の H100 GPU を持つ単一ノードで Llama 3.2 1B を効率的にトレーニングします。

主な概念:

TRL (トランスフォーマー強化学習):強化学習と教師あり微調整を使用して言語モデルをトレーニングするためのツールを提供するライブラリです。
DeepSpeed ZeRO ステージ 3: 大規模なモデルのトレーニングを可能にするために、GPU 全体でモデルパラメーター、グラデーション、およびオプティマイザーの状態をパーティション分割するメモリ最適化手法。
AI ランタイム: トレーニングワークロード用の GPU リソースを自動的にプロビジョニングおよびスケーリングする Databricks マネージド GPU コンピューティング。

詳細については、「 AI ランタイム」を参照してください。

必要条件

このノートブックには、次のものが必要です。

AI ランタイム: ノートブックでは、分散トレーニングに 8 個の H100 GPU を備えた Databricks AI ランタイムが使用されます。クラスター構成は必要ありません。
Unity カタログ: モデルチェックポイントを格納し、トレーニング済みのモデルを登録するための Unity カタログカタログとスキーマ。
HuggingFace トークン: Databricks シークレットに格納されている HuggingFace アクセストークンで、基本モデルとデータセットをダウンロードします。
Python パッケージ: AI ランタイムは、最も必要なライブラリをプレインストールします。以下のセットアップセクションでは、 deepspeedがインストールされます。

サーバーレス GPU コンピューティングに接続する

このノートブックには、サーバーレス GPU コンピューティングが必要です。接続するには

右上にあるノートブックのコンピューティングセレクターをクリックし、[ サーバーレス GPU] を選択します。
右側にある [環境] ボタンをクリックします。
アクセラレータとして 8xH100 を選択します。
このノートブックの例を実行するために必要なすべてのライブラリが含まれている右側のパネルから AI v5 環境を選択します。
[適用] をクリックします。

トレーニング機能は、分散トレーニング用に 8 個の H100 GPU を自動的にプロビジョニングします。

必要なパッケージをインストールする

AI ランタイムには、必要なライブラリの大部分が既にプレインストールされています。この例では、 deepspeedのみをインストールする必要があります。

%pip install deepspeed==0.19.1
%restart_python

Unity カタログと環境変数を構成する

モデルチェックポイントを格納し、トレーニング済みのモデルを登録するための Unity カタログの場所を設定します。ノートブックでは、クエリパラメーターを使用して次の構成が行われます。

カタログとスキーマ: モデルとチェックポイントを整理するための Unity カタログ名前空間
モデル名: Unity カタログに登録されているモデルの名前
ボリューム: トレーニング中にモデルチェックポイントを格納するための Unity カタログボリューム

また、この構成では、Databricks シークレットから HuggingFace トークンを取得し、トレーニングメトリックを追跡するための MLflow 実験を設定します。

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

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

DeepSpeed ZeRO ステージ 3 の構成を作成する

DeepSpeed ZeRO (ゼロ冗長オプティマイザー) ステージ 3 では、GPU あたりのメモリ消費量を減らすために、すべての GPU のモデルパラメーター、グラデーション、およびオプティマイザーの状態がパーティション分割されます。これにより、1 つの GPU のメモリに収まらない大規模なモデルのトレーニングが可能になります。

主要な構成設定:

bf16 対応: bfloat16 精度を使用して、トレーニングの高速化とメモリ使用量の削減を実現
ステージ 3 の最適化: GPU 間ですべてのモデルの状態をパーティション分割する
CPU オフロードなし: H100 ハードウェアで最大限のパフォーマンスを実現するために、GPU 上のすべてのデータを保持します
オーバーラップ通信: 勾配通信と計算を重ねて効率を高める

def create_deepspeed_config(stage: int):
    """Create a DeepSpeed ZeRO configuration for single-node 8xH100 training."""

    deepspeed_config = {
        "fp16": {
            "enabled": False
        },
        "bf16": {
            "enabled": True
        },
        "zero_optimization": {
            "stage": stage,
            "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": 0,
            "stage3_max_live_parameters": 1e9,
            "stage3_max_reuse_distance": 1e9,
            "stage3_gather_16bit_weights_on_model_save": True
        },
        "gradient_accumulation_steps": 1,
        "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
zero_stage = 3
deepspeed_config = create_deepspeed_config(zero_stage)
print(f"⚙️  DeepSpeed ZeRO Stage {zero_stage} configuration created")

トレーニングパラメーターを定義する

監視対象の微調整パラメーターを構成します。

モデル: Llama 3.2 1B Instruct、H100 GPU 用のコンパクトなモデル
データセット: 会話型 AI トレーニング用の TRL ライブラリの Capybara データセット
バッチサイズ: デバイスあたり 2 のバッチサイズと 4 ステップの勾配累積により、実効バッチサイズを 64 にします。
学習率: コサインスケジューラとウォームアップを使用した 2e-4
トレーニング手順: デモンストレーション用の 60 ステップ (完全なトレーニングの場合は増加)

この構成では、bfloat16 の精度とグラデーションのチェックポイント処理を使用して、メモリ使用量を最適化します。

def create_training_config():
    """Create training configuration for TRL SFT."""

    # Model and dataset configuration (not part of TrainingArguments)
    model_config = {
        "model_name": "meta-llama/Llama-3.2-1B-Instruct",  # Small Llama model suitable for 8xH100
        "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": 1,
        "learning_rate": 2e-4,
        "max_steps": 60,   # TO DO remove when fine-tuning on full dataset. Demo purposes only.
        "logging_steps": 10,
        "save_steps": 30,
        "eval_steps": 30,
        "eval_strategy": "steps",
        "warmup_steps": 10,
        "lr_scheduler_type": "cosine",
        "gradient_checkpointing": False,
        "fp16": False,
        "bf16": True,
        "optim": "adamw_torch",
        "remove_unused_columns": False,
        "run_name": f"llama3.2-1b-fft-zero3",
        "report_to": "mlflow",
        "save_total_limit": 2,
        "load_best_model_at_end": True,
        "metric_for_best_model": "eval_loss",
        "greater_is_better": False,
    }

    return model_config, training_args_config

# Create training configuration
model_config, training_args_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']}")

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

@distributed ライブラリの serverless_gpu デコレーターを使用すると、Databricks AI ランタイムで GPU ワークロードを実行できます。デコレーターは 8 個の H100 GPU をプロビジョニングし、分散トレーニングのセットアップを自動的に処理します。

キーパラメーター:

gpus=8: 分散トレーニング用に 8 個の GPU を要求する
gpu_type='H100': H100 GPU ハードウェアを指定します

トレーニング関数:

HuggingFace から基本モデルとトークナイザーを読み込みます
会話型 AI のチャット書式設定を設定する
トレーニングデータセットを読み込みます
DeepSpeed 最適化を使用して TRL SFTTrainer を初期化します。
モデルをトレーニングし、チェックポイントを保存する
トレーニング結果と MLflow 実行 ID を返します

詳細については、 AI ランタイム API のドキュメントを参照してください。

from serverless_gpu import distributed

mlflow.set_experiment(MLFLOW_EXPERIMENT_NAME)
@distributed(
    gpus=8,
    gpu_type='H100',
)
def run_distributed_trl_sft():
    """
    Distributed TRL SFT training function using AI Runtime.

    This function will be executed on the H100 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
    from trl import SFTTrainer, SFTConfig

    # 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 H100 GPU...")

    try:
        # Load tokenizer
        print(f"📥 Loading tokenizer from 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

        # 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 = SFTConfig(
            **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_config["model_name"],
            args=training_args,
            train_dataset=dataset["train"],
            eval_dataset=dataset["test"] if "test" in dataset else None,
            processing_class=tokenizer,
        )

        # 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() を呼び出して、トレーニング関数を実行します。これにより、AI ランタイムリソースがプロビジョニングされ、DeepSpeed 最適化を使用して 8 つの H100 GPU 間でトレーニングが実行され、結果が返されます。

トレーニングプロセス:

8 個の H100 GPU を自動的にプロビジョニングする
HuggingFace からモデルとデータセットをダウンロードします
完全な微調整を使用してモデルをトレーニングする
チェックポイントを Unity カタログボリュームに保存します
メトリックを MLflow に記録する
トレーニングの状態、最終的な損失、および MLflow の実行 ID を返します

# 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']}")

微調整されたモデルを保存し、推論をテストする

このオプションの手順では、微調整されたモデルを読み込み、サンプルプロンプトでテストして結果を確認します。

プロセス:

保存したモデルを読み込む
サンプルの会話プロンプトを使用してモデルをテストする

def save_and_load_trained_model():
    """Load the fully fine-tuned model from the Unity Catalog volume."""

    import torch
    from transformers import AutoModelForCausalLM, AutoTokenizer

    print(f"📥 Loading fine-tuned model from: {training_args_config['output_dir']}")

    # Load the fully fine-tuned model directly from the checkpoint directory
    model = AutoModelForCausalLM.from_pretrained(
        training_args_config['output_dir'],
        torch_dtype=torch.bfloat16,
        trust_remote_code=True,
        device_map={"":0}
    )
    tokenizer = AutoTokenizer.from_pretrained(training_args_config['output_dir'], trust_remote_code=True)

    print("✅ Model loaded successfully!")
    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
test_trained_model(model, tokenizer)

Unity カタログにモデルを登録する

微調整されたモデルを MLflow に記録し、デプロイと提供のために Unity カタログに登録します。モデルは次の方法でログに記録されます。

モデルとトークナイザー: 推論に必要な両方のコンポーネント
タスクの種類: 会話型 AI の llm/v1/chat として構成
入力例: テスト用のサンプルチャットメッセージ形式
Unity カタログの登録: 構成済みのカタログとスキーマにモデルを自動的に登録します

登録されると、モデルをモデルサービスエンドポイントにデプロイするか、バッチ推論に使用できます。

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

次のステップ

ノートブックの例

AI ランタイムを使用して Llama 3.2 1B を微調整する

ノートブックを入手

フィードバック

このページはお役に立ちましたか?

Last updated on 2026-08-18

AI ランタイムを使用して Llama 3.2 1B を微調整する

必要条件

サーバーレス GPU コンピューティングに接続する

必要なパッケージをインストールする

Unity カタログと環境変数を構成する

DeepSpeed ZeRO ステージ 3 の構成を作成する

トレーニング パラメーターを定義する

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

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

微調整されたモデルを保存し、推論をテストする

Unity カタログにモデルを登録する

次のステップ

ノートブックの例

AI ランタイムを使用して Llama 3.2 1B を微調整する

フィードバック

その他のリソース

トレーニングパラメーターを定義する

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