巨大気象データと戦う ― サロゲートモデル学習を高速化する圧縮技術

気象予測AIの学習では、大規模な再解析データを繰り返し読み込む必要があり、GPUの計算性能よりもデータI/Oがボトルネックとなることがあります。本講演では、日本領域再解析データ ClimCORE を対象に、誤差許容圧縮を活用して学習パイプラインを高速化する取り組みを紹介します。気象データ特有の課題や、科学技術データ向け圧縮技術 cuSZ を用いた高速化について、HPCと機械学習の両面からお話しします。

Transcript

1. 巨大気象データと戦う ― サロゲートモデル学習を高速化する圧縮技術 Kevin Xu 1) 2) ◯ 福田 圭祐1) 前田 新一1) 1) 2) ㈱ Preferred Networks 東京大学 @ GPU UNITE 勉強会 #5: HPC × AI 最前線
2. 2 本発表は、2026年並列／分散／協調処理に関するサマー・ワークショップ（SWoPP 2026）における発表 「気象サロゲートモデル学習パイプラインにおける誤差許容圧縮を用いたI/O最適化」 の内容を含んでいます
3. 背景と課題 3
4. 機械学習とは 4 何らかの関数（中身不明）を考える
5. 機械学習とは 5 入力 x と 出力 y の組が多数あるとき、これらからFを近似的に再現する
6. 機械学習とは 6 関数は、複雑なものでも良い “neko” データの量と計算パワー次第で、どれだけ複雑な関数も近似できる（かもしれない）
7. 機械学習とは 7 数値シミュレーションの１ステップも、関数とみなせる データの量と計算パワー次第で、どれだけ複雑な関数も近似できる（かもしれない）
8. 8 気象シミュレーションの複数ステップも、関数とみなせる データの量と計算パワー次第で、どれだけ複雑な関数も近似できる（かもしれない） ※ 機械学習は、物理法則を学習するのではなく、外から見たI/Oの写像だけを学習している
9. Flops are cheap, bandwidth are expensive 9 "AI and Memory Wall",Gholami et al., arXiv, (2024)
10. 気象予報モデル 10 • 数値予報モデル： 物理法則に基づく時間発展方程式などを数値的に解く • サロゲート/機械学習 (ML) モデル： ERA5などの再解析データを用いた学習により、ニューラルネットワークで近似する GraphCast (Science 2023), PanguWeather (Nature 2023) など モデル 現在の気象場 x(t) 気温・風・湿度・気圧など 予測された気象場 y(t+1)
11. NNで気象サロゲートモデルを作ると何が嬉しいか？ 11 • Flops are cheap, bandwidth are expensive ◦ 計算機のBytes-to-FLOPs (B/F) 比はどんどん悪くなっている ◦ 流体・気象などの数値計算は、メモリバンド幅律速であることが多い ◦ →計算機の進化の恩恵を受けられない • 複雑なシミュレーションの入出力関係をNNで近似することで、現代GPUの膨大な FLOPSを使いやすい計算へ変換できる。 • 原理にこだわらずに学習できるので、数値計算 ＋ 観測データの両方を合わせた現 象を学習できる ◦ （純粋な数値計算より予測精度がよいかもしれない） "AI and Memory Wall",Gholami et al., arXiv, (2024)
12. 気象サロゲートモデル 12 サロゲートモデルは、数値予報モデルと比較して高速かつ同等かそれ以上の精度 （上から順に） IFS HRES （数値予報） 数時間 約10,000倍 RMSE [K] 1.4秒 0.1秒 IFS (数値予報) スパコン数百ノード Pangu V100×1 Pangu (ML) Pangu H100×1 図：全球予測の実行時間比較 （24時間先予測） GraphCast (ML) better リードタイム［ hours］ 図：全球予測の精度比較 （気温@850hPa, 2020年, 全球平均）
13. 参考：Matlantis™ 材料探索用の汎用原子レベルシミュレーター 業界の課題 • 従来材料の持続可能性の限界 • 新素材探索にかかる莫大な時間と コスト PFNは、持続可能な未来を実現する新しい電池材料、半導体、合成燃料向け触媒、潤滑剤などの新素材の探索を従来 の1万倍以上高速化する汎用原子レベルシミュレーターMatlantis™をENEOSと共同で開発。グループ会社の Matlantis株式会社がクラウドサービスとして国内外150以上の企業・団体に提供しています。 96元素のあらゆる組み合わせで分子、 結晶など幅広い材料の種類に対応 未知の材料の物性等も従来の1万倍以上 （最大7,000万倍*1）の速度でブラウザ 上でシミュレーション 6,300万以上の構造からなるMatlantisの訓練データ の生成には、1台のGPUで処理すると3,200年かかる 計算量が費やされています。*2 詳細: https://matlantis.com/ 触媒 電池 半導体 合金 潤滑剤 セラミック 吸着剤 分離膜 サステナビリティに貢献する多様な材料 の探索を高速化 *1: 3,000原子でのR2SCANレベルのDFT計算時間推定値と比較した場合 *2: Matlantisの機械学習原子間ポテンシャル「PFP」は、PFNのスーパーコンピュータおよび国 立研究開発法人産業技術総合研究所のAI橋渡しクラウド（ABCI） を用いて開発されました。 13
14. 課題：ストレージI/OによるGPU待機 14 学習時には多変数・高解像度データのストレージI/Oが律速となる Flops are cheap, bandwidth are expensive (again) CPU GPU • I/O（N iter） 学習 (N-1) idle I/O（N+1 iter） 学習 (N) idle 高性能なGPU計算機では、学習時間 << ストレージI/O となり、GPUが待機する 例：Miyabi-G共有ファイルシステム（総帯域1.0 TB/s, 1,310ノード）の1ノードあたり約0.76 GB/sに おいて、ERA5全球の入出力データ（float32、約2.8 GB）では、1iterのI/Oは約3.7秒を要すると想定 される。一方で学習の計算時間は0.5秒未満（H100）。 • ①GPUの演算性能向上 ②気象データの高解像度化 により、I/O律速はさらに深刻化する
15. 科学データ向け誤差許容圧縮 15 指定した誤差範囲内で大幅に圧縮し、ストレージ容量やI/O時間を削減するa ただし、データの展開 (解凍) 処理と転送処理がボトルネックとなる • • CPU圧縮器（SZ3, ZFP, など） • • 圧縮器が逐次処理のため低速 展開データのCPU-GPU転送が必要 展開 I/O CPU CPU-GPU転送 GPU 学習 idle 図：CPU展開の実行タイムライン • • GPU圧縮器（cuSZ, nvCOMPなど） CPU 展開は高速だが、圧縮率が低い 展開と学習がGPU資源を奪い合う GPU I/O 展開 学習 idle 図：GPU展開の実行タイムライン
16. 提案手法 16
17. 提案手法の全体像 17 • • • SZ圧縮アルゴリズムをベースとし、学習データのCPU展開・GPU展開の両パターンを実装 性能モデルを構築し、必要な圧縮率とCPU展開・GPU展開を選択可能 GH200において、実際に誤差許容圧縮器を用いた学習パイプラインを構築 1. 2. 3. 4. 【I/O】圧縮データをストレージから読む 【展開】Grace CPU上でSZ3のデコード処理 【CPU-GPU間転送】展開済みデータをNVLink-C2CでH100 GPUへ供給 【学習】GPUは学習に専念 転送 圧縮データ I/O CPU 1F 7C 10 FF … I/O 展開 学習 CPU GPU 図：CPU展開の場合のデータフロー 展開 GPU 転 送 学習 図：実行タイムライン
18. 圧縮器の要件 18 図：最適な実行タイムライン 展開速度の要件： • • GPU展開では展開時間が必ず露出する CPU展開では展開が学習計算より速けれ ば完全に隠蔽できる CPU 生データ I/O 学習 GPU idle 圧縮なし（ベースライン） 圧縮率の要件： • 圧縮後のI/O時間が学習時間を下回る 例：ClimCOREデータ データ 520 MB/iter、I/O帯域 0.76 GB/s、学 習 0.1 s/iter の場合 → CPU展開では、圧縮率 約7倍以上・ 展開速度 約5.2 GB/s以上が必要 CPU GPU 圧縮データ I/O 展開 学習 GPU展開（提案手法１） CPU GPU 圧縮データ I/O 展開 転送 学習 CPU展開（提案手法２）
19. CPU展開の実装：パイプライン化・並列化 19 • • 課題：逐次的なCPU展開では必要なデータ供給性能を満たせない。 解決策：依存関係のないサンプル間並列 & 変数単位の展開処理を並列化する • 使用したライブラリ：NVIDIA DALI、SZ3 I/O（N+1番目） I/O（N+2番目） 展開（N-1番1変数目） I/O（N+4番目） I/O（N+3番目) 展開（N+2番1変数目） 転送 展開（N-1番2変数目） 転送 展開（N+1番1変数目） CPU 展開 転送 転送 展開（N番1変数目） 展開（N+2番2変数目） 転送 展開（N番2変数目） GPU N-3番目：学習 N-2番目：学習 展開（N+3番2変数目） 転送 展開（N+1番2変数目） 転送 I/O 展開 展開（N+3番1変数目） 転送 N-1番目：学習 N番目：学習
20. GPU展開の実装：cuSZの圧縮率の向上 20 • • 課題：GPU圧縮・展開器（cuSZ）は高速だが、圧縮率が低い 解決策：エントロピー符号化のみを学習ベースに置換 ◦ • ハフマン符号化のかわりに rANS を採用し、シンボル出現の確率分布を軽量NN により学習した。 結果：ベンチマーク（CESM-ATM）で、追加コスト約 100flops/byte未満で、 圧縮率の1.2倍の向上を達成 Huffman 符号化 → 速度と圧縮率の両立 使用したライブラリ: FZGPUModules 生データ 予測 (spline) 量子化 軽量NN分布推定 + rANS 予測はNNに置き換えない NN予測モデル 図：提案手法のデータフロー 圧縮データ
21. 性能モデル 21 • • 各方式の1反復あたりレイテンシを、データサイズ・I/O帯域・圧縮率・展開速度で定式化 性能モデルにより、システム環境に応じて必要な圧縮率と展開方式（CPU / GPU）を選択 レイテンシ 学習に露出する時間 非圧縮 - GPU展開 CPU展開 I/O隠蔽の条件（圧縮率） 条件を満たせば０ （ ）
22. 実験・評価 22
23. 評価環境：Miyabi（NVIDIA GH200 ） 23 GRACE CPU NVLink C2C 72Core, 3.0GHz 450 GB/s Hopper GPU HBM 989.5 TFLOP/s (BF16) 4.0 TB/s （片方向） 提案手法２「CPU展開」を用いてGPUアイドル時間を排除できる性能要件を満たす • 高性能CPUによる高速な展開 • CPU–GPU間の高速なデータ転送 → NVIDIA GH200 Grace Hopper Superchip（Miyabi-G）はこれらの要件を満たす
24. 実験設定 24 Q1. 圧縮パイプラインで学習はどれだけ高速化するか? Q2. 予測精度はどれだけ維持されるか? 表：データセット ERA5 ClimCORE 対象領域 全球 日本近郊 水平解像度 25km（721×1440） 5km（661×817） 鉛直方向 13層 17層 変数 上空5変数 + 地表4変数 上空6変数 + 地表4変数 リードタイム 24h 1h
25. 学習時間の性能モデル評価 25 図：性能モデル（横軸: I/O 帯域幅、縦軸: iteration time [ms]） • ERA5では圧縮なし(750 ms/iter)に対して 3.4倍 (220 ms/iter)* *全球MLモデルの学習が10.9時間で完了する • ClimCOREでは圧縮なし(518 ms/iter)に対して 5.2倍(100 ms/iter) の高速化を実現可能
26. 結果１：学習時間の実測値 26 • • • Miyabi の共有ファイルシステムは全ノードで帯域を共有（総帯域 1.0 TB/s） 実測時は他ジョブの I/O 負荷が低く、実効帯域は想定を大きく上回った このため非圧縮でも I/O 律速が顕在化せず、圧縮による高速化は限定的だった 表：学習速度と実効 I/O帯域の実測値（ Miyabi、16 ノード） ノード 数 データセッ ト 非圧縮： 実効帯域* [GB/s] 非圧縮： 学習速度 [ms/iter] 圧縮： 実効帯域 * [GB/s] 圧縮： 学習速度 [s/iter] 高速化 16 ERA5 2.65 GB/s 215.9 ms 0.88 GB/s 212.4 ms 1.02倍 16 ClimCORE 2.44 GB/s 176.6 ms 0.51 GB/s 112.3 ms 1.57倍 ※ 参考：Miyabi共有ファイルシステムは、カタログスペック総バンド幅1.0TB/s, 1ノードあたり約0.76 GB/s
27. 結果２：予測精度 27 • • 圧縮で失われる小スケール成分は、本質的に予測不可能なため、誤差が許容されると 考えられる 5km・1h の高解像度データ（ClimCORE）でも同等精度を確認 表：ERA5（25km, 24h先予測） 手法 Z500 T850 W700 U850 T2M TSL U10M V10M 非圧縮 64.9 0.82 6.95e-5 1.58 0.85 80.4 1.03 1.06 圧縮 64.3 0.83 6.92e-5 1.56 0.84 80.0 1.02 1.05 表：ClimCORE（5km, 1h先予測） 手法 Z500 T850 W700 U850 Tsfc PSEA Usfc Vsfc 非圧縮 2.25 0.227 8.79e-5 0.528 0.415 0.276 0.465 0.499 圧縮 2.41 0.226 8.64e-5 0.529 0.393 0.284 0.452 0.492
28. 今後の取り組みとまとめ 28
29. 今後の取り組み 29 • 予測精度のより詳細な評価 ：圧縮パイプラインが予測精度に及ぼす影響について、台 風や集中豪雨などの個別事例に基づく検証を行う。 • より高密度・高I/O要件のデータへの適用：3次元・多層・高頻度の海洋モデルな ど、圧縮・展開性能への要求がさらに厳しい領域へ手法を展開する。 • 他の計算機システムへの適用：PFNの自社クラスタ等の環境では、ファイルシステム のバンド幅が比較的低く、より高い効果が期待できる。 • 予測精度のさらなる向上：台風予測に適した損失関数の設計や、モデルアーキテク チャの改良を通じて、精度の一層の向上を図る。
30. まとめ 30 • 気象サロゲートモデル学習向けに、誤差許容圧縮を用いた I/O最適な学習パイプラインを 提案 ◦ ◦ • 隠蔽に必要な圧縮率・展開速度の要件を解析モデルで定式化 SZ圧縮器について、ハードウェアの特性に応じた2種類（CPU/GPU）の展開方法を改良 Miyabi-G (GH200) においては、性能モデル上では5.2倍、実測では1.5倍の性能向上が得 られた ◦ 今回の実験時にはシステム負荷が低かったと考えられ、提案手法の恩恵が薄れた。混雑時や他シ ステム・高解像度なデータでは効果が期待できると考えられる • 今回の実験設定では予測精度への影響はほとんどない（RMSEは非圧縮と同等） 謝辞： • 本研究に使用したデータ（RRJ-ClimCORE 試供版データ）は、ClimCORE より提供を受けた。
31. Making the real world computable