複数の A10 GPU 上で Unsloth を使用して Llama-3.2-3B を分散微調整する

このノートブックはLlama-3.2-3Bを微調整する方法を示しています複数の A10 GPU にわたる分散トレーニングを使用した大規模言語モデル。これは、最適化された効率的なファイン チューニングのためのUnslothライブラリと、分散トレーニング オーケストレーションのためのserverless_gpuライブラリを組み合わせたものです。

ノートブックには次の内容が記載されています。

• 4つのA10 GPUにわたる分散トレーニングの構成
• LoRA アダプターを使用した Llama-3.2-3B のロードと微調整
• FineTome-100k データセットからのトレーニング データの処理
• 教師ありファインチューニング (SFT) とMLflowトラッキングを使用したトレーニング
• アダプターをマージし、Unity Catalogにモデルを登録する

分散トレーニングでは、モデルの品質を維持しながら複数の GPU 間で計算を並列化することで、トレーニング時間を大幅に短縮します。

要件: A10 アクセラレータを備えたサーバレス GPU コンピュート

このノートブックには、A10 アクセラレータによる GPU コンピュートが必要です。 環境パネルでアクセラレータとして A10 を選択し、 [適用] をクリックします。

注: コンピュート プロビジョニングには最大 8 分かかる場合があります。 分散トレーニングでは、トレーニング機能の実行時に 4 つの A10 GPU が自動的にプロビジョニングされます。

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

CUDA 12.4 およびPyTorch 2.6.0 をサポートする Unsloth ライブラリを、分散トレーニング用の Accelerate、追加ユーティリティ用の unsloth_zoo、およびエクスペリメント追跡用のMLflowとともにインストールします。 インストール後、Python ランタイムが再起動し、新しいパッケージがロードされます。

Python

%pip install unsloth[cu124-torch260]==2025.9.6
%pip install accelerate==1.7.0
%pip install unsloth_zoo==2025.9.8
%pip install mlflow>=3.6
%restart_python

Unity Catalogとモデル設定を構成する

ノートブック ウィジェットを使用してUnity Catalog場所とモデル構成を定義し、簡単にカスタマイズできます。 構成には以下が含まれます。

• Unity Catalog名前空間 (カタログ、スキーマ、モデル名、ボリューム)
• ベースモデルの選択 (Llama-3.2-3B-Instructアンスロートより
• チェックポイントをUnity Catalogボリュームに保存するための出力ディレクトリ
• トレーニングデータセット（FineTome-100k）

Python

dbutils.widgets.text("uc_catalog", "main")
dbutils.widgets.text("uc_schema", "default")
dbutils.widgets.text("uc_model_name", "llama-3_2-3b")
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}")

# Model selection - Choose based on your compute constraints
MODEL_NAME = "unsloth/Llama-3.2-3B-Instruct" # or choose "unsloth/Llama-3.2-1B-Instruct"
OUTPUT_DIR = f"/Volumes/{UC_CATALOG}/{UC_SCHEMA}/{UC_VOLUME}/{UC_MODEL_NAME}" # Save checkpoint to UC Volume
DATASET_NAME = "mlabonne/FineTome-100k"

print(f"MODEL_NAME: {MODEL_NAME}")
print(f"OUTPUT_DIR: {OUTPUT_DIR}")
print(f"DATASET_NAME: {DATASET_NAME}")

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

マルチ GPU トレーニングを有効にするには、 @distributed(gpus=4, gpu_type='a10', remote=True)で装飾されたトレーニング関数を作成します。この関数は、トレーニング ワークフロー全体をカプセル化します。

• ベースLlama-3.2-3Bのロードモデルとトークナイザー
• LoRA アダプターを適用して効率的なファインチューニングを実現する
• チャット テンプレートを使用した FineTome-100k データセットの処理
• 分散トレーニング設定でSFTTrainerを構成する
• MLflowトラッキングによるモデルのトレーニング
• トレーニング済みのアダプターとトークナイザーをUnity Catalogボリュームに保存する

@distributedデコレータは、4 つの A10 GPU にわたる GPU プロビジョニングと分散トレーニング オーケストレーションを自動的に処理します。

Python

from serverless_gpu import distributed
from serverless_gpu import runtime as rt

