そして、AF2のオープンソースソフトウェアであるOpenFoldを富岳へ実装し、CPU超並列環境におけるスループット性の改善として、検索タスクにおける大量のタンパク質の一斉処理技術を開発したとする。1万個のタンパク質の入力アミノ酸配列で評価を行った結果、高速化前と比較して、平均8.5倍の高速化(スループット増加)が達成されたという。

  • 検索タスクにおける高速化手法の概要。検索タスクの実行順序および同時実行数が最適化され、メモリ不足時には計算資源量を自動で再割り当てする機構が構築された。

    検索タスクにおける高速化手法の概要。検索タスクの実行順序および同時実行数が最適化され、メモリ不足時には計算資源量を自動で再割り当てする機構が構築された。(出所:理研Webサイト)

さらにモデル計算タスクにおいては、性能低下を回避するため、バッチ行列積の入力全体をあらかじめ並べ替える工夫を施した上、アテンション機構など、一部のGPU向け実装をCPUに移植して最適化を実施。その結果、最適化前に比べて平均1.3倍の高速化が実現されたとする。

  • モデル計算タスクにおける高速化手法の概要。

    モデル計算タスクにおける高速化手法の概要。(出所:理研Webサイト)

これらの高速化が施された富岳実装版OpenFoldを用いて、1万件の入力アミノ酸配列の評価を行ったところ、推論手順全体で最適化前と比較して6.3倍の高スループット性を達成したとのこと。これは、仮に富岳の全系を用いた場合、1時間あたり約120万件のアミノ酸配列を処理できることになり、富岳がタンパク質立体構造の大規模高速推論に対して実用的なCPUシステムであることが示されたとしている。

理研 R-CCSのHPC／AI駆動型医薬プラットフォーム部門においては、富岳への創薬DXプラットフォームの構築が進められている。今回開発された富岳実装版OpenFoldは、創薬DXプラットフォームにおいて、標的タンパク質の構造を高速に推論する重要な要素技術となることが期待できるとした。

理研は現在、富岳のソフトウェアやアプリケーションの成果を、Armアーキテクチャベースの商業クラウド上でも直接利用できる「バーチャル富岳」環境の整備を進めている最中だ。今後、富岳実装版OpenFoldも同様にクラウド環境で利用可能にし、タンパク質の構造多様性を考慮した推論ができるよう改良を進める予定としている。