シングルGPUでHugging Face モデルをファインチューニングする

この記事では、1 つの GPU で Hugging Face transformers ライブラリを使用して Hugging Face モデルをファインチューンする方法について説明します。また、レイクハウスからデータを読み込み、モデルを MLflow にログ記録するための Databricks 固有の推奨事項も含まれているため、Databricks でモデルを使用および管理できます。

Hugging Face transformers ライブラリには、 Transformers モデルの読み込みとファインチューニングを可能にする Trainer ユーティリティと Auto Model クラスが用意されています。

これらのツールは、簡単な変更で次のタスクに使用できます。

ファインチューンするモデルの読み込み。
Hugging Face Transformersトレーナーユーティリティの構成を構築します。
1 つの GPU でトレーニングを実行する。

Hugging Face Transformersとはを参照してください。

要件

ドライバーに 1 つの GPU を搭載した単一ノードクラスター。
Databricks Runtime 13.0 ML以降の GPU バージョン。
- このファインチューニングの例では、 Databricks Runtime 13.0 ML以降に含まれている🤗Transformers、🤗データセット、および🤗評価パッケージが必要です。
MLflow 2.3.
Transformersを使用してモデルをファインチューンするために準備およびロードされたデータ。

Hugging Face データセットをトークン化する

Hugging Face Transformers モデルでは、ダウンロードされたデータのテキストではなく、トークン化された入力が必要です。基本モデルとの互換性を確保するには、基本モデルから読み込まれたオートトークナイザーを使用します。 Hugging Face datasets を使用すると、トレーニングデータとテストデータの両方に一貫してトークナイザーを直接適用できます。

例：

from transformers import AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained(base_model)
def tokenize_function(examples):
    return tokenizer(examples["text"], padding=False, truncation=True)

train_test_tokenized = train_test_dataset.map(tokenize_function, batched=True)

トレーニング構成を設定する

Hugging Face トレーニング構成ツールを使用して、トレーナーを構成できます。トレーナークラスでは、ユーザーが以下を提供する必要があります。

メトリクス
基本モデル
トレーニング構成

コンピュート Trainer デフォルトの loss メトリクスに加えて、評価メトリクスを設定できます。次の例は、 accuracy をメトリクスとして追加する方法を示しています。

import numpy as np
import evaluate
metric = evaluate.load("accuracy")
def compute_metrics(eval_pred):
    logits, labels = eval_pred
    predictions = np.argmax(logits, axis=-1)
    return metric.compute(predictions=predictions, references=labels)

NLPの自動モデルクラスを使用して、タスクに適したモデルをロードします。

テキスト分類の場合は、 AutoModelForSequenceClassification を使用して、テキスト分類の基本モデルを読み込みます。モデルを作成するときに、データセットの準備中に作成されたクラスの数とラベルマッピングを指定します。

from transformers import AutoModelForSequenceClassification
model = AutoModelForSequenceClassification.from_pretrained(
        base_model,
        num_labels=len(label2id),
        label2id=label2id,
        id2label=id2label
        )

次に、トレーニング構成を作成します。 TrainingArguments クラスを使用すると、出力ディレクトリ、評価方法、学習率、およびその他のパラメーターを指定できます。

from transformers import TrainingArguments, Trainer
training_args = TrainingArguments(output_dir=training_output_dir, evaluation_strategy="epoch")

データコレーターを使用すると、トレーニングデータセットと評価データセットの入力がバッチ処理されます。DataCollatorWithPadding は、テキスト分類のベースラインパフォーマンスが良好です。

from transformers import DataCollatorWithPadding
data_collator = DataCollatorWithPadding(tokenizer)

これらのパラメーターをすべて構築したら、 Trainerを作成できます。

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=train_test_dataset["train"],
    eval_dataset=train_test_dataset["test"],
    compute_metrics=compute_metrics,
    data_collator=data_collator,
)

トレーニングして MLflowに記録する

Hugging Face MLflow と適切に連携し、 MLflowCallback を使用してモデルのトレーニング中にメトリクスを自動的にログに記録します。ただし、トレーニング済みのモデルを自分でログに記録する必要があります。

