Storage Requirements for AIの紹介

# はじめに

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

* [Storage Requirements for AI](https://www.snia.org/educational-library/storage-requirements-ai-2024)
    * Compute, Memory, and Storage Summit 2024にて発表されたもの

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

# 概要

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

# 背景

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

# AIのライフサイクル

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/2960620/07f3a9f6-caf3-7c13-fa83-c7ebc81c7731.png)

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

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

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

# Training & Tuning の概要

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/2960620/60d90ffe-aef8-a75f-4432-a6bd945168d2.png)

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

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

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

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

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

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

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/2960620/c5f2a6c4-cf59-a667-8bfa-f18cd8b8d33c.png)

| 列 | 意味 |
|---|---|
|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](https://blog.roboflow.com/what-is-resnet-50/)というモデルに対し、DLIOベンチマークを実行した際のI/OをトレースすることでどのようなI/Oアクセスを行っているかを解析した結果を示す。（NFSとS3で実施）

:::note info
Resnet50はディープニューラルネットワークのトレーニングに関連する課題に対処するために設計された一連のモデル（よくわかっていない）。  
DLIOはGPUをエミュレートするAIストレージベンチマークであり、以下のような特徴がある。
* GPU不要でストレージパフォーマンスを測定可能
* 学習データセットは1024～150KBのイメージテンソルを含むTFRecord
* 学習データはTensorflowのデータローダを使用して読み込まれる
:::

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/2960620/0c21a8f4-9fd6-93e8-7966-ab6fb021a36f.png)

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

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/2960620/c80b63a3-7638-3281-d274-44b27707c177.png)

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

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

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

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/2960620/6e58e529-1bc0-c81c-9728-b576a590c439.png)

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

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

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

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/2960620/0d50b283-d959-bac8-837d-886e54963993.png)

| 列 | 意味 |
|---|---|
|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の帯域幅が必要

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

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/2960620/35c4a2e9-acf0-adfe-1d07-07eb7c998caa.png)

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

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

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

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/2960620/28813283-131f-fc6f-f559-b23f4d2e3739.png)

| 列 | 意味 |
|---|---|
|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クラスタとストレージ

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/2960620/46f64b9a-8cac-ad51-d709-d35f7cc10212.png)

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

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

# まとめ

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/2960620/c082a42f-e254-5401-f509-d32d99d87c1d.png)

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