事前学習環境構築　～シェルスクリプトによる再現性の担保～

jaked97

2024/10/18に公開

 はじめにGENIAC 松尾研 LLM開発プロジェクト　Team Kumaの加藤純です。
今回、Tanukiモデルの事前学習環境について、必要なライブラリの調査、スクリプトの構築を担当しました。今回はその担当作業にあたっての取り組み内容を報告します。
報告にあたり、本記事で紹介している環境構築スクリプトの実演も含む勉強会を以下の日時で開催しました。
開催日時：2024年10月13日（日）17:00 - 19:00
テーマ：小型LlamaモデルのTransformer Engineを用いた事前学習の環境構築
また、勉強会の様子は動画としてアップロードされています。
この勉強会では、環境構築スクリプトの説明や実際の環境構築の様子に加えて、Transformer Engineを利用したfp8での事前学習についてもデモンストレーションしています。本記事と合わせてご活用ください。

 本記事の概要本記事では、Tanukiモデルの事前学習環境を自動的に構築できるスクリプトについて解説します。
モデルの学習は実行する環境が異なると結果が変わったり、意図しないエラーが発生したりすることがあります。そのためプロジェクトメンバー全員が同一の環境でモデルを学習できるよう、学習環境構築について再現性を担保することが非常に重要です。
こうした課題に対応するために、GENIAC 松尾研 LLM開発プロジェクトでは、誰でもスクリプトを実行するだけで同じ環境を簡単に構築できる仕組みを開発しました。
本記事ではこの環境構築スクリプトについて、以下の2つのセクションに分けて詳しく解説します。

Tanukiモデルで利用した環境構築スクリプトの概要
作成したスクリプトの具体的な内容、動作、設定について詳しく解説します。


環境構築スクリプト作成の取り組み
スクリプト作成の過程や直面した課題、今後の改善点について解説します。


 1. Tanukiモデルで利用した環境構築スクリプトの概要
 1.1 前提松尾研の提供環境（GCP）での構築です。
vCPU：208

GPU：H100 × 8

メインメモリ：1,872GB

/home領域の各ユーザの利用可能なデータサイズ：30GB

共有ディスク領域で各サブチームが使用可能なサイズ：30TB

OS：Ubuntu 20.04

NVIDIA-SMI：535.104.05

Driver Version：535.104.05

CUDA Verion：12.2
構築スクリプトのベースは以下を参考に作成しました。
matsuolab/ucllm_nedo_prod/train/README-train_gcp_honban_multi_node_multi_gpu.md
llm-jp/Megatron-LM/scripts/tsubame/install.sh
また、以下環境においても構築スクリプトが正常動作することを確認しています。（勉強会での動作実演より）


項目
Azure
GCP
Windows10 Pro + WSL2


OS
Ubuntu 20.04.6 LTS
Ubuntu 22.04.3 LTS
Ubuntu 22.04.5 LTS

GPU
GH100（fp8/bf16）
A100-PCIE-40GB（bf16のみ）
GeForce RTX 4090（fp8/bf16）・RTX 3090（bf16のみ）

GPUメモリ
94GB
40GB
24GB

NVIDIA-SMI
535.183.06
535.129.03
565.51.01

Driver Version
535.183.06
535.129.03
565.90

CUDA Version
なし
12.2
12.1

gccおよびg++
9.4.0
11.4.0
11.4.0

GNU Make
4.2.1
4.3
4.3


 1.2 動作方法https://github.com/matsuolab/nedo_project_code の team_hatakeyama_phase2 ブランチをcloneして、share-jk_pretrain_env.sh を実行します。
# GitHubリポジトリのクローン
git clone <https://github.com/matsuolab/nedo_project_code.git>

# リポジトリのディレクトリに移動
cd nedo_project_code

# 指定されたブランチをチェックアウト
git checkout team_hatakeyama_phase2

# スクリプトの格納ディレクトリに移動
cd team_hatakeyama_phase2/hatakeyama/pretrain

# スクリプトを実行
. share-jk_pretrain_env.sh
注意： minicondaを利用してPython実行の仮想環境を構築します。もしminicondaをインストールしている場合はアインストールしてください。minicondaのインストールも含めてshare-jk_pretrain_env.shで実行します。

 1.3 スクリプト抜粋解説github上の環境構築用スクリプトshare-jk_pretrain_env.shについて、実行時※のポイントを解説します。
