「学習」のアジリティを高める SageMaker Training Job

ツール導入前の課題

2022 年の 11 月頃、内製の OCR エンジンである NineOCR をリリースし、弊社プロダクト内での運用が開始しました。このリリースによって、リアルな課題が継続的にフィードバックされるようになりました。また、リリースまでの開発・リリース後の改善がスムーズに実現できるように NineOCR の開発メンバーも増員されました。

メンバーの増員・課題の増加という状況の変化が起きたものの、開発や改善のための検証は EC2 の GPU インスタンス（ml.g4dn.xlarge や ml.g4dn.metal）を数台起動して各メンバーが作業を行う方式を採っていました。この開発方式は、クリティカルではないものの以下の課題が顕在化していました。

1. パラメータチューニング等で一時的に GPU リソースを増やしたいときに、各種コストを加味して躊躇う場面がある
2. 常に GPU リソースを使うとは限らず、必要以上のコストを支払っている場面がある

1 について、必要に応じてインスタンスを増やすことがストレートな解決策ですが、非常に簡素ではあるものの権限を保有するメンバーへの依頼が必要であり作成に若干のリードタイムがあります。また、インスタンスを作成すると棚卸し作業や OS の更新等で一定のメンテナンスコストがかかります。中長期的に利用する見立てがあればインスタンスを増やすのが良い選択肢ですが、一時的な用途の場合は作成のコストが上回ると感じる場面がありました。

2 について、一般に GPU を積んだ EC2 インスタンスの維持コストは比較的高価（例えば東京リージョンにおける g4dn.metal の年間維持コストは 2025 年 2 月 5 日現在の価格で約 1,384 万円）です。開発では常に GPU リソースを利用するとは限らず、データ分析や評価、API 仕様の修正等を行う時間もあるため、GPU リソースを使い切れていないと感じる場面がありました。

どのような状態を目指していたか

上述した２つの課題を解決できる状態を目指していました。

1. 必要なときに必要なだけのリソースを確保でき、インフラストラクチャレベルでのメンテナンスが不要なサービスを利用している状態
2. 使用していない間はインスタンスを停止する等、リソース管理の仕組みが導入されている状態

既に関連するリソースが AWS に存在していることから利用するクラウド環境は AWS を前提として考えていました。サービスとしては ECS, EKS, SageMaker 等の選択肢がありますが、GPU インスタンスを大きな手間なく扱うことができる点から SageMaker の利用を早い段階で決めました。しかし、ひとえに Amazon SageMaker Training を利用するしても、AWS 公式ブログ記事にもある通り、いくつか実装パターンがあります。

1. AWS が提供しているコンテナイメージを拡張する方法
2. 独自のコンテナイメージに SageMaker Training Toolkit をインストールする方法
3. スクラッチでコンテナイメージを作成する方法

上記の選択肢からどの実装パターンを採用するかをいくつかの観点で比較しました。

選定にあたって重要視したのは、以下の 2 点です。

1. ローカル（EC2）のホストでも、コードをほとんど変更せずに実行可能な構成にできるか
2. 共通処理（ファイルマウント・ログ・エラーハンドリング・S3 への保存等）を実装をスキップできるか

1 について補足します。深層学習のモデル開発は小規模に実験してコンセプトを確かめるフェーズと、大規模に実験して高精度化を狙うフェーズに大別され、これらを繰り返して試行錯誤するフローと考えています。小規模に実験する際は、テストやデバッグをしながら高速に実装を進めたいです。また、一部の設定（学習率スケジューラの挙動・バッチサイズ等）は実際に学習しながら決定したいです。このフェーズではアジリティを担保するために EC2 等のローカル環境で開発を行うのが望ましいと考えます。一方で大規模に実験するフェーズでは多少のアジリティは犠牲にして実験を複数同時に実行できること・強力なインスタンスが利用できることが望ましいと考えます。試行錯誤を高速に行うために、ローカル環境で小規模に実験する際と大規模に実験する際でコードの変更が極力不要な構成にしたく、この観点を重視しています。

2 については共通処理をマネージドサービスに委ねることで開発のボリュームが減って検証のアジリティが高まることを期待し、重視しています。

まず、それぞれの構成が、重要視している観点を充足するかを机上で調査しました。

1. AWS が提供しているコンテナイメージを拡張する方法

AWS が提供しているイメージは、任意の Python バージョン、深層学習フレームワークのバージョンが提供されているとは限りません。ローカル環境と利用可能なイメージのバージョンに差異がある場合、ローカル環境側のバージョンを変更しなければならないため、ローカルでの実行はハードルが高い場合があります。一方で、提供されたコンテナイメージは、エラー処理やログ出力に加え複数 GPU での学習やマネージドスポットトレーニングにも対応しており、共通処理の実装コストは非常に小さく済みます。