MLflow の実行でトレーニングをラップします。これにより、トークナイザーとトレーニング済みのモデルから Transformers パイプラインが構築され、ローカルディスクに書き込まれます。最後に、 mlflow.transformers.log_model を使用してモデルを MLflow に記録します。

from transformers import pipeline

with mlflow.start_run() as run:
  trainer.train()
  trainer.save_model(model_output_dir)
  pipe = pipeline("text-classification", model=AutoModelForSequenceClassification.from_pretrained(model_output_dir), batch_size=1, tokenizer=tokenizer)
  model_info = mlflow.transformers.log_model(
        transformers_model=pipe,
        artifact_path="classification",
        input_example="Hi there!",
    )

パイプラインを作成する必要がない場合は、トレーニングで使用されるコンポーネントをディクショナリに送信できます。

model_info = mlflow.transformers.log_model(
  transformers_model={"model": trainer.model, "tokenizer": tokenizer},
  task="text-classification",
  artifact_path="text_classifier",
  input_example=["MLflow is great!", "MLflow on Databricks is awesome!"],
)

推論のためにモデルを読み込む

モデルがログに記録され、準備ができたら、推論のためにモデルを読み込むことは、MLflow でラップされた事前トレーニング済みモデルを読み込むことと同じです。

logged_model = "runs:/{run_id}/{model_artifact_path}".format(run_id=run.info.run_id, model_artifact_path=model_artifact_path)

# Load model as a Spark UDF. Override result_type if the model does not return double values.
loaded_model_udf = mlflow.pyfunc.spark_udf(spark, model_uri=logged_model, result_type='string')

test = test.select(test.text, test.label, loaded_model_udf(test.text).alias("prediction"))
display(test)

一般的なCUDAエラーのトラブルシューティング

このセクションでは、一般的なCUDAエラーとその解決方法に関するガイダンスについて説明します。

OutOfMemoryError: CUDA のメモリ不足

大規模なモデルをトレーニングする場合、発生する可能性のある一般的なエラーは、CUDA のメモリ不足エラーです。

例：

OutOfMemoryError: CUDA out of memory. Tried to allocate 20.00 MiB (GPU 0; 14.76 GiB total capacity; 666.34 MiB already allocated; 17.75 MiB free; 720.00 MiB reserved in total by PyTorch) If reserved memory is >> allocated memory try setting max_split_size_mb to avoid fragmentation.  See documentation for Memory Management and PYTORCH_CUDA_ALLOC_CONF.

このエラーを解決するには、次の推奨事項を試してください。

トレーニングのバッチサイズを小さくします。トレーニング引数の per_device_train_batch_size 値を小さくすることができます。
精度の低いトレーニングを使用します。 fp16=True は [トレーニング引数] で設定できます。
トレーニング引数でgradient_accumulation_stepsを使用して、バッチ全体のサイズを効果的に増やします。
8ビットのアダムオプティマイザを使用します。
トレーニングの前に GPU メモリをクリーンアップします。場合によっては、GPUメモリが未使用のコードで占有されている可能性があります。
```
from numba import cuda
device = cuda.get_current_device()
device.reset()
```

CUDA カーネルエラー

トレーニングを実行すると、CUDAカーネルエラーが発生する場合があります。

例：

CUDA kernel errors might be asynchronously reported at some other API call, so the stacktrace below might be incorrect.

For debugging, consider passing CUDA_LAUNCH_BLOCKING=1.

トラブルシューティングを行うには:

CPUでコードを実行して、エラーが再現可能かどうかを確認してください。
別のオプションは、 CUDA_LAUNCH_BLOCKING=1を設定してより良いトレースバックを取得することです。
```
import os
os.environ["CUDA_LAUNCH_BLOCKING"] = "1"
```

ノートブック: 1 つの GPUでのテキスト分類のファインチューニング

コード例をすばやく使用するために、このサンプルノートブックでは、テキスト分類用にモデルをファインチューンするためのエンドツーエンドの例を提供します。この記事の以降のセクションでは、 DatabricksでHugging Face を使用してファインチューンする方法について詳しく説明します。

テキスト分類Hugging Faceモデルのファインチューニングノートブック

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