Storage Requirements for AIの紹介

はじめに

本記事ではSNIA（Storage Networking Industry Association）の以下のプレゼンテーション発表の内容について記載します。

• Storage Requirements for AI
  • Compute, Memory, and Storage Summit 2024にて発表されたもの

※自分向けメモとして雑多に書いています。そのため内容を簡潔にまとめるということは出来ていませんので、読みづらいかもしれませんがご了承ください。
※この記事内の画像はすべて公開済みの上記プレゼンテーション資料より引用したものです。

概要

• GPUが脚光を浴びることが多い一方で、AIインフラにおいてストレージが果たす重要な役割を認識することが不可欠である
• データの準備から事前学習、微調整、チェックポイント、推論に至るまで、AIのライフサイクル全体を通じて、ストレージは極めて重要である
• 本プレゼンテーションでは、具体例を用いてこれらの各段階を掘り下げ、特にパフォーマンスに関連する重要な要件を特定する

背景

• 現代のDeep Learningでは何百万もの行列演算が必要となる
• これらの計算を合理的な時間内に行うには並列処理が必須である
• 並列処理のためGPUの需要が高まっており取り合い状態になっている
• AIデータセンターではGPUの利用率を最大化することが求められる

AIのライフサイクル

このスライドでは、AIのライフサイクルを以下の3つのフェーズに分類し、それぞれのタスクおよびストレージに求められる要件について述べている。

• Data Preparation
  • 生データを分散処理フレームワークなどによってAIモデルに適した形式に変換する
  • オープンデータ形式で構造化データと非構造化データを保存し、パフォーマンスの高いアクセスを提供することが重要である
• Training & Tuning
  • GPUに学習データを供給し続けるための高いRead帯域幅が必要となる
  • 学習中にモデルのチェックポイントを素早く保存するための高いWrite帯域幅が必要となる
• Inference
  • ストレージからモデルデータを素早く読み取る必要がある
  • バッチ推論の場合はGPUにデータを供給し続けるための高いRead帯域幅が必要となる

本プレゼンテーションでは、Training & Tuningを中心に説明する。

Training & Tuning の概要

このスライドでは、Training & Tuningについてどのような処理が行われるかを解説している。

• 学習データはストレージから読み込まれ、比較的少数のランダム化されたバッチにパッケージングされる
• バッチを使ってモデルに推論させた結果と正解データを比較して損失スコアを計算する
• スコアをもとにモデルの重み（パラメータ）を調整することでモデルの精度を上げていく
• モデルが安定した精度に収束するまで繰り返す
• 障害に備えて定期的にチェックポイントを保存する

Training & Tuningに必要なストレージのRead性能要件

Training & Tuningについてストレージに求められる性能要件を解説する。

• GPUの利用率を最大化することが目的の場合、GPUベンチマークを使用して様々なモデルに必要なピークパフォーマンスから、それを維持するために必要なストレージのRead性能を逆算することが合理的なアプローチである
  • ベンチマークはMLCommonsやMLPerf Trainingを使うのが最適

NVIDIA H100 80GB GPUを例として、各モデルのベンチマークから算出した必要となるストレージのRead帯域幅は以下である。

列	意味
Name	モデル名
Model Size	モデルサイズ（パラメータ数）
Training Dataset	学習データセットの詳細
Input Size	各学習データサンプルのサイズ
# H100 GPUs	GPU数
Model Throughput	1秒間に学習できたデータ量
Training Data Read BandWidth	Model Throughputを維持するのに必要なRead帯域幅

• モデルによって様々だが40GB/s以上のRead帯域幅が必要になるケースもある
• 必ずしもパラメータ数が大きいモデルがストレージ性能要件も高いとは限らない（Input Sizeの方が影響する）

I/Oパターンはアクセスするストレージによって変わる。
例として、Resnet50というモデルに対し、DLIOベンチマークを実行した際のI/OをトレースすることでどのようなI/Oアクセスを行っているかを解析した結果を示す。（NFSとS3で実施）

Resnet50はディープニューラルネットワークのトレーニングに関連する課題に対処するために設計された一連のモデル（よくわかっていない）。
DLIOはGPUをエミュレートするAIストレージベンチマークであり、以下のような特徴がある。

