Support stage1-2 MoE resume on g0050

scripts/resume_stage1_2_pretrain_8192_8gpu_mbs14.sh

@@ -0,0 +1,119 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+# Host path layout on g0050 after /ssd was replaced by /data.
+HOST_WORKSPACE="${HOST_WORKSPACE:-/data/workspace/yi}"
+CONTAINER_WORKSPACE="${CONTAINER_WORKSPACE:-/ssd/workspace/yi}"
+
+REPO_ROOT="${REPO_ROOT:-${CONTAINER_WORKSPACE}/laoyao_2b_moe}"
+HOST_REPO_ROOT="${HOST_REPO_ROOT:-${HOST_WORKSPACE}/laoyao_2b_moe}"
+IMAGE="${IMAGE:-laoyao/nemo-megatron:26.06-flashattn4}"
+RUN_NAME="${RUN_NAME:-stage1_2_8192_8gpu_dp8_mbs14_full_recompute_general55_code15_math15_science15}"
+RUN_DIR="${RUN_DIR:-${REPO_ROOT}/runs/${RUN_NAME}}"
+CKPT_DIR="${CKPT_DIR:-${RUN_DIR}/checkpoints}"
+TENSORBOARD_DIR="${TENSORBOARD_DIR:-${RUN_DIR}/tensorboard}"
+LOAD_DIR="${LOAD_DIR:-${CKPT_DIR}}"
+DATA_MANIFEST="${DATA_MANIFEST:-${CONTAINER_WORKSPACE}/laoyao_2b_moe_pretraining_dataset/megatron_bridge/stage2_8192_by_category_general55_code15_math15_science15_100b_v1/stage1_2_weighted_manifest.json}"
+VALIDATION_PREFIX="${VALIDATION_PREFIX:-${REPO_ROOT}/dataset/val/megatron_8192_glm52/heldout_2p8k_text_document}"
+TRAIN_ITERS="${TRAIN_ITERS:-281245}"
+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}"
+LR="${LR:-3.0e-4}"
+MIN_LR="${MIN_LR:-5.0e-6}"
+WARMUP_FRACTION="${WARMUP_FRACTION:-0.01}"
+LOG_FILE="${LOG_FILE:-/tmp/laoyao_stage1_2_8192_8gpu_mbs14.log}"
+CONTAINER_NAME="${CONTAINER_NAME:-laoyao_stage1_2_pretrain}"
+
+HOST_DATA_MANIFEST="${DATA_MANIFEST/${CONTAINER_WORKSPACE}/${HOST_WORKSPACE}}"
+HOST_LOAD_DIR="${LOAD_DIR/${CONTAINER_WORKSPACE}/${HOST_WORKSPACE}}"
+HOST_VALIDATION_PREFIX="${VALIDATION_PREFIX/${CONTAINER_WORKSPACE}/${HOST_WORKSPACE}}"
+HOST_CKPT_DIR="${CKPT_DIR/${CONTAINER_WORKSPACE}/${HOST_WORKSPACE}}"
+HOST_TENSORBOARD_DIR="${TENSORBOARD_DIR/${CONTAINER_WORKSPACE}/${HOST_WORKSPACE}}"
+
+if [[ ! -d "${HOST_REPO_ROOT}" ]]; then
+  echo "ERROR: host repo dir not found: ${HOST_REPO_ROOT}" >&2
+  exit 1
+fi
+if [[ ! -f "${HOST_DATA_MANIFEST}" ]]; then
+  echo "ERROR: host data manifest not found: ${HOST_DATA_MANIFEST}" >&2
+  exit 1
+fi
+if [[ ! -d "${HOST_LOAD_DIR}" ]]; then
+  echo "ERROR: host load checkpoint dir not found: ${HOST_LOAD_DIR}" >&2
+  exit 1
+fi
+if [[ ! -f "${HOST_VALIDATION_PREFIX}.idx" ]]; then
+  echo "WARNING: host validation prefix not found: ${HOST_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 "${HOST_CKPT_DIR}" "${HOST_TENSORBOARD_DIR}"
+
+echo "Starting Laoyao 2B MoE stage1-2 pretraining"
+echo "host_workspace: ${HOST_WORKSPACE}"
+echo "container_workspace: ${CONTAINER_WORKSPACE}"
+echo "repo: ${REPO_ROOT}"
+echo "data_manifest: ${DATA_MANIFEST}"
+echo "save_checkpoint_dir: ${CKPT_DIR}"
+echo "load_checkpoint_dir: ${LOAD_DIR}"
+echo "load_checkpoint_iteration: $(cat "${HOST_LOAD_DIR}/latest_checkpointed_iteration.txt" 2>/dev/null || echo unknown)"
+echo "train_iters: ${TRAIN_ITERS}"
+echo "log_file: ${LOG_FILE}"
+
+cd "${HOST_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 "${HOST_WORKSPACE}:${CONTAINER_WORKSPACE}" \
+  -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}' \
+      --lr '${LR}' \
+      --min-lr '${MIN_LR}' \
+      --warmup-fraction '${WARMUP_FRACTION}' \
+      --save-dir '${CKPT_DIR}' \
+      --load-dir '${LOAD_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"