インスタンスサイズを上げたのに実行時間が変わらない？SageMaker notebook instancesでの機械学習のボトルネックをCloudWatchで特定した話

自己紹介

NTTドコモ データプラットフォーム部(以下DP部)の江口です。普段はドコモが保有する顧客基盤データによる価値創出のためのパートナー企業向けデータプロダクトの企画/開発を行っております。その中でドコモデータを用いた生成モデルの学習を行うことになり、GPUを用いた機械学習モデルの構築をSageMaker notebook instances上で行いながら、実験中のリソース状況をモニタリングする必要性がでてきました。

通常、SageMakerで本格的なモデル学習を行う際は、管理やコストの観点から「SageMaker Training Jobs」を使用するのが一般的です。しかし、今回のフェーズはモデル構築の初期段階であり、ハイパーパラメータの調整やコードの試行錯誤を高速に繰り返す必要がありました。

そこで、対話的な実行が可能でハンドリングが良い「SageMaker notebook instances」をあえて実験環境として採用しました。しかし、notebook上での学習は手軽な反面、リソースの使用状況が見えづらくなるという課題もあります。その課題への一つの対応策を本記事にまとめました。

なお、本記事の掲載の取組については、支援メンバーであるAMBLの濵田さんに詳細検討・実装を進めてもらっており、以下は、濵田さんに執筆いただいています。

はじめに

皆さんこんにちは。NTTドコモ データプラットフォーム部でデータプロダクト開発を担当しているAMBL株式会社の濵田幸杜です。
本記事は、アナリストとして仕事をしてきた私が、業務ではじめて機械学習モデルの学習を進める中で直面した 「パフォーマンスの壁」 と、その解決の鍵となった Amazon CloudWatch についてお話しします。本記事では実験環境として、取り回しが容易なAmazon SageMaker notebook instancesを使用しております。

インスタンスサイズを上げても実行時間が変わらない

私は現在業務でとある機械学習モデルの構築・学習を行っています。
開発当初は、少量のデータ を用いて比較的安価なインスタンスで学習を行っていました。

その後、精度向上や本番適用を見据えて学習データ量を徐々に増やしたため、処理時間の短縮を目的として、インスタンスサイズを段階的に上げていきました。

ちなみに当たり前の話ではありますが、インスタンスサイズを上げると以下のように時間当たりの料金も性能に比例して上がっていきます。

インスタンスごとの性能と料金（スタンダードインスタンス）

インスタンスごとの性能と料金（高速コンピューティング）

しかし、ある時点から奇妙な現象が起きました。

「あれ？ インスタンスサイズを大きくしたのに、実行時間が変わらない……？」

「性能を上げる＝インスタンスサイズを大きくする」
機械学習モデリングをはじめたての私は単純にそう考えていましたが、コストだけが上がり、パフォーマンスが頭打ちになったことで、他に工夫できる点はないか考えました。

Amazon SageMakerにおけるml.g5シリーズなどの一部インスタンスの場合、例えば2xlargeから4xlargeへサイズを上げても、搭載されているGPUは数もスペックも変わらず、CPUの数とメモリだけが増えるケースがあることをこのタイミングで知ることになりましたが、そもそものボトルネックがどこにあるのかを可視化したくなりました。

ボトルネックの特定：何を調査すべきか

闇雲にインスタンスを変えるのではなく、「何が原因なのか」を特定するため、以下の観点で調査を行うことにしました。それにより今後の計算リソースの適切な選択ができる知識を得られるのではないかと考えました。

• CPUとGPUの連携状況（データ供給の遅延はないか？）

  • GPUの計算速度は非常に高速ですが、そのGPUにデータを渡すための「前処理」を行っているのはCPUです。「CPUによるデータ供給が遅すぎて、GPUが待ちぼうけ（手待ち状態）になっていないか？」を確認します。

• GPUリソースの使用効率（余力は残っていないか？）

• 高性能なGPUに変更しても、一度に処理させるデータ量（バッチサイズ）が少なければ、宝の持ち腐れになります。「GPUの計算コア（Compute）やメモリ容量（VRAM）がスカスカの状態（低負荷）になっていないか？」を確認し、設定を上げる余地があるかを判断します。

これらのリソース状況は、notebookの画面を見ているだけでは分かりません。
また、nvidia-smi 等のコマンドでは「その瞬間」の値は確認することができますが、今回は長時間のパフォーマンス推移を把握したいと考えました。
そこで、インスタンス内部のパフォーマンス推移を簡易に可視化するために Amazon CloudWatch を活用することにしました。

『CloudWatch』とは？

AWSの公式ドキュメントには、CloudWatchについて以下のように記述されています。

Amazon CloudWatch は、Amazon Web Services (AWS) リソースと AWS で実行されているアプリケーションをリアルタイムでモニタリングする機能を備え、アプリケーションのパフォーマンス、運用状態、リソース使用率をシステム全体で把握できるさまざまなツールが用意されています。

私が実際に使用してみた感覚も含めて噛み砕いて言うと、要するに 「AWSを使ったときの裏側がどうなってるか確認ツール」 のようなものです。

CPU使用率やGPU使用率、メモリ状況などを詳細にトラッキングすることができます。
トラッキング指標は様々ですが、先の仮説の検証のためには以下の指標をピックアップします。