※本スクリプトの実行時間は約50分です
Transformer Engine（tev）、Flash Attention（fav）、Pythonのバージョン、ダウンロードするminicondaのバージョンを指定します。
tev=1.7
fav=2.5.7
python_Mver="3"
python_mver="10"
conda_ver="py${python_Mver}${python_mver}_23.10.0-1"
llm-jp/Megatron-LMリポジトリのtsubameブランチをgit cloneし、コードの一部を修正します。
requirements.txtからmpi4pyをコメントアウトし、numpyのバージョンを固定します。
# mpi4py：condaでインストールするためコメントアウト
# numpy : v1.26.0をインストール
sed -i \\
	-e "/^mpi4py/s/^/# /g" \\
	-e "s/^numpy$/numpy==1.26.4/g" \\
	$DIR/Megatron-LM/requirements.txt
datasetsディレクトリのMakefileにあるpython3-configのパスを修正します。
# python3-config：パス修正（llmjp_tsubameのユーザパス残存のため）
sed -i 's@/gs/fs/tga-bayes-crest/taishi/pyenv/versions/3.10.14/bin/@@'  $DIR/Megatron-LM/megatron/core/datasets/Makefile
trainingディレクトリのglobal_var.pyにあるwandbのログ記録部分について、マルチノードに対応させるために修正します。
# wandb：マルチノードログ対応
sed -i \\
	-e "/if getattr(args, 'wandb_project', '') and args.rank == (args.world_size - 1):/i \\    local_last_rank = args.rank%torch.cuda.device_count() - (torch.cuda.device_count() -1)" \\
	-e "s/if getattr(args, 'wandb_project', '') and args.rank == (args.world_size - 1):/if getattr(args, 'wandb_project', '') and local_last_rank==0:/" \\
	-e 's/exp_name = args.wandb_name + "-" + now/exp_name = args.wandb_name + "-" + now + "_" + str(args.rank)/' \\
	$DIR/Megatron-LM/megatron/training/global_vars.py
conda環境の構築場所を指定し、各環境変数※を設定します。（※activate⇔deactivateでオンオフ可能）
conda create -p $DIR/miniconda3/envs/$env_name python=${python_Mver}.${python_mver} -y

mkdir -p $DIR/miniconda3/envs/$env_name/etc/conda/activate.d && \\
echo 'export ORIGINAL_LD_LIBRARY_PATH=$LD_LIBRARY_PATH' > $DIR/miniconda3/envs/$env_name/etc/conda/activate.d/edit_environment_variable.sh && \\
echo 'export LD_LIBRARY_PATH="'$DIR'/miniconda3/envs/'$env_name'/lib:$LD_LIBRARY_PATH"' >> $DIR/miniconda3/envs/$env_name/etc/conda/activate.d/edit_environment_variable.sh && \\

# 以下をコメントアウトするとCUDNN_PATHが通らない状態に
echo 'export CUDNN_PATH="'$DIR'/miniconda3/envs/'$env_name'/lib/python'${python_Mver}.${python_mver}'/site-packages/nvidia/cudnn"' >> $DIR/miniconda3/envs/$env_name/etc/conda/activate.d/edit_environment_variable.sh && \\
chmod +x $DIR/miniconda3/envs/$env_name/etc/conda/activate.d/edit_environment_variable.sh

mkdir -p $DIR/miniconda3/envs/$env_name/etc/conda/deactivate.d && \\
echo 'export LD_LIBRARY_PATH=$ORIGINAL_LD_LIBRARY_PATH' > $DIR/miniconda3/envs/$env_name/etc/conda/deactivate.d/rollback_environment_variable.sh && \\
echo 'unset ORIGINAL_LD_LIBRARY_PATH' >> $DIR/miniconda3/envs/$env_name/etc/conda/deactivate.d/rollback_environment_variable.sh && \\

# 以下をコメントアウトするとCUDNN_PATHが通らない状態に
echo 'unset CUDNN_PATH' >> $DIR/miniconda3/envs/$env_name/etc/conda/deactivate.d/rollback_environment_variable.sh && \\
chmod +x $DIR/miniconda3/envs/$env_name/etc/conda/deactivate.d/rollback_environment_variable.sh
CUDA関連のソフトウェア・ライブラリ、pipのバージョンアップ、pythonライブラリをインストールします。
echo "########## conda install : cuda-toolkit, cudnn, nccl, mpi4py ##########"
conda install -c nvidia/label/cuda-12.1.0 cuda-toolkit -y
conda install -c conda-forge cudnn=8.9.7 nccl=2.18.3 mpi4py=3.1.6 -y

echo "########## pip install : pip & requirements ##########"
pip install pip==24.0
cd $DIR/Megatron-LM && pip install -r requirements.txt
apex、Transformer Engine、Flash Attentionライブラリをインストールします。
echo "########## pip install : apex ##########"
cd $DIR/Megatron-LM/apex && pip install -v --disable-pip-version-check --no-cache-dir --no-build-isolation --config-settings "--build-option=--cpp_ext" --config-settings "--build-option=--cuda_ext" ./

