Databricksサーバレス GPU で Mosaic LLM Foundry を使用してLlama 3.1 8B を微調整する

このノートブックでは、 Databricksサーバレス GPU でMosaic LLM Foundryを使用してLlama 3.1 8B モデルを微調整する方法を示します。 LLM Foundry は、分散トレーニング戦略をサポートする大規模言語モデルのトレーニング、ファインチューニング、評価、およびデプロイのためのコードベースです。

ノートブックでは以下を使用します:

• Mosaic LLM Foundry : FSDP、効率的なデータ読み込み、 MLflow統合のサポートが組み込まれた、LLM のトレーニングと微調整のためのフレームワーク
• FSDP (Fully Sharded Data Parallel) : モデルの欠点、勾配、オプティマイザーの状態を GPU 全体に分散します
• Databricksサーバレス GPU : 自動的にプロビジョニングと GPU コンピュート リソースを管理
• Unity Catalog : モデルのチェックポイントと登録するトレーニング済みモデルを保存します
• MLflow : 体験の追跡とトレーニング メトリクス

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

分散トレーニング用のMosaic LLM Foundryと依存関係をインストールします。

• llm-foundry: LLMトレーニングとファインチューニングのためのコアフレームワーク
• mlflow: 体験追跡とモデルレジストリ
• flash-attention: より速いトレーニングのための最適化された注意の実装
• hf_transfer: Hugging Faceからのモデルのダウンロードが高速化
• yamlmagic: ノートブックのセルで YAML 構成を有効にします

Python

%pip install llm-foundry[gpu]==0.20.0
%pip install mlflow==3.6
%pip install matplotlib==3.10.0
%pip install --force-reinstall --no-cache-dir --no-deps "https://github.com/Dao-AILab/flash-attention/releases/download/v2.7.4.post1/flash_attn-2.7.4.post1+cu12torch2.6cxx11abiFALSE-cp312-cp312-linux_x86_64.whl"
%pip install hf_transfer
%pip install git+https://github.com/josejg/yamlmagic.git

Python環境を再起動します

新しくインストールされたパッケージが利用可能であることを確認するために、Python カーネルを再起動します。

Python

dbutils.library.restartPython()

モデル保存用のUnity Catalogパスを設定する

モデルのチェックポイントを保存し、トレーニング済みモデルを登録するためのUnity Catalog場所を設定します。 この構成では、コードを編集せずにカスタマイズできるクエリを使用します。

Python

dbutils.widgets.text("uc_catalog", "main")
dbutils.widgets.text("uc_schema", "default")
dbutils.widgets.text("uc_model_name", "llama3_1-8b")
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")

MLFLOW_EXPERIMENT_NAME = '/Workspace/Shared/llm-foundry-sgc' # TODO: update this name

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"EXPERIMENT_NAME: {MLFLOW_EXPERIMENT_NAME}")

# Model selection - Choose based on your compute constraints
OUTPUT_DIR = f"/Volumes/{UC_CATALOG}/{UC_SCHEMA}/{UC_VOLUME}/{UC_MODEL_NAME}" # Save checkpoint to UC Volume

print(f"OUTPUT_DIR: {OUTPUT_DIR}")

YAMLを使用してトレーニング構成を定義する

YAML 形式からファインチューニング設定を読み込みます。 構成では以下を指定します。

• モデルアーキテクチャと事前学習済み重み（Llama 3.1 8B）
• 分散トレーニングのためのFSDP設定
• トレーニングハイパーパラメータ（学習率、バッチサイズ、オプティマイザー）
• データセット構成 (mosaicml/dolly_hhrlhf)
• MLflow のログ記録とモデルのチェックポイント
• モニタリングと最適化のためのコールバック

Python

%load_ext yamlmagic

Python

%%yaml config
seed: 17
model:
  name: hf_causal_lm
  pretrained: true
  init_device: mixed
  use_auth_token: true
  use_flash_attention_2: true
  pretrained_model_name_or_path: meta-llama/Llama-3.1-8B
loggers:
  mlflow:
    resume: true
    tracking_uri: databricks
    rename_metrics:
      time/token: time/num_tokens
      lr-DecoupledLionW/group0: learning_rate
    log_system_metrics: true
    experiment_name: "mlflow_experiment_name"
    run_name: llama3_8b-finetune
    model_registry_uri: databricks-uc
    model_registry_prefix: main.linyuan
