[論文レビュー] High-level python abstractions for optimal checkpointing in inversion problems

Navjot Kukreja, Jan Hückelheim+5 more|arXiv (Cornell University)|Jan 12, 2018

Seismic Imaging and Inversion Techniques参考文献 13被引用数 16

ひとこと要約

本稿では、偏微分方程式（PDE）の有限差分DSLであるDevitoに、Revolveチェックポイントアルゴリズムを統合する高水準なPython API、pyRevolveを提示する。この手法により、大規模な逆問題における最適なメモリ使用が可能となり、正確性を損なわずにメモリ使用量を削減できる。コード生成とRevolveの最適チェックポイントスケジューリングを組み合わせることで、正確性を保ちつつメモリフットプリントを削減し、計算効率と正しさを維持したまま、より大きな地震波動逆問題を解くことが可能になる。

ABSTRACT

Inversion and PDE-constrained optimization problems often rely on solving the adjoint problem to calculate the gradient of the objec- tive function. This requires storing large amounts of intermediate data, setting a limit to the largest problem that might be solved with a given amount of memory available. Checkpointing is an approach that can reduce the amount of memory required by redoing parts of the computation instead of storing intermediate results. The Revolve checkpointing algorithm o ers an optimal schedule that trades computational cost for smaller memory footprints. Integrat- ing Revolve into a modern python HPC code and combining it with code generation is not straightforward. We present an API that makes checkpointing accessible from a DSL-based code generation environment along with some initial performance gures with a focus on seismic applications.

研究の動機と目的

• PDE制約付き最適化のための高性能Pythonコードへの最適チェックポイントの統合を簡素化すること。
• 科学的HPCワークロードにおけるチェックポイントを用いた前向きおよびアドジョイントシミュレーションの手動管理の複雑さに対処すること。
• 最適チェックポイントによりメモリ要件を削減することで、地震画像処理におけるより大規模な逆問題を可能にすること。
• 数学的勾配テストとベンチマークを用いて、実装の正しさとパフォーマンスを検証すること。
• 今後の拡張の基盤を築くこと、特に並列チェックポイントとマルチステージディスクベースのチェックポイントを想定する。

提案手法

• 著者らは、Revolveアルゴリズムと連携して最適チェックポイントスケジュールを生成する高水準なPython API、pyRevolveを実装した。
• このAPIは、有限差分時間領域PDEソルバーのためのドメイン特化言語であるDevitoに統合され、チェックポイント付きの前向きおよびアドジョイントシミュレーションの自動コード生成を可能にした。
• このアプローチは、Devitoにおける記号的抽象化を用いて波動方程式を定義し、Revolveの最適スケジュールに基づいて自動的にチェックポイントロジックを挿入する。
• 実装は、完全メモリ保存とチェックポイント戦略の両方をサポートしており、プロファイルツールを用いてメモリ使用量と実行時間を測定した。
• 正しさの検証のため、1615ステップの230×230×230グリッド上で勾配テストを実施し、チェックポイントありとなしの結果を比較した。
• ピークメモリとソリューションまでの総実行時間を複数回の実行で測定するために、Pythonのmemory_profilerとtimeモジュールを活用した。

実験結果

リサーチクエスチョン

• RQ1 高水準なPython抽象化は、科学的HPCアプリケーションにおける最適チェックポイントの統合を効果的に簡素化できるか？
• RQ2 メモリ使用量と実行時間の観点から、チェックポイントベースの勾配計算は、完全メモリ保存と比べてどのように異なるか？
• RQ3 チェックポイントのオーバーヘッドは、並列実行におけるスケーラビリティとロードバランシングにどの程度影響を及えるか？
• RQ4 提案されたAPIは、顕著なメモリフットプリントの削減を実現しながらも、数値的正確性を維持できるか？
• RQ5 並列化と動的時間ステップ法の観点から、現在のチェックポイント戦略の実用的限界は何か？

主な発見

• 3次元地震波動逆問題において、実装は勾配を正確に再現した。完全メモリ保存とチェックポイント実行の両方で得られた結果は完全に一致し、数値的正確性が確認された。
• ピークメモリ使用量は、完全保存の約80 GBから、チェックポイントを用いることで最低で約10 GBまで削減され、固定メモリ制限内でのより大規模なシミュレーションが可能になった。
• 理論的に最適であるにもかかわらず、チェックポイント実装は完全保存より遅く実行された。これは、繰り返しのPythonからC関数呼び出しとnumpyのディープコピーによるオーバーヘッドが原因であった。
• チェックポイントを各時間ステップに保存した場合でさえ、完全保存ケースと同一のメモリ使用量であったが、性能劣化が顕著に現れた。
• 結果から、現在の実装のパフォーマンスは、とりわけ細粒度の時間ステップ処理において、Pythonの実行モデルが制限要因であることが示された。
• 著者らは、並列チェックポイントとオンラインで適応するチェックポイントの今後の研究の必要性を指摘した。これは、複雑なシミュレーションにおける動的時間ステップ法をサポートするためである。