echo "########## pip install : Transfomer Engine ##########"
cd $DIR/Megatron-LM/TransformerEngine && pip install .

echo "########## pip install : Flash Attention ##########"
cd $DIR && pip install --no-build-isolation flash-attn==$fav
Transformer Engine、Flash Attentionが正常にインストールできたかどうかを確認します。
echo "===== import check : transformer_engine & flash_attn ====="
python -c "from transformer_engine import pytorch" && echo "----- import OK -----" || echo "----- import NG -----"

 2. 環境構築スクリプト作成の取り組み
 2.1 スクリプト作成までの流れ環境構築スクリプト作成は以下のステップで進めました。
松尾研から提供された事前学習用の構築コマンドをシェルスクリプトにまとめ、自動化と再現性の向上を図りました。
学習効率を高めるため、Transformer Engineの公式ドキュメントを参照し、必要なライブラリの追加および不要なライブラリの削除を実施しました。このプロセスは環境の最適化に重要な役割を果たしました。
先行事例として参考にしたllmjpのinstall.shと比較し、ライブラリの有無やバージョンの差異を確認しました。また、インストールの順序を慎重に検討し、エラーを回避しながら安定した環境を構築しました。

 2.2 スクリプト作成における苦労およびエラー解消方法スクリプト作成で最も困難だった点は、Transformer Engineの構築例が非常に少なく、特に日本語の情報がほとんどなかったことです。そのため、エラーが発生するたびに自力で解決する必要がありました。特に、Transformer EngineとFlash Attentionが正常に動作するバージョンを見つけるために、さまざまな組み合わせを試行しなければなりませんでした。また、再現性を確保するためには、多くの方にテストを実施してもらい、スクリプトが正しく動作することを確認する必要があったため、スクリプト作成後の検証作業にも多くの時間がかかりました。
以下、スクリプト作成時に発生したエラーおよびその解消方法を記載します。
エラー発生時の解消方法
Python3.10でのmpi4pyライブラリのインストールエラー
requirements.txtの通りにPyPIからmpi4pyをインストールすると、Pythonバージョンによりエラーが発生します。
参考：mpi4py installation failed with Python 3.9 and 3.10 #335
調査の結果、condaからmpi4pyをインストールすることで、エラーを回避することができました。
conda install -c conda-forge cudnn=8.9.7 nccl=2.18.3 mpi4py=3.1.6 -y
CUDAソフトウェアのインストール後の学習コードの動作エラー
cuda-toolkitのバージョンとtorchのcudaバージョンに差異がある場合、学習コード実行時に動作エラーが発生します。
ただし、バージョンを揃えるように設定した場合でもエラーが発生することがあります。
調査の結果、cuda-toolkitをcondaインストールする場合に、ツールのチャンネルを明示的に指定（c nvidia/label/cuda-12.1.0）することで、エラーを回避することができました。
conda install -c nvidia/label/cuda-12.1.0 cuda-toolkit -y
Transformer Engineの動作エラー
Transformer EngineとFlash Attentionのバージョンや各ライブラリのインストール順番によって、動作（import）エラーが発生します。
多くの組合せで試した結果、以下の順番およびバージョンの組み合わせの場合に、エラーを回避することができました。


順番
ライブラリ
バージョン


1
requirements.txt
ー

2
apex
0.1

3
Transformer Engine
1.7

4
Flash Attention
2.5.7


CuDNNライブラリのPATHが設定されていない場合、学習コード実行時に動作エラーが発生します。
調査の結果、condaからインストールしたcudnnライブラリのパスをCUDNN_PATHに設定することで、エラーを回避することができました。
echo 'export ORIGINAL_LD_LIBRARY_PATH=$LD_LIBRARY_PATH' > $DIR/miniconda3/envs/$env_name/etc/conda/activate.d/edit_environment_variable.sh && \\
echo 'export LD_LIBRARY_PATH="'$DIR'/miniconda3/envs/'$env_name'/lib:$LD_LIBRARY_PATH"' >> $DIR/miniconda3/envs/$env_name/etc/conda/activate.d/edit_environment_variable.sh && \\
echo 'export CUDNN_PATH="'$DIR'/miniconda3/envs/'$env_name'/lib/python'${python_Mver}.${python_mver}'/site-packages/nvidia/cudnn"' >> $DIR/miniconda3/envs/$env_name/etc/conda/activate.d/edit_environment_variable.sh && \\
chmod +x $DIR/miniconda3/envs/$env_name/etc/conda/activate.d/edit_environment_variable.sh

