refine training scripts (#95)

Author: odashi

pretrain/installers/v4-upstream-megatron-abci/README.md

@@ -5,7 +5,7 @@
 ABCI 3.0上で以下のコマンドを実行し、`<env_install_path>`に環境を構築できる
 
 ```bash
-cd pretrain/installers/v5-megatron-abci/
+cd pretrain/installers/v4-upstream-megatron-abci/
 bash run_setup.sh <env_install_path>
 ```
 

pretrain/scripts/v4-upstream-training-template/README.md

@@ -13,7 +13,7 @@ ABCI 3.0 上で Megatron-LM を利用した LLM-jp v5 用の学習スクリプ
 
 ```bash
 cd $EXP_DIR
-git clone git@github.com:llm-jp/scripts.git
+git clone https://github.com/llm-jp/scripts.git
 ```
 
 次に、 [pretrain/installers/v5-megatron-abci](../../installers/v5-megatron-abci/README.md) を利用し、`$EXP_DIR/env` に環境を構築する。
@@ -47,20 +47,22 @@ cp -r scripts/pretrain/task_template/ $EXP_DIR/tasks/$TASK_NAME
 
 ```bash
 cd $EXP_DIR/scripts/pretrain/$TRAINING_SCRIPT_DIR/
-bash run_train.sh <RESERVATION_ID> <EXPERIMENT_ID> <EXPERIMENT_DIR> <TASK_NAME> <WANDB_PROJECT> <NUM_NODES>
+bash run_train.sh <GROUP_ID> <RESERVATION_ID> <JOB_NAME> <EXPERIMENT_DIR> <TASK_NAME> <WANDB_PROJECT> <NUM_NODES> <WALLTIME>
 
 # Example:
-bash run_train.sh R0123456789 0123 /path/to/0123_experiment task_name 0123_experiment 32
+bash run_train.sh gcg51557 R0123456789 0123_pretrain /path/to/0123_experiment task_name 0123_experiment 32 720:00:00
 ```
 
 CLIからは以下の引数を指定する
 
+- `<GROUP_ID>`: ABCI グループ ID
 - `<RESERVATION_ID>`: ABCI の予約キュー ID
-- `<EXPERIMENT_ID>`: 実験の識別子 (e.g. `0123`)
-- `<EXPERIMENT_DIR>`: 実験ディレクトリのパス (e.g. `/home/ach17726fj/experiments/0123_experiment`)
+- `<JOB_NAME>`: ジョブ名 (e.g., `0123_pretrain`)
+- `<EXPERIMENT_DIR>`: 実験ディレクトリのパス (e.g. `/path/to/0123_experiment`)
 - `<TASK_NAME>`: タスクディレクトリ名 (e.g. `task_name`)
 - `<WANDB_PROJECT>`: WandB に記録するプロジェクト名 (e.g. `0123_experiment`)
 - `<NUM_NODES>`: 使用するノード数 (e.g. `32`)
+- `<WALLTIME>`: ジョブの制限時間 (e.g., `720:00:00`)
 
 ### Training Configuration
 
@@ -70,5 +72,3 @@ CLIからは以下の引数を指定する
   - Megatron-LM の `pretrain_gpt.py` に渡す引数をこのファイル内の変数に定義する
 - `train_data.sh`: 学習データのパス及び利用するトークン数などを定義するスクリプト
   - Megatron-LM の `--train-data` 引数に渡す値をこのファイル内の `$TRAIN_DATA_PATH` 変数に定義する
-- `train_iters.txt`: 学習イテレーション数を定義するファイル
-  - 学習するイテレーション数を記載し、他には何も記載しない

pretrain/scripts/v4-upstream-training-template/convert/qsub_convert.sh

@@ -0,0 +1,153 @@
+#!/bin/bash
+
+# Predefined variables:
+# * EXPERIMENT_DIR: Experiment directory 
+# * TASK_NAME: Name of the task
+# * ITER: Target iteration number
+# * TOKENIZER_DIR: Directory of the tokenizer model
+
+cd ${PBS_O_WORKDIR}
+
+TASK_DIR=${EXPERIMENT_DIR}/tasks/${TASK_NAME}
+JOB_ID=${PBS_JOBID%%.*}
+
+mkdir -p ${TASK_DIR}/logs
+LOGFILE=${TASK_DIR}/logs/convert-${JOB_ID}.out
+ERRFILE=${TASK_DIR}/logs/convert-${JOB_ID}.err
+exec > ${LOGFILE} 2> ${ERRFILE}
+
+set -eu -o pipefail
+
+ENV_DIR=${EXPERIMENT_DIR}/env
+SCRIPT_DIR=${EXPERIMENT_DIR}/scripts
+
+# Load common environment variables
+source ${ENV_DIR}/scripts/environment.sh
+
+# Load modules
+source /etc/profile.d/modules.sh
+module load cuda/${PRETRAIN_CUDA_VERSION}/${PRETRAIN_CUDA_VERSION}.${PRETRAIN_CUDA_VERSION_PATCH}
+module load cudnn/${PRETRAIN_CUDNN_VERSION}/${PRETRAIN_CUDNN_VERSION_WITH_PATCH}
+module load hpcx/${PRETRAIN_HPCX_VERSION}
+module load nccl/${PRETRAIN_NCCL_VERSION}/${PRETRAIN_NCCL_VERSION_WITH_PATCH}
+# For logging
+module list
+
+# Load Python venv
+source ${ENV_DIR}/venv/bin/activate
+
+## Debug/logging flags
+export LOGLEVEL=INFO
+export NCCL_DEBUG=WARN
+export NCCL_DEBUG_SUBSYS=WARN
+export PYTHONFAULTHANDLER=1
+export CUDA_DEVICE_MAX_CONNECTIONS=1
+export CUDA_LAUNCH_BLOCKING=0
+export CUDNN_LOGDEST_DBG=stderr
+export CUDNN_LOGERR_DBG=1
+
+export MASTER_ADDR=$(head -n 1 ${PBS_NODEFILE} | hostname -f)
+export MASTER_PORT=$((10000 + RANDOM % 1000))
+echo "hostname: ${MASTER_ADDR}"
+
+ITER_NAME=iter_$(printf %07d ${ITER})  # iter_0123456
+
+MEGATRON_PATH=${ENV_DIR}/src/Megatron-LM
+OUTPUT_DIR=${TASK_DIR}/checkpoints_hf/${ITER_NAME}
+
+# Setup working directory
+TEMP_DIR=$(mktemp -d "${TASK_DIR}/tmp_converter_${JOB_ID}_XXXXXX")
+echo "TEMP_DIR=${TEMP_DIR}"
+function rm_tempdir {
+    if [ -e ${TEMP_DIR} ]; then
+        echo "Removing temporary directory: ${TEMP_DIR}"
+        rm -rf ${TEMP_DIR}
+        echo "Done removing"
+    fi
+}
+trap rm_tempdir EXIT
+trap 'trap - EXIT; rm_tempdir; exit 1' INT PIPE TERM
+
+########
+# Step 1: Convert `torch_dist` format to `torch`
+# This process requires to launch the trainer script with the same parallelism configs.
+########
+echo "Start converting: torch_dist --> torch"
+
+# Prepare source model at specific iteration
+mkdir ${TEMP_DIR}/torch_dist
+echo ${ITER} > ${TEMP_DIR}/torch_dist/latest_checkpointed_iteration.txt
+ln -s ${TASK_DIR}/checkpoints/${ITER_NAME} ${TEMP_DIR}/torch_dist/${ITER_NAME}
+
+# Training data: TRAIN_DATA_PATH
+source ${TASK_DIR}/train_data.sh
+
+# Synthesize all model params: ALL_PARAMS
+# Requires TRAIN_ITERS and TRAIN_DATA_PATH
+source ${TASK_DIR}/params.sh
+
+# Add params for model conversion
+ALL_PARAMS+=(
+    --load ${TEMP_DIR}/torch_dist
+    --ckpt-convert-format torch
+    --ckpt-convert-save ${TEMP_DIR}
+)
+
+echo "ALL_PARAMS: ${ALL_PARAMS[@]}"
+
+NUM_NODES=$(wc -l < ${PBS_NODEFILE})
+NUM_GPUS_PER_NODE=8
+NUM_GPUS=$((${NUM_NODES} * ${NUM_GPUS_PER_NODE}))
+echo "nnodes: ${NUM_NODES}; ngpus: ${NUM_GPUS}"
+echo NUM_NODES=${NUM_NODES}
+echo NUM_GPUS_PER_NODE=${NUM_GPUS_PER_NODE}
+echo NUM_GPUS=${NUM_GPUS}
+
+# Launch trainer script to convert the checkpoint
+mpirun \
+    --display-allocation \
+    --report-bindings \
+    --oversubscribe \
+    -np ${NUM_GPUS} \
+    --npernode ${NUM_GPUS_PER_NODE} \
+    -bind-to none \
+    -map-by slot \
+    python \
+        ${MEGATRON_PATH}/pretrain_gpt.py \
+        ${ALL_PARAMS[@]}
+
+echo "Files created by the Step 1:"
+find ${TEMP_DIR}/torch | sort
+
+########
+# Step 2: Convert `torch` to `Hugging Face`
+########
+
+echo "Start converting: torch --> hf"
+
+python ${MEGATRON_PATH}/tools/checkpoint/convert.py \
+    --model-type GPT \
+    --loader mcore \
+    --saver llmjp4_hf \
+    --load-dir ${TEMP_DIR}/torch \
+    --save-dir ${OUTPUT_DIR} \
+    --hf-tokenizer-path ${TOKENIZER_DIR} \
+    --save-dtype bfloat16 \
+    --loader-transformer-impl transformer_engine \
+    --megatron-path ${MEGATRON_PATH}
+
+echo "Files created by the Step 2:"
+find ${OUTPUT_DIR} | sort
+
+########
+# Step 3: Replace tokenizer model
+########
+
+echo "Start replacing tokenizer"
+
+cp ${TOKENIZER_DIR}/* ${OUTPUT_DIR}
+
+echo "Final model files:"
+find ${OUTPUT_DIR} | sort
+
+echo "Done processing"

pretrain/scripts/v4-upstream-training-template/convert/run_convert.sh

@@ -0,0 +1,53 @@
+#!/bin/bash
+
+set -eu -o pipefail
+
+if [ $# -ne 6 ]; then
+    >&2 echo "Usage: $0 <RESERVATION_ID> <EXPERIMENT_ID> <EXPERIMENT_DIR> <TASK_NAME> <TOKENIZER_DIR> <NUM_NODES>"
+    >&2 echo "Example: $0 R0123456789 0123 /path/to/0123_experiment task_name /path/to/tokenizer 1"
+    exit 1
+fi
+
+# NOTE(odashi):
+# Some variables are not used, but maintained for compatibility with training script.
+RESERVATION_ID=$1; shift
+EXPERIMENT_ID=$1; shift
+EXPERIMENT_DIR=$1; shift
+TASK_NAME=$1; shift
+TOKENIZER_DIR=$1; shift
+NUM_NODES=$1; shift
+
+# This directory
+SCRIPT_ROOT=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd )
+
+TASK_DIR=${EXPERIMENT_DIR}/tasks/${TASK_NAME}
+LAST_ITER=$(cat ${TASK_DIR}/checkpoints/latest_checkpointed_iteration.txt)
+
+dependency=()
+
+for iter in $(seq 1000 1000 ${LAST_ITER}); do
+    if [ ! -e ${TASK_DIR}/checkpoints/iter_$(printf '%07d' ${iter}) ]; then
+        #echo "Skip iter=${iter}: Source model does not exist."
+        continue
+    fi
+    if [ -e ${TASK_DIR}/checkpoints_hf/iter_$(printf '%07d' ${iter})/tokenizer.json ]; then
+        #echo "Skip iter=${iter}: Converted model already exists."
+        continue
+    fi
+
+    # NOTE(odashi): RTYPE=rt_HG doesn't work for 8B models.
+    job_id=$(qsub \
+        ${dependency[@]} \
+        -P gcg51557 \
+        -q ${RESERVATION_ID} \
+        -N ${EXPERIMENT_ID}_convert \
+        -l select=${NUM_NODES},walltime=6:00:00 \
+        -v RTYPE=rt_HF,EXPERIMENT_DIR=${EXPERIMENT_DIR},TASK_NAME=${TASK_NAME},ITER=${iter},TOKENIZER_DIR=${TOKENIZER_DIR} \
+        -o /dev/null \
+        -e /dev/null \
+        -m n \
+        ${SCRIPT_ROOT}/qsub_convert.sh
+    )
+    echo "Submitted iter=${iter}: job_id=${job_id}"
+    #dependency=(-W depend=afterany:${job_id})
+done

pretrain/scripts/v4-upstream-training-template/pretrain/qsub_train.sh

@@ -5,6 +5,8 @@
 # * TASK_NAME: Name of the task
 # * WANDB_PROJECT: W&B project name
 
+set -eu -o pipefail
+
 cd ${PBS_O_WORKDIR}
 
 TASK_DIR=${EXPERIMENT_DIR}/tasks/${TASK_NAME}
@@ -15,8 +17,6 @@ LOGFILE=${TASK_DIR}/logs/pretrain-${JOB_ID}.out
 ERRFILE=${TASK_DIR}/logs/pretrain-${JOB_ID}.err
 exec > ${LOGFILE} 2> ${ERRFILE}
 
-set -eu -o pipefail
-
 ENV_DIR=${EXPERIMENT_DIR}/env
 SCRIPT_DIR=${EXPERIMENT_DIR}/scripts
 
@@ -55,21 +55,19 @@ echo "hostname: ${MASTER_ADDR}"
 NUM_NODES=$(wc -l < ${PBS_NODEFILE})
 NUM_GPUS_PER_NODE=8
 NUM_GPUS=$((${NUM_NODES} * ${NUM_GPUS_PER_NODE}))
-echo "nnodes: ${NUM_NODES}; ngpus: ${NUM_GPUS}"
 echo NUM_NODES=${NUM_NODES}
 echo NUM_GPUS_PER_NODE=${NUM_GPUS_PER_NODE}
 echo NUM_GPUS=${NUM_GPUS}
 
+# For logging
+echo "PBS_NODEFILE:"
 cat ${PBS_NODEFILE}
 
-# Training steps
-TRAIN_ITERS=$(cat ${TASK_DIR}/train_iters.txt)
-
-# Training data: TRAIN_DATA_PATH
+# Load training data: TRAIN_DATA_PATH
 source ${TASK_DIR}/train_data.sh
 
-# Synthesize all model params: ALL_PARAMS
-# Requires TRAIN_ITERS and TRAIN_DATA_PATH
+# Load model params: ALL_PARAMS
+# Requires TRAIN_DATA_PATH
 source ${TASK_DIR}/params.sh
 
 # Add logging params
@@ -82,15 +80,30 @@ ALL_PARAMS+=(
 )
 
 # Add Checkpointing params
+BASE_CHECKPOINT_DIR=${TASK_DIR}/base_checkpoints
 TASK_CHECKPOINT_DIR=${TASK_DIR}/checkpoints
+
+if [ -e ${TASK_CHECKPOINT_DIR}/latest_checkpointed_iteration.txt ]; then
+    echo "Resume from the last checkpoint in this task"
+    LOAD_DIR=${TASK_CHECKPOINT_DIR}
+elif [ -e ${BASE_CHECKPOINT_DIR}/latest_checkpointed_iteration.txt ]; then
+    echo "Start from the base checkpoint"
+    LOAD_DIR=${BASE_CHECKPOINT_DIR}
+else
+    echo "Start from scratch"
+    LOAD_DIR=${TASK_CHECKPOINT_DIR}
+fi
+
 ALL_PARAMS+=(
-    --load ${TASK_CHECKPOINT_DIR}
+    --load ${LOAD_DIR}
     --save ${TASK_CHECKPOINT_DIR}
     --save-interval 1000
 )
 
+# For logging
 echo "ALL_PARAMS: ${ALL_PARAMS[@]}"
 
+echo "Start training..."
 mpirun \
     --display-allocation \
     --report-bindings \
@@ -102,3 +115,5 @@ mpirun \
     python \
         ${ENV_DIR}/src/Megatron-LM/pretrain_gpt.py \
         ${ALL_PARAMS[@]}
+
+echo "Training completed successfully."

pretrain/scripts/v4-upstream-training-template/pretrain/run_train.sh

@@ -2,31 +2,30 @@
 
 set -eu -o pipefail
 
-if [ $# -ne 6 ]; then
-    >&2 echo "Usage: $0 <RESERVATION_ID> <EXPERIMENT_ID> <EXPERIMENT_DIR> <TASK_NAME> <WANDB_PROJECT> <NUM_NODES>"
-    >&2 echo "Example: $0 R0123456789 0123 /path/to/0123_experiment task_name 0123_experiment 32"
+if [ $# -ne 8 ]; then
+    >&2 echo "Usage: $0 <GROUP_ID> <RESERVATION_ID> <JOB_NAME> <EXPERIMENT_DIR> <TASK_NAME> <WANDB_PROJECT> <NUM_NODES> <WALLTIME>"
+    >&2 echo "Example: $0 gcg51557 R0123456789 0123 /path/to/0123_experiment task_name 0123_experiment 32 720:00:00"
     exit 1
 fi
 
+GROUP_ID=$1; shift
 RESERVATION_ID=$1; shift
-EXPERIMENT_ID=$1; shift
+JOB_NAME=$1; shift
 EXPERIMENT_DIR=$1; shift
 TASK_NAME=$1; shift
 WANDB_PROJECT=$1; shift
 NUM_NODES=$1; shift
+WALLTIME=$1; shift
 
 # This directory
 SCRIPT_ROOT=$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd )
 
-WALLTIME=720:00:00 # 30 days
-# WALLTIME=01:00:00 # 1 hour
-
 qsub \
-    -P gcg51557 \
+    -P ${GROUP_ID} \
     -q ${RESERVATION_ID} \
-    -N ${EXPERIMENT_ID}_pretrain \
+    -N ${JOB_NAME} \
     -l select=${NUM_NODES},walltime=${WALLTIME} \
-    -v RTYPE=rt_HF,EXPERIMENT_DIR=${EXPERIMENT_DIR},TASK_NAME=${TASK_NAME},WANDB_PROJECT=${WANDB_PROJECT} \
+    -v RTYPE=rt_HF,USE_SSH=1,EXPERIMENT_DIR=${EXPERIMENT_DIR},TASK_NAME=${TASK_NAME},WANDB_PROJECT=${WANDB_PROJECT} \
     -o /dev/null \
     -e /dev/null \
     -m n \

pretrain/scripts/v4-upstream-training-template/task_template/params.sh

@@ -47,9 +47,9 @@ ALL_PARAMS+=(
 
 # Scheduler
 ALL_PARAMS+=(
-    --train-iters ${TRAIN_ITERS}
+    --train-iters 100000
     --lr-warmup-iters 2000
-    --lr-decay-iters ${TRAIN_ITERS}
+    --lr-decay-iters 100000
     --lr-decay-style cosine
     --eval-interval 999999999
     --eval-iters 0