データセットの読み込みにはEBS、その成果物にはEFSを使うとよい　文献から学ぶクラウド機械学習のベストプラクティス

文献についての紹介

服部圭悟氏（以下、服部）： 本日は、文献『Cost-efficient and scalable ML-experiments in AWS with spot-instances, Kubernetes and Horovod』について紹介したいと思います。

実はこちら、Twitterのほうで紹介してもらったのを読んだのですが、すごくよいなと思って。ABEJA Platformというプラットフォーム事業をやっているのですが、その経験を踏まえてもすごくいいところがあると思ったので、その経験で補足できるところを補足しながら紹介していきたいと思います。

前提としてAWSの各種サービスとKubernetesのサービスを使っています。できるかぎり補足しようとは思うんですけれども、足りないところはご了承ください。

それからプラットフォームを作ることを前提として置いているので、例えばなにか自分1人でやっているプロジェクトだとかPoCとかでトライ＆エラーをたくさん回すようなフェーズだと、ここまでのものはたぶん必要ないと思うので、そのへんも踏まえてお聞きください。

私は服部と申します。機械学習をずっとやっていて、3〜4年ぐらい前からMLOpsをやり始めています。MLOpsが大好きなので日々キャッチアップをしています。

LINEとABEJAではインフラすごく強い人がいっぱいいたので、彼らに頼って生きてきています（笑）。なので、今日もAWSのサービスのこととか聞かれてもちょっと困るので、そのへんは強い人、誰かが助けてくれるといいなと思っています。よろしくお願いします。

今日持ち帰れるもの

まず機械学習に関わる仕事について紹介したいと思います。機械学習は大きく分けると2つあります。機械学習のモデルを作る学習周りの作業と、そのモデルをプロダクトで実運用する運用のところの作業です。今回の話は学習のプラットフォームの話です。

学習のプラットフォームでもけっこういろいろとサブタスクがいっぱいありますし、運用のところとかもすごいサブタスクがいっぱいがあって、お酒を飲みながらみなさんとお話ししたいなと思っているのですが、今日はちゃんと文献の紹介をしたいと思います。

最初に結論を言います。今回この文献からは、これらのことが持ち帰れるので、これらだけ持って帰って適当にアーキテクチャを組んでも、そこそこいいものができると思います。

まず、学習のデータセットはEBSにもって、学習の成果物はEFSに保存するようなかたち……両方ともインスタンスにマウントする感じですね。で使うとよいよという話。そしてリソースの管理はSpot Instancesを使うとコストも抑えられていいよという話。

それから機械学習のジョブはどうやって回すかというと、Kubernetes Jobを使ってやるといいよって話があります。最後に、分散学習はHorovodを使うと変更も少なく簡単にできるよという話があります。

これらについてもう少し深い話をしていきたいと思います。

S3、EBS、EFS

はじめにデータの話をします。最初にAWSのストレージの特徴を少し紹介したいと思います。ここではDBは扱いません。

大きく分けると3つあります。S3、EBS、EFSの3つですね。S3が一番安くて、みなさん馴染みがあるものだと思います。容量を気にしないで使えるストレージで堅牢性と可用性が非常に高いので、とりあえずファイルをここに置いておくと安心・安全というかたちで使えます。

S3はオブジェクトストレージになっていて、ファイル単位にエンドポイントがあります。なので、ファイル単位でRead、Write、そして更新をやる感じになっています。

次がEBSです。EBSはブロックストレージになっていて、値段はS3よりは高くてEFSよりは安いです。ブロックストレージなので、容量は事前に、10G（ギガ）使うとか20G使うとか定義するかたちになります。それをマウントするかたちで使うのが一般的です。

Availability Zone（AZ）が1つの中でしか動かなくて、インスタンスも1つからしかマウントできなません。実際はちょっと違うんですけど、ここではそういう説明をさせてください。

そしてEFSです。EFSはほかの2つのストレージに比べると高いです。ただ、S3と同じように容量を気しないで使うことができます。Multi-AZで使えるのと、マルチインスタンスからアタッチできるというところがあって、けっこう便利に使えます。

データの読み込みにはEBS

というのを踏まえた上で、機械学習の実際の作業ではデータをどう使うかというところを少し整理したいと思います。大きく分けると「蓄積」「読み込み」「書き出し」ですね。

蓄積は学習に使うデータとかがポンポン蓄積していくということをすると思います。例えばアノテーション作業をして、アノテーションのデータもポンポン貯めていくと。実際に機械学習を回すときは、その貯めたデータを読み込んで使う、成果物はどこかに書き出すようなかたちになっていると思います。

パイプラインとかを組めば、前処理・後処理のパイプラインでそれぞれ成果物を保存する・読み込むところが発生すると思いますし、例えば「100エポック回すときの中間結果を取りたいから、10エポックごとに結果を吐き出す」ような使い方もあるでしょう。けれども大きく分けると、やっぱり「蓄積」「読み込み」「書き出し」という作業になります。

この文献では、データセットの読み込みにはEBSを使って、その成果物、アーティファクトにはEFSを使うとよいという話をしています。

ここからは、我々のABEJAでの経験を踏まえた上でのよいところを考察しながら紹介したいと思います。データセットにEBSを使うことのすごくいい点は、学習における総実行時間が短縮できるところですね。

何を言っているかというと、機械学習でとくに深層学習だとすごい大量のデータを扱うので、ファイルの読み込みが非常に多く発生します。「すべてのファイルを舐めたあとに、次のエポックでまた舐める」みたいなことが行なわれます。

データがメモリに全部乗ればいいのですが、メモリどころかディスクにも乗らないみたいなこともよくあって、どうしてもやっぱりディスクに乗らないので、例えばS3をマウントして使うとか、もしくは都度S3のファイルを参照しながら学習を行うという選択肢になります。