2. 独自のコンテナイメージに SageMaker Training Toolkit をインストールする方法

ローカルの OS と合ったイメージを選択し、バージョン管理ツール経由でライブラリをインストールすることでローカルとコンテナイメージの環境差異をほとんどなくすことができます。共通処理は SageMaker Training Toolkit が多くを担うため、お作法に従った記述を強いられるものの実装コストを抑えることができます。

3. スクラッチでコンテナイメージを作成する方法

ローカルの OS と合ったイメージを選択し、バージョン管理ツール経由でライブラリをインストールすることでローカルとコンテナイメージの環境差異をほとんどなくすことができます。共通処理をサポートする機能はないので全て独自で実装する必要があります。

まとめ

1 の手法はローカルホストとは環境差異が発生する可能性が高く、3 の手法は共通処理をすべて実装する必要があり開発のコストが高くなるため、2 の方法を採用しました。

改善したかった課題はどれくらい解決されたか

SageMaker Training Job の採用によって柔軟にインスタンスタイプを選択しつつ、複数の学習を同時に実行できるようになりました。当初課題に感じていた「パラメータチューニング等で一時的に GPU リソースを増やしたいときに、各種コストを加味して躊躇う場面がある」といった課題は概ね解決されたと考えています。 また、普段使いの作業用インスタンスと SageMaker Training Job を使い分けることで「常に GPU リソースを使うとは限らず、必要以上のコストを支払っている場面がある」という課題に関しても概ね解決されたと考えています。

どのような成果が得られたか

研究員のメンバーが新規参画した際は、検証用の比較的軽量な GPU インスタンスを用意しつつ、大規模に実験したい場合は SageMaker Training Job を利用するというスタイルが根付きつつあります。大規模な GPU インスタンスをメンバーごとに用意する場合と比べて大きくコストが削減されており、レバレッジの効く成果が得られています。

また、弊社は 2024 年 8 月頃に文書画像をデータ化する Vision-Language モデル "Viola" を作成し、リリースしました。 当時は Vision-Language モデルに関する知見が社内に少なかったため、技術不確実性が高いプロジェクトとして開始しました。このとき、初期フェーズではコストをかけすぎずに（従量課金かつ安いインスタンスを使って）不確実性を潰し、ある程度価値創出の見込みが立ったら（大規模なインスタンスを使って）大きく投資する、という「小さく始め、大きく育てる」戦略を採用しました。これを実現できたのは SageMaker Training Job の機能を活用したことが大きいです。 新規に機械学習のプロジェクトを立ち上げる際に強力な基盤として利用できたことからも大きな成果を得ることが出来たと考えています。

上述した通り、機械学習の検証プロセスは小規模な実験と大規模な実験を行ったり来たりします。 クラウド環境での動作を前提とする設計を採用すると、小規模な実験を行いたいときにクラウド環境での長い実行リードタイムが強いペインとなり、普及展開が困難になると考えていました。

上記背景から、手元の環境で動くコードをほとんど変更せずにクラウド環境で実行可能な設計にすることは必須の要件と捉えており、この構成を実現することに苦労しました。

機械学習における「学習」は比較的ピーク性のあるタスクであり、クラウドマネージドサービスの恩恵（従量課金・スケーラブル）を受けやすい性質を持つため、一度は検討の価値があると考えています。

私は、「学習」は一般的なバッチジョブとは異なり、試しに数分実行したいシーンが頻繁にあると思っています。例えばバッチサイズは一度学習を実行してから決めることが多いです。学習率もしばらく学習を行ってメトリクスを確認しながら調整することがあります。このようなユースケースは、完全にクラウド化してしまうと迅速な対応が困難になる可能性があります。

上記の例は普遍的ではないかもしれませんが、実際に学習を行うメンバーと連携したうえで課題に沿った適切なツールを選定することが重要であると考えています。今回の意思決定がどなたかの参考になれば大変幸いです。

機械学習モデルの継続的な改善を可能にする体制を整えています。 今回の内容とは別軸ですが NineOCR の学習を効率化するために、「Feature Store の構築」を行いました。 また、実験管理ツールの MLFlow が SageMaker に統合されたマネージドサービスとして利用可能になったため、利用を開始しました。いくつかのモデルは SageMaker Processing を用いてデータセットの作成・評価を行っています。

周辺プロセスを SageMaker と連携しより多くのトイルを削減することで「学習」に限らず機械学習ワークフロー全体のアジリティを高めていきたいと考えています。