Google Colaboratory Pro+で24時間を超える実行を実現する方法

大規模言語モデルの学習などにはGPUが必須になってくるため、Google Colaboratory(以下colab)を利用している方も多いかと思います。
colabは課金したバージョンでも24時間しか接続ができないため長時間の学習には不向きです。
とはいえGCPやAWSなどのクラウドサービスのGPUよりは経済的であり、さらに環境構築もほぼないため運用が楽です。
ここではHuggingfaceのTransformerモデルをファインチューンするときに、colabの24時間制限の中でうまく学習を継続させる方法を共有いたします。
それと余談ではありますがcolabが割安であるということも述べます。

24時間を超えて学習を継続させる方法

モデルの学習をチェックポイントで保存し、24時間たったら再接続を行なってチェックポイントから学習を再開します。
モデルの学習をチェックポイントで保存するためには、
以下のようにTrainingArgumentsのoutput_dirにモデルの保存先を指定します。
なおHuggingFaceの公式ドキュメントはこちら。

from transformers import TrainingArguments

training_args = TrainingArguments(
    output_dir = {チェックポイントでモデルを保存したいディレクトリのパス},
    save_steps=50 # 何ステップごとにチェックポイントで保存するか
)

この指定したディレクトリに、学習途中のモデルのチェックポイントが保存されます。
チェックポイントを保存する頻度はsave_stepsで指定することができます。
このsave_stepsはデフォルトでは500になっているのですが、
あまり低頻度だとチェックポイント排出間の時間が長くなってしまい、
colabの接続が切れる前の最後のチェックポイントから接続が切れるまでの時間の学習が無駄になってしまいます。
かといってsave_stepsを小さな値にすると頻繁にチェックポイントが保存されてしまってストレージを圧迫してしまいます(一つのチェックポイントフォルダは数GB)。
最新のチェックポイントが保存される時には過去のチェックポイントは消去されるのですが、ゴミ箱に残りっぱなしになっているのでストレージは圧迫されます。
colabでファインチューニングする際にはGoogle Driveをマウントしていると思うので、Google Driveをアップグレードするのも手です。
まあファインチューニングするモデルによって最適なsave_stepsは異なってくると思うので、この部分は最初に手探りするしかないでしょう。

保存したチェックポイントをモデルに読み込ませるには以下のようにします。

model = AutoModel.from_pretrained(checkpoint_dir+"/checkpoint-{チェックポイントのステップ数}")

さて、チェックポイントから学習を再開させるには以下のようにもう一工夫が必要です。

from transformers import Trainer

trainer=Trainer(
    args=args,
    model=model,
    tokenizer=tokenizer,
    train_dataset=train_dataset,
    eval_dataset=eval_dataset
)

# resume_from_checkpointでチェックポイントから学習を再開するか指定する
trainer.train(resume_from_checkpoint=True)

これでcolabの24時間制限が来ても、再接続を行なって学習を継続させることができます。

上記のtrainerではチェックポイントがあることを前提にしていますが、
チェックポイントがない状態からファインチューンを始めて、いちいちハードコード的に書き換えずに済ますにはおおよそ以下のようになります。

import os
from transformers import Trainer, TrainingArguments 

checkpoint_dir =　{チェックポイントでモデルを保存したいディレクトリのパス}
MODEL = {HuggingFaceのモデル}

training_args = TrainingArguments(
    output_dir = checkpoint_dir,
    save_steps=50 # 何ステップごとにチェックポイントで保存するか
)

if os.path.exists(checkpoint_dir):
 # チェックポイントからモデルを読み込み 
  model = AutoModelForMultipleChoice.from_pretrained(checkpoint_dir+"/checkpoint-{チェックポイントのステップ数}")
else:
  # HuggingFaceからモデルを読み込む
  model = AutoModelForMultipleChoice.from_pretrained(MODEL)

tokenizer=AutoTokenizer.from_pretrained(MODEL)

# 2回目のcolab接続からは、最初の接続で加工したデータセットをpickleなどで読み込むと早い
train_dataset = ・・・
eval_dataset = ・・・

trainer=Trainer(
    args=args,
    model=model,
    tokenizer=tokenizer,
    train_dataset=train_dataset,
    eval_dataset=eval_dataset
)

if from_checkpoint:
  # チェックポイントから学習を再開
  trainer.train(resume_from_checkpoint=True)
else:
  # 初めから学習を開始
  trainer.train()

データセットに関して上記のコメントにも書きましたが、
2回目からのcolab接続では、初回の接続で作成したデータセットをpickleでdumpしておいたものを読み込むほうが早いです。
そもそもGPUを使ってなお数十時間以上を要する大規模な学習では、データセットの加工でも相当な時間を要するはずです。
colabの再接続でカーネルがリセットされた時にも再利用できるように、一度作ったデータセットは保存しておきましょう。
データセットの加工でcolabの計算Unitを消費してしまうのはもったいないです。
なお、pickleでdumpする方法とdumpしたものをロードする方法はこちら。

Google Colaboratoryは安い

colabの課金形態には、毎月の定期課金と、好きなタイミングで100ユニットか500ユニットを課金できる"Pay As You Go"があります。

colabでGPUのNVIDIA A100を使うと1時間で13.06コンピューティングユニットを消費します。
colabの一回の課金では500コンピューティングユニットが付与されるため、38時間強しか利用できないことになります。
もし何日にもわたって学習を継続させるためには、二日に一回は500コンピューティングユニットを追加購入する必要があります。
colabの最上級の課金であるColaboratory Pro+(500コンピューティングユニットと速いGPUなどを使える)は税込5767円であり、初めはお高く感じます。

しかし実際にはクラウドサービスと比較するとかなり割安です。
例えば、細かい設定は置いておいて、同じGoogleのGCPでNVIDIA A100を使うとなると、およそ1時間で4ドルほどかかります。
colabでは1時間でおよそ150円程度(=5767 * 13.06/500)にしかならないことを考えると、かなりcolabが経済的かがわかると思います（しかもcolabは円建てであるので円安の影響も受けない！）。

colab、おすすめです。