116 lines
3.7 KiB
Bash
116 lines
3.7 KiB
Bash
#!/usr/bin/env bash
|
||
# =============================================================================
|
||
# launch_korean_3b.sh — 8-GPU FP8 pretraining launcher for 3B Korean LLM
|
||
#
|
||
# Usage:
|
||
# bash scripts/launch_korean_3b.sh # full run (~60B tokens)
|
||
# bash scripts/launch_korean_3b.sh --max_steps 50 # quick benchmark
|
||
# bash scripts/launch_korean_3b.sh --resume checkpoints/korean_3b_fp8_run1/checkpoint-XXXXX
|
||
#
|
||
# Effective batch size: 8 (local) × 8 GPU × 4 (grad_accum) × 4096 (seq_len)
|
||
# = 1,048,576 tokens / step
|
||
# =============================================================================
|
||
set -euo pipefail
|
||
|
||
RUN_NAME="${RUN_NAME:-korean_3b_fp8_run1}"
|
||
CONFIG="${CONFIG:-configs/3b_pretrain.yaml}"
|
||
CKPT_DIR="checkpoints/${RUN_NAME}"
|
||
LOG_FILE="${CKPT_DIR}/train.log"
|
||
NPROC=8
|
||
MASTER_PORT="${MASTER_PORT:-29502}"
|
||
|
||
MAX_STEPS=57000
|
||
BATCH_SIZE=4
|
||
GRAD_ACCUM=8
|
||
LR=1.5e-4
|
||
WARMUP_STEPS=2000
|
||
SEED=42
|
||
|
||
EXTRA_ARGS="$@"
|
||
|
||
# ---- B200 / NVSwitch NCCL tuning -------------------------------------------
|
||
export NCCL_IB_DISABLE=1
|
||
export NCCL_ALGO=Ring
|
||
export NCCL_PROTO=Simple
|
||
export NCCL_MIN_NCHANNELS=16
|
||
export NCCL_MAX_NCHANNELS=16
|
||
export NCCL_BUFFSIZE=67108864
|
||
export OMP_NUM_THREADS=4
|
||
export MKL_NUM_THREADS=4
|
||
|
||
# cd FIRST — 이후 상대경로 체크가 프로젝트 루트 기준으로 동작
|
||
cd "$(dirname "$0")/.."
|
||
|
||
# TRAIN_DATA fallback: cd 이후에 상대경로 체크
|
||
if [[ -f "data/merged_3b_train.bin" ]]; then
|
||
TRAIN_DATA="${TRAIN_DATA:-data/merged_3b_train.bin}"
|
||
echo "Using merged training data: data/merged_3b_train.bin"
|
||
elif [[ -f "data/korean_train.bin" ]]; then
|
||
TRAIN_DATA="${TRAIN_DATA:-data/korean_train.bin}"
|
||
echo "Using fallback training data: data/korean_train.bin"
|
||
else
|
||
echo "ERROR: No training data found (data/merged_3b_train.bin or data/korean_train.bin)"
|
||
exit 1
|
||
fi
|
||
|
||
# VAL_DATA fallback: cd 이후에 상대경로 체크
|
||
VAL_DATA="${VAL_DATA:-data/merged_3b_val.bin}"
|
||
if [[ ! -f "${VAL_DATA}" ]]; then
|
||
VAL_DATA="data/korean_val.bin"
|
||
fi
|
||
|
||
if [[ ! -f "${TRAIN_DATA}" ]]; then
|
||
echo "ERROR: Training data not found: ${TRAIN_DATA}"
|
||
exit 1
|
||
fi
|
||
if [[ ! -f "${VAL_DATA}" ]]; then
|
||
echo "ERROR: Validation data not found: ${VAL_DATA}"
|
||
exit 1
|
||
fi
|
||
|
||
mkdir -p "${CKPT_DIR}"
|
||
|
||
echo "=================================================================="
|
||
echo " Run name : ${RUN_NAME}"
|
||
echo " Config : ${CONFIG}"
|
||
echo " Train data : ${TRAIN_DATA}"
|
||
echo " CKPT dir : ${CKPT_DIR}"
|
||
echo " Max steps : ${MAX_STEPS}"
|
||
echo " LR : ${LR}"
|
||
echo " Batch size : ${BATCH_SIZE} (local) × ${NPROC} GPU × ${GRAD_ACCUM} grad_accum"
|
||
echo " Started : $(date)"
|
||
echo "=================================================================="
|
||
|
||
export PYTHONWARNINGS="ignore::UserWarning:torch.library"
|
||
|
||
torchrun \
|
||
--nproc_per_node=${NPROC} \
|
||
--master_port=${MASTER_PORT} \
|
||
train/pretrain.py \
|
||
--config "${CONFIG}" \
|
||
--train_data "${TRAIN_DATA}" \
|
||
--val_data "${VAL_DATA}" \
|
||
--checkpoint_dir "${CKPT_DIR}" \
|
||
--log_file "${LOG_FILE}" \
|
||
--max_steps ${MAX_STEPS} \
|
||
--batch_size ${BATCH_SIZE} \
|
||
--lr ${LR} \
|
||
--grad_accum ${GRAD_ACCUM} \
|
||
--warmup_steps ${WARMUP_STEPS} \
|
||
--seed ${SEED} \
|
||
${EXTRA_ARGS} \
|
||
2>&1 | grep -v "UserWarning" \
|
||
| grep -v "Warning only once" \
|
||
| grep -v "Overriding a previously" \
|
||
| grep -v "dispatch key:" \
|
||
| grep -v "previous kernel:" \
|
||
| grep -v "new kernel:" \
|
||
| grep -v "operator: flash_attn" \
|
||
| grep -v "registered at /usr/local" \
|
||
| grep -v "self.m.impl" \
|
||
| tee -a "${LOG_FILE}"
|
||
|
||
echo "=================================================================="
|
||
echo " Done : $(date)"
|
||
echo "=================================================================="
|