@distributed(gpus=4, gpu_type='a10', remote=True)
def run_train():
    from datasets import load_dataset
    import logging
    import mlflow
    import torch

    # IMPORTANT: import unsloth BEFORE trl
    from unsloth import FastLanguageModel, is_bfloat16_supported
    from unsloth.chat_templates import get_chat_template, standardize_sharegpt, train_on_responses_only

    from trl import SFTTrainer
    from transformers import TrainingArguments, DataCollatorForSeq2Seq
    from transformers.integrations import MLflowCallback

    max_seq_length = 2048 # Choose any!
    dtype = None # None for auto detection. Float16 for Tesla T4, V100, Bfloat16 for Ampere+
    load_in_4bit = False # Use 4bit quantization to reduce memory usage. Can be False.

    # Load model and tokenizer
    model, tokenizer = FastLanguageModel.from_pretrained(
        model_name=MODEL_NAME,
        max_seq_length=max_seq_length,
        dtype=dtype,
        load_in_4bit=load_in_4bit,
        # token = "hf_...", # use one if using gated models like meta-llama/Llama-2-7b-hf
    )

    model = FastLanguageModel.get_peft_model(
        model,
        r=16, # Choose any number > 0 ! Suggested 8, 16, 32, 64, 128
        target_modules=["q_proj", "k_proj", "v_proj", "o_proj",
                        "gate_proj", "up_proj", "down_proj",],
        lora_alpha=16,
        lora_dropout=0, # Supports any, but = 0 is optimized
        bias="none",    # Supports any, but = "none" is optimized
        # [NEW] "unsloth" uses 30% less VRAM, fits 2x larger batch sizes!
        use_gradient_checkpointing="unsloth", # True or "unsloth" for very long context
        random_state=3407,
        use_rslora=False,  # We support rank stabilized LoRA
        loftq_config=None, # And LoftQ
    )

    # Process data
    tokenizer = get_chat_template(
        tokenizer,
        chat_template="llama-3.1",
    )

    def formatting_prompts_func(examples):
        convos = examples["conversations"]
        texts = [tokenizer.apply_chat_template(convo, tokenize = False, add_generation_prompt = False) for convo in convos]
        return { "text" : texts, }


    dataset = load_dataset(DATASET_NAME, split="train")

    dataset = standardize_sharegpt(dataset)
    dataset = dataset.map(formatting_prompts_func, batched=True,)

    trainer = SFTTrainer(
        model = model,
        tokenizer = tokenizer,
        train_dataset = dataset,
        dataset_text_field = "text",
        max_seq_length = max_seq_length,
        data_collator = DataCollatorForSeq2Seq(tokenizer = tokenizer),
        dataset_num_proc = 6,
        packing = False, # Can make training 5x faster for short sequences.
        args = TrainingArguments(
            per_device_train_batch_size = 2,
            gradient_accumulation_steps = 4,
            warmup_steps = 5,
            # num_train_epochs = 1, # Set this for 1 full training run.
            max_steps = 25,
            learning_rate = 2e-4,
            fp16 = not is_bfloat16_supported(),
            bf16 = is_bfloat16_supported(),
            logging_steps = 1,
            optim = "adamw_8bit",
            weight_decay = 0.01,
            lr_scheduler_type = "linear",
            seed = 3407,
            output_dir = OUTPUT_DIR,
            report_to = "mlflow", # Use MLflow to track model metrics,
            run_name = f"{MODEL_NAME}-finetune-unsloth",
        ),
    )
    trainer = train_on_responses_only(
        trainer,
        instruction_part="<|start_header_id|>user<|end_header_id|>\n\n",
        response_part="<|start_header_id|>assistant<|end_header_id|>\n\n",
        num_proc=1
    )

    trainer.train()

    # Save model
    if rt.get_global_rank() == 0:
        logging.info("\nSaving trained model...")
        trainer.save_model(OUTPUT_DIR)
        logging.info("✓ LoRA adapters saved - use with base model for inference")
        tokenizer.save_pretrained(OUTPUT_DIR)
        logging.info("✓ Tokenizer saved with model")
        logging.info(f"\n🎉 All artifacts saved to: {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

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

4 つの A10 GPU にわたる分散トレーニング機能を起動します。.distributed()メソッドは GPU をプロビジョニングし、トレーニング ワークロードを分散し、追跡用の MLflow 実行 ID を返します。このステップは、プロビジョニング コンピュート リソースとトレーニング ループを実行するため、数分かかる場合があります。

Python

run_id = run_train.distributed()[0]

モデル登録と展開戦略

分散トレーニングが完了したら、本番運用用に微調整されたモデルを登録します。

• MLflow Tracking - 体験追跡用の記録済みモデル アーティファクト、トレーニング メトリクス、およびメタデータ
• Unity Catalog - 集中ガバナンス、アクセス制御、リネージ追跡のためのモデルを登録する
• モデルのバージョン管理 - 自動バージョン管理により、モデルのライフサイクル管理とロールバック機能が有効になります。
• メタデータ - 完全なモデル情報により再現性とコンプライアンスが保証されます

アダプターをマージしてUnity Catalogに登録する

トレーニング済みのLoRAアダプターをロードし、ベースとなるLlama-3.2-3Bとマージします。モデルの重みを設定し、最終モデルをUnity Catalogに登録します。 このプロセス:

• チェックポイントディレクトリからベースモデルとトレーニング済みのLoRAアダプタをロードします。
• アダプタの重みをベースモデルにマージして、単一の展開可能なモデルを作成します。
• 適切なメタデータとともにマージされたモデルをMLflowに記録します
• ガバナンスとデプロイメントのために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

# 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=run_id):
    model_info = mlflow.transformers.log_model(
        transformers_model=components,
        name="model",
        task=task,
        registered_model_name=full_model_name,
        metadata={
            "task": task,
            "pretrained_model_name": MODEL_NAME,
            "databricks_model_family": "Llama3.2",
        },
    )

print(f"✓ Model successfully registered in Unity Catalog: {full_model_name}")
print(f"✓ MLflow model URI: {model_info.model_uri}")
print(f"✓ Model version: {model_info.registered_model_version}")

# Print deployment information
print(f"\n📦 Model Registration Complete!")
print(f"Unity Catalog Path: {full_model_name}")
print(f"Optimization: Liger Kernels + LoRA")

次のステップ

微調整されたモデルはUnity Catalogに登録され、デプロイできる状態になりました。 分散トレーニングとモデルサービングの詳細については、以下をご覧ください。

• バッチまたはリアルタイム推論用のモデルをデプロイする
• モデルサービングエンドポイントの作成と管理
• Unsloth ドキュメント

サンプルノートブック

複数の A10 GPU 上で Unsloth を使用して Llama-3.2-3B を分散微調整する

Open notebook in new tab