#例：CUDNN_PATH="＄HOME/miniconda3/envs/【環境名】/lib/python3.10/site-packages/nvidia/cudnn"


 2.3 スクリプトの今後の改善案環境構築スクリプトの共有化という目標は達成できましたが、いくつかの改善の余地があります。これらは時間の制約により実現できなかった部分です。
インストールと動作検証の統合
現在の環境構築スクリプトは、事前学習環境のセットアップに重点を置いていますが、その環境が正しく動作しているかを確認するためには、別途スクリプトを実行する必要があります。理想は環境構築とサンプルデータを用いた簡単な学習テストを1つのスクリプトにまとめることです。動作確認まで行って初めて環境構築が成功したと言えるため、インストールが成功しただけでなく、実際の動作が確認できる仕組みが理想的と考えます。
fp8利用の明示化
fp8はディープラーニングの効率を大幅に向上させる技術です。通常、計算精度として使用されるfp16よりも少ないビット数で計算を行うことで、メモリ消費を削減し、トレーニング速度を向上させます。しかし現在のスクリプトではfp8が使用されているかどうかがログから確認できません。
今回はTFLOPSのパフォーマンスが16ビットから8ビットに移行した際に速度が向上したことを根拠にfp8が使用されていると判断しましたが、今後は明示的にfp8の利用がログに記録されるようにすることが理想と考えます。
Transformer Engineのオンオフ
現在のスクリプトでは、Transformer Engineをオフ（--transformer-impl local）にした状態でLlamaモデルの学習を実行すると以下のエラーが発生します。
AssertionError: (RMSNorm) is not supported in FusedLayerNorm when instantiating FusedLayerNorm
現状ではTransformer Engineの使用を前提としているため大きな問題にはなりませんが、今後はTransformer Engineを利用できない環境での学習を考え、Transformer Engineをオン・オフできるようにすることが望ましいと考えます。
Docker版環境構築
異なる環境で一貫した動作を保証するために、Docker版の環境構築スクリプトも望ましいと考えます。

 おわりにGENIAC 松尾研 LLM開発プロジェクトでは実行環境の重要性を認識し、環境構築の専任担当として責任を持って取り組んできました。特に、Transformer Engineに関する情報が限られている中で、デバッグを繰り返しながらTanukiモデルの環境を無事に構築できたことに非常に満足しています。
公開したスクリプトが「実行できない」と言われることがないように、高い再現性を確保するため、さまざまな環境で何度もテストを繰り返しました。その結果、今回公開した環境構築スクリプトの品質には自信を持っています。
最後に、この構築スクリプトの検証にご協力いただいた関係者の皆様に、深く感謝申し上げます。皆様のご協力のおかげで、再現性のある堅実な環境構築スクリプトを提供することができました。
!この成果は、NEDO（国立研究開発法人新エネルギー・産業技術総合開発機構）の助成事業「ポスト５Ｇ情報通信システム基盤強化研究開発事業」（JPNP20017）の結果得られたものです。

項目	Azure	GCP	Windows10 Pro + WSL2
OS	Ubuntu 20.04.6 LTS	Ubuntu 22.04.3 LTS	Ubuntu 22.04.5 LTS
GPU	GH100（fp8/bf16）	A100-PCIE-40GB（bf16のみ）	GeForce RTX 4090（fp8/bf16）・RTX 3090（bf16のみ）
GPUメモリ	94GB	40GB	24GB
NVIDIA-SMI	535.183.06	535.129.03	565.51.01
Driver Version	535.183.06	535.129.03	565.90
CUDA Version	なし	12.2	12.1
gccおよびg++	9.4.0	11.4.0	11.4.0
GNU Make	4.2.1	4.3	4.3

順番	ライブラリ	バージョン
1	requirements.txt	ー
2	apex	0.1
3	Transformer Engine	1.7
4	Flash Attention	2.5.7

東大松尾・岩澤研究室 | LLM開発プロジェクト[GENIAC]Publication

東京大学松尾・岩澤研究室が運営する松尾研LLMコミュニティのLLM開発プロジェクト[GENIAC] の開発記録、情報発信になります。各種リンクはこちら linktr.ee/matsuolab_community

はじめに

本記事の概要

1. Tanukiモデルで利用した環境構築スクリプトの概要

1.1 前提

1.2 動作方法

1.3 スクリプト抜粋解説

2. 環境構築スクリプト作成の取り組み

2.1 スクリプト作成までの流れ

2.2 スクリプト作成における苦労およびエラー解消方法

2.3 スクリプトの今後の改善案

おわりに

Discussion