📑

事前学習の方法(10/27日勉強会公開用)

2024/10/27に公開

はじめに

この記事の内容は、以下の勉強会で報告したものになります。

開催日時:2024年10月27日(日)17:00 - 19:00

テーマ小型LlamaモデルのMegatron-LMを用いた事前学習と継続事前学習

この勉強会では、トークナイザー学習、事前学習や継続事前学習についても報告を行っています。本記事と合わせてご活用ください。

本記事では、TanukiI-8B の事前学習と継続事前学習を行うための説明を行います。

0. TanukiI-8Bとは

  • 日本語に特化した大規模言語モデル(LLM)で、対話能力や長文生成において優れた性能を示すモデルです。

https://zenn.dev/matsuolab/articles/24d1cdeb3a2d11

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

    https://zenn.dev/matsuolab/articles/ce8642479a448f

2. 学習コード公開

学習コード

team-hatakeyama-phase2/Tanuki_pretraining

学習済みの事前学習モデル

Drive
team-hatakeyama-phase2/Tanuki-8B-base-v1.0

TP、PP変換したものも上げたい

3. 学習の方法

3.1 トークナイズ

Tanuki_pretraining/tokenize/ を使用する
詳細は README を参照

トークナイズを行う学習コーパスを tokenize/input_jsonl に格納

tokenize_config.yaml の中で入力したデータを指定

実行シェルスクリプトを、実行

3.2 事前学習

Tanuki_pretraining/training/ を使用する
詳細はREADMEを参照

  1. トークナイズを行った学習コーパスを training/training_corpus に格納

  2. 以下を実行 (実行場所に注意)

Tanuki_pretraining$ ./training/8B-A100_4_pretraining_initial_v2.sh 

3.3 継続事前学習

Tanuki_pretraining/training/ を使用する
詳細はREADMEを参照

  1. 事前学習を行ったチェックポイントをcheckpoints に格納
      (自身で学習した場合はcheckpoints にある)
  2. 追加するコーパスをトークナイズし training/training_corpus に格納
  3. Tanuki_pretraining$ ./training/8B-A100_4_pretraining_continue_v2.sh の中のCHECKPOINT_LOAD_DIR を継続事前学習の元のcheckpoints を指定する。
# model config
TOKENIZER_MODEL="./training/tokernizer/tokenizer_scale200.model"
CHECKPOINT_LOAD_DIR="./training/checkpoints/8b-iter0126000_tp1_pp4"
CHECKPOINT_SAVE_DIR="./training/checkpoints/Llama-3-8/tp${TENSOR_PARALLEL_SIZE}-pp${PIPELINE_PARALLEL_SIZE}-ct${CONTEXT_PARALLEL_SIZE}-LR${LR}-MINLR${MIN_LR}-WD${WEIGHT_DECAY}-WARMUP${LR_WARMUP_STEPS}"

この時、TP、PPが異なっていると、エラーになるので注意

また、Megatron形式である必要がある。
Hugging Fase 形式の場合は変換が必要になる
変更に必要なツールは tools の中に格納されている。

  1. 以下を実行 (実行場所に注意)
Tanuki_pretraining$ ./training/8B-A100_4_pretraining_initial_v2.sh 

おわりに

GENIAC 松尾研 LLM開発プロジェクトでは実行環境の重要性を認識し、環境構築の専任担当として責任を持って取り組んできました。特に、Transformer Engineに関する情報が限られている中で、デバッグを繰り返しながらTanukiモデルの環境を無事に構築できたことに非常に満足しています。

公開したスクリプトが「実行できない」と言われることがないように、高い再現性を確保するため、さまざまな環境で何度もテストを繰り返しました。その結果、今回公開した環境構築スクリプトの品質には自信を持っています。

最後に、この構築スクリプトの検証にご協力いただいた関係者の皆様に、深く感謝申し上げます。皆様のご協力のおかげで、再現性のある堅実な環境構築スクリプトを提供することができました。

この成果は、NEDO(国立研究開発法人新エネルギー・産業技術総合開発機構)の助成事業「ポスト5G情報通信システム基盤強化研究開発事業」(JPNP20017)の結果得られたものです。

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

Discussion