97 lines
3.4 KiB
Bash
Executable File
97 lines
3.4 KiB
Bash
Executable File
#!/usr/bin/env bash
|
|
set -euo pipefail
|
|
|
|
REPO_ROOT="${REPO_ROOT:-/ssd/workspace/yi/laoyao_2b_moe}"
|
|
IMAGE="${IMAGE:-laoyao/nemo-megatron:26.06-flashattn4}"
|
|
RUN_NAME="${RUN_NAME:-pretrain_8192_8gpu_dp8_mbs14_full_recompute_weighted_heldoutval_resume10000}"
|
|
RUN_DIR="${RUN_DIR:-${REPO_ROOT}/runs/${RUN_NAME}}"
|
|
CKPT_DIR="${CKPT_DIR:-${RUN_DIR}/checkpoints}"
|
|
TENSORBOARD_DIR="${TENSORBOARD_DIR:-${RUN_DIR}/tensorboard}"
|
|
DATA_MANIFEST="${DATA_MANIFEST:-/ssd/workspace/yi/laoyao_2b_moe_pretraining_dataset/megatron_bridge/pretrain_8192_glm52_direct_v1/manifest.json}"
|
|
VALIDATION_PREFIX="${VALIDATION_PREFIX:-${REPO_ROOT}/dataset/val/megatron_8192_glm52/heldout_2p8k_text_document}"
|
|
TRAIN_ITERS="${TRAIN_ITERS:-184209}"
|
|
SEQ_LENGTH="${SEQ_LENGTH:-8192}"
|
|
MICRO_BATCH_SIZE="${MICRO_BATCH_SIZE:-14}"
|
|
GLOBAL_BATCH_SIZE="${GLOBAL_BATCH_SIZE:-112}"
|
|
SAVE_INTERVAL="${SAVE_INTERVAL:-2500}"
|
|
EVAL_INTERVAL="${EVAL_INTERVAL:-15000}"
|
|
EVAL_ITERS="${EVAL_ITERS:-10}"
|
|
DATASET_WORKERS="${DATASET_WORKERS:-4}"
|
|
LOG_FILE="${LOG_FILE:-/tmp/laoyao_pretrain_8192_8gpu_dp8_mbs14_full_recompute_weighted_heldoutval_resume10000_resume.log}"
|
|
CONTAINER_NAME="${CONTAINER_NAME:-laoyao_pretrain_resume}"
|
|
|
|
if [[ ! -d "${REPO_ROOT}" ]]; then
|
|
echo "ERROR: repo dir not found: ${REPO_ROOT}" >&2
|
|
exit 1
|
|
fi
|
|
if [[ ! -f "${DATA_MANIFEST}" ]]; then
|
|
echo "ERROR: data manifest not found: ${DATA_MANIFEST}" >&2
|
|
exit 1
|
|
fi
|
|
if [[ ! -d "${CKPT_DIR}" ]]; then
|
|
echo "ERROR: checkpoint dir not found: ${CKPT_DIR}" >&2
|
|
exit 1
|
|
fi
|
|
if [[ ! -f "${VALIDATION_PREFIX}.idx" ]]; then
|
|
echo "WARNING: validation prefix not found: ${VALIDATION_PREFIX}" >&2
|
|
fi
|
|
|
|
if docker ps --format '{{.Names}}' | grep -Fxq "${CONTAINER_NAME}"; then
|
|
echo "ERROR: container already running: ${CONTAINER_NAME}" >&2
|
|
exit 1
|
|
fi
|
|
|
|
mkdir -p "${CKPT_DIR}" "${TENSORBOARD_DIR}"
|
|
|
|
echo "Resuming Laoyao 2B MoE training"
|
|
echo "repo: ${REPO_ROOT}"
|
|
echo "checkpoint_dir: ${CKPT_DIR}"
|
|
echo "latest_checkpointed_iteration: $(cat "${CKPT_DIR}/latest_checkpointed_iteration.txt" 2>/dev/null || echo unknown)"
|
|
echo "log_file: ${LOG_FILE}"
|
|
|
|
cd "${REPO_ROOT}"
|
|
|
|
nohup docker run --rm --name "${CONTAINER_NAME}" \
|
|
--gpus all \
|
|
--ipc=host \
|
|
--network=host \
|
|
--ulimit memlock=-1 \
|
|
--ulimit stack=67108864 \
|
|
-e CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 \
|
|
-v /ssd:/ssd \
|
|
-w "${REPO_ROOT}" \
|
|
"${IMAGE}" \
|
|
bash -lc "torchrun --nproc_per_node=8 \
|
|
training/megatron_bridge/laoyao_2b_moe_pretrain.py \
|
|
--data-manifest '${DATA_MANIFEST}' \
|
|
--validation-prefix '${VALIDATION_PREFIX}' \
|
|
--seq-length '${SEQ_LENGTH}' \
|
|
--train-iters '${TRAIN_ITERS}' \
|
|
--micro-batch-size '${MICRO_BATCH_SIZE}' \
|
|
--global-batch-size '${GLOBAL_BATCH_SIZE}' \
|
|
--tensor-parallel 1 \
|
|
--pipeline-parallel 1 \
|
|
--expert-parallel 1 \
|
|
--context-parallel 1 \
|
|
--split 999,1,0 \
|
|
--dataset-workers '${DATASET_WORKERS}' \
|
|
--save-dir '${CKPT_DIR}' \
|
|
--load-dir '${CKPT_DIR}' \
|
|
--tensorboard-dir '${TENSORBOARD_DIR}' \
|
|
--save-interval '${SAVE_INTERVAL}' \
|
|
--keep-last-checkpoints 10 \
|
|
--log-interval 10 \
|
|
--eval-interval '${EVAL_INTERVAL}' \
|
|
--eval-iters '${EVAL_ITERS}' \
|
|
--use-distributed-optimizer \
|
|
--overlap-grad-reduce \
|
|
--overlap-param-gather \
|
|
--recompute-granularity full \
|
|
--recompute-method uniform \
|
|
--recompute-num-layers 1" \
|
|
> "${LOG_FILE}" 2>&1 &
|
|
|
|
echo "Launched. Check with:"
|
|
echo " tail -f ${LOG_FILE}"
|
|
echo " docker ps"
|