[論文レビュー] DistStat.jl: Towards Unified Programming for High-Performance Statistical Computing Environments in Julia

Seyoon Ko, Hua Zhou+2 more|arXiv (Cornell University)|Oct 30, 2020

Gene expression and cancer classification参考文献 42被引用数 4

ひとこと要約

DistStat.jl は、分散配列抽象化を介して CPU クラスタとマルチ GPU 環境を統合し、高パフォーマンスな統計計算を実現する Julia パッケージである。これにより、統計アルゴリズムの透過的かつスケーラブルな実行が可能となり、大規模なゲノム解析が加速される。AWS 上の 20 ノードの仮想クラスタで 400,000 人の UK Biobank データに対して 500,000 変数の Cox モデルを 50 分未塔でフィッティングしたが、これは CPU では同等の Python 実装より 20–30% 速く、GPU でも最適化されていないカーネルでも同等の性能を発揮した。

ABSTRACT

The demand for high-performance computing (HPC) is ever-increasing for everyday statistical computing purposes. The downside is that we need to write specialized code for each HPC environment. CPU-level parallelization needs to be explicitly coded for effective use of multiple nodes in cluster supercomputing environments. Acceleration via graphics processing units (GPUs) requires to write kernel code. The Julia software package DistStat.jl implements a data structure for distributed arrays that work on both multi-node CPU clusters and multi-GPU environments transparently. This package paves a way to developing high-performance statistical software in various HPC environments simultaneously. As a demonstration of the transparency and scalability of the package, we provide applications to large-scale nonnegative matrix factorization, multidimensional scaling, and $\ell_1$-regularized Cox proportional hazards model on an 8-GPU workstation and a 720-CPU-core virtual cluster in Amazon Web Services (AWS) cloud. As a case in point, we analyze the on-set of type-2 diabetes from the UK Biobank with 400,000 subjects and 500,000 single nucleotide polymorphisms using the $\ell_1$-regularized Cox proportional hazards model. Fitting a half-million-variate regression model took less than 50 minutes on AWS.

研究の動機と目的

• 統計計算における、マルチ CPU クラスタやマルチ GPU システムといった異なるハイパフォーマンスコンピューティング（HPC）環境向けに、別々の低レベルコードを書くという課題に対処すること。
• CPU クラスタ、マルチ GPU ワークステーション、クラウドベースの仮想クラスタといった異種 HPC 環境を統一的にプログラミング可能にするため、1 つの高レベルインターフェースを提供すること。
• 非負値行列分解、多次元スケーリング、ℓ₁-正則化付き Cox レジッションといった広く使われる統計最適化アルゴリズムを、多様な HPC プラットフォームで効率的かつスケーラブルに実行できること。
• Julia の多重ディ patch とネイティブ並列処理が、特に大規模データに対して、既存の Python ベースの HPC ツールを上回る性能を発揮できることを示すこと。
• 全ゲノム関連研究のようなテラバイトからペタバイト規模のデータセットを、統一的かつ高パフォーマンスなコードで分析できることを支援すること。

提案手法

• パッケージは、CPU 配列や CuArrays といった下位の配列型に抽象化された分散配列データ構造を実装しており、対象の HPC 環境に自動的に適応する。
• Julia の多重ディ patch と組み込みのベクトル化機能を活用し、分散データ上で要素演算や線形代数演算を高レベルで表現力豊かに記述できる構文を提供する。
• 通信レイヤーとして MPI を使用しており、クラウドベースの仮想クラスタやオンプレミスクラスタを含む、あらゆる MPI 対応環境への移植性を実現する。
• CPU と GPU の実行パスを透過的に切り替える仕組みを備えており、ユーザーはコードの変更ではなく設定の切り替えでバックエンドを切り替えられる。
• GPU 加速のため、CUDA.jl と統合し、ℓ₁-正則化付き回帰における行列-ベクトル積のような計算集約的処理に CUDA カーネルを活用する。
• AVX 指令セットを CPU で、CUDA カーネルを GPU で使用した、P_{(n+1)}δ のような操作に最適化されたメモリ効率の良いカーネルを実装している。

実験結果

リサーチクエスチョン

• RQ1 マルチ CPU クラスタとマルチ GPU システムを含む異種 HPC 環境において、1 つの高レベルプログラミングインターフェースが統計計算を統合可能か。
• RQ2 Julia ベースの分散配列ライブラリのパフォーマンスは、統計最適化用途における既存の Python ベースの HPC ツールと比べてどうか。
• RQ3 400,000 人の被験者と 500,000 個の SNP を含む、多変量生存モデルを用いた、テラスケールのゲノムデータセットに対して、統一されたコードベースがどの程度スケーラブルか。
• RQ4 統計アルゴリズムのための GPU と CPU の別々のコードベースを書く・保守する負担を、この抽象化がどの程度軽減できるか。
• RQ5 MPI と GPU 加速を組み合わせた高レベル言語である Julia が、大規模統計推論においてどの程度のパフォーマンス上限に達するか。

主な発見

• 400,000 × 500,000 のゲノムデータセットに対して、ℓ₁-正則化付き Cox 比例ハザードモデルが、20 ノードの AWS 仮想クラスタ（720 コア）上で 50 分未塔でフィッティングされた。
• 同じ解析で、CPU では Python ベースの dist_stat 実装より 20–30% 速く、スレッド制御の優位性と Julia の柔軟なスレッドモデルによるものであった。
• 8 GPU のワークステーションでは、最適化されていない GPU カーネルを使用しても dist_stat と同等のパフォーマンスを達成した。これは GPU 加速の大きな潜在的価値を示している。
• 本パッケージにより、全ゲノムデータにおける最大規模の多変量生存モデル解析が成功裏に実施され、前例に比べて被験者数を倍増させた。
• 分散配列抽象化により、CPU と GPU バックエンド間の切り替えが最小限の設定変更で透明に可能となり、統一的 HPC プログラミングの実現可能性が裏付けられた。
• AVX と CUDA を活用した P_{(n+1)}δ のメモリ効率の良いカーネル実装により、特に CPU および GPU 最適化された線形代数演算で顕著なパフォーマンス向上が得られた。