そう考えたときに何が使えるかという選択肢を、S3・EBS・EFSで考えると、まずS3はオブジェクトストレージなので、1個1個ファイルを指定して読み込むファイルローダを書かなきゃいけないんですけど、たぶんみなさん（機械学習エンジニア）はローカルで作業するときにそんな使い方しないですよね。

「データフォルダがあって、データフォルダの下にデータを置いて、そこを参照する」みたいな使い方をすると思います。

goofysみたいなものを使って無理やりS3をマウントすることもできなくはないですけど、不安定でけっこうマウントが外れるんですよね。学習の途中でマウント外れたときの扱いとかを制御をいちいち書くのはけっこうつらい。

残ったEBS・EFSだとどっちを使うのがよいかというと、EFSは非常に読み書きが遅いので、消去法的にEBSを使うという話になります。

EBSの辛いところ

「EBS、これ、こういうふうに使えたら本当はもっとよかったのにな」と思うところもあります。例えばMulti-AZで計算機クラスタを組んで、耐障害性を上げるためにMulti-AZを使って、そこの上で学習ジョブをいっぱい走らせたいとき。

例えばResNet50やResNet100などネットワークのデザインが異なる学習ジョブを作ったりハイパーパラメータの値を変えたりして「同じ学習データの中で一番性能のいいモデルを吐き出したい」という用途があります。1つのEBSに学習データを貯めてすべての学習ジョブからReadOnlyでマウントできたら、シンプルな構成でお金も節約できると思うんですけど、EBSはSingle-AZなのでMulti-AZな計算クラスタからは使えません。

KubernetesのPersistent VolumeでReadWriteManyでマウントできたらよかったのになってところがちょっとつらいところですね。

文献（Rosebud AI）ではどういうふうにしたかというと、EBSをジョブごとに作成して、データセットをS3からスナップショットしてくるという使い方をしています。EBSはReadOnlyでマウントしていて、ジョブが完了したらそいつを破棄するのが基本的なオペレーションになっています。

再利用するというオプションもあって、その場合は次のジョブですぐそのボリュームをマウントして動かします。そうするとデータのスナップショットの時間が省略できるので、すぐにジョブが走るというメリットがあります。

ABEJA PlatformではEFSを採用

ABEJA Platformではどうしたかというと、EBSではなくEFSを使っています。

どうやってデータをダウンロードしているかというと、EFSってMulti-Attachができるので、複数のジョブでマウントできるんですね。「保存先にファイルが存在したらダウンロード処理は実行しない」という制御を1つ入れるだけで、複数のジョブがそれぞれマルチスレッドでダウンロードするみたいなことができるんです。そうするとダウンロード効率がよくなるだろうと、こういうふうにしています。

ただ、やっぱりEFSはファイル読み込みがめちゃくちゃ遅いので、EBSにしようかという話をしているところでした。

EFSのいい点と悪い点

続きましてEFSを使う理由です。最初に挙げられるのは、機械学習では成果物のサイズは予測できないという点です。モデルのサイズってパラメータ次第なのでデザインに使ったパラメータでモデルのサイズは予測できるんですけど、そんなことする人はいないでしょうし。

試行錯誤をする段階でディスク容量を気にするというのは、本質的じゃないなというので、アーキテクチャのデザインとしては体験がよくないなというのが、1つあります。例えば、成果物としてログファイルを含めるとか、ログの内容を何にするかとか、いろいろと自由度もあるので。

もう1つは、Disk Fullです。Disk Fullってことは死ぬってことで、つまり成果物が得られないことになるので、それはもう機械学習エンジニアにとっては本当に発狂ものだというところで、その意味でもやっぱりデザインとしてはよくないと。例えば1つのインスタンスに複数のジョブが走る場合に、自分のジョブが誰か知らないやつのジョブがDisk Fullになったせいで巻き添えで死ぬってなったら、それも発狂ポイントだろうと。

逆にEFSを使ってつらい点は何かというと、気をつけないとクラウド破産するところですね。何を言っているかというと、機械学習エンジニアがディスク容量を気にしないでいいようにデザインしたばかりに、彼らは本当にディスク容量を気にしないんですよ。ランニングコストを気しないのでプラットフォーマーのことは気にせず、バンバン使えるから使うと。どんどん中間ファイルを吐き出すみたいなこともします。

例えば自然言語処理で言語モデルとかBERTとかで言語モデルとか作っていると、数十Gのモデルができて、それが何個も試行錯誤で増えて、かつ10エポックごとにモデルを吐き出すみたいになってくると、もう一気に数百G……T（テラ）みたいなかたちになってしまいます。これを放っておくと本当ランニングコストにダイレクトに効いてくるので、つらいところがあります。

文献ではどうしているかというと、成果物をEFSに置いています。チェックポイントを設けて、中間出力をEFSに置くという運用をしています。これはなぜかというと、何かの理由でサービスが落ちたりとかして実験が中断したときに、中断したところから実験を再開できるメリットがあるからです。

さっきEFSがMulti-Attachできると言いましたが、TensorBoardから見たときもすごくよくて、各ジョブが吐き出した結果をTensorBoardが参照してすべてのジョブの俯瞰ができるってところが、1つメリットとしてあります。

ABEJA Platformではどうしたかというと、文献のやったことに加えてJupyter Notebookのマウントをできるようにしています。各ジョブの結果も見れるし、なんなら各ジョブのコードも見れるし、使ったデータも見れるようになっています。

それからランニングコストを下げるために、EFSを定期的に掃除するエージェントを走らせています。しかし、これによって成果物がロストしたら元も子もないので、学習ジョブが完了したらその成果物を圧縮してS3にアップロードするエージェントもいます。