• GPU不要でストレージパフォーマンスを測定可能
• 学習データセットは1024～150KBのイメージテンソルを含むTFRecord
• 学習データはTensorflowのデータローダを使用して読み込まれる

• NFSで実施した結果
  • 左上は64KB-256KBの安定したIOストリームを示す
  • 右側はバッチ当たりの学習データ数が変化してもIOサイズの分布は変化していないことを示す
  • 左下はデータが順番に読まれている（シーケンシャルReadである）ことを示している
  • GPUを最大限活用するには、64～256KBのIOシーケンシャルReadに対して、6.1GB/sのRead帯域幅を達成する必要がある（6.1GB/sは前のベンチマークの結果から）

• S3で実施した場合
  • 左上は20～50MBの大きなIOが発生していることを示す
  • 右側はNFS同様
  • 左下はランダムアクセスであることを示している

利用するストレージによってアクセスパターンが変わるため、AIストレージは多様なアクセスパターンに対して高いRead帯域幅を確保する必要がある。

チェックポイント書き出しの概要

このスライドでは、Training & Tuningのチェックポイントの書き出しについて解説している。

• 大規模なAIモデル学習には数日～数週間かかる場合がある
• その間、学習データが処理されるにつれてモデルの重みは常に変化する
• 定期的に永続的なストレージにチェックポイントとして保存することでデータ損失を防ぐ
• チェックポイントは通常、1つ以上のファイルとして保存され、それぞれが単一のライターによって順番に書き込まれる
• 学習は通常、チェックポイント中に一時停止され、GPUの使用率が低下するため、チェックポイントの書き出しは迅速に完了することが重要

チェックポイント書き出しに必要なWrite帯域幅の推定

列	意味
Model Parameters	モデルパラメータ数
Checkpoint Size	チェックポイントの容量（1パラメータあたり14バイトで推定）
Total Write BW (GBps) Needed to Checkpoint within Time Limit	制限時間内に書き出しが完了するために必要なWrite帯域幅

• モデルパラメータ数と、何分以内に書き出す必要があるかによって異なる
• 例：175Bのパラメータの場合2450GBのチェックポイントの書き出しが発生、これが360秒（2時間のうちの5%）で書き出す場合は6.8GB/sの帯域幅が必要

チェックポイントのリストア

このスライドでは、Training & Tuningのチェックポイントのリストアについて解説している。

• 1つのチェックポイントが例えば3つのGPUでロードされ、モデルの同じ部分を保持する
• チェックポイントファイルは通常複数のGPUに並列で読み取られる
• チェックポイントがすべてのGPUに復元されるまでモデルのトレーニングを再開できない

チェックポイントのリストアに必要なRead帯域幅の推定

列	意味
Model Parameters	モデルパラメータ数
Checkpoint Size	チェックポイントの容量（1パラメータあたり14バイトで推定）
Total Read BW (GBps) Needed to Restore Checkpoint within 5 Minutes	5分以内にリストアが完了するために必要なRead帯域幅

• 例：175Bのモデルを16のGPUに5分以内にロードする場合、2.18GB/sのRead帯域幅が必要

GPUクラスタとストレージ

このスライドでは、GPUクラスタを考えた場合に必要となるストレージについて述べている。

• これまでは単一サーバ・単一モデルの場合を考えてきた
• 現実的には、GPUクラスタではAIライフサイクルの様々な段階の、多数のAIワークロードをホストしている可能性がある
• 分散スケジューラがクラスタ内のサーバにジョブを割り当てるため、どこに配置されるか関係なく平等にデータアクセスできる必要がある
• 単一のNameSpaceから、AIアクセスパターンが大きく異なる複数のワークロードのパフォーマンスニーズを満たす必要がある、かつスケールアウトできることがストレージとして求められる

まとめ

• ストレージはAIライフサイクルのすべてに関与し、要件はライフサイクルの各段階で異なる
• 適切なAIストレージソリューションを選択するには、予想されるワークロードやサービスが求める要件を洞察することが必要である