初めまして、JX通信社のMLチームでインターンをしている田中です。これまで物体検知や深層距離学習を使った画像分類、自然言語の分類や生成などのAI作成、MLOpsなどをやっています。今回、MLエンジニアのヨンテさんのもとで文書生成AIのハイパーパラメータ調整の並列化に挑戦しました。
Hydraで書かれたコードをVertex AIでハイパーパラメータ調整するためのサンプルコードはgithubにて公開しています！ ぜひ、一度試してみてください！😊
背景
JX通信社のMLチームでは「力を使うべき場所に注力しよう」の理念のもとに、機械学習用のテンプレートコードを作成しており、その中ではPyTorch LightningやHydraなどが使用されています。
理念やテンプレートコードについてはヨンテさんが書いた解説記事があるので読んでみてください。
tech.jxpress.net tech.jxpress.net
文書生成AIの学習は、１回の学習に数時間程度かかります。
ハイパーパラメータ調整の学習を一つのマシーンで直列で回していると、図１の上の例ように１回の学習時間×学習の試行回数以上の時間がかかり、数日以上かかってしまいます。
JX通信社で大事にしているバリューの一つにSpeedがあり、その理念をML側からも支えるために、学習をより早く終わらせ、事業の開発スピードを早めることを目指しています。
そこで、ハイパーパラメータ調整の際は、たくさんの実験を並列で回したいです 。このような並列学習は、フルマネージドMLサービスであるVertex AIを使用すれば簡単に実現できます(図１下の例)。
一方で、Vertex AI Trainingのハイパーパラメータ調整機能と、JX通信社のテンプレートコードで用いられているHydraの相性が悪く、そのまま使えなかったので、工夫して動くようにしました！
問題点
ハイパーパラメータ調整はVertex AI側がコマンドライン引数を用いてハイパーパラメータを学習コンテナに渡し学習が行われ、学習が終わった後に最適化の評価値を専用のライブラリを使用しVertex AIに伝えることで、ハイパーパラメータの最適化が行われます。
しかし、以下のようにVertex AIが渡すコマンドライン引数の形式とHydraが対応するコマンドライン引数の形式が違うためうまくできませんでした。
Vertex AIの公式ドキュメントに記載されている形式
argparse を推奨
python3 -m my_trainer --learning_rate learning-rate-in-this-trial
⚠️ 注意
自分が調査した限り、実際渡されている形式は上の形式ではなく
--learning_rate=learning-rate-in-this-trialでした。
推奨されている
argparseだとこの形式でも受け取れます。
Hydraの形式
python3 my_trainer.py learning_rate=learning-rate-in-this-trial
解決方法
解決方法はいくつか考えられます。
Vertex AIがHydra形式に対応するのを待つ
→ 一番良いが、対応されるか分からない
Vertex AIのハイパーパラメータ調整機能を使わず、オープンソースのハイパーパラメータ自動最適化フレームワークOptunaを使用する
→ 並列ではなく直列でHydra+Optunaで使用しており、一つのインスタンスで完結するため、さまざまな環境で動かすことができるので便利。しかし、同一マシーンではなく並列で回す際はSQLのサーバーを立てる必要があり大変。(Optunaの並列について)
argparseなど対応しているコマンドライン引数解析ライブラリを使う→ 推奨されているが、パーサーの定義が大変、設定ファイルを使用できない
Vertex AIが渡す引数の形式を変換する処理を挟んで実行する
→ 強引だが、導入コストが一番低い
今回は導入コストが一番低い4を採用しました。
Hydra形式へ変換し、本来の学習コードを実行するスクリプトの作成
まずは、Vertex AIからコンテナに渡される入力(コマンドライン引数)の変換を考えます。今回、コマンドライン引数の変換にはshell scriptを使いました。
shell scriptでのコマンドライン引数の受け取りですが、一番目の引数が欲しい時は $1 で取り、n番目は $nで取れます。今回は引数の個数は可変なため $@で引数全体をそのまま文字列として受け取ります。
次に受け取った文字列を正規表現で置換を行いHydra形式に変換します。
shell scriptなのでsedコマンドを使用して正規表現の置換を行います。
以下のようなコマンドにすることで、引数の形式については
--key=value と --key value 両方とも受け取れるようにします。
sed -r 's/--([^= ]*)[= ]([^ ]*)/\1=\2/g'
sedコマンドはVimの置換と同様に /で区切りs/正規表現/置換/フラグと言う形式で(expr)が \1, \2と対応しています。そして最後の gは複数回置換するフラグです。また、sedコマンドは最短一致に対応しておらず、マッチしたいものの後ろにある文字以外を繰り返すという処理で対処しています。( [^ ]* のところなど)
- 実行例
❯ echo "--lr=0.001 --batch-size 64" | sed -r 's/--([^= ]*)[= ]([^ ]*)/\1=\2/g' lr=0.001 batch-size=64
最後に変換した引数を使って本来の学習コードを実行できるようにします。
$(expr)はexprを実行した結果を文字列として受け取れるので、それを用いて
python train.py $(echo $@ | sed -r 's/--([^= ]*)[= ]([^ ]*)/\1=\2/g') # train.pyはHydraを用いたコード
とすると、変換した引数を本来の学習コードの実行コマンドの後ろに追加できます。
あとはこれを train.shなどでファイルを作成し、Vertex AIの設定ファイルの実行コマンドに ./train.shなどと記載すれば完成です。設定ファイルの実行コマンドの記載場所は HyperparameterTuningJob.trialJobSpec.workerPoolSpec.containerSpec.commandです(参考)。
HydraのコードでVertex AIのハイパーパラメータ調整ができた
このような工夫を行った結果、Vertex AIでハイパーパラメータ調整を並列に計算することができました(図2~4)。図3,4には、複数の色の線が同時に存在しており、これは複数インスタンスが並行して学習していることを意味してます。
これにより、これまで数日かかってしまう学習も、数時間で終えることができました！
まとめ
Vertex AIのハイパーパラメータ調整ジョブの時、スクリプトを挟み正規表現でパラメータの形式を変換することで、Vertex AIから渡されるハイパーパラメータをHydraが受け取れる形式に変換することができました。
これにより、Hydraを使用した既存のコードを使用しながら並列でハイパーパラメータ調整できるようになり、学習時間の短縮ができました。 今回の内容について、sample codeをgithubに公開しています。このコードが皆さまの役に立てれば幸いです。
今回の方法以外により良い方法があるかもしれません。私自身まだまだ、勉強中なので改善点などアドバイス頂けられたら幸いです。
補足：チームでR&Dに取り組む工夫
このブログでご紹介しましたが、JX通信社ではPytorch Lightningをベースとした、テンプレートコードを用いて学習を行うことで、チームで効率よく実験ができるようにしています。
興味があれば以下のブログも読んでみてください
tech.jxpress.net tech.jxpress.net
我々とともに挑戦する仲間を求めています
我々とともに成長しながら、より良い社会のためのMLを開発したい仲間を社員・インターン問わず積極的に募集しています！また、MLエンジニアはもちろん、あらゆる職種のエンジニアを求めています！
正社員、インターン、そして副業・復業として体験的に働くことで、当社のカルチャーや働き方等を知っていただいたうえで正式入社を見極めていただくことが可能な「おためし入社」制度などもあります！ほんの少しでも興味を持たれた方はこちらを覗いてみてください！
*1:⚠️ Vertex AIを利用するさらなる利点
ハイパーパラメータ調整をする方法は主に以下の２つに分けられます。
①実験前に予め、ハイパーパラメータを決めて実験を行い、その中で最も性能の良いパラメータを採用する方法 (Grid searchなどが有名)
②とあるハイパーパラメータを与え実験し、その結果を元に更に次のハイパーパラメータ決めて実験を行い、を繰り返すことで最適なハイパーパラメータを探索する方法 (ベイズ推定などが有名)
①の場合は探索するハイパーパラメータは実験前から決まっているため、すべての実験を一度に並行で回せばよく、図１の下の例の１回目の実験で終わります。しかし、場当たり的にハイパーパラメータを選んでいるため、無駄な計算も多くなってしまい、最適なハイパーパラメータを得るためには膨大なコストがかかる可能性があります。
一方、②の方法では、戦略的にハイパーパラメータを選んでいるため最適なパラメーターをコスパよく高精度に探索することができます。しかし、前の学習結果に依存して次の実験で探索するハイパーパラメータが決まるため、全ての学習を一度に並行で回す事はできません。
並行学習を直列でつなぐことができれば、並行学習とベイス最適化の両方の恩恵を受けることができますが、一般的には実装が非常に大変です。しかし、Vertex AIを利用することで並行学習の直列化をほとんど苦労することなく実施することができます (図1 下の例)。
