#!/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"
