マルチ GPU サーバレス コンピュートで Axolotl を使用して Olmo3 7B を微調整する
このノートブックでは、 サーバーレス GPU コンピュートでAxolotl を使用してOlmo3 7B Instruct モデルを 微調整する方法を示します。DatabricksAxolotl は、QLoRA (量子化低ランク適応) を使用したLLMポストトレーニング用の高性能フレームワークを提供し、マルチ GPU インフラストラクチャでの効率的なファインチューニングを可能にします。 トレーニングされたモデルは MLflow に記録され、デプロイメントのために Unity Catalog に登録されます。
必要な依存関係をインストールする
Flash アテンションをサポートする Axolotl、エクスペリメント トラッキング用のMLflow 、互換性のあるバージョンのトランスフォーマーと最適化ライブラリをインストールします。 cut-cross-entropyパッケージは、大規模言語モデルのメモリ効率の高い損失計算を提供します。
Python
%pip install -U packaging setuptools wheel ninja
%pip install mlflow>=3.6
%pip install --no-build-isolation axolotl[flash-attn]>=0.12.0
%pip install transformers==4.57.3
%pip uninstall -y awq autoawq
%pip install "cut-cross-entropy[transformers] @ git+https://github.com/axolotl-ai-cloud/ml-cross-entropy.git@f4b5712"
dbutils.library.restartPython()
Retrieve HuggingFace ウイルス
Databricks シークレットから HuggingFace 認証トークンを取得します。このトークンは、HuggingFace Hub から Olmo3 7B ベース モデルをダウンロードするために必要です。
Python
HF_TOKEN = dbutils.secrets.get(scope="sgc-nightly-notebook", key="hf_token")
トレーニングの設定
olmo3-7b-qlora.yaml の例に基づいて Axolotl トレーニング構成を設定します。主な変更点は次のとおりです。
- エクスペリメント追跡のためのMLflow統合
- チェックポイント保存用のUnity Catalogボリュームパス
- より広範な GPU 互換性のために Flash Attention の代わりに SDPA (Scaled Dot Attention) を採用
Unity Catalogパスを定義する
モデル チェックポイントを保存するためのUnity Catalog場所を指定するウィジェットを作成します。 出力ディレクトリは、カタログ、スキーマ、ボリューム、およびモデル名を完全修飾パスに結合します。
Python
dbutils.widgets.text("uc_catalog", "main")
dbutils.widgets.text("uc_schema", "default")
dbutils.widgets.text("uc_volume", "checkpoints")
dbutils.widgets.text("model", "openai/gpt-oss-20b")
UC_CATALOG = dbutils.widgets.get("uc_catalog")
UC_SCHEMA = dbutils.widgets.get("uc_schema")
UC_VOLUME = dbutils.widgets.get("uc_volume")
UC_MODEL_NAME = dbutils.widgets.get("model")
print(f"UC_CATALOG: {UC_CATALOG}")
print(f"UC_SCHEMA: {UC_SCHEMA}")
print(f"UC_VOLUME: {UC_VOLUME}")
print(f"UC_MODEL_NAME: {UC_MODEL_NAME}")
OUTPUT_DIR = f"/Volumes/{UC_CATALOG}/{UC_SCHEMA}/{UC_VOLUME}/{UC_MODEL_NAME}"
print(f"OUTPUT_DIR: {OUTPUT_DIR}")
Python
import os
os.environ['AXOLOTL_DO_NOT_TRACK'] = '1'
テレメトリを無効にする
環境変数を設定して、Axolotl の使用状況追跡を無効にします。
Axolotl の設定を作成する
Axolotl のDictDefault形式を使用して完全なトレーニング構成を定義します。これには、モデル設定 (4 ビット量子化の QLoRA)、データセット構成 (Alpaca 形式)、LoRA ハイパーパラメーター (ランク 32、アルファ 16)、トレーニング ポイント (1 エポック、バッチ サイズ 2、勾配累積 4)、およびエクスペリメント トラッキングのためのMLflow統合が含まれます。
Python
from axolotl.cli.config import load_cfg
from axolotl.utils.dict import DictDefault
# Config is based on with some changes to fit GPU types
# https://raw.githubusercontent.com/axolotl-ai-cloud/axolotl/main/examples/olmo3/olmo3-7b-qlora.yaml
# Axolotl provides full control and transparency over model and training configuration
config = DictDefault(
base_model="allenai/Olmo-3-7B-Instruct-SFT",
plugins=[
"axolotl.integrations.cut_cross_entropy.CutCrossEntropyPlugin"
],
load_in_8bit=False,
load_in_4bit=True,
datasets=[
{
"path": "fozziethebeat/alpaca_messages_2k_test",
"type": "chat_template"
}
],
dataset_prepared_path="last_run_prepared",
val_set_size=0.1,
output_dir=OUTPUT_DIR,
adapter="qlora",
lora_model_dir=None,
sequence_len=2048,
sample_packing=True,
lora_r=32,
lora_alpha=16,
lora_dropout=0.05,
lora_target_linear=True,
lora_target_modules=[
"gate_proj",
"down_proj",
"up_proj",
"q_proj",
"v_proj",
"k_proj",
"o_proj"
],
wandb_project=None,
wandb_entity=None,
wandb_watch=None,
wandb_name=None,
wandb_log_model=None,
gradient_accumulation_steps=4,
micro_batch_size=2,
num_epochs=1,
optimizer="adamw_bnb_8bit",
lr_scheduler="cosine",
learning_rate=0.0002,
bf16="auto",
tf32=False,
gradient_checkpointing=True,
resume_from_checkpoint=None,
logging_steps=1,
flash_attention=False,
warmup_ratio=0.1,
evals_per_epoch=1,
saves_per_epoch=1,
# Eval dataset is too small
eval_sample_packing=False,
# Write metrics to MLflow
use_mlflow=True,
mlflow_tracking_uri="databricks",
mlflow_run_name="olmo3-7b-qlora-axolotl",
hf_mlflow_log_artifacts=False,
wandb_mode="disabled",
attn_implementation="sdpa",
sdpa_attention=True,
save_first_step=True,
device_map=None,
)
Python
from axolotl.utils import set_pytorch_cuda_alloc_conf
set_pytorch_cuda_alloc_conf()
PyTorch CUDA メモリ割り当てを構成する
マルチ GPU セットアップでの効率的なトレーニングのために GPU メモリ管理を最適化します。
サーバレスGPUコンピュートで分散トレーニングを実行
サーバレス GPU APIの@distributedデコレータを使用して、Axolotl トレーニング ジョブを 8 つの H100 GPU に分散します。 デコレータはマルチ GPU オーケストレーションを処理し、手動でクラスターをセットアップすることなく、トレーニング機能を分散環境で実行できるようにします。
Python
from serverless_gpu.launcher import distributed
from serverless_gpu.compute import GPUType
@distributed(gpus=8, gpu_type=GPUType.H100)
def run_train(cfg: DictDefault):
import os
os.environ['HF_TOKEN'] = HF_TOKEN
from axolotl.common.datasets import load_datasets
# Load, parse and tokenize the datasets to be formatted with qwen3 chat template
# Drop long samples from the dataset that overflow the max sequence length
# validates the configuration
cfg = load_cfg(cfg)
dataset_meta = load_datasets(cfg=cfg)
from axolotl.train import train
# just train the first 16 steps for demo.
# This is sufficient to align the model as we've used packing to maximize the trainable samples per step.
cfg.max_steps = 16
model, tokenizer, trainer = train(cfg=cfg, dataset_meta=dataset_meta)
import mlflow
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
Python
result = run_train.distributed(config)
トレーニングジョブを実行する
分散トレーニング ジョブを起動します。この関数はデータセットをロードし、構成を検証し、16 ステップのモデルをトレーニングし、追跡用のMLflow実行 ID を返します。
Python
run_id = result[0]
print(run_id)
MLflow実行IDの抽出
モデルの登録と体験追跡のためにトレーニング結果からMLflow実行 ID を取得します。
調整したモデルをUnity Catalogに登録する
トレーニング済みの LoRA アダプターを読み込み、ベースモデルとマージし、結合したモデルをMLflow経由でUnity Catalogに登録します。 これにより、モデルを展開および推論に使用できるようになります。
注: このステップでは、モデル チェックポイントをロードするために H100 GPU コンピュートが必要です。 小さい GPU で実行すると、CUDA のメモリ不足エラーが発生する可能性があります。
Python
from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
try:
from transformers.activations import NewGELUActivation, PytorchGELUTanh, GELUActivation
except ImportError:
from transformers.activations import NewGELUActivation, GELUTanh as PytorchGELUTanh, GELUActivation
from peft import PeftModel
import mlflow
from mlflow import transformers as mlflow_transformers
import torch
HF_MODEL_NAME = "allenai/Olmo-3-7B-Instruct-SFT"
torch.cuda.empty_cache()
# 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(
HF_MODEL_NAME,
trust_remote_code=True
)
# Load tokenizer
tokenizer = AutoTokenizer.from_pretrained(HF_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}")
text_gen_pipe = pipeline(
task="text-generation",
model=peft_model,
tokenizer=tokenizer,
)
input_example = ["Hello, world!"]
with mlflow.start_run(run_id=run_id):
model_info = mlflow.transformers.log_model(
transformers_model=text_gen_pipe, # 🚨 pass the pipeline, not just the model
artifact_path="model",
input_example=input_example,
# optional: save_pretrained=False for reference-only PEFT logging
# save_pretrained=False,
)
# Start MLflow run and log model
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")