• cpu_usage_active（CPU使用率）：全体の使用率だけでなく、特定のコアだけが張り付いていないかなども考慮します。
• utilization_gpu（GPU使用率）：この値が低い場合、インスタンスサイズが大きくてもうまく使えていないことになります。
• utilization_memory（GPUメモリ使用率）：GPU上のメモリ（VRAM）がどの程度使われているかを示します。この値が低い場合、バッチサイズを上げるなどして、一度により多くのデータを処理できる余地があることを示唆します。

これらを可視化し、インスタンスの中で何が起きているかを把握します。

モニタリング方法

モニタリングに必要なものは以下の二つです。
* モニタリング指標を記述したjsonファイル * 実行しているnotebookのhostname

1つ目のjsonファイルは以下のように記述します（命名は何でもいいですが、私はconfig.jsonと命名していました）

{
  "agent": {
    "metrics_collection_interval": 60,
    "run_as_user": "root"
  },
  "metrics": {
    "metrics_collected": {
      "cpu": {
        "measurement": [
          "cpu_usage_active"
        ],
        "totalcpu": true
      },
      "nvidia_gpu": {
        "measurement": [
          "utilization_gpu",
          "utilization_memory"
        ]
      }
    }
  }
}

2つ目のhostnameはnotebook上で!hostnameと入力すると出力されるものです

hostname出力例

これらを用意した上で以下のコードをnotebook上で実行します。

!sudo /opt/aws/amazon-cloudwatch-agent/bin/amazon-cloudwatch-agent-ctl \
-a fetch-config \
-m ec2 \
-c file:config.json \
-s

これらの準備が完了したら、いよいよCloudWatchの画面に遷移します。

実際にモニタリングしてみた

GPU搭載インスタンスのml.g5.2xlarge内で機械学習を行った際に、以下の指標（メトリクス）をダッシュボードで確認しました。

• cpu_usage_active（CPU使用率）
  
• utilization_gpu（GPU使用率）
  
• utilization_memory（GPUメモリ使用率）
  

CloudWatchの画面への遷移は以下の通りです

まずはホーム画面の左上の検索画面から「CloudWatch」を入力し、選択します。

CloudWatch検索画面

次に画面左のメトリクスから「すべてのメトリクス」を選択します。

メトリクス選択

最後に、notebook上で確認したhostnameを入力すると、jsonファイルで指定したメトリクスがダッシュボード上で確認できるようになります。

※hostnameはnotebookを再起動する度に変更されるため、モニタリングする際はnotebook起動後に必ずコードを実行してください

実際のダッシュボード画面は以下の通りです。
青線：cpu_usage_active
黄線：utilization_gpu
黒線：utilization_memory

ダッシュボード例

上記のようにGPUとCPUの使用率やメモリの使用率をモニタリングすることができます。
また、今回私は使用しませんでしたが、アラームを設定して閾値を超えた場合に通知を送信するようにしたり、ログ管理でログデータの収集をしたりすることができます。

実際の改善事項

私が実行していたコードでパフォーマンスが頭打ちになっていた原因はGPU使用率が低く、GPUメモリにも余裕があることでした。

これは、GPUの計算能力を十分にいかせていないということです。
そのため実験に用いていた機械学習アルゴリズムの中で、PyTorch由来の以下2つのハイパーパラメータを変更し、「バッチあたりの計算量の向上」と「データ供給量のパイプライン化」を図りました。

• バッチサイズ（batch_size）を大きくする
  • utilization_memoryが低く、GPUのメモリに余裕がありました。そこで、一度にGPUに渡すデータ量（バッチサイズ）を増やすことで、GPUの計算リソースをより効率的に活用するようにしました。
  • ただし、バッチサイズを大きくしすぎると「CUDA Out of Memoryエラー」が発生するため、メモリ使用率を見ながら調整を行いました。
• 並列読み込み数（num_workers）を増やす
  • 「見かけ上のCPU使用率は低い（余裕があるように見える）のに、実は特定コアがパンクしていてGPUへのデータ供給が滞っている」 というケースは頻繁に起こります。そのため、num_workers を増やして処理をマルチプロセス化し、複数のコアに負荷を分散させることでボトルネックを解消しました。
  • CPU使用率は 「全コアの平均値」 であるため、数値の解釈には注意しながら調整を行いました。

改善結果

これらの変更を加えた結果、実際にGPU使用率が大幅に上昇しました。

コード改善前ダッシュボード

コード改善後ダッシュボード

また、実行時間も大幅に改善されました。
（GPU搭載インスタンスのml.g5.2xlargeで約4h30mから約3h00mへ短縮）

※改善するために行った「バッチサイズを大きくする」は、学習の挙動が変わる可能性があり、実際に採用する場合は学習率の調整が必要となる可能性があります。今回は計算リソースのパフォーマンス最適化を目的としており、この目的は達成できました。

おわりに

今回の経験で、ただコードを書くだけでなく、「確保したリソースがどう使われているか」を可視化することの重要性を改めて痛感しました。
ただ闇雲にインスタンスサイズを上げるのではなく、そのインスタンスサイズで最大限のパフォーマンスを発揮できているかをCloudWatchを使って確認することで、コストパフォーマンスを最大化することができます。

この記事を読んだ皆さんも、まずは「計測」を意識して、CloudWatchを活用してみてはいかがでしょうか？

参考資料

• Amazon SageMaker AI の料金
  https://aws.amazon.com/jp/sagemaker-ai/pricing/
  
• AWS CloudWatchについて
  https://docs.aws.amazon.com/ja_jp/AmazonCloudWatch/latest/monitoring/WhatIsCloudWatch.html