callbacks:
  lr_monitor: {}
  run_timeout:
    timeout: 7200
  scheduled_gc:
    batch_interval: 1000
  speed_monitor:
    window_size: 10
  memory_monitor: {}
  runtime_estimator: {}
  hf_checkpointer:
    save_folder: "dbfs:/Volumes/main/sgc/checkpoints/llama3_1-8b-hf"
    save_interval: "1ep"
    precision: "bfloat16"
    overwrite: true

    mlflow_registered_model_name: "main.sgc.llama3_1_8b_full_ft"
    mlflow_logging_config:
      task: "llm/v1/completions"
      metadata:
        pretrained_model_name: "meta-llama/Llama-3.1-8B-Instruct"
optimizer:
  lr: 5.0e-07
  name: decoupled_lionw
  betas:
  - 0.9
  - 0.95
  weight_decay: 0
precision: amp_bf16
scheduler:
  name: linear_decay_with_warmup
  alpha_f: 0
  t_warmup: 10ba
tokenizer:
  name: meta-llama/Llama-3.1-8B
  kwargs:
    model_max_length: 1024
algorithms:
  gradient_clipping:
    clipping_type: norm
    clipping_threshold: 1
autoresume: false
log_config: false
fsdp_config:
  verbose: false
  mixed_precision: PURE
  state_dict_type: sharded
  limit_all_gathers: true
  sharding_strategy: FULL_SHARD
  activation_cpu_offload: false
  activation_checkpointing: true
  activation_checkpointing_reentrant: false
max_seq_len: 1024
save_folder: "output_folder"
dist_timeout: 600
max_duration: 20ba
progress_bar: false
train_loader:
  name: finetuning
  dataset:
    split: test
    hf_name: mosaicml/dolly_hhrlhf
    shuffle: true
    safe_load: true
    max_seq_len: 1024
    packing_ratio: auto
    target_prompts: none
    target_responses: all
    allow_pad_trimming: false
    decoder_only_format: true
  timeout: 0
  drop_last: false
  pin_memory: true
  num_workers: 8
  prefetch_factor: 2
  persistent_workers: true
eval_interval: 1
save_interval: 1h
log_to_console: true
save_overwrite: true
python_log_level: debug
save_weights_only: false
console_log_interval: 10ba
device_eval_batch_size: 1
global_train_batch_size: 32
device_train_microbatch_size: 1
save_num_checkpoints_to_keep: 1

Python

config["loggers"]["mlflow"]["experiment_name"] = MLFLOW_EXPERIMENT_NAME
config["save_folder"] = OUTPUT_DIR
config["callbacks"]["hf_checkpointer"]["save_folder"] = OUTPUT_DIR
config["callbacks"]["hf_checkpointer"]["mlflow_registered_model_name"] = f"{UC_CATALOG}.{UC_SCHEMA}.{UC_MODEL_NAME}"

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

このセルは、 @distributedデコレータを使用して 8 つの H100 GPU 上で実行されるトレーニング関数を定義します。機能:

• モデルアクセス用にHugging Faceを構成します
• 高速モデルダウンロードを可能にする hf_transfer
• YAML設定でLLM Foundry train()関数を呼び出します
• エクスペリメントを追跡するためのMLflow実行 ID を返します。

@distributedデコレータは自動的にプロビジョニング サーバレス GPU コンピュートを実行し、分散トレーニング オーケストレーションを処理します。

Python

from serverless_gpu import distributed
from serverless_gpu import runtime as sgc_runtime
from llmfoundry.command_utils.train import train
from omegaconf import DictConfig
import mlflow
import yaml
from huggingface_hub import constants

HF_TOKEN = dbutils.secrets.get(scope="sgc-nightly-notebook", key="hf_token")

@distributed(gpus=8, gpu_type='H100')
def run_llm_foundry():
    import os
    import logging
    os.environ["HUGGING_FACE_HUB_TOKEN"] = HF_TOKEN
    constants.HF_HUB_ENABLE_HF_TRANSFER = True
    train(DictConfig(config))

    logging.info("\n✓ Training completed successfully!")

    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上でトレーニング関数を実行します。この関数はMLflow実行 ID を返します。この ID は、 MLflow UI でメトリクスの追跡、ログの表示、トレーニング済みモデルへのアクセスに使用できます。

Python

mlflow_run_id = run_llm_foundry.distributed()[0]
print(mlflow_run_id)

次のステップ

• マルチGPUおよびマルチノード分散トレーニング
• サーバレス GPU コンピュートのベスト プラクティス
• サーバーレス GPU コンピュートの問題のトラブルシューティング
• Mosaic LLM Foundry ドキュメント
• Unity Catalogモデルレジストリ

サンプルノートブック

Databricksサーバレス GPU で Mosaic LLM Foundry を使用してLlama 3.1 8B を微調整する

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