富岳へmesh-tensorflowをインストールする

富岳でモデル並列による分散トレーニングを試すには、 mesh-tensorflow を使うのが一番手っ取り早いだろうと、インストールしてみた。

富岳計算ノード側で用意されている TensorFlow2.2.0 の venv環境を作って pip install mesh_tensorflow を実行したところ、どうも TensorBoard プラグイン情報を書き込むときに権限がなく、完了できなかった。

そこで、TensorFlow含めてソースコードから環境を構築した。

私のアカウントは3月末で計算ノードが使用できなくなるが、4月最初の保守ですでに計算ノードOSや富士通コンパイラのバージョンアップが行われる予定になっており、このとおりには実行できない可能性がある。

以下、手順である。

---

TensorFlow 2.2.0 をソースコードからインストールする

2022年3月24日時点での手順は以下の Qiita 記事にまとめている。

• 富岳上にTensorFlow2.2.0/PyTorch1.7.0共存環境を構築する

mesh-tensorflow 0.1.16 をソースコードからインストール

pip install mesh_tensorflow でもインストールできるが、念の為こちらも富士通＋理化学研究所がベンチマークで使用したソースコードを使ってインストールする。

• ML Commons MLPerf(TM) HPC v1.0 results
• v1.0 Results
• https://github.com/mlcommons/hpc_results_v1.0

ここではv1.0 Resultsのcodeリンク先にあるGitHubリポジトリを使用する。

毎年ベンチマークを実行しているようなので、今年もおそらく同様のリポジトリが公開されるかもしれない。

• ログインノードにSSH接続
• 以下のコマンドを実行

cd
mkdir -p projects
cd projects
git clone https://github.com/mlcommons/hpc_results_v1.0.git
cd hpc_results_v1.0
git checkout main

${HOME}/.local/venv にTensorFlowがインストールされているvenv環境 tf220_pt170 が存在する前提とする。

• 計算ノードに会話ジョブでログイン
• 以下のコマンドを実行

export ENV_NAME="tf220_pt170"
export LD_LIBRARY_PATH=${HOME}/.local/lib:.:${LD_LIBRARY_PATH}
export PATH=${HOME}/.local/bin:.:${PATH}

cd ${HOME}/.local/venv/${ENV_NAME}
source bin/activate

export TCSDS_PATH=/opt/FJSVxtclanga/tcsds-1.2.34
export CC="${TCSDS_PATH}/bin/fcc -Nclang -Kfast -Knolargepage -lpthread"
export CXX="${TCSDS_PATH}/bin/FCC -Nclang -Kfast -Knolargepage -lpthread"

cd ~/projects/hpc_results_v1.0/Fujitsu+RIKEN/benchmarks/cosmoflow/implementations/implementation_fugaku/setup/mpi4py
python3 setup.py install
cd ~/projects/hpc_results_v1.0/Fujitsu+RIKEN/benchmarks/cosmoflow/implementations/implementation_fugaku/mesh
python3 setup.py install

---

mesh-tensorflow のプログラムを実行するには、複数ノードで実行することになるので、mpiexec 経由になる。
このため、現実的には計算ノードへ投入する通常ジョブを作成することになる。
通常ジョブジョブスクリプトはベンチマークの 実行コード あたりを参考に作成のこと。

なお、プリポストノード群の一部にGPUが搭載されたノードが存在するが、こちらはTensorFlow含めPyPI管理下のパッケージがそのままpipインストールできる。富士通コンパイラで作成されたバイナリではないが、こちらを使用することで動作可能である。GPUノード上にvenv環境を構築してmesh/examples/mnist.pyをテスト実行するバッチサンプルは次の通り：

#!/bin/bash
#SBATCH -p ppsq
#SBATCH --time=24:00:00
#SBATCH --gpus=1
#SBATCH --gpus-per-node=1
#SBATCH --gpus-per-task=1
#SBATCH --mail-type="ALL"
#SBATCH --mail-user="[e-mail]"
#SBATCH -J "installMesh"
#SBATCH --comment="mesh-tensorflow install for gpu node"

###
# install mesh-tensorflow for gpu node
###
export ENV_NAME="gpu"
mkdir -p ${HOME}/.local/avx512/venv
python3 -m venv ${ENV_NAME}
cd ${HOME}/.local/avx512/venv/${ENV_NAME}
source bin/activate
export LD_LIBRARY_PATH=${LD_LIBRARY_PATH}:/usr/local/cuda-11.2/lib64
export PATH=${PATH}:/usr/local/cuda-11.2/bin:.
pip3 install tensorflow-gpu
pip3 install mesh-tensorflow

#####
# test (mnist.py)
#####
mkdir -p ${HOME}/projects
https://github.com/tensorflow/mesh.git
cd ${HOME}/projects/mesh/examples
rm -rf gpu
mkdir -p gpu
mkdir -p gpu/data
mkdir -p gpu/model
date
python3 mnist.py --model_dir gpu/model --data_dir gpu/data
date
#####
# EoF
#####

環境構築の点では mesh-tensorflow は、GPU環境で実行するほうが簡